From f75f581a57d0c5f4ff0b635d106e95f79a8bf15c Mon Sep 17 00:00:00 2001 From: sufubao Date: Fri, 13 Nov 2026 18:54:39 +0800 Subject: [PATCH 1/6] test: slim unit tests --- unit_tests/README.md | 26 + .../common/att_resources/manager_fixtures.py | 45 - .../common/att_resources/model_helpers.py | 342 ---- .../att_resources/test_backend_startup.py | 374 ---- .../att_resources/test_bounded_state.py | 440 ----- .../common/att_resources/test_capacity.py | 1113 ----------- .../att_resources/test_compute_binding.py | 147 -- .../att_resources/test_dsv4_integration.py | 287 --- .../common/att_resources/test_dsv4_state.py | 505 ----- .../att_resources/test_fp8_compute_binding.py | 201 -- .../test_fp8_e8m0_packed_mla_state.py | 235 --- .../att_resources/test_full_att_cpu_kernel.py | 412 ----- .../test_full_att_dp_transfer.py | 221 --- .../att_resources/test_full_att_transfer.py | 338 ---- .../common/att_resources/test_full_state.py | 283 --- .../test_full_storage_initialization.py | 120 -- .../att_resources/test_gqa_layout_identity.py | 114 -- .../att_resources/test_gqa_pd_kernel.py | 283 --- .../common/att_resources/test_gqa_transfer.py | 192 -- .../test_hybrid_layer_mapping.py | 130 -- .../common/att_resources/test_layer_groups.py | 164 -- .../test_layer_state_selection.py | 216 --- .../att_resources/test_manager_helpers.py | 316 ---- .../att_resources/test_model_calibration.py | 261 --- .../test_model_config_normalization.py | 235 --- .../test_model_initialization.py | 185 -- .../att_resources/test_model_resources.py | 214 --- .../test_model_storage_definition.py | 276 --- .../att_resources/test_mtp_state_commit.py | 114 -- .../att_resources/test_network_transport.py | 320 ---- .../att_resources/test_no_attention_draft.py | 84 - .../att_resources/test_page_pool_storage.py | 167 -- .../att_resources/test_paged_storage.py | 22 - .../test_quantized_full_state.py | 388 ---- .../test_quantized_paged_storage.py | 37 - .../att_resources/test_role_transfers.py | 149 -- .../att_resources/test_runtime_transport.py | 571 ------ .../att_resources/test_state_creation.py | 235 --- .../att_resources/test_state_identity.py | 356 ---- .../att_resources/test_state_queries.py | 189 -- .../common/att_resources/test_tp_config.py | 223 --- .../att_resources/test_transfer_views.py | 233 --- .../test_validate_model_cache.py | 41 - .../test_window_compressor_state.py | 195 -- unit_tests/common/basemodel/__init__.py | 0 .../basemodel/attention/fa3/test_neo.py | 119 -- .../flashinfer/test_decode_graph_state.py | 100 - .../flashinfer/test_decode_wrapper_init.py | 132 -- .../flashinfer/test_mla_cuda_graph_plan.py | 178 -- .../attention/flashinfer/test_mla_nope.py | 106 -- .../attention/flashinfer/test_workspace.py | 56 - .../basemodel/attention/linear/test_gdn.py | 438 ----- .../basemodel/attention/test_dsv4_flashmla.py | 377 ---- .../attention/test_fp8_logical_layer.py | 117 -- .../basemodel/attention/test_fp8_page_size.py | 45 - .../basemodel/attention/test_lazy_import.py | 248 --- .../attention/test_state_lifetime.py | 174 -- .../common/basemodel/test_att_sink_weight.py | 37 - .../basemodel/test_cuda_graph_autotune.py | 130 -- .../basemodel/test_cuda_graph_layout.py | 133 -- .../common/basemodel/test_hidden_collector.py | 193 -- .../common/basemodel/test_model_input.py | 212 --- .../common/basemodel/test_model_output.py | 344 ---- .../basemodel/test_moe_route_info_manager.py | 263 --- .../common/basemodel/test_mtp_manager.py | 133 -- .../common/basemodel/test_overlap_utils.py | 213 --- .../test_prefill_cuda_graph_state.py | 81 - .../gqa/mtp_diverse/test_mtp_diverse.py | 196 -- .../mtp_diverse/test_mtp_diverse_autotune.py | 301 --- .../gqa/mtp_diverse/test_stage2_autotune.py | 220 --- .../gqa/test_flash_decoding_autotune.py | 270 --- .../decode_att/int4kv/test_int4kv_autotune.py | 295 --- .../test_int8kv_flash_decoding_diverse.py | 139 -- ...st_int8kv_flash_decoding_diverse_stage1.py | 188 -- ...st_int8kv_flash_decoding_diverse_stage2.py | 192 -- ...st_int8kv_flash_decoding_diverse_stage3.py | 35 - .../decode_att/int8kv/test_normal_autotune.py | 300 --- .../test_context_attention_fwd_neo.py | 340 ---- .../test_context_flashattention_nopad.py | 104 -- .../test_context_flashattention_nopad1.py | 196 -- .../att/prefill_att/test_fa3_neo.py | 561 ------ .../kv_copy/test_mla_destindex_copy_kv.py | 42 - .../kv_copy/test_ppl_int4kv_copy_kv.py | 62 - .../kv_copy/test_ppl_int8kv_copy_kv.py | 86 - .../linear_att/test_acceptance_fusion.py | 272 --- .../linear_att/test_causal_conv1d_mtp.py | 1026 +---------- .../test_fused_recurrent_strided.py | 71 - .../linear_att/test_kda_autotune.py | 287 --- .../test_mtp_fused_recurrent_autotune.py | 374 ---- .../test_mtp_fused_recurrent_equiv.py | 167 -- .../linear_att/test_mtp_state_params.py | 65 - .../linear_att/test_replayssm.py | 561 +----- .../linear_att/test_replayssm_execution.py | 530 ------ .../linear_att/test_ssm_autotune_memory.py | 70 - .../triton_kernel/mhc/test_autotune.py | 105 -- .../basemodel/triton_kernel/mhc/test_mhc.py | 275 --- .../decode_att/test_gqa_flash_decoding.py | 112 -- .../triton_kernel/test_add_in_place.py | 31 - .../test_apply_constraint_mask.py | 104 -- .../triton_kernel/test_apply_invalid_token.py | 50 - .../triton_kernel/test_atomic_event.py | 28 - .../triton_kernel/test_diverse_utils.py | 48 - .../triton_kernel/test_dynamic_mtp_utils.py | 360 ---- .../basemodel/triton_kernel/test_fa3_utils.py | 82 - .../triton_kernel/test_gen_decode_params.py | 26 - .../test_gen_mtp_prefill_params.py | 44 - .../triton_kernel/test_gen_nsa_ks_ke.py | 87 - .../triton_kernel/test_gen_prefill_params.py | 43 - .../triton_kernel/test_gen_sampling_params.py | 86 - .../basemodel/triton_kernel/test_mtp_utils.py | 107 -- .../triton_kernel/test_multimodal_emb.py | 54 - .../triton_kernel/test_qk_rmsnorm_fused.py | 64 - .../test_redundancy_topk_ids_repair.py | 151 -- .../triton_kernel/test_repack_kv_index.py | 50 - .../triton_kernel/test_routing_capture.py | 137 -- .../test_runtime_scalar_specialization.py | 146 -- .../test_select_kv_index_from_req.py | 75 - .../triton_kernel/test_sp_pad_kernel.py | 36 - unit_tests/common/dsv4_test_utils.py | 10 - .../fused_moe/test_activation_config.py | 132 -- .../test_append_shared_expert_topk.py | 222 --- unit_tests/common/fused_moe/test_deepep.py | 311 ---- .../test_deepep_runtime_specialization.py | 61 - .../fused_moe/test_deepgemm_activation.py | 96 - .../fused_moe/test_ep_reduce_decode_output.py | 36 - .../fused_moe/test_grouped_fused_moe.py | 331 ---- .../test_grouped_fused_moe_ep_sms.py | 11 - .../fused_moe/test_grouped_fused_moe_speed.py | 51 - .../common/fused_moe/test_grouped_topk.py | 100 - .../test_moe_silu_and_mul_mix_quant_ep.py | 196 -- .../common/fused_moe/test_softmax_topk.py | 126 +- .../kv_trans_kernel/test_kv_trans_v2.py | 123 -- .../kv_trans_kernel/test_nixl_kv_trans.py | 141 -- .../common/quantization/test_no_quant.py | 43 - .../test_ue8m0_fp8_quantization.py | 513 ------ unit_tests/common/test_autotuner.py | 486 ----- .../common/test_cache_copy_control_flow.py | 98 - .../common/test_deepseek4_paged_cache.py | 1628 ----------------- .../test_hybrid_req_manager_pd_sharing.py | 145 -- unit_tests/common/test_linear_state.py | 334 ---- .../common/test_linear_state_contract.py | 241 --- unit_tests/common/test_mem_manager_ipc.py | 110 -- .../test_pinned_memory_non_blocking_copy.py | 39 - .../common/test_pinned_memory_stream.py | 162 -- .../common/test_req_manager_hold_page.py | 27 - unit_tests/common/test_req_manager_page.py | 67 - .../models/deepseek2/test_rope_repeat.py | 25 - .../triton_kernel/test_hadamard_transform.py | 83 - .../test_topk_index_to_mem_index.py | 24 - .../models/deepseek_v4/test_c128_workspace.py | 548 ------ unit_tests/models/deepseek_v4/test_kernels.py | 321 ---- unit_tests/models/deepseek_v4/test_model.py | 300 --- .../deepseek_v4/test_shared_expert_overlap.py | 247 --- .../models/deepseek_v4/test_tokenizer.py | 79 - unit_tests/models/deepseek_v4/test_vision.py | 350 ---- .../deepseek_v4/test_vision_integration.py | 321 ---- .../models/deepseek_v4/test_workspace.py | 97 - unit_tests/models/glm5_next/test_cache.py | 270 --- .../glm5_next/test_compact_topk_modes.py | 65 - .../glm5_next/test_direct_nsa_indices.py | 257 --- .../glm5_next/test_fp8_compile_cache.py | 94 - unit_tests/models/glm5_next/test_fp8_kv.py | 1073 ----------- .../models/glm5_next/test_fused_mtp_norm.py | 126 -- .../models/glm5_next/test_indexer_topk.py | 161 -- .../glm5_next/test_indexer_workspace.py | 78 - .../models/glm5_next/test_kda_conv_qkv.py | 186 -- unit_tests/models/glm5_next/test_kernels.py | 665 ------- .../models/glm5_next/test_kpool_gather.py | 141 -- .../test_kpool_runtime_specialization.py | 200 -- unit_tests/models/glm5_next/test_layers.py | 451 ----- unit_tests/models/glm5_next/test_mtp.py | 308 ---- .../glm5_next/test_mtp_prefill_cache.py | 278 --- unit_tests/models/glm5_next/test_pd_cache.py | 443 ----- .../glm5_next/test_prefill_workspace.py | 116 -- .../glm5_next/test_token_major_mla_bmm.py | 92 - unit_tests/models/glm5_next/test_visual.py | 216 --- unit_tests/models/phi3/test_rotary_emb.py | 50 - .../qwen2_vl/test_get_mrope_position_ids.py | 48 - .../models/qwen2_vl/test_infer_struct.py | 68 - unit_tests/models/qwen2_vl/test_mrope.py | 66 - .../models/qwen2_vl/test_rotary_pos_emb.py | 55 - .../models/qwen3-vl/test_deepstack_emb.py | 82 - .../qwen3next/test_shared_expert_gate.py | 22 - .../models/sensenova_loop/test_gdn_state.py | 151 -- unit_tests/models/sensenova_loop/test_loop.py | 376 ---- .../models/sensenova_loop/test_loop_mhc.py | 107 -- .../test_multimodal_tokenizer_positions.py | 98 - unit_tests/models/test_neo_registry.py | 30 - unit_tests/models/test_post_layer_output.py | 112 -- .../models/test_qwen3_dspark_model_output.py | 517 ------ .../models/test_qwen3_eagle_model_input.py | 100 - unit_tests/models/test_qwen_vl_tokenizer.py | 69 - unit_tests/models/test_yarn_rotary.py | 142 -- .../vit/test_flash_attention_forward.py | 139 -- .../core/objs/test_out_token_circlequeue.py | 127 -- unit_tests/server/core/objs/test_req.py | 127 -- .../server/core/objs/test_sampling_params.py | 104 -- unit_tests/server/core/objs/test_shm_array.py | 74 - .../server/core/objs/test_shm_req_manager.py | 13 - .../core/objs/test_stop_sequence_output.py | 951 ---------- .../core/objs/test_token_ids_validation.py | 109 -- .../test_multimodal_token_id_range.py | 61 - unit_tests/server/grammar_helpers.py | 154 -- unit_tests/server/httpserver/test_grammar.py | 340 ---- .../server/httpserver/test_pd_continuation.py | 93 - .../httpserver/test_pd_generate_error.py | 506 ----- .../httpserver/test_pd_kv_status_channel.py | 109 -- .../test_pd_master_cached_tokens.py | 769 -------- .../httpserver/test_pd_master_discovery.py | 371 ---- .../httpserver/test_pd_split_connections.py | 294 --- .../httpserver/test_pd_websocket_lifecycle.py | 81 - .../test_running_request_lifecycle.py | 565 ------ .../test_cache_placement_controller.py | 367 ---- .../test_multimodal_cache_matching.py | 120 -- unit_tests/server/pd_master_fault_launcher.py | 94 - unit_tests/server/pd_master_test_utils.py | 188 -- .../router/dynamic_prompt/test_radix_cache.py | 722 -------- .../mode_backend/test_dp_big_pages.py | 286 --- .../test_dp_overlap_spec_engine.py | 491 ----- .../mode_backend/test_generic_pre_process.py | 176 -- .../mode_backend/test_multi_level_kv_cache.py | 507 ----- .../mode_backend/test_nccl_control_channel.py | 151 -- .../test_paged_side_allocations.py | 253 --- .../mode_backend/test_pd_dynamic_split.py | 244 --- .../mode_backend/test_pd_manager_startup.py | 174 -- .../mode_backend/test_pd_notify.py | 223 --- .../mode_backend/test_pd_page_io.py | 53 - .../mode_backend/test_pd_prefill_abort.py | 199 -- .../mode_backend/test_pd_token_indices.py | 126 -- .../mode_backend/test_pd_transfer_progress.py | 176 -- .../test_pd_transfer_readiness.py | 38 - .../test_prefill_queue_strategy.py | 401 ---- .../mode_backend/test_rl_backend_ops.py | 57 - .../mtp_speculative/test_dflash.py | 180 -- .../mtp_speculative/test_dflash2.py | 117 -- .../mtp_speculative/test_dspark.py | 129 -- .../mtp_speculative/test_eagle_no_att.py | 93 - .../mtp_speculative/test_eagle_overlap.py | 363 ---- .../mtp_speculative/test_eagle_with_att.py | 213 --- .../mtp_speculative/test_planner.py | 964 ---------- .../mtp_speculative/test_vanilla_no_att.py | 211 --- .../mtp_speculative/test_vanilla_overlap.py | 147 -- .../mtp_speculative/test_vanilla_prefill.py | 237 --- .../structured_output/test_error_output.py | 66 - .../model_infer/structured_output/test_mtp.py | 444 ----- .../structured_output/test_overlap.py | 152 -- .../structured_output/test_reasoning.py | 194 -- .../structured_output/test_sampling.py | 320 ---- .../structured_output/test_state.py | 407 ----- .../model_infer/test_cpu_resource_layout.py | 513 ------ .../test_hybrid_cpu_state_cache.py | 362 ---- .../model_infer/test_multimodal_chunking.py | 404 ---- .../req_queue/test_pd_queue_selection.py | 108 -- unit_tests/server/router/test_kv_capacity.py | 62 - .../server/router/test_profiler_service.py | 53 - unit_tests/server/test_api_responses.py | 313 ---- unit_tests/server/test_api_start_deepep.py | 45 - unit_tests/server/test_application_stop.py | 200 -- .../server/test_deepseek_dsml_stream_args.py | 72 - .../test_deepseek_tool_parser_stream.py | 143 -- unit_tests/server/test_detokenization.py | 114 -- unit_tests/server/test_dp_request_capacity.py | 109 -- .../test_function_call_common_prefix.py | 67 - .../server/test_generation_error_response.py | 91 - unit_tests/server/test_glm47_tool_parser.py | 127 -- unit_tests/server/test_hybrid_start_args.py | 138 -- unit_tests/server/test_hypercorn_config.py | 26 - unit_tests/server/test_mtp_start_args.py | 55 - .../server/test_multimodal_chunk_args.py | 15 - unit_tests/server/test_pd_cache_aware.py | 321 ---- unit_tests/server/test_pd_local_selection.py | 919 ---------- .../server/test_pd_master_haproxy_proxy.py | 434 ----- .../test_pd_master_launcher_integration.py | 97 - unit_tests/server/test_pd_master_mode.py | 612 ------- .../server/test_pd_master_multiprocess_e2e.py | 496 ----- .../test_pd_master_multiprocess_launcher.py | 399 ---- unit_tests/server/test_pd_master_req_ids.py | 149 -- .../server/test_pd_master_topology_api.py | 96 - unit_tests/server/test_pd_node_stats.py | 321 ---- unit_tests/server/test_pd_start_args.py | 311 ---- .../server/test_prefill_queue_start_args.py | 27 - .../server/test_prompt_cache_tree_service.py | 228 --- .../server/test_qwen3_coder_stream_fc.py | 151 -- .../server/test_reasoning_token_usage.py | 222 --- .../test_req_id_generator_performance.py | 131 -- unit_tests/server/test_ssm_state_mode.py | 165 -- unit_tests/server/test_tokenizer.py | 177 -- .../server/test_vocab_sampling_start_args.py | 57 - unit_tests/test_format_out.py | 47 - unit_tests/test_static_benchmark_paged_kv.py | 155 -- unit_tests/utils/test_atexit_exit_modes.py | 188 -- unit_tests/utils/test_config_utils.py | 53 - unit_tests/utils/test_custom_kernel_utils.py | 35 - unit_tests/utils/test_envs_utils.py | 179 -- unit_tests/utils/test_log_utils.py | 42 - .../utils/test_multimodal_chunk_utils.py | 68 - unit_tests/utils/test_service_shm_cleanup.py | 280 --- unit_tests/utils/test_sgl_utils.py | 288 --- unit_tests/utils/test_shm_utils.py | 88 - unit_tests/utils/test_speculative_utils.py | 413 ----- unit_tests/utils/test_start_utils.py | 319 ---- .../utils/test_startup_process_lifecycle.py | 135 -- 302 files changed, 128 insertions(+), 64389 deletions(-) create mode 100644 unit_tests/README.md delete mode 100644 unit_tests/common/att_resources/manager_fixtures.py delete mode 100644 unit_tests/common/att_resources/model_helpers.py delete mode 100644 unit_tests/common/att_resources/test_backend_startup.py delete mode 100644 unit_tests/common/att_resources/test_bounded_state.py delete mode 100644 unit_tests/common/att_resources/test_capacity.py delete mode 100644 unit_tests/common/att_resources/test_compute_binding.py delete mode 100644 unit_tests/common/att_resources/test_dsv4_integration.py delete mode 100644 unit_tests/common/att_resources/test_dsv4_state.py delete mode 100644 unit_tests/common/att_resources/test_fp8_compute_binding.py delete mode 100644 unit_tests/common/att_resources/test_fp8_e8m0_packed_mla_state.py delete mode 100644 unit_tests/common/att_resources/test_full_att_cpu_kernel.py delete mode 100644 unit_tests/common/att_resources/test_full_att_dp_transfer.py delete mode 100644 unit_tests/common/att_resources/test_full_att_transfer.py delete mode 100644 unit_tests/common/att_resources/test_full_state.py delete mode 100644 unit_tests/common/att_resources/test_full_storage_initialization.py delete mode 100644 unit_tests/common/att_resources/test_gqa_layout_identity.py delete mode 100644 unit_tests/common/att_resources/test_gqa_pd_kernel.py delete mode 100644 unit_tests/common/att_resources/test_gqa_transfer.py delete mode 100644 unit_tests/common/att_resources/test_hybrid_layer_mapping.py delete mode 100644 unit_tests/common/att_resources/test_layer_groups.py delete mode 100644 unit_tests/common/att_resources/test_layer_state_selection.py delete mode 100644 unit_tests/common/att_resources/test_manager_helpers.py delete mode 100644 unit_tests/common/att_resources/test_model_calibration.py delete mode 100644 unit_tests/common/att_resources/test_model_config_normalization.py delete mode 100644 unit_tests/common/att_resources/test_model_initialization.py delete mode 100644 unit_tests/common/att_resources/test_model_resources.py delete mode 100644 unit_tests/common/att_resources/test_model_storage_definition.py delete mode 100644 unit_tests/common/att_resources/test_mtp_state_commit.py delete mode 100644 unit_tests/common/att_resources/test_network_transport.py delete mode 100644 unit_tests/common/att_resources/test_no_attention_draft.py delete mode 100644 unit_tests/common/att_resources/test_page_pool_storage.py delete mode 100644 unit_tests/common/att_resources/test_paged_storage.py delete mode 100644 unit_tests/common/att_resources/test_quantized_full_state.py delete mode 100644 unit_tests/common/att_resources/test_quantized_paged_storage.py delete mode 100644 unit_tests/common/att_resources/test_role_transfers.py delete mode 100644 unit_tests/common/att_resources/test_runtime_transport.py delete mode 100644 unit_tests/common/att_resources/test_state_creation.py delete mode 100644 unit_tests/common/att_resources/test_state_identity.py delete mode 100644 unit_tests/common/att_resources/test_state_queries.py delete mode 100644 unit_tests/common/att_resources/test_tp_config.py delete mode 100644 unit_tests/common/att_resources/test_transfer_views.py delete mode 100644 unit_tests/common/att_resources/test_validate_model_cache.py delete mode 100644 unit_tests/common/att_resources/test_window_compressor_state.py delete mode 100644 unit_tests/common/basemodel/__init__.py delete mode 100644 unit_tests/common/basemodel/attention/fa3/test_neo.py delete mode 100644 unit_tests/common/basemodel/attention/flashinfer/test_decode_graph_state.py delete mode 100644 unit_tests/common/basemodel/attention/flashinfer/test_decode_wrapper_init.py delete mode 100644 unit_tests/common/basemodel/attention/flashinfer/test_mla_cuda_graph_plan.py delete mode 100644 unit_tests/common/basemodel/attention/flashinfer/test_mla_nope.py delete mode 100644 unit_tests/common/basemodel/attention/flashinfer/test_workspace.py delete mode 100644 unit_tests/common/basemodel/attention/linear/test_gdn.py delete mode 100644 unit_tests/common/basemodel/attention/test_dsv4_flashmla.py delete mode 100644 unit_tests/common/basemodel/attention/test_fp8_logical_layer.py delete mode 100644 unit_tests/common/basemodel/attention/test_fp8_page_size.py delete mode 100644 unit_tests/common/basemodel/attention/test_lazy_import.py delete mode 100644 unit_tests/common/basemodel/attention/test_state_lifetime.py delete mode 100644 unit_tests/common/basemodel/test_att_sink_weight.py delete mode 100644 unit_tests/common/basemodel/test_cuda_graph_autotune.py delete mode 100644 unit_tests/common/basemodel/test_cuda_graph_layout.py delete mode 100644 unit_tests/common/basemodel/test_hidden_collector.py delete mode 100644 unit_tests/common/basemodel/test_model_input.py delete mode 100644 unit_tests/common/basemodel/test_model_output.py delete mode 100644 unit_tests/common/basemodel/test_moe_route_info_manager.py delete mode 100644 unit_tests/common/basemodel/test_mtp_manager.py delete mode 100644 unit_tests/common/basemodel/test_overlap_utils.py delete mode 100644 unit_tests/common/basemodel/test_prefill_cuda_graph_state.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_mtp_diverse.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_mtp_diverse_autotune.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_stage2_autotune.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/test_flash_decoding_autotune.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/att/decode_att/int4kv/test_int4kv_autotune.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage1.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage2.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage3.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_normal_autotune.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_attention_fwd_neo.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_flashattention_nopad.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_flashattention_nopad1.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_fa3_neo.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/kv_copy/test_mla_destindex_copy_kv.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/kv_copy/test_ppl_int4kv_copy_kv.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/kv_copy/test_ppl_int8kv_copy_kv.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/linear_att/test_fused_recurrent_strided.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/linear_att/test_kda_autotune.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_autotune.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_equiv.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_state_params.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune_memory.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/mhc/test_autotune.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/mhc/test_mhc.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/mla_att/decode_att/test_gqa_flash_decoding.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/test_add_in_place.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/test_apply_constraint_mask.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/test_apply_invalid_token.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/test_atomic_event.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/test_diverse_utils.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/test_dynamic_mtp_utils.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/test_fa3_utils.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/test_gen_decode_params.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/test_gen_mtp_prefill_params.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/test_gen_nsa_ks_ke.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/test_gen_prefill_params.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/test_gen_sampling_params.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/test_mtp_utils.py delete mode 100755 unit_tests/common/basemodel/triton_kernel/test_multimodal_emb.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/test_qk_rmsnorm_fused.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/test_redundancy_topk_ids_repair.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/test_repack_kv_index.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/test_routing_capture.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/test_runtime_scalar_specialization.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/test_select_kv_index_from_req.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/test_sp_pad_kernel.py delete mode 100644 unit_tests/common/fused_moe/test_activation_config.py delete mode 100644 unit_tests/common/fused_moe/test_append_shared_expert_topk.py delete mode 100644 unit_tests/common/fused_moe/test_deepep.py delete mode 100644 unit_tests/common/fused_moe/test_deepep_runtime_specialization.py delete mode 100644 unit_tests/common/fused_moe/test_deepgemm_activation.py delete mode 100644 unit_tests/common/fused_moe/test_ep_reduce_decode_output.py delete mode 100644 unit_tests/common/fused_moe/test_grouped_fused_moe.py delete mode 100644 unit_tests/common/fused_moe/test_grouped_fused_moe_ep_sms.py delete mode 100644 unit_tests/common/fused_moe/test_grouped_fused_moe_speed.py delete mode 100755 unit_tests/common/fused_moe/test_grouped_topk.py delete mode 100644 unit_tests/common/fused_moe/test_moe_silu_and_mul_mix_quant_ep.py delete mode 100644 unit_tests/common/kv_trans_kernel/test_kv_trans_v2.py delete mode 100644 unit_tests/common/kv_trans_kernel/test_nixl_kv_trans.py delete mode 100644 unit_tests/common/quantization/test_no_quant.py delete mode 100644 unit_tests/common/quantization/test_ue8m0_fp8_quantization.py delete mode 100644 unit_tests/common/test_autotuner.py delete mode 100644 unit_tests/common/test_cache_copy_control_flow.py delete mode 100644 unit_tests/common/test_hybrid_req_manager_pd_sharing.py delete mode 100644 unit_tests/common/test_linear_state.py delete mode 100644 unit_tests/common/test_linear_state_contract.py delete mode 100644 unit_tests/common/test_mem_manager_ipc.py delete mode 100644 unit_tests/common/test_pinned_memory_non_blocking_copy.py delete mode 100644 unit_tests/common/test_req_manager_hold_page.py delete mode 100644 unit_tests/models/deepseek2/test_rope_repeat.py delete mode 100644 unit_tests/models/deepseek3_2/triton_kernel/test_hadamard_transform.py delete mode 100644 unit_tests/models/deepseek3_2/triton_kernel/test_topk_index_to_mem_index.py delete mode 100644 unit_tests/models/deepseek_v4/test_c128_workspace.py delete mode 100644 unit_tests/models/deepseek_v4/test_kernels.py delete mode 100644 unit_tests/models/deepseek_v4/test_model.py delete mode 100644 unit_tests/models/deepseek_v4/test_shared_expert_overlap.py delete mode 100644 unit_tests/models/deepseek_v4/test_tokenizer.py delete mode 100644 unit_tests/models/deepseek_v4/test_vision.py delete mode 100644 unit_tests/models/deepseek_v4/test_vision_integration.py delete mode 100644 unit_tests/models/deepseek_v4/test_workspace.py delete mode 100644 unit_tests/models/glm5_next/test_cache.py delete mode 100644 unit_tests/models/glm5_next/test_compact_topk_modes.py delete mode 100644 unit_tests/models/glm5_next/test_direct_nsa_indices.py delete mode 100644 unit_tests/models/glm5_next/test_fp8_compile_cache.py delete mode 100644 unit_tests/models/glm5_next/test_fp8_kv.py delete mode 100644 unit_tests/models/glm5_next/test_fused_mtp_norm.py delete mode 100644 unit_tests/models/glm5_next/test_indexer_topk.py delete mode 100644 unit_tests/models/glm5_next/test_indexer_workspace.py delete mode 100644 unit_tests/models/glm5_next/test_kda_conv_qkv.py delete mode 100644 unit_tests/models/glm5_next/test_kernels.py delete mode 100644 unit_tests/models/glm5_next/test_kpool_gather.py delete mode 100644 unit_tests/models/glm5_next/test_kpool_runtime_specialization.py delete mode 100644 unit_tests/models/glm5_next/test_layers.py delete mode 100644 unit_tests/models/glm5_next/test_mtp.py delete mode 100644 unit_tests/models/glm5_next/test_mtp_prefill_cache.py delete mode 100644 unit_tests/models/glm5_next/test_pd_cache.py delete mode 100644 unit_tests/models/glm5_next/test_prefill_workspace.py delete mode 100644 unit_tests/models/glm5_next/test_token_major_mla_bmm.py delete mode 100644 unit_tests/models/glm5_next/test_visual.py delete mode 100644 unit_tests/models/phi3/test_rotary_emb.py delete mode 100644 unit_tests/models/qwen2_vl/test_get_mrope_position_ids.py delete mode 100644 unit_tests/models/qwen2_vl/test_infer_struct.py delete mode 100644 unit_tests/models/qwen2_vl/test_mrope.py delete mode 100644 unit_tests/models/qwen2_vl/test_rotary_pos_emb.py delete mode 100644 unit_tests/models/qwen3-vl/test_deepstack_emb.py delete mode 100644 unit_tests/models/qwen3next/test_shared_expert_gate.py delete mode 100644 unit_tests/models/sensenova_loop/test_gdn_state.py delete mode 100644 unit_tests/models/sensenova_loop/test_loop.py delete mode 100644 unit_tests/models/sensenova_loop/test_loop_mhc.py delete mode 100644 unit_tests/models/test_multimodal_tokenizer_positions.py delete mode 100644 unit_tests/models/test_neo_registry.py delete mode 100644 unit_tests/models/test_qwen3_dspark_model_output.py delete mode 100644 unit_tests/models/test_qwen3_eagle_model_input.py delete mode 100644 unit_tests/models/test_qwen_vl_tokenizer.py delete mode 100644 unit_tests/models/test_yarn_rotary.py delete mode 100644 unit_tests/models/vit/test_flash_attention_forward.py delete mode 100644 unit_tests/server/core/objs/test_out_token_circlequeue.py delete mode 100644 unit_tests/server/core/objs/test_req.py delete mode 100644 unit_tests/server/core/objs/test_shm_array.py delete mode 100644 unit_tests/server/core/objs/test_token_ids_validation.py delete mode 100644 unit_tests/server/embed_cache/test_multimodal_token_id_range.py delete mode 100644 unit_tests/server/grammar_helpers.py delete mode 100644 unit_tests/server/httpserver/test_grammar.py delete mode 100644 unit_tests/server/httpserver/test_pd_continuation.py delete mode 100644 unit_tests/server/httpserver/test_pd_generate_error.py delete mode 100644 unit_tests/server/httpserver/test_pd_kv_status_channel.py delete mode 100644 unit_tests/server/httpserver/test_pd_master_cached_tokens.py delete mode 100644 unit_tests/server/httpserver/test_pd_master_discovery.py delete mode 100644 unit_tests/server/httpserver/test_pd_split_connections.py delete mode 100644 unit_tests/server/httpserver/test_pd_websocket_lifecycle.py delete mode 100644 unit_tests/server/httpserver/test_running_request_lifecycle.py delete mode 100644 unit_tests/server/multi_level_kv_cache/test_cache_placement_controller.py delete mode 100644 unit_tests/server/multi_level_kv_cache/test_multimodal_cache_matching.py delete mode 100644 unit_tests/server/pd_master_fault_launcher.py delete mode 100644 unit_tests/server/pd_master_test_utils.py delete mode 100644 unit_tests/server/router/dynamic_prompt/test_radix_cache.py delete mode 100644 unit_tests/server/router/model_infer/mode_backend/test_dp_big_pages.py delete mode 100644 unit_tests/server/router/model_infer/mode_backend/test_dp_overlap_spec_engine.py delete mode 100644 unit_tests/server/router/model_infer/mode_backend/test_generic_pre_process.py delete mode 100644 unit_tests/server/router/model_infer/mode_backend/test_multi_level_kv_cache.py delete 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_paged_side_allocations.py delete mode 100644 unit_tests/server/router/model_infer/mode_backend/test_pd_dynamic_split.py delete mode 100644 unit_tests/server/router/model_infer/mode_backend/test_pd_manager_startup.py delete mode 100644 unit_tests/server/router/model_infer/mode_backend/test_pd_notify.py delete mode 100644 unit_tests/server/router/model_infer/mode_backend/test_pd_page_io.py delete mode 100644 unit_tests/server/router/model_infer/mode_backend/test_pd_prefill_abort.py delete mode 100644 unit_tests/server/router/model_infer/mode_backend/test_pd_token_indices.py delete mode 100644 unit_tests/server/router/model_infer/mode_backend/test_pd_transfer_progress.py delete mode 100644 unit_tests/server/router/model_infer/mode_backend/test_pd_transfer_readiness.py delete mode 100644 unit_tests/server/router/model_infer/mode_backend/test_prefill_queue_strategy.py delete mode 100644 unit_tests/server/router/model_infer/mode_backend/test_rl_backend_ops.py delete mode 100644 unit_tests/server/router/model_infer/mtp_speculative/test_dflash.py delete mode 100644 unit_tests/server/router/model_infer/mtp_speculative/test_dflash2.py delete mode 100644 unit_tests/server/router/model_infer/mtp_speculative/test_dspark.py delete mode 100644 unit_tests/server/router/model_infer/mtp_speculative/test_eagle_no_att.py delete mode 100644 unit_tests/server/router/model_infer/mtp_speculative/test_eagle_overlap.py delete mode 100644 unit_tests/server/router/model_infer/mtp_speculative/test_eagle_with_att.py delete mode 100644 unit_tests/server/router/model_infer/mtp_speculative/test_planner.py delete mode 100644 unit_tests/server/router/model_infer/mtp_speculative/test_vanilla_no_att.py delete mode 100644 unit_tests/server/router/model_infer/mtp_speculative/test_vanilla_overlap.py delete mode 100644 unit_tests/server/router/model_infer/mtp_speculative/test_vanilla_prefill.py delete mode 100644 unit_tests/server/router/model_infer/structured_output/test_error_output.py delete mode 100644 unit_tests/server/router/model_infer/structured_output/test_mtp.py delete mode 100644 unit_tests/server/router/model_infer/structured_output/test_overlap.py delete mode 100644 unit_tests/server/router/model_infer/structured_output/test_reasoning.py delete mode 100644 unit_tests/server/router/model_infer/structured_output/test_sampling.py delete mode 100644 unit_tests/server/router/model_infer/structured_output/test_state.py delete mode 100644 unit_tests/server/router/model_infer/test_cpu_resource_layout.py delete mode 100644 unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py delete mode 100644 unit_tests/server/router/model_infer/test_multimodal_chunking.py delete mode 100644 unit_tests/server/router/req_queue/test_pd_queue_selection.py delete mode 100644 unit_tests/server/router/test_kv_capacity.py delete mode 100644 unit_tests/server/router/test_profiler_service.py delete mode 100644 unit_tests/server/test_api_responses.py delete mode 100644 unit_tests/server/test_api_start_deepep.py delete mode 100644 unit_tests/server/test_application_stop.py delete mode 100644 unit_tests/server/test_deepseek_dsml_stream_args.py delete mode 100644 unit_tests/server/test_deepseek_tool_parser_stream.py delete mode 100644 unit_tests/server/test_detokenization.py delete mode 100644 unit_tests/server/test_dp_request_capacity.py delete mode 100644 unit_tests/server/test_function_call_common_prefix.py delete mode 100644 unit_tests/server/test_generation_error_response.py delete mode 100644 unit_tests/server/test_glm47_tool_parser.py delete mode 100644 unit_tests/server/test_hybrid_start_args.py delete mode 100644 unit_tests/server/test_hypercorn_config.py delete mode 100644 unit_tests/server/test_mtp_start_args.py delete mode 100644 unit_tests/server/test_multimodal_chunk_args.py delete mode 100644 unit_tests/server/test_pd_cache_aware.py delete mode 100644 unit_tests/server/test_pd_local_selection.py delete mode 100644 unit_tests/server/test_pd_master_haproxy_proxy.py delete mode 100644 unit_tests/server/test_pd_master_launcher_integration.py delete mode 100644 unit_tests/server/test_pd_master_mode.py delete mode 100644 unit_tests/server/test_pd_master_multiprocess_e2e.py delete mode 100644 unit_tests/server/test_pd_master_multiprocess_launcher.py delete mode 100644 unit_tests/server/test_pd_master_req_ids.py delete mode 100644 unit_tests/server/test_pd_master_topology_api.py delete mode 100644 unit_tests/server/test_pd_node_stats.py delete mode 100644 unit_tests/server/test_pd_start_args.py delete mode 100644 unit_tests/server/test_prefill_queue_start_args.py delete mode 100644 unit_tests/server/test_prompt_cache_tree_service.py delete mode 100644 unit_tests/server/test_qwen3_coder_stream_fc.py delete mode 100644 unit_tests/server/test_reasoning_token_usage.py delete mode 100644 unit_tests/server/test_req_id_generator_performance.py delete mode 100644 unit_tests/server/test_ssm_state_mode.py delete mode 100644 unit_tests/server/test_tokenizer.py delete mode 100644 unit_tests/server/test_vocab_sampling_start_args.py delete mode 100644 unit_tests/test_format_out.py delete mode 100644 unit_tests/test_static_benchmark_paged_kv.py delete mode 100644 unit_tests/utils/test_atexit_exit_modes.py delete mode 100644 unit_tests/utils/test_config_utils.py delete mode 100644 unit_tests/utils/test_custom_kernel_utils.py delete mode 100644 unit_tests/utils/test_envs_utils.py delete mode 100644 unit_tests/utils/test_log_utils.py delete mode 100644 unit_tests/utils/test_multimodal_chunk_utils.py delete mode 100644 unit_tests/utils/test_service_shm_cleanup.py delete mode 100644 unit_tests/utils/test_sgl_utils.py delete mode 100644 unit_tests/utils/test_shm_utils.py delete mode 100644 unit_tests/utils/test_speculative_utils.py delete mode 100644 unit_tests/utils/test_start_utils.py delete mode 100644 unit_tests/utils/test_startup_process_lifecycle.py diff --git a/unit_tests/README.md b/unit_tests/README.md new file mode 100644 index 0000000000..a6ad0a6682 --- /dev/null +++ b/unit_tests/README.md @@ -0,0 +1,26 @@ +# Unit tests + +Keep tests that can catch an incorrect result, broken public contract, resource +lifetime bug, or a previously observed failure. Test the production entry point +with small inputs; mock external dependencies only as needed. + +- Add a regression case to an existing test when it exercises the same behavior. + Parameterize inputs instead of copying setup and assertions into another test. +- For kernels, compare outputs and mutated state against an independent reference. + Include relevant dtype, layout, padding, and boundary cases without taking the + Cartesian product of unrelated options. +- Do not test class inheritance, private method inventories, or dataclass field + declarations. Exercise the behavior those implementations support. +- A CPU reimplementation tested only against its own expected output does not + validate the CUDA path. Run the real implementation or explicitly skip when its + prerequisites are unavailable. +- Keep timing loops and performance reports out of unit tests. Benchmarks and + service-level tests belong under `test/`; ordinary unit tests need assertions. +- When removing a test, check whether another test covers the same failure mode. + Similar setup alone does not make cache ownership, transfer boundaries, or + asynchronous failure tests redundant. + +Run affected tests with `python -m pytest ` in an environment with their +required dependencies. For changed Python files, run +`pre-commit run --files `. Report skipped hardware/provider checks separately +from passing tests. diff --git a/unit_tests/common/att_resources/manager_fixtures.py b/unit_tests/common/att_resources/manager_fixtures.py deleted file mode 100644 index cfcf1ad8bf..0000000000 --- a/unit_tests/common/att_resources/manager_fixtures.py +++ /dev/null @@ -1,45 +0,0 @@ -"""Concrete main/draft composition used to exercise manager transfer operations.""" - -from __future__ import annotations - -from typing import TYPE_CHECKING -from dataclasses import replace -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 -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 - -if TYPE_CHECKING: - import torch - - -class MainDraftTestResourceManager(AttentionStateManager): - 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( - 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.register_state("draft0", draft_state) - self.register_state("main", full_state(device=device)) - prepared_full_state = MLAFullAttentionState( - layout_config=MLAConfig( - logical_layer_groups=self.main_gqa_fp32_h1_d4_state.config.logical_layer_groups, - latent_dim=self.main_gqa_fp32_h1_d4_state.config.head_dim, - latent_dtype=self.main_gqa_fp32_h1_d4_state.config.dtype, - ), - page_size=self.main_gqa_fp32_h1_d4_state.page_size, - device=device, - ) - prepared_full_state.allocate_storage(self.main_gqa_fp32_h1_d4_state.capacity) - self.register_state("draft0", prepared_full_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_registration() diff --git a/unit_tests/common/att_resources/model_helpers.py b/unit_tests/common/att_resources/model_helpers.py deleted file mode 100644 index 86e62cb671..0000000000 --- a/unit_tests/common/att_resources/model_helpers.py +++ /dev/null @@ -1,342 +0,0 @@ -"""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.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, -) - - -if TYPE_CHECKING: - from lightllm.models.qwen3next.config import Qwen3NextResourceConfig - from lightllm.common.attention_state.allocators import FullAttKVAllocator - from collections.abc import Mapping, Sequence - from lightllm.common.attention_state.manager import AttentionStateManager - 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, - config: ( - GQAAttentionConfig - | MLAConfig - | GroupQuantizedAttentionConfig - | StaticFP8AttentionConfig - | PackedMLAAttentionConfig - | CompressedAttentionConfig - | FP8E8M0PackedMLAConfig - | 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, - ) - 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, - ) - 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.allocate_storage(token_capacity) - 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( - sorted( - (state for state in attributes.values() if isinstance(state, state_class)), - key=lambda state: state.state_name, - ) - ) - - -def main_linear_state(resources: AttentionStateManager): - (state,) = (state for state in resources.get_bounded_attention_states("main") if hasattr(state, "conv_state")) - return state - - -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"))) - 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_attention_states() - - -def create_full_attention_roles( - 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" - states = [] - for start, end in ((0, boundary), (boundary, len(layers))): - if start == end: - states.append(None) - continue - 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, - 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.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.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") - ) - return fields - return { - 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, - ssm_state_mode: str = "native", - ssm_replay_config: dict | None = None, -) -> AttentionStateManager: - """测试按既定几何参数组装通用管理器，不依赖旧模型管理器的构造/绑定流程。""" - 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.bounded_att import WindowCompressorStateConfig, create_window_compressor_state - from lightllm.models.glm5_next.config import Glm5NextResourceConfig - - 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( - 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", - ssm_state_mode=ssm_state_mode, - ssm_replay_config=ssm_replay_config, - gate_lower_bound=config.gate_lower_bound if glm else -5.0, - ) - bounded_state.initialize_page_pools( - 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)): - if state is None: - continue - manager.register_state(role, state) - if glm: - window = create_window_compressor_state( - config=WindowCompressorStateConfig( - logical_layer_ids=tuple(layer for group in state.config.logical_layer_groups for layer in group), - 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, - ), - request_capacity=request_capacity, - mtp_step=mtp_step, - device="cuda", - ) - 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.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_backend_startup.py b/unit_tests/common/att_resources/test_backend_startup.py deleted file mode 100644 index 457fdfa1c3..0000000000 --- a/unit_tests/common/att_resources/test_backend_startup.py +++ /dev/null @@ -1,374 +0,0 @@ -"""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.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: - 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) - request_allocations = [] - manager_creations = [] - from lightllm.common.attention_state import AttentionStateManager - - def create_empty_manager(request_capacity: int): - 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.AttentionStateManager", create_empty_manager) - - 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 已创建完成，可通过缓存查询取得资源。 - for item in models: - 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_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 - - def warmup(model: TpPartBaseModel): - assert model.backends_bound - assert all(item.req_manager is mapping for item in models) - 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)]} - 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 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.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): - assert all(item.req_manager is mapping for item in models) - 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 - 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 - - 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_state_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: - 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), - ) - for item in models: - item.bind_req_manager(req_manager=req_manager) - for index, item in enumerate(models): - 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_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) - 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 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 []) - 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) - 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"), - 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_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( - torch.tensor([0], dtype=torch.int32, device="cuda"), - torch.tensor([i + 2], 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.logical_layer_groups) == 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_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] - 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_runtime_to_pages( - req_indices=[0], - state_token_lens=[4], - page_kind="big", - 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_runtime_from_pages( - req_indices=[0], - state_token_lens=[4], - page_kind="big", - 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_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_state.py b/unit_tests/common/att_resources/test_bounded_state.py deleted file mode 100644 index e2c08b1461..0000000000 --- a/unit_tests/common/att_resources/test_bounded_state.py +++ /dev/null @@ -1,440 +0,0 @@ -from __future__ import annotations - -from typing import TYPE_CHECKING -from dataclasses import replace - -import pytest -import torch - -from lightllm.common.attention_state import AttentionStateManager -from lightllm.common.attention_state.bounded_att import ( - GDNPagesPool, - GDNState, - KDAPagesPool, - KDAState, - ConvSSMStateConfig, -) -from lightllm.common.attention_state.types import ( - AxisMapping, - StateTransferLayout, - TransferRegion, -) - -if TYPE_CHECKING: - from lightllm.common.attention_state.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", - 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, - ) - state = cls( - config=config, - request_capacity=3, - mtp_step=mtp_step, - device=device, - ) - 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 - - -def page_layout(state: KDAState | GDNState): - cfg = state.config - 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( - state_name=state.state_name, - 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_mappings=( - AxisMapping( - axis="channel", - local_start=0, - external_start=0, - length=cfg.conv_channels, - ), - AxisMapping( - axis="window", - local_start=0, - external_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_mappings=( - AxisMapping( - axis="head", - local_start=0, - external_start=0, - length=cfg.num_v_heads, - ), - AxisMapping( - axis="key_dim", - local_start=0, - external_start=0, - length=cfg.key_dim, - ), - AxisMapping( - axis="value_dim", - local_start=0, - external_start=0, - length=cfg.value_dim, - ), - ), - ), - ), - ) - - -@pytest.mark.parametrize("cls,pool_cls", [(KDAState, KDAPagesPool), (GDNState, GDNPagesPool)]) -@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], - 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( - 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, - ) - 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() - 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_runtime_to_pages( - req_indices=[0], - state_token_lens=[8], - page_kind=page_kind, - page_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_runtime_from_pages( - req_indices=[1], - state_token_lens=[8], - page_kind=page_kind, - 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]) - 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=config, - request_capacity=3, - mtp_step=2, - device="cuda", - ) - 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), - ) - state.conv_state[:, 0, :, :2].copy_(conv) - state.ssm_state[:, 0].copy_(ssm) - for page_kind in ("big", "small"): - state.save_bounded_att_runtime_to_pages( - req_indices=[0], - state_token_lens=[8], - page_kind=page_kind, - page_slot_ids=[0], - ) - state.reset_bounded_att_state([0]) - for req, page_kind in enumerate(("big", "small"), start=1): - state.restore_bounded_att_runtime_from_pages( - req_indices=[req], - state_token_lens=[8], - page_kind=page_kind, - page_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"): - 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, device="cuda") - state.conv_state[:, 0].fill_(7) - state.ssm_state[:, 0:3].fill_(11) - state.save_bounded_att_runtime_to_pages( - req_indices=[0], - state_token_lens=[6], - page_kind="small", - page_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), - 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, device="cuda") - source.conv_state[:, 0].fill_(13) - source.ssm_state[:, 0].fill_(29) - source.save_bounded_att_runtime_to_pages( - req_indices=[0], - state_token_lens=[8], - page_kind="big", - page_slot_ids=[0], - ) - layout = page_layout(source) - buffer = torch.full((3, layout.byte_end), 255, dtype=torch.uint8, pin_memory=True) - source.write_state_pages_to_cpu( - page_kind="big", - page_slot_ids=[0], - cpu_page_buffers=buffer, - cpu_page_indices=[2], - layout=layout, - ) - torch.cuda.synchronize() - assert (buffer[:2] == 255).all() and (buffer[2, :16] == 255).all() - target = make_state(cls, device="cuda") - 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_runtime_from_pages( - req_indices=[2], - state_token_lens=[8], - page_kind="big", - page_slot_ids=[1], - ) - assert (target.conv_state[:, 2, :, :2] == 13).all() - assert (target.ssm_state[:, 6] == 29).all() - # 页池只保存快照数据；前缀关联由调用方维护，读回的快照可以再次写出。 - target.read_state_pages_from_cpu( - cpu_page_buffers=buffer, - cpu_page_indices=[2], - big_page_slot_ids=[0], - layout=layout, - ) - target.write_state_pages_to_cpu( - page_kind="big", - page_slot_ids=[0], - cpu_page_buffers=buffer, - cpu_page_indices=[0], - layout=layout, - ) - torch.cuda.synchronize() - assert torch.equal(buffer[0], buffer[2]) - - -@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], - 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( - size=(3, layout.byte_end), - fill_value=254, - dtype=torch.uint8, - device=device, - ) - page = allocation[1] - source.write_to_pd_page( - 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) - target.read_from_pd_page( - 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() - 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", ["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_runtime_state_from_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, 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) - 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_pages_from_dp( - 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]) - 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 deleted file mode 100644 index b6dff3dc0e..0000000000 --- a/unit_tests/common/att_resources/test_capacity.py +++ /dev/null @@ -1,1113 +0,0 @@ -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 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.bounded_att import WindowCompressorState, WindowCompressorStateConfig -from lightllm.server.core.objs.start_args_type import StartArgs -from lightllm.common.attention_state.full_att import ( - FullAttentionState, - GQAFullAttentionState, - GQAAttentionConfig, - MLAConfig, - create_gqa_state, - create_mla_state, - GroupQuantizedAttentionConfig, - GroupQuantizedFullAttentionState, - MLAFullAttentionState, - PackedMLAFullAttentionState, - StaticFP8FullAttentionState, -) - -if TYPE_CHECKING: - from easydict import EasyDict - - -def definition(kind: str): - if kind == "mla": - return create_mla_state( - 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, - format="packed_group128_fp8" if kind == "packed_fp8" else "packed_bf16", - ), - 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 - 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 - - -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=template, - config=template.config, - token_capacity=20, - device="cpu", - ) - actual = allocated_bytes(state, token_only=True) - 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_operator().get_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 = AttentionStateManager(request_capacity=1) - manager.register_state("main", state) - 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_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_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) - monkeypatch.setattr( - "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 = 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_registration() - if expected is None: - with pytest.raises(MemoryError, match="HOLD"): - resources.initialize_storage( - 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_storage( - args=args, - token_capacity=None, - mem_fraction=1.0, - device="cpu", - ) - else: - resources.initialize_storage( - args=args, - token_capacity=None, - mem_fraction=1.0, - device="cpu", - ) - assert resources.full_att_allocator.size == expected - assert resources.get_pd_operator().get_page_bytes() == manager.get_pd_operator().get_page_bytes() - - -def test_shared_allocator_accounts_for_distinct_state_stores(): - main, draft = definition("int4"), definition("static_head") - 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( - template=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]) -@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], - 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, - ssm_dtype=torch.float32, - ) - tail = WindowCompressorState( - config=WindowCompressorStateConfig( - logical_layer_ids=(1,), - window_size=4, - compress_stride=4, - target="indexer", - candidate_dim=3, - dtype=torch.bfloat16, - ), - request_capacity=3, - mtp_step=3, - device="cuda", - ) - reference = state_cls( - config=config, - request_capacity=3, - mtp_step=3, - device="cuda", - ) - 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 - ) - 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_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_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.get_pd_operator().write_page( - 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.buffer.zero_() - manager.get_pd_operator().read_page( - 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.buffer[:, 0]) - for position in range(4, 8): - 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( - 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] - - -@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 = AttentionStateManager(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), - ): - if state_cls is MLAFullAttentionState: - config = MLAConfig( - 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, - ) - 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_operator().get_page_bytes() == 72 - monkeypatch.setattr( - "lightllm.common.attention_state.manager.FullAttKVAllocator", partial(FullAttKVAllocator, pin_memory=False) - ) - manager.allocate_storage( - token_capacity=4, - page_size=1, - big_page_tokens=4, - device="cuda", - ) - 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_operator().get_page_bytes()), dtype=torch.uint8, device="cuda") - page = pages[1] - 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.get_pd_operator().write_page( - 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.get_pd_operator().read_page( - 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() - - -@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( - logical_layer_groups=((0,),), - global_head_num=1, - head_dim=128, - scale_dtype=torch.float32, - bits=8, - group_size=128, - ) - manager = AttentionStateManager(request_capacity=1) - prepared_full_state = AlignedFieldsState(config, page_size=1, device="cuda") - 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.allocate_storage(8) - manager.register_state("draft0", prepared_full_state) - manager.full_att_allocator = FullAttKVAllocator(8, 1, pin_memory=False) - 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_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() - 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() - assert get_layouts() is layouts - pages = torch.full( - 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.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)] - 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.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.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.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.get_pd_operator().read_page( - 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]) - 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]) -@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, - 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 - for region in super().get_pd_transfer_regions() - ) - - manager = AttentionStateManager(request_capacity=1) - if full_page_tokens is not None: - pd_args.pd_kv_page_size = full_page_tokens - config = GQAAttentionConfig( - logical_layer_groups=((0,),), - global_head_num=1, - head_dim=1, - dtype=torch.float32, - ) - manager.register_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, - ), - request_capacity=1, - device="cuda", - ) - 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_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.pd_operator.logger.warning", warnings.append) - expected = max(80, (full_page_tokens or 0) * 8) - 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.get_pd_operator().write_page( - 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.get_pd_operator().read_page( - 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) - if warnings: - assert "64 bytes" in warnings[0] and "using 80 bytes per transfer page" in warnings[0] - # 相同 manager 的容量查询命中缓存，不重复打印 warning。 - assert manager.get_pd_operator().get_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 = AttentionStateManager(request_capacity=1) - config = GQAAttentionConfig( - logical_layer_groups=((0,),), - global_head_num=1, - head_dim=3, - dtype=torch.bfloat16, - ) - 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 - ) - assert get_page_bytes() == 48 - monkeypatch.setattr( - "lightllm.common.attention_state.manager.FullAttKVAllocator", partial(FullAttKVAllocator, pin_memory=False) - ) - 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_cache_operator().get_page_layouts() - if transfer == "cpu" - else manager.get_pd_operator().get_pd_full_att_page_layouts() - ) - assert layouts["main_gqa_bf16_h1_d3_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 = AttentionStateManager(request_capacity=1) - config = GQAAttentionConfig( - logical_layer_groups=((0,),), - global_head_num=1, - head_dim=1, - dtype=torch.float32, - ) - full_state = GQAFullAttentionState(config, page_size=1, device="meta") - manager.register_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, - ) - state = cls(config, request_capacity=1, device="meta") - manager.register_state(role, state) - if with_tail: - state = WindowCompressorState( - config=WindowCompressorStateConfig( - logical_layer_ids=(0,), - window_size=4, - compress_stride=4, - target="indexer", - candidate_dim=1, - dtype=torch.float32, - ), - request_capacity=1, - mtp_step=0, - device="meta", - ) - 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_registration() - 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_h1_d1_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) - 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) - for role in ("main", "draft0"): - 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( - 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( - config=bounded_config, - request_capacity=3, - mtp_step=2, - device="meta", - ) - 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_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 - } - 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_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): - 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) - if region.field == "ssm" - else region - for region in super().get_cpu_transfer_regions() - ) - - manager = AttentionStateManager(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, - ), - request_capacity=1, - device="meta", - ) - manager.register_state("main", state) - manager.finish_state_registration() - # Conv 占 6 字节，SSM 位于 [72, 76)，State 与物理页终点均对齐到 80。 - 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 - - -@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=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 = 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_registration() - 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) - 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( - size=(3, manager.get_cpu_cache_operator().get_page_bytes()), - fill_value=165, - dtype=torch.uint8, - pin_memory=True, - ) - 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, - ) - torch.cuda.synchronize() - assert (pages[1, data_end:gqa_start] == 165).all() - assert (pages[[0, 2]] == 165).all() - 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, - ) - 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]) -@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 = 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)): - 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, - ) - state = state_cls( - config=config, - request_capacity=1, - mtp_step=2, - device="cuda", - ) - state.initialize_page_pools(big_page_capacity=2, small_page_capacity=1, pin_memory=True) - manager.register_state(role, state) - manager.finish_state_registration() - states = manager.get_bounded_attention_states() - 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_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) - state.ssm_state.fill_(i + 7) - source_slots = manager.bounded_att_big_page_allocator.alloc(1) - manager.save_bounded_att_runtime_to_pages( - req_indices=[0], - state_token_lens=[8], - page_kind="big", - page_slot_ids=source_slots, - ) - 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], - 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]) - target_slots = manager.bounded_att_big_page_allocator.alloc(1) - manager.get_cpu_cache_operator().read_req( - 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, - ) - 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_token", "pd_runtime"]) -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 = AttentionStateManager(request_capacity=1) - config = GQAAttentionConfig( - logical_layer_groups=((0,),), - global_head_num=1, - head_dim=3, - 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, - ), - request_capacity=1, - device="cpu", - ), - ) - manager.finish_state_registration() - 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_h1_d3_state - before_layout = before[state.state_name] - monkeypatch.setattr( - "lightllm.common.attention_state.manager.FullAttKVAllocator", partial(FullAttKVAllocator, pin_memory=False) - ) - manager.allocate_storage( - token_capacity=8, - page_size=4, - big_page_tokens=8, - device="cpu", - ) - after = getter() - 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_h1_d3_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() - ) - 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()) - 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): - @deepcopy_result - @lru_cache(maxsize=None) - 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 = GQAAttentionConfig( - logical_layer_groups=((0,),), - global_head_num=1, - head_dim=2, - dtype=torch.float32, - ) - manager = AttentionStateManager(request_capacity=1) - prepared_full_state = ReversedCpuRegionsState(config, page_size=4, device="cuda") - 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.allocate_storage(8) - 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() - ) - 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.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 - - 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, - device="cpu" if transfer == "cpu" else "cuda", - pin_memory=transfer == "cpu", - ) - if transfer == "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.get_cpu_cache_operator().read_req( - 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.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.get_pd_operator().read_page( - 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 deleted file mode 100644 index 85988d5edd..0000000000 --- a/unit_tests/common/att_resources/test_compute_binding.py +++ /dev/null @@ -1,147 +0,0 @@ -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 unit_tests.common.att_resources.model_helpers import replace_logical_layer_groups - -from dataclasses import replace -from types import SimpleNamespace - -import pytest -import torch - -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 ( - 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=template, - config=replace_logical_layer_groups(template.config, ((3,), (8,))), - token_capacity=8, - device="cuda", - ) - for storage_buffer in full_state_test_buffers(state).values(): - storage_buffer.zero_() - 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_registration() - state = getattr(manager, f"{role}_{template.get_state_type()}_state") - infer = InferStateInfo() - infer.att_state_manager = manager - infer.is_draft_model = role == "draft0" - 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.latent_dim if kind == "packed_glm" else state.config.head_dim - 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( - 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( - 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() - 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 state.get_storage_layer_index(logical_layer_id=8) == 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=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( - size=(2, 1, 576), - dtype=torch.bfloat16, - device="cuda", - ) - state.write_full_att_kv(0, history_ids, history) - 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, - 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( - 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_dsv4_integration.py b/unit_tests/common/att_resources/test_dsv4_integration.py deleted file mode 100644 index badc2adf0c..0000000000 --- a/unit_tests/common/att_resources/test_dsv4_integration.py +++ /dev/null @@ -1,287 +0,0 @@ -"""通过模型工厂和生产 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", - 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() - 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.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 = [ - 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 deleted file mode 100644 index ae2ba27179..0000000000 --- a/unit_tests/common/att_resources/test_dsv4_state.py +++ /dev/null @@ -1,505 +0,0 @@ -"""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.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 create_compressed_att_state - - definition = make_full(ratio=4) - 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 - 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_cache_operator().get_page_layouts() - assert bounded[-1].state_name not in layouts - 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.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"), - 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.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 - - -@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.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.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 deleted file mode 100644 index ec5a247576..0000000000 --- a/unit_tests/common/att_resources/test_fp8_compute_binding.py +++ /dev/null @@ -1,201 +0,0 @@ -"""FP8 backends obtain calibration from the role's matching storage scheme.""" - -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 unit_tests.common.att_resources.model_helpers import replace_logical_layer_groups - -from typing import TYPE_CHECKING -from dataclasses import replace -from types import SimpleNamespace -import weakref - -import pytest -import torch - -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 -from unit_tests.common.att_resources.test_capacity import definition - -if TYPE_CHECKING: - from collections.abc import Sequence - - -@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=template, - config=replace_logical_layer_groups(template.config, ((3,), (8,))), - token_capacity=8, - device="cpu", - scales=scales, - q_scales=q_scales, - ) - ) - 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") - plain = rebuild_full_state( - template=plain, - config=plain.config, - token_capacity=8, - device="cpu", - ) - manager.register_state(role, plain) - other = rebuild_full_state( - template=template, - config=replace(template.config, logical_layer_groups=((20,), (21,)), global_head_num=4), - token_capacity=8, - device="cpu", - scales=torch.full((2, 8 if kind == "static_head" else 2), 99.0), - q_scales=torch.full((2, 4), 99.0) if kind == "static_head" else None, - ) - manager.register_state(role, other) - 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_state_manager = manager - infer.is_draft_model = role == "draft0" - 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="cpu") - 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="cpu") - - def allocate(shape: Sequence[int], dtype: torch.dtype, device: str | torch.device | int): - return torch.empty(shape, dtype=dtype, device="cpu") - - if prefill: - att.prefill_att( - q=q, - k=k, - v=v, - att_control=AttControl(logical_layer_id=8), - alloc_func=allocate, - ) - else: - att.decode_att( - q=q, - k=k, - v=v, - att_control=AttControl(logical_layer_id=8), - 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.skipif(not torch.cuda.is_available(), reason="CUDA required") -@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.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 - 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, - scale_mode=scale_mode, - ), - 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.allocate_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(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) - 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_fp8_e8m0_packed_mla_state.py b/unit_tests/common/att_resources/test_fp8_e8m0_packed_mla_state.py deleted file mode 100644 index 8e2ee61214..0000000000 --- a/unit_tests/common/att_resources/test_fp8_e8m0_packed_mla_state.py +++ /dev/null @@ -1,235 +0,0 @@ -"""独立 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, - PackedMLALogicalConfig, - FullAttentionState, - PackedMLAFullAttentionState, - create_packed_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_packed_mla_state( - config=PackedMLALogicalConfig( - logical_layer_groups=((0, 4), (8,)), - latent_dim=512, - format="packed_group64_e8m0", - tp_rank=tp_rank, - tp_world_size=2, - ), - page_size=4, - device=device, - ) - 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 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 - 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="latent 576"): - create_packed_mla_state( - config=PackedMLALogicalConfig( - logical_layer_groups=((0,),), - latent_dim=512, - format="packed_group128_fp8", - ), - device="meta", - ) - with pytest.raises(AssertionError, match="512-dimensional"): - create_packed_mla_state( - config=PackedMLALogicalConfig( - logical_layer_groups=((0,),), - latent_dim=576, - format="packed_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.packed_mla.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_cpu_kernel.py b/unit_tests/common/att_resources/test_full_att_cpu_kernel.py deleted file mode 100644 index 0197b17daa..0000000000 --- a/unit_tests/common/att_resources/test_full_att_cpu_kernel.py +++ /dev/null @@ -1,412 +0,0 @@ -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.attention_state.types import AxisMapping -from lightllm.common.attention_state.full_att import GQAAttentionConfig, 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 = GQAAttentionConfig( - 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") -@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( - 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] - 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( - 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)) - 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( - AxisMapping( - axis="head", - local_start=1, - external_start=0, - length=1, - ) - if item.axis == "head" - else item - for item in region.axis_mappings - ) - 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") - 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( - 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() - 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( - 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( - 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: - 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( - token_indices=ids, - cpu_page_buffers=pages, - cpu_page_indices=[0], - layout=layout, - ) - else: - 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") -@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( - 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() - 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.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) - with pytest.raises(AssertionError, match="token indices must reside on CUDA"): - 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, - 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_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_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_page( - 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_dp_transfer.py b/unit_tests/common/att_resources/test_full_att_dp_transfer.py deleted file mode 100644 index 78bc924405..0000000000 --- a/unit_tests/common/att_resources/test_full_att_dp_transfer.py +++ /dev/null @@ -1,221 +0,0 @@ -"""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.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 - - -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]) -@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_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_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() diff --git a/unit_tests/common/att_resources/test_full_att_transfer.py b/unit_tests/common/att_resources/test_full_att_transfer.py deleted file mode 100644 index 8c11087762..0000000000 --- a/unit_tests/common/att_resources/test_full_att_transfer.py +++ /dev/null @@ -1,338 +0,0 @@ -"""各 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.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 - -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": "packed", - "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( - state_name=source.state_name, - 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_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() - - # 比较整行原始字节：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( - 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) - 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_page( - token_indices=source_ids, - pd_page_buffers=pages, - pd_page_index=1, - layout=layout, - ) - state.read_from_pd_page( - 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( - state_name=state.state_name, - 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( - 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=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): - 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=token_indices, - cpu_page_buffers=pages[:, 16:-16], - cpu_page_indices=page_indices, - layout=layout, - cpu_page_ready_flags=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 deleted file mode 100644 index 1339ebb4b0..0000000000 --- a/unit_tests/common/att_resources/test_full_state.py +++ /dev/null @@ -1,283 +0,0 @@ -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.attention_state.full_att import ( - FullAttentionState, - GQAFullAttentionState, - MLAFullAttentionState, - GQAAttentionConfig, - MLAConfig, -) -from lightllm.common.attention_state.types import ( - 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, - ) - 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 - - -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_mappings=( - AxisMapping( - axis="head", - local_start=0, - external_start=0, - length=1, - ), - AxisMapping( - axis="dim", - local_start=0, - external_start=0, - length=4, - ), - ), - ) - ) - offset += 128 - return StateTransferLayout( - state_name=state.state_name, - 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( - 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_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_page( - pd_page_buffers=pages, - token_indices=ids, - pd_page_index=1, - layout=layout, - ) - return - state.write_to_pd_page( - 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_page( - 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() - - -@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( - 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( - 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() - - -@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( - size=(layout.byte_end + 256,), - fill_value=53, - dtype=torch.uint8, - device="cuda", - ) - with pytest.raises(RuntimeError, match="shape"): - 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, - 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 = GQAAttentionConfig( - 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.allocate_storage(12) - target = GQAFullAttentionState(config, page_size=4, device=device) - 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_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_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, logical_layer_groups=config.logical_layer_groups[:1]), page_size=4, device=device - ) - 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_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 deleted file mode 100644 index 26d4425ac7..0000000000 --- a/unit_tests/common/att_resources/test_full_storage_initialization.py +++ /dev/null @@ -1,120 +0,0 @@ -"""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.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 - -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.usefixtures("runtime"), -] - - -@pytest.mark.parametrize("kind,encoding,model_cls", CASES) -def test_model_storage_write_and_copy_preserve_all_fields( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, - kind: str, - encoding: str, - model_cls: type[TpPartBaseModel], -): - _, (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") - is_mla = isinstance(state, (MLAFullAttentionState, PackedMLAFullAttentionState)) - dim = config.latent_dim if is_mla else config.head_dim - kv = torch.randn( - size=(2, 1 if is_mla 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() - 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( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, - mode: str, - kind: str, - encoding: str, - model_cls: type[TpPartBaseModel], -): - _, (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 - 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)) - 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_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_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_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_layout_identity.py b/unit_tests/common/att_resources/test_gqa_layout_identity.py deleted file mode 100644 index f4c3174f44..0000000000 --- a/unit_tests/common/att_resources/test_gqa_layout_identity.py +++ /dev/null @@ -1,114 +0,0 @@ -from dataclasses import replace -from types import SimpleNamespace - -import pytest -import torch - -from lightllm.common.attention_state import AttentionStateManager -from lightllm.common.attention_state.full_att import GQAAttentionConfig, create_gqa_state -from lightllm.common.basemodel.attention.base_att import AttControl - - -@pytest.mark.parametrize("kv_type", ["None", "int4kv", "int8kv", "fp8kv_spt", "fp8kv_sph"]) -def test_gqa_layouts_coexist_and_match_across_tp_and_registration_order(kv_type, pd_args, monkeypatch): - config = GQAAttentionConfig(logical_layer_groups=((0,),), global_head_num=8, head_dim=128, dtype=torch.bfloat16) - configs = [ - config, - replace(config, logical_layer_groups=((1,),), global_head_num=16), - replace(config, logical_layer_groups=((2,),), head_dim=256), - ] - source, target = (AttentionStateManager(request_capacity=0) for _ in range(2)) - for manager, items, tp in ((source, configs, 1), (target, reversed(configs), 2)): - for item in items: - state = create_gqa_state(replace(item, tp_rank=tp - 1, tp_world_size=tp), kv_type=kv_type, device="meta") - manager.register_state("main", state) - manager.finish_state_registration() - states = source.get_full_attention_states("main") - assert len({state.state_name for state in states}) == 3 - for layer in range(3): - state = source.get_state("main", "gqa", layer) - peer = target.get_state("main", "gqa", layer) - assert state.state_name == peer.state_name - assert state.config.global_head_num == configs[layer].global_head_num - assert state.config.head_dim == configs[layer].head_dim - for getter in ( - lambda m: m.get_cpu_cache_operator().get_page_layouts(), - lambda m: m.get_pd_operator().get_pd_full_att_page_layouts(), - ): - layouts, peer_layouts = getter(source), getter(target) - assert set(layouts) == set(peer_layouts) == {state.state_name for state in states} - ordered = sorted(layouts.values(), key=lambda layout: layout.byte_start) - assert all(a.byte_end <= b.byte_start for a, b in zip(ordered, ordered[1:])) - copied = [] - for state in states: - monkeypatch.setattr( - state, - "copy_tokens_from_dp", - lambda *, remote_dp_full_att_state, **_: copied.append(remote_dp_full_att_state), - ) - source.get_dp_shared_operator().copy_from_dp( - remote_att_manager=target, - src_token_indices=torch.tensor([0]), - dst_token_indices=torch.tensor([1]), - src_big_page_ids=[], - dst_big_page_ids=[], - src_req_idx=0, - dst_req_idx=1, - ) - assert copied == [getattr(target, state.state_name) for state in states] - - -@pytest.mark.parametrize("prefill", [False, True]) -def test_fa3_reads_descale_for_each_head_layout(prefill, monkeypatch): - from lightllm.common.basemodel.attention.fa3 import fp8 as fa3_fp8 - - manager = AttentionStateManager(request_capacity=2) - for layer, heads in enumerate((2, 4)): - config = GQAAttentionConfig( - logical_layer_groups=((layer,),), global_head_num=heads, head_dim=128, dtype=torch.bfloat16 - ) - state = create_gqa_state(config, kv_type="fp8kv_sph", device="cpu") - state.scales = torch.arange(1, 2 * heads + 1, dtype=torch.float32).view(1, -1) - state.q_scales = torch.arange(1, heads + 1, dtype=torch.float32).view(1, heads) - manager.register_state("main", state) - manager.finish_state_registration() - # prefill 必须按 b_seq_len 取 batch，不能读这个未设置的请求 batch。 - infer = SimpleNamespace( - att_state_manager=manager, - att_state_role="main", - batch_size=None, - b_seq_len=torch.ones(2, dtype=torch.int32), - max_q_seq_len=1, - ) - cls = fa3_fp8.Fp8Fa3PrefillAttState if prefill else fa3_fp8.Fp8Fa3DecodeAttState - parent = fa3_fp8.Fa3PrefillAttState if prefill else fa3_fp8.Fa3DecodeAttState - monkeypatch.setattr(parent, "init_state", lambda self: None) - monkeypatch.setattr(fa3_fp8, "q_per_head_static_fp8_quant", lambda q, scales: q) - observed = {} - - def attend(**kwargs): - observed.update(kwargs) - return kwargs["q"] - - monkeypatch.setattr(fa3_fp8, "flash_attn_with_kvcache", attend) - att = cls(backend=None, _infer_state_weak_ref=lambda: infer) - att.b_att_seq_len = torch.ones(3, dtype=torch.int32) - att.page_table = att.cu_seqlens_q = att.cu_seqlens_k = None - att.decode_max_q_seq_len = 1 - att.causal = True - att.init_state() - batch_size = infer.b_seq_len.shape[0] if prefill else att.b_att_seq_len.shape[0] - for state in manager.get_full_attention_states("main"): - heads = state.config.head_num - layer = state.config.logical_layer_groups[0][0] - q = k = v = torch.zeros((2, heads, state.config.head_dim)) - observed.clear() - run = att.prefill_att if prefill else att.decode_att - run(q=q, k=k, v=v, att_control=AttControl(logical_layer_id=layer)) - row = state.scales[0] - k_descale = observed["k_descale"] - v_descale = observed["v_descale"] - assert k_descale.data_ptr() == row[:heads].data_ptr() - assert v_descale.data_ptr() == row[heads:].data_ptr() - torch.testing.assert_close(k_descale, row[:heads].expand(batch_size, heads)) - torch.testing.assert_close(v_descale, row[heads:].expand(batch_size, heads)) diff --git a/unit_tests/common/att_resources/test_gqa_pd_kernel.py b/unit_tests/common/att_resources/test_gqa_pd_kernel.py deleted file mode 100644 index 2870ca126e..0000000000 --- a/unit_tests/common/att_resources/test_gqa_pd_kernel.py +++ /dev/null @@ -1,283 +0,0 @@ -"""GQA PD 直接搬运：原始位模式、物理页选择、TP 切片及有效 token 范围。""" - -from __future__ import annotations - -from dataclasses import replace -from typing import TYPE_CHECKING - -import pytest -import torch - -from lightllm.common.attention_state.full_att import GQAAttentionConfig, GQAFullAttentionState -from lightllm.common.attention_state.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=GQAAttentionConfig( - logical_layer_groups=((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.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): - 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( - 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), - 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_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=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_page( - 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=GQAAttentionConfig( - 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", - ) - state.allocate_storage(token_capacity=8) - state.k_storage.fill_(19) - state.v_storage.fill_(19) - regions = state.get_pd_transfer_regions() - layout = 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, - ) - 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_page( - token_indices=ids, - pd_page_buffers=pages, - pd_page_index=1, - layout=layout, - ) - state.read_from_pd_page( - 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=GQAAttentionConfig( - logical_layer_groups=((0,), (2,)), - global_head_num=2, - head_dim=128, - dtype=torch.bfloat16, - ), - page_size=4, - device="cuda:1", - ) - state.allocate_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( - state_name=state.state_name, - 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_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_page( - 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 deleted file mode 100644 index ac19dc39ab..0000000000 --- a/unit_tests/common/att_resources/test_gqa_transfer.py +++ /dev/null @@ -1,192 +0,0 @@ -"""GQA 共用搬运 kernel 在受限 grid 下完整覆盖多轮、跨页和尾块数据。""" - -from typing import TYPE_CHECKING - -import pytest -import torch - -from lightllm.common.attention_state.full_att import GQAAttentionConfig, GQAFullAttentionState -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: - 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 = GQAAttentionConfig( - 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.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( - state_name=source.state_name, - 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() - - -@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.attention_state.full_att import gqa_transfer - - pd_args.pd_kv_page_size = 2 - config = GQAAttentionConfig( - 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.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: - 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( - state_name=source.state_name, - 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() diff --git a/unit_tests/common/att_resources/test_hybrid_layer_mapping.py b/unit_tests/common/att_resources/test_hybrid_layer_mapping.py deleted file mode 100644 index fb9ce7ba5b..0000000000 --- a/unit_tests/common/att_resources/test_hybrid_layer_mapping.py +++ /dev/null @@ -1,130 +0,0 @@ -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("mode", ["native", "replay"]) -@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, - mode: str, -) -> 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", - mtp_step=2, - ssm_state_mode=mode, - ssm_replay_config={"cache_len": 8} if mode == "replay" else {}, - ) - 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", - ) - if mode == "native": - stride = bounded_state.mtp_step + 1 - assert bounded_state.ssm_state.shape[1] == bounded_state.request_capacity * stride - if mode == "replay": - assert bounded_state.capacity == 8 - 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 deleted file mode 100644 index 0f55fe5986..0000000000 --- a/unit_tests/common/att_resources/test_layer_groups.py +++ /dev/null @@ -1,164 +0,0 @@ -"""逻辑层共享物理 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.attention_state import AttentionStateManager -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, - replace_logical_layer_groups, -) -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_logical_layer_groups(template.config, ((0, 4), (8, 11))) - state = rebuild_full_state( - template=template, - config=config, - token_capacity=8, - device="cpu", - ) - 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) - - 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] - - 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(kind: str, groups: tuple) -> None: - with pytest.raises(AssertionError): - replace_logical_layer_groups(definition(kind).config, groups) - - -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,) = 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，保持四份物理层存储。 - assert state.config.logical_layer_groups == ((10,), (11,), (12, 14), (13, 15)) - 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") -@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=template, - config=replace_logical_layer_groups(template.config, ((0, 4), (8,))), - token_capacity=8, - device="cuda", - ) - target = rebuild_full_state( - template=template, - config=replace_logical_layer_groups(template.config, ((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 = 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") - if transfer == "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_cache_operator().get_page_layouts(), - ) - 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_cache_operator().get_page_layouts(), - ) - else: - 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_operator().get_pd_full_att_page_layouts(), - ) - 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_operator().get_pd_full_att_page_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_layer_state_selection.py b/unit_tests/common/att_resources/test_layer_state_selection.py deleted file mode 100644 index 52337199a7..0000000000 --- a/unit_tests/common/att_resources/test_layer_state_selection.py +++ /dev/null @@ -1,216 +0,0 @@ -"""按角色、层归属及方案查询资源，不在推理层保存绑定状态。""" - -from contextlib import nullcontext -from types import SimpleNamespace - -import pytest -import torch - -from lightllm.common.attention_state import AttentionStateManager -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_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 - - -def test_creation_does_not_bind_layers(monkeypatch: pytest.MonkeyPatch) -> None: - 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.max_req_num = 0 - model.torch_memory_saver = SimpleNamespace(region=lambda **_: nullcontext()) - state = GQAFullAttentionState( - layout_config=GQAAttentionConfig( - logical_layer_groups=((0,),), - global_head_num=1, - head_dim=128, - dtype=torch.bfloat16, - ), - device="meta", - ) - monkeypatch.setattr( - LlamaTpPartModel, - "create_full_att_states", - classmethod(lambda cls, **kwargs: [state]), - ) - - model.prepare_attention_states() - - assert model.att_state_manager.main_gqa_bf16_h1_d128_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 = 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", - ), - ) - 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_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 = [] - for logical_layer_id in range(4): - layer = TransformerLayerInfer.__new__(TransformerLayerInfer) - layer.layer_num_ = logical_layer_id - model.layers_infer.append(layer) - - model.prepare_attention_states() - - assert all(not hasattr(layer, "full_att_state_name") for layer in model.layers_infer) - manager.finish_state_registration() - state = manager.get_state(role, "gqa", 0) - assert state.state_name == f"{role}_gqa_bf16_h1_d128_state" - assert manager.get_state(role, "gqa", 2) is state - assert manager.get_state(role, "gqa", 3).state_name == f"{role}_gqa_fp16_h1_d128_state" - - -@pytest.mark.parametrize("finalized", [False, True]) -def test_prefix_and_logical_layer_select_quantized_state(finalized: bool) -> None: - from dataclasses import replace - - manager = AttentionStateManager(request_capacity=1) - plain = definition("gqa") - quantized = definition("int8") - quantized = type(quantized)( - layout_config=replace_logical_layer_groups(quantized.config, ((10, 12),)), - device="meta", - ) - 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 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_state("main", "gqa", 99) - - -def test_selection_requires_finished_creation_and_reuses_cache() -> None: - manager = AttentionStateManager(request_capacity=1) - plain = definition("gqa") - manager.register_state("main", plain) - logical_layer_id = plain.config.logical_layer_groups[0][0] - 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 = AttentionStateManager(request_capacity=1) - plain = definition("gqa") - 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_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 = 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_state("main", state.get_state_type(), 3) is state - - -def test_mla_selection_excludes_same_layer_tail() -> None: - from dataclasses import replace - - manager = AttentionStateManager(request_capacity=3) - tail = make_window() - template = definition("mla") - state = type(template)( - layout_config=replace_logical_layer_groups(template.config, ((3, 7),)), - device="meta", - ) - 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_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 - ) - 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) - for model_name in ("draft1", "draft10"): - manager.register_state(model_name, definition("gqa")) - manager.finish_state_registration() - assert manager.get_state(role, "gqa", 0).state_name == f"{role}_gqa_bf16_h2_d16_state" - - -@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_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})) - state = type(template)(config=config, request_capacity=3, device="cpu") - manager.register_state("main", state) - states.append(state) - manager.finish_state_registration() - logical_layer_id = template.config.logical_layer_ids[0] - for state in states: - 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_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 deleted file mode 100644 index 9405b46c89..0000000000 --- a/unit_tests/common/att_resources/test_manager_helpers.py +++ /dev/null @@ -1,316 +0,0 @@ -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.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 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 - -if TYPE_CHECKING: - 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"): - return MainDraftTestResourceManager(device) - - -def layouts_for(manager: AttentionStateManager, kind: str = "cpu"): - layouts = [] - if kind != "att_state": - layouts += [layout_for(state) for state in (manager.main_gqa_fp32_h1_d4_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) - ] - - offset = 0 - result = {} - for layout in layouts: - state_end = offset + layout.byte_end - result[layout.state_name] = 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: AttentionStateManager): - for i, state in enumerate((manager.main_gqa_fp32_h1_d4_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_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) - state.ssm_state[:, 3:6].fill_(47 + i) - slots = manager.bounded_att_big_page_allocator.alloc(2) - manager.save_bounded_att_runtime_to_pages( - req_indices=[0, 1], - state_token_lens=[4, 8], - page_kind="big", - page_slot_ids=slots, - ) - if manager.main_kda_conv_fp32_ssm_fp32_state.conv_state.is_cuda: - torch.cuda.synchronize() - 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(device="cuda") - slots = init_source(manager) - manager.restore_bounded_att_runtime_from_pages( - req_indices=[2], - state_token_lens=[4], - page_kind="big", - 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() - assert (state.ssm_state[:, 6] == 31 + i).all() - - def fail(*args: object, **kwargs: object): - raise RuntimeError("simulated copy failure") - - 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_runtime_to_pages( - req_indices=[0, 1], - state_token_lens=[4, 8], - page_kind="big", - page_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 - - -@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() - calls = [] - - def record_lengths( - req_indices: list[int], - state_token_lens: list[int], - page_kind: PageKind, - page_slot_ids: list[int], - ) -> None: - 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) - getattr(manager, operation)( - req_indices=[0, 2], - state_token_lens=[17, 29], - page_kind="big", - page_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() - 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_state_pages_from_cpu", lambda **kwargs: None) - restored = [] - - def record_restore( - req_indices: list[int], - state_token_lens: list[int], - page_kind: PageKind, - page_slot_ids: list[int], - ) -> None: - restored.append((req_indices, state_token_lens, page_kind, page_slot_ids)) - - monkeypatch.setattr(manager, "restore_bounded_att_runtime_from_pages", record_restore) - # 只加载绝对位置 [64, 70)，本批长度 6；恢复位置必须是 70，且末尾碎页使用 LOAD。 - 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"), - 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_page", "read_from_pd_page"]) -def test_pd_bounded_only_forwards_prefix_length(operation: str, monkeypatch: pytest.MonkeyPatch): - manager = composition() - 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.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", - 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") -@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"), composition("cuda") - 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_runtime_to_pages( - req_indices=[page_count - 1], - state_token_lens=[token_count], - page_kind="small", - page_slot_ids=small_slots, - ) - torch.cuda.synchronize() - source.get_cpu_cache_operator().write_req( - 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() - - # 保留末槽给 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.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"), - state_token_len=token_count, - 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.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") - big_slots = init_source(source) - small_slots = source.bounded_att_small_page_allocator.alloc(1) - source.save_bounded_att_runtime_to_pages( - req_indices=[1], - state_token_lens=[6], - page_kind="small", - page_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( - 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.get_cpu_cache_operator().write_req( - 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() - 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.get_cpu_cache_operator().write_req( - token_indices=token_ids, - cpu_page_buffers=pages, - cpu_page_indices=page_ids, - layouts=layouts, - cpu_page_ready_flags=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 deleted file mode 100644 index 9544084bb6..0000000000 --- a/unit_tests/common/att_resources/test_model_calibration.py +++ /dev/null @@ -1,261 +0,0 @@ -"""统一标定阶段不依赖 State 注册顺序，不在布局定义阶段加载参数。""" - -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.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 -from lightllm.server.core.objs.start_args_type import StartArgs - -if TYPE_CHECKING: - from pathlib import Path - from easydict import EasyDict - - -@pytest.mark.parametrize("encoding", ["fp8kv_sph", "fp8kv_spt"]) -def test_calibration_binds_all_model_states_after_creation(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 - - # 反向创建、注册，标定仍按 main 和 draft 的数值顺序绑定。 - manager = AttentionStateManager(request_capacity=1) - for index in reversed(range(len(models))): - model = models[index] - 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( - 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.state_name is None - assert state.scales is None and state.q_scales is None - 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_state_role) - 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.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": - 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.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) - 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 = AttentionStateManager(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=GQAAttentionConfig( - 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.register_state(role, state) - if encoding != "None": - expected.append((state, row)) - with pytest.raises(AssertionError, match="State registration"): - manager.load_full_att_calibration(args=args) - 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) - 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.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)) - 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.attention_state.full_att import GQAAttentionConfig, create_gqa_state - - state = create_gqa_state( - config=GQAAttentionConfig( - logical_layer_groups=((0, 1), (2,)), - global_head_num=2, - head_dim=16, - dtype=torch.bfloat16, - ), - kv_type=encoding, - device="cpu", - ) - manager = AttentionStateManager(request_capacity=1) - manager.register_state("main", state) - manager.finish_state_registration() - with pytest.raises(AssertionError, match="calibration must be bound"): - state.allocate_storage(4) - if encoding == "fp8kv_sph": - 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.allocate_storage(4) - assert state.k_storage.shape[0] == 2 diff --git a/unit_tests/common/att_resources/test_model_config_normalization.py b/unit_tests/common/att_resources/test_model_config_normalization.py deleted file mode 100644 index 6e3dfbc92c..0000000000 --- a/unit_tests/common/att_resources/test_model_config_normalization.py +++ /dev/null @@ -1,235 +0,0 @@ -"""模型启动和 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] - - -@pytest.mark.parametrize( - "model_type,text_model_type,wrapper", - [ - ("internvl_chat", kind, "llm_config") - for kind in ("phi3", "internlm2", "llama", "qwen2", "deepseek_v2", "qwen3", "qwen3_moe") - ] - + [ - ("neo_chat", "qwen3", "llm_config"), - ("neo_chat", "qwen3_moe", "llm_config"), - ("qwen3_omni_moe", "qwen3_moe", "thinker_config"), - ("qwen3_vl", "qwen3", "text_config"), - ("llama", "llama", None), - ], -) -def test_multimodal_cpu_layout_matches_runtime(model_type, text_model_type, wrapper, tmp_path, monkeypatch): - from dataclasses import asdict - from lightllm.common.attention_state import AttentionStateManager - from lightllm.models import get_model_class - from lightllm.utils.cpu_cache_layout import build_cpu_cache_layouts - from lightllm.utils.envs_utils import get_env_start_args - from lightllm.utils.llm_utils import get_llm_model_class - - text = dict( - model_type=text_model_type, - num_hidden_layers=2, - num_attention_heads=4, - num_key_value_heads=2, - hidden_size=256, - kv_lora_rank=64, - qk_rope_head_dim=32, - ) - raw = ( - text - if wrapper is None - else { - "model_type": model_type, - wrapper: {"text_config": text} if wrapper == "thinker_config" else text, - } - ) - original = deepcopy(raw) - (tmp_path / "config.json").write_text(json.dumps(raw)) - args = StartArgs(model_dir=str(tmp_path), data_type="bfloat16", enable_cpu_cache=True, cpu_cache_token_page_size=8) - monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(asdict(args))) - get_env_start_args.cache_clear() - get_llm_model_class.cache_clear() - try: - layouts = build_cpu_cache_layouts() - cls = get_model_class(raw) - model = cls.__new__(cls) - model.weight_dir_, model.finetune_config = str(tmp_path), None - model._init_config() - normalized = cls.normalize_model_config(raw) - assert model.config == normalized - assert cls.normalize_model_config(normalized) == normalized - assert raw == original - manager = AttentionStateManager(request_capacity=0) - for state in cls.create_full_att_states( - model_config=model.config, - args=args, - model_logical_layer_ids=(0, 1), - ): - manager.register_state("main", state) - manager.finish_state_registration() - assert layouts and layouts == manager.get_cpu_cache_operator().get_page_layouts() - finally: - get_llm_model_class.cache_clear() - get_env_start_args.cache_clear() diff --git a/unit_tests/common/att_resources/test_model_initialization.py b/unit_tests/common/att_resources/test_model_initialization.py deleted file mode 100644 index 09e88b4a93..0000000000 --- a/unit_tests/common/att_resources/test_model_initialization.py +++ /dev/null @@ -1,185 +0,0 @@ -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.attention_state.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_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 - 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_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) - 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.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_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( - 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_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_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") - 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 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 - assert resources.bounded_att_small_page_allocator.size == 2 - assert not hasattr(model, "mtp_state_cursors") - linear_state = main_linear_state(resources) - assert linear_state.mtp_state_index is not None - assert linear_state.mtp_enabled == bool(mtp) - if mode == "prefill": - 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 - 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 - assert draft_state.capacity == full_state.capacity - 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 deleted file mode 100644 index a571817fae..0000000000 --- a/unit_tests/common/att_resources/test_model_resources.py +++ /dev/null @@ -1,214 +0,0 @@ -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 - -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.attention_state.allocators import FullAttKVAllocator -from unit_tests.common.att_resources.model_helpers import create_full_attention_roles -from lightllm.common.attention_state.full_att import ( - GQAFullAttentionState, - GQAAttentionConfig, -) -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 -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) - 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_head_dim=512 if glm else 16, - global_linear_k_heads=gk, - global_linear_v_heads=gv, - head_linear_k_dim=8, - head_linear_v_dim=8, - conv_kernel_size=3, - conv_state_dtype=torch.bfloat16, - ssm_state_dtype=torch.float32, - full_attention_interval=4 if glm else 2, - all_layer_num=4, - ) - if glm: - full = PackedMLAFullAttentionState( - PackedMLAAttentionConfig( - 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, - ) - full.allocate_storage(16) - else: - full = GQAFullAttentionState( - GQAAttentionConfig( - 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.allocate_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, - full_state=main, - full_allocator=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=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 - 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_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_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_runtime_to_pages( - req_indices=[0], - state_token_lens=[8], - page_kind="big", - 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())) - 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.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_cache_operator().get_page_layouts(), - big_page_slot_ids=[0], - ) - torch.cuda.synchronize() - for rank, (resources, full) in enumerate(targets): - if glm: - resources.main_compressor_indexer_w4_s4_bf16_state.buffer.fill_(13) - slot = resources.bounded_att_big_page_allocator.alloc(1)[0] - resources.get_cpu_cache_operator().read_req( - cpu_page_buffers=buffer, - cpu_page_indices=[1], - token_indices=torch.arange( - start=8, - end=16, - dtype=torch.long, - device="cuda", - ), - state_token_len=8, - layouts=resources.get_cpu_cache_operator().get_page_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: - # 无重叠窗口在组边界不依赖旧历史，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() - - -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( - glm=True, - tp=1, - rank=0, - mtp=0, - ) - 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 deleted file mode 100644 index 70970756df..0000000000 --- a/unit_tests/common/att_resources/test_model_storage_definition.py +++ /dev/null @@ -1,276 +0,0 @@ -"""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.cpu_cache_layout import build_cpu_cache_layouts -from lightllm.utils.envs_utils import get_env_start_args -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: - 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, - ), - ( - 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, - ), -] - - -@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_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_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_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_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") - - -@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 = 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_registration() - args = StartArgs( - page_size=4, - run_mode="prefill", - pd_kv_page_size=8, - pd_kv_page_num=3, - ) - 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 - 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.attention_state.manager.get_current_device_id", lambda: 3) - monkeypatch.setattr( - "lightllm.common.attention_state.manager.FullAttKVAllocator", partial(FullAttKVAllocator, pin_memory=False) - ) - template = rebuild_full_state( - template=state, - config=state.config, - token_capacity=0, - device="meta", - ) - 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, - device="cpu", - ) - assert resources.full_att_allocator.size == 28 - assert resources.get_pd_operator().get_page_bytes() == reference.get_pd_operator().get_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.attention_state.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_state_manager.finish_state_registration() - runtime_layouts = model.att_state_manager.get_cpu_cache_operator().get_page_layouts() - assert set(cpu_layouts) == set(runtime_layouts) == {"main_gqa_bf16_h8_d16_state", "main_gqa_fp16_h8_d16_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 deleted file mode 100644 index ea114d2ba9..0000000000 --- a/unit_tests/common/att_resources/test_mtp_state_commit.py +++ /dev/null @@ -1,114 +0,0 @@ -"""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.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 = 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_registration() - 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"), - torch.tensor([3], dtype=torch.int32, device="cuda"), - role="main", - ) - 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_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"), - torch.tensor([1], dtype=torch.int32, device="cuda"), - role="draft0", - ) - 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") -def test_commit_without_recurrent_state_and_invalid_role(): - 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"), - 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"), - 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"), - 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_state_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=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] - assert accepted_index.tolist() == expected_mask - 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 deleted file mode 100644 index f66767b9f4..0000000000 --- a/unit_tests/common/att_resources/test_network_transport.py +++ /dev/null @@ -1,320 +0,0 @@ -"""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.attention_state.manager import AttentionStateManager - 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 - - -@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: AttentionStateManager, - 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_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 - - -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( - 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( - 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 - - transporter = NixlKVTransporter(os.getpid(), device, buffer) - connection.send( - ( - "metadata", - PDAgentMetadata( - agent_name=transporter.agent_name, - agent_metadata=transporter.agent_metadata, - num_pages=transporter.num_pages, - page_reg_desc=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_compressor_indexer_w4_s4_bf16_state.buffer) - ): - tensor.copy_(torch.arange(tensor.numel(), device=device).reshape(tensor.shape) % (23 + i * 7)) - 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): - copy_page( - 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), - task_kind=kind, - req_idx=1, - state_token_len=task.end_kv_index, - write_to_pd_page=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=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) - 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=transporter, - stage="done", - pending=pending, - key=request.get_key(), - ) - assert done.get_key() == request.get_key() - if backend == "nccl": - 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( - 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), - task_kind=kind, - req_idx=2, - state_token_len=request.end_kv_index, - write_to_pd_page=False, - ) - copy_stream.synchronize() - 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() - 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() - 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 deleted file mode 100644 index d7d285493b..0000000000 --- a/unit_tests/common/att_resources/test_no_attention_draft.py +++ /dev/null @@ -1,84 +0,0 @@ -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.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 -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 = 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") - 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_state_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_attention_states() - resources.finish_state_registration() - draft._init_att_backend() - draft._init_att_backend1() - 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)) - - 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_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 - 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"]) -def test_attention_draft_prepares_its_own_state_before_shared_allocation(mode: str): - draft, resources = draft_model(mode) - 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.state_name == f"draft0_{state.get_state_type()}_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 deleted file mode 100644 index 032f7674eb..0000000000 --- a/unit_tests/common/att_resources/test_page_pool_storage.py +++ /dev/null @@ -1,167 +0,0 @@ -from __future__ import annotations - -import gc -from typing import TYPE_CHECKING - -import pytest -import torch -import torch.multiprocessing as mp - -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 - 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( - 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) - 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( - 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) - 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 deleted file mode 100644 index e146261f11..0000000000 --- a/unit_tests/common/att_resources/test_paged_storage.py +++ /dev/null @@ -1,22 +0,0 @@ -from __future__ import annotations - -import pytest - -from lightllm.common.attention_state.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 deleted file mode 100644 index de13eb28fc..0000000000 --- a/unit_tests/common/att_resources/test_quantized_full_state.py +++ /dev/null @@ -1,388 +0,0 @@ -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 typing import TYPE_CHECKING -from math import prod - -import pytest -import torch - -from lightllm.common.attention_state.full_att import ( - GroupQuantizedAttentionConfig, - GroupQuantizedFullAttentionState, - StaticFP8AttentionConfig, - StaticFP8FullAttentionState, - PackedMLAAttentionConfig, - PackedMLALogicalConfig, - PackedMLAFullAttentionState, -) -from lightllm.common.attention_state.types import ( - AxisMapping, - StateTransferLayout, - TransferRegion, -) - -if TYPE_CHECKING: - from collections.abc import Sequence - from lightllm.common.attention_state.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( - logical_layer_groups=((1,), (3,)), - global_head_num=2, - head_dim=16, - scale_dtype=torch.float16, - bits=int(kind[-1]), - ) - prepared_full_state = GroupQuantizedFullAttentionState(config, page_size=4, device=device) - prepared_full_state.allocate_storage(8) - return prepared_full_state - if kind.startswith("fp8_"): - per_head = kind == "fp8_head" - config = StaticFP8AttentionConfig( - logical_layer_groups=((1,), (3,)), - global_head_num=2, - 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 - ) - 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.allocate_storage(8) - return prepared_full_state - fp8 = kind == "packed_fp8" - config = PackedMLAAttentionConfig( - 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) - 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_mappings=tuple( - AxisMapping( - axis=axis, - local_start=0, - external_start=0, - length=dim, - ) - for axis, dim in zip(field.axes, shape) - if axis not in ("layer", "token") - ), - ) - ) - offset += prod(shape) * itemsize - return StateTransferLayout( - state_name=state.state_name, - 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=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( - 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( - 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_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_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_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_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_page( - 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) - - 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_dp( - src_token_indices=src_indices, - dst_token_indices=dst_indices, - remote_dp_full_att_state=source, - ) - return - target.copy_tokens_from_dp( - src_token_indices=src_indices, - dst_token_indices=dst_indices, - remote_dp_full_att_state=source, - ) - assert_token_fields_equal( - source=source, - source_ids=[0, 1], - target=target, - target_ids=[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.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") - state.write_full_att_kv(1, ids, kv) - if isinstance(state, GroupQuantizedFullAttentionState): - 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) - 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( - actual=dequant, - expected=kv.float(), - rtol=0.13, - atol=0.03, - ) - else: - view = state.get_layer_view(1) - assert (view.indexer_k == 29).all() - if state.config.dtype != torch.float8_e4m3fn: - 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() - - -@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", - ) - 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 torch.equal(state.scales, source.scales) - state.device = torch.device("cpu") - 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)) diff --git a/unit_tests/common/att_resources/test_quantized_paged_storage.py b/unit_tests/common/att_resources/test_quantized_paged_storage.py deleted file mode 100644 index ca97cc7b70..0000000000 --- a/unit_tests/common/att_resources/test_quantized_paged_storage.py +++ /dev/null @@ -1,37 +0,0 @@ -from unit_tests.common.att_resources.model_helpers import full_state_test_buffers - -import torch - -from lightllm.common.attention_state.full_att import ( - GroupQuantizedAttentionConfig, - GroupQuantizedFullAttentionState, -) - - -def test_quantized_kv_buffers_include_the_full_hold_page(): - allocations = [] - for bits in (8, 4): - state = GroupQuantizedFullAttentionState( - GroupQuantizedAttentionConfig( - logical_layer_groups=((0,), (1,), (2,)), - global_head_num=2, - head_dim=16, - scale_dtype=torch.float16, - bits=bits, - ), - page_size=4, - device="cpu", - ) - state.allocate_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 deleted file mode 100644 index 67c2fe8c2e..0000000000 --- a/unit_tests/common/att_resources/test_role_transfers.py +++ /dev/null @@ -1,149 +0,0 @@ -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_and_dp_windows( - glm: bool, - pd_args: EasyDict, - monkeypatch: pytest.MonkeyPatch, -): - pd_args.pd_kv_page_size = 128 - 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 = source.get_full_attention_states(role)[0] - 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}_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_operator().get_page_bytes(), dtype=torch.uint8, device="cuda") - for kind, getter in ( - ("kv", "get_pd_full_att_page_layouts"), - ("att_state", "get_pd_bounded_att_page_layouts"), - ): - source.get_pd_operator().write_page( - pd_page_buffers=page.unsqueeze(0), - 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} - ), - ) - 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} - ), - ) - 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.get_dp_shared_operator().copy_from_dp( - remote_att_manager=source, - src_token_indices=src_indices, - dst_token_indices=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 = target.get_full_attention_states(role)[0] - 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}_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 (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( - req_indices=[0], - state_token_lens=[8], - page_kind="big", - page_slot_ids=[slot], - ) - torch.cuda.synchronize() - 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_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.get_cpu_cache_operator().read_req( - cpu_page_buffers=cpu, - cpu_page_indices=[0], - token_indices=torch.arange( - start=8, - end=16, - dtype=torch.long, - device="cuda", - ), - state_token_len=8, - layouts=target.get_cpu_cache_operator().get_page_layouts(), - big_page_slot_ids=[dst_slot], - req_idx=1, - ) - for index, role in enumerate(("main", "draft0")): - state = target.get_full_attention_states(role)[0] - 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: - # 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 deleted file mode 100644 index 1371cfe47e..0000000000 --- a/unit_tests/common/att_resources/test_runtime_transport.py +++ /dev/null @@ -1,571 +0,0 @@ -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 - -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 -from contextlib import nullcontext -import pickle -import traceback -import uuid - -import pytest -import torch -import torch.multiprocessing as mp - -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 -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 = AttentionStateManager(request_capacity=2) - manager.register_state("main", state) - manager.full_att_allocator = FullAttKVAllocator(8, 4, pin_memory=False) - 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()): - 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( - 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_page=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_page=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_page=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_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( - token_indices=source_indices, - cpu_page_buffers=cpu_pages, - cpu_page_indices=[0], - layouts=managers[0].get_cpu_cache_operator().get_page_layouts(), - ) - else: - 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_cache_operator().get_page_layouts(), - ) - if device == "cuda": - torch.cuda.synchronize() - 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_cache_operator().get_page_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].get_dp_shared_operator().copy_from_dp( - remote_att_manager=managers[0], - 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].get_dp_shared_operator().copy_from_dp( - remote_att_manager=managers[0], - 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, - ) - assert_token_fields_equal( - source=src, - source_ids=[1, 3], - target=dst, - target_ids=[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=glm, - tp=1, - rank=0, - mtp=2, - ) - with torch.cuda.device(0): - target, target_full = make_resources( - glm=glm, - tp=1, - 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) - 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_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): - 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_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, - ) - 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_state_manager=target), - dp_att_state_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=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 - 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 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") -@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( - 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(): - storage_buffer[:, 4].zero_() - main_linear_state(resources).conv_state.fill_(11) - 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_() - pool.ssm_state.zero_() - slot = resources.bounded_att_big_page_allocator.alloc(1) - with torch.cuda.device(source_device): - resources.save_bounded_att_runtime_to_pages( - req_indices=[0], - state_token_lens=[8], - page_kind="big", - page_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=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)) - 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() - 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() - 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( - 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): - 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( - 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"), - task_kind=kind, - req_idx=1, - state_token_len=8, - write_to_pd_page=True, - ) - copy_page( - 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"), - task_kind=kind, - req_idx=2, - state_token_len=8, - 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()): - 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 - - -@pytest.mark.parametrize("kind", ["pd", "dp"]) -def test_ipc_readers_accept_different_pickle_lengths(kind: str, monkeypatch: pytest.MonkeyPatch) -> None: - """各读者独立导出；模拟 IPC 计数变化引起的长度差异，并验证补齐后仍可逐份读取。""" - resources = SimpleNamespace( - 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 = [] - 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) - for kind in ipc_kinds: - resources, mapping = shared_att_resources_utils.load_att_resources(0, kind=kind) - 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 - if kind == "dp": - 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.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": - 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") - ) - for window in window_states(resources): - window.reset_bounded_att_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: - 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 deleted file mode 100644 index 6eb8fa1be8..0000000000 --- a/unit_tests/common/att_resources/test_state_creation.py +++ /dev/null @@ -1,235 +0,0 @@ -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 - -import pytest -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, -) -from lightllm.common.attention_state.bounded_att import ( - BoundedAttentionState, - ConvSSMStateConfig, - create_gdn_state, - create_kda_state, -) -from lightllm.common.attention_state.bounded_att import WindowCompressorStateConfig, create_window_compressor_state - - -@pytest.mark.parametrize( - "kv_type, expected_type, storage_dtype, width", - [ - ("None", "gqa_bf16_h8_d128", torch.bfloat16, 128), - ("int4kv", "gqa_int4_group16_bf16_scale_h8_d128", torch.int8, 64), - ("int8kv", "gqa_int8_group16_bf16_scale_h8_d128", torch.int8, 128), - ("fp8kv_spt", "gqa_fp8_static_per_tensor_h8_d128", torch.uint8, 128), - ("fp8kv_sph", "gqa_fp8_static_per_head_h8_d128", 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 = GQAAttentionConfig( - 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, - ) - assert state.get_state_type() == expected_type - 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 - 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 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) - 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) - 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_group128"), - ], -) -def test_mla_creation_computes_packed_width_from_logical_dimensions( - packed: bool, - kv_type: str, - kv_dim: int, - row_width: int, - expected_type: str, -): - 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 - 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 - 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 torch.bfloat16) - if packed: - # 保持旧算子读取的物理布局：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 - - -@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, - ) - state = create_state( - config=config, - request_capacity=3, - mtp_step=2, - device="cpu", - ) - 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,) - assert state.big_pages_pool is None and state.small_pages_pool is None - - -def test_kpool_creation_preserves_pool_and_mtp_dimensions(): - config = WindowCompressorStateConfig( - logical_layer_ids=(3, 7), - window_size=4, - compress_stride=4, - target="indexer", - candidate_dim=128, - dtype=torch.bfloat16, - ) - 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.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(): - config = GQAAttentionConfig( - 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") - packed_config = PackedMLALogicalConfig( - 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") - 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", - ) - - -@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 deleted file mode 100644 index 856b4ef4af..0000000000 --- a/unit_tests/common/att_resources/test_state_identity.py +++ /dev/null @@ -1,356 +0,0 @@ -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 unit_tests.common.att_resources.model_helpers import replace_logical_layer_groups - -from typing import TYPE_CHECKING -from dataclasses import replace - -import pytest -import torch - -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 - -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, "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 - - 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.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 - assert state.get_token_bytes() == token_bytes - 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) - - -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.register_state(role, definition("gqa")) - with pytest.raises(AssertionError, match="already bound"): - 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_h2_d16_state,) - assert manager.get_full_attention_states("draft10") == (manager.draft10_gqa_bf16_h2_d16_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.attention_state.full_att import GQAFullAttentionState - - class CustomGQAState(GQAFullAttentionState): - def get_state_type(self) -> str: - return "gqa_custom_fp32" - - state = full_state(CustomGQAState) - 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.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_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_window_compressor_state import make_window - - 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_window() - storage = state.buffer - config = state.config - 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.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": - assert state.conv_state is storage and state.big_pages_pool is pool - else: - 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) - 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(): - 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_state_role == "draft35" - draft.mtp_previous_draft_models.append(object()) - with pytest.raises(AssertionError, match="at most 36"): - _ = draft.att_state_role - - -@pytest.mark.parametrize("role", ["draft", "draft36", "draft-1", "draft01"]) -def test_register_state_rejects_invalid_draft_numbers(role: str): - manager = AttentionStateManager(request_capacity=1) - with pytest.raises(AssertionError, match="unknown model role"): - manager.register_state(role, definition("gqa")) - - -def allocated_state(kind: str, device: str = "cpu"): - template = definition(kind) - 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") -@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.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(): - 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 = 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_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 = 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_registration() - 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_cache_operator().get_page_layouts(), - ) - 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_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_cache_operator().get_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 = 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_logical_layer_groups(template.config, ((4,), (6,)))) - manager.register_state("main", quant) - if overlap: - overlapping = allocated_state("static_head") - manager.register_state("main", overlapping) - manager.finish_state_registration() - assert manager.main_gqa_bf16_h2_d16_state is main - assert manager.draft0_gqa_bf16_h2_d16_state is draft - assert manager.main_gqa_int8_group8_bf16_scale_h2_d16_state is quant - assert manager.get_full_attention_states("draft0") == (draft,) - expected = [main, quant] - if overlap: - assert manager.main_gqa_fp8_static_per_head_h2_d16_state is overlapping - expected.append(overlapping) - 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") -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() == "gqa_bf16_h2_d16" - assert quant.get_state_type() == "gqa_int8_group8_bf16_scale_h2_d16" - 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 = AttentionStateManager(request_capacity=2) - 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") - ) - 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_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.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=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.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=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( - 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) - 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 deleted file mode 100644 index 0d96541acf..0000000000 --- a/unit_tests/common/att_resources/test_state_queries.py +++ /dev/null @@ -1,189 +0,0 @@ -"""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.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 - - -def test_created_state_queries_keep_the_same_objects_through_storage_initialization(monkeypatch: pytest.MonkeyPatch): - args = StartArgs(page_size=4) - definitions = [] - 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 + "_") - ) - ) - for query in ( - manager.get_full_attention_states, - manager.get_bounded_attention_states, - ): - with pytest.raises(AssertionError, match="State registration"): - query() - assert manager._get_states.cache_info().currsize == 0 - - cached = [] - allocate_storage = AttentionStateManager.allocate_storage - - 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) - allocate_storage(manager, *args, **kwargs) - - monkeypatch.setattr(AttentionStateManager, "allocate_storage", check_created_before_allocation) - monkeypatch.setattr( - "lightllm.common.attention_state.manager.FullAttKVAllocator", partial(FullAttKVAllocator, pin_memory=False) - ) - manager.finish_state_registration() - manager.initialize_storage( - 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) - 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) - 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_h1_d4_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,) - - 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 + 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() - - -def test_finished_state_creation_allows_unallocated_storage(): - manager = AttentionStateManager(request_capacity=1) - manager.register_state("main", definition("gqa")) - manager.full_att_allocator = FullAttKVAllocator(0, 4, pin_memory=False) - manager.finish_state_registration() - 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.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"): - 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.state_name) - 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.state_name) for state in source_states - ) - for state in source.get_full_attention_states(): - 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 - source = MainDraftTestResourceManager() - 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.get_dp_shared_operator().create_ipc_copy() - return - 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 - for state in source.get_bounded_attention_states(): - 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 - - -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.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 - 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, - ): - 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 + 2 - - -def test_clearing_query_cache_releases_manager_and_states(): - manager = MainDraftTestResourceManager() - manager.get_full_attention_states() - manager.get_bounded_attention_states() - manager_ref = weakref.ref(manager) - 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 - 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_tp_config.py b/unit_tests/common/att_resources/test_tp_config.py deleted file mode 100644 index 1fd660fa8d..0000000000 --- a/unit_tests/common/att_resources/test_tp_config.py +++ /dev/null @@ -1,223 +0,0 @@ -"""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 unit_tests.common.att_resources.model_helpers import replace_logical_layer_groups - -from typing import TYPE_CHECKING -import pickle -from dataclasses import replace - -import pytest -import torch - -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 ( - GQAAttentionConfig, - MLAConfig, - GQAFullAttentionState, - MLAFullAttentionState, -) - -if TYPE_CHECKING: - from collections.abc import Sequence - from easydict import EasyDict - 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 = AttentionStateManager(request_capacity=1) - manager.register_state("main", state) - manager.finish_state_registration() - return manager.get_cpu_cache_operator().get_page_layouts()[state.state_name] - - -@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 - if state_cls is MLAFullAttentionState: - config = MLAConfig( - logical_layer_groups=((0,),), - latent_dim=4, - latent_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(): - 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.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(): - field.tensor.copy_(fields[name][:, :, start:end]) - 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.allocate_storage(4) - layout = cpu_layout(target) - assert layout.byte_end == reference_layout.byte_end - 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]) - - -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, - ) - 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( - 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 = GQAAttentionConfig( - 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)( - config=other_config, - request_capacity=source.request_capacity, - mtp_step=source.mtp_step, - device="meta", - ) - else: - other = rebuild_full_state(source, config=other_config) - 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,) - assert query() == (source, other) - - 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_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 - ) - 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 - ) - 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 deleted file mode 100644 index e3fa612818..0000000000 --- a/unit_tests/common/att_resources/test_transfer_views.py +++ /dev/null @@ -1,233 +0,0 @@ -"""传输视图整理必须保留存储别名、位模式、未参与区域和独立寻址轴。""" - -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( - "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( - "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 - - -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/att_resources/test_validate_model_cache.py b/unit_tests/common/att_resources/test_validate_model_cache.py deleted file mode 100644 index 42508cd0ee..0000000000 --- a/unit_tests/common/att_resources/test_validate_model_cache.py +++ /dev/null @@ -1,41 +0,0 @@ -from types import SimpleNamespace - -import pytest - -from lightllm.common.attention_state import AttentionStateManager -from lightllm.common.attention_state.full_att import create_compressed_att_state -from lightllm.models.deepseek_v4.config import DSV4AttentionConfig -from test.att_resources import validate_model_cache - - -@pytest.mark.parametrize("draft_count,bad_role", [(0, None), (2, None), (2, "main"), (2, "draft1")]) -def test_validator_checks_all_compressed_states(monkeypatch, draft_count, bad_role): - resources = AttentionStateManager(request_capacity=0) - resources.full_att_allocator = SimpleNamespace(size=512) - req_manager = object() - models = [] - for index, role in enumerate(["main"] + [f"draft{i}" for i in range(draft_count)]): - config = DSV4AttentionConfig( - logical_layer_ids=(index * 2, index * 2 + 1), - compress_ratios=(4, 128), - head_dim=512, - rope_head_dim=64, - indexer_head_dim=128, - sliding_window=128, - ) - for ratio in (4, 128): - state = create_compressed_att_state(config.create_compressed_config(ratio), device="meta") - state.allocate_storage(256 if role == bad_role and ratio == 4 else 512) - resources.register_state(role, state) - models.append(SimpleNamespace(att_state_manager=resources, att_state_role=role, req_manager=req_manager)) - resources.finish_state_registration() - monkeypatch.setattr(validate_model_cache, "StaticBenchmarkExecutor", lambda **kwargs: None) - monkeypatch.setattr(validate_model_cache, "TokenSource", lambda args: None) - # Stop at input validation after checking resources, before any GPU execution. - args = SimpleNamespace(input_len=255, cpu_cache_token_page_size=256) - if bad_role is not None: - with pytest.raises(AssertionError): - validate_model_cache.validate(models[0], args, models[1:]) - else: - with pytest.raises(ValueError, match="input_len must contain complete CPU pages"): - validate_model_cache.validate(models[0], args, models[1:]) diff --git a/unit_tests/common/att_resources/test_window_compressor_state.py b/unit_tests/common/att_resources/test_window_compressor_state.py deleted file mode 100644 index 106743b126..0000000000 --- a/unit_tests/common/att_resources/test_window_compressor_state.py +++ /dev/null @@ -1,195 +0,0 @@ -"""统一窗口状态：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_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_cache_operator().get_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/common/basemodel/__init__.py b/unit_tests/common/basemodel/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/unit_tests/common/basemodel/attention/fa3/test_neo.py b/unit_tests/common/basemodel/attention/fa3/test_neo.py deleted file mode 100644 index 9ef88130cf..0000000000 --- a/unit_tests/common/basemodel/attention/fa3/test_neo.py +++ /dev/null @@ -1,119 +0,0 @@ -import importlib.util -import sys -from types import ModuleType, SimpleNamespace - -import pytest -import torch - -from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.common.basemodel.attention.fa3 import fp, neo - - -def _load_neo(monkeypatch, attention_func): - interface = None - if attention_func is not None: - interface = ModuleType("flash_attn_interface") - interface.flash_attn_with_kvcache = attention_func - monkeypatch.setitem(sys.modules, "flash_attn_interface", interface) - name = "lightllm.common.basemodel.attention.fa3._neo_import_test" - spec = importlib.util.spec_from_file_location(name, neo.__file__) - module = importlib.util.module_from_spec(spec) - monkeypatch.setitem(sys.modules, name, module) - spec.loader.exec_module(module) - return module - - -def test_recognizes_neo_interface_without_replacing_regular_fa3(monkeypatch): - def neo_attention(q, image_token_end=None): - return q - - regular_attention = fp.flash_attn_with_kvcache - module = _load_neo(monkeypatch, neo_attention) - - assert module.HAS_FLASH_ATTN_INTERFACE - assert module.flash_attn_with_kvcache_neo is neo_attention - assert fp.flash_attn_with_kvcache is regular_attention - - -@pytest.mark.parametrize("interface_kind", ["missing", "regular", "legacy_tag"]) -def test_unavailable_neo_interface_provides_source_install_instructions(monkeypatch, interface_kind): - def regular_attention(q): - return q - - def legacy_attention(q, image_token_tag=None): - return q - - attention_func = {"missing": None, "regular": regular_attention, "legacy_tag": legacy_attention}[interface_kind] - with pytest.warns(UserWarning, match="Neo FA3 is unavailable") as warnings: - module = _load_neo(monkeypatch, attention_func) - - assert not module.HAS_FLASH_ATTN_INTERFACE - assert module.flash_attn_with_kvcache_neo is None - message = str(warnings[0].message) - assert "https://github.com/WANDY666/flash-attention.git" in message - assert "cd flash-attention/hopper" in message - assert "FLASH_ATTENTION_FORCE_BUILD=TRUE" in message - assert "--no-deps --force-reinstall" in message - assert "--llm_prefill_att_backend triton" in message - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="no CUDA") -@pytest.mark.parametrize("page_size", [1, 16, 64]) -def test_prefill_initializes_paged_kv_metadata_and_forwards_image_ends(monkeypatch, page_size): - def indices(values): - return torch.tensor(values, dtype=torch.int32, device="cuda") - - image_ends = indices([2 * page_size, 2 * page_size, 0]) - # Each request uses two nonadjacent physical pages, in a different order. - req_to_token_indexs = indices( - [ - list(range(2 * page_size, 3 * page_size)) + list(range(page_size)), - list(range(3 * page_size, 4 * page_size)) + list(range(page_size, 2 * page_size)), - ] - ) - infer_state = InferStateInfo() - infer_state.batch_size = 2 - infer_state.max_kv_seq_len = 2 * page_size - infer_state.max_q_seq_len = 2 - infer_state.input_ids = torch.zeros(3, dtype=torch.int64, device="cuda") - infer_state.b_req_idx = indices([1, 0]) - infer_state.b_seq_len = indices([2 * page_size, 2 * page_size - 1]) - infer_state.b1_cu_q_seq_len = indices([0, 2, 3]).long() - infer_state.b1_cu_kv_seq_len = indices([0, 2 * page_size, 4 * page_size - 1]).long() - infer_state.b_image_token_end = image_ends - infer_state.req_manager = SimpleNamespace(req_to_token_indexs=req_to_token_indexs) - backend = SimpleNamespace(uses_causal_attention=lambda: True, infer_page_size=page_size) - state = neo.NeoFa3AttBackend.create_att_prefill_state(backend, infer_state) - state.init_state() - - q = torch.randn(3, 2, 16, device="cuda") - k = torch.randn(4 * page_size, 1, 16, device="cuda") - v = torch.randn_like(k) - output = torch.empty_like(q) - - def neo_attention(*, image_token_end, **kwargs): - assert image_token_end is image_ends - assert kwargs["q"] is q - assert kwargs["k_cache"].shape == kwargs["v_cache"].shape == (4, page_size, 1, 16) - torch.testing.assert_close(kwargs["page_table"], indices([[3, 1], [2, 0]])) - token_indices = req_to_token_indexs[infer_state.b_req_idx] - torch.testing.assert_close(kwargs["k_cache"][kwargs["page_table"]].flatten(1, 2), k[token_indices]) - torch.testing.assert_close(kwargs["v_cache"][kwargs["page_table"]].flatten(1, 2), v[token_indices]) - torch.testing.assert_close(kwargs["cu_seqlens_q"], indices([0, 2, 3])) - torch.testing.assert_close(kwargs["cu_seqlens_k_new"], indices([0, 2 * page_size, 4 * page_size - 1])) - assert kwargs["cache_seqlens"] is infer_state.b_seq_len - assert kwargs["max_seqlen_q"] == 2 - assert kwargs["causal"] is True - assert "image_token_tag" not in kwargs - return output - - monkeypatch.setattr(neo, "HAS_FLASH_ATTN_INTERFACE", True) - monkeypatch.setattr(neo, "flash_attn_with_kvcache_neo", neo_attention) - assert state.prefill_att(q, k, v) is output - - -def test_prefill_missing_extension_reports_install_instructions(monkeypatch): - monkeypatch.setattr(neo, "HAS_FLASH_ATTN_INTERFACE", False) - state = neo.NeoFa3PrefillAttState() - with pytest.raises(ImportError, match="WANDY666/flash-attention"): - state.prefill_att(None, None, None) diff --git a/unit_tests/common/basemodel/attention/flashinfer/test_decode_graph_state.py b/unit_tests/common/basemodel/attention/flashinfer/test_decode_graph_state.py deleted file mode 100644 index 190d8f1162..0000000000 --- a/unit_tests/common/basemodel/attention/flashinfer/test_decode_graph_state.py +++ /dev/null @@ -1,100 +0,0 @@ -import sys -import weakref -from types import ModuleType, SimpleNamespace - -import torch - -from lightllm.common.basemodel.attention.flashinfer.fp import FlashInferDecodeAttState -from lightllm.common.basemodel.attention.flashinfer import mla -from lightllm.common.basemodel.infer_struct import InferStateInfo - - -def test_fp_graph_copy_refreshes_plan_with_runtime_length_and_preserves_buffers(monkeypatch): - plan_calls = [] - flashinfer = ModuleType("flashinfer") - decode = ModuleType("flashinfer.decode") - decode.fast_decode_plan = lambda wrapper, **kwargs: plan_calls.append((wrapper, kwargs)) - flashinfer.decode = decode - monkeypatch.setitem(sys.modules, "flashinfer", flashinfer) - monkeypatch.setitem(sys.modules, "flashinfer.decode", decode) - backend = SimpleNamespace( - infer_page_size=16, - tp_q_head_num=4, - tp_kv_head_num=2, - head_dim=8, - q_data_type=torch.bfloat16, - kv_data_type=torch.bfloat16, - ) - captured_owner = InferStateInfo() - captured_owner.batch_size = 2 - runtime_owner = InferStateInfo() - runtime_owner.max_kv_seq_len = 17 - captured = FlashInferDecodeAttState( - backend=backend, - _infer_state_weak_ref=weakref.ref(captured_owner), - kv_starts=torch.tensor([0, 4, 8], dtype=torch.int32), - kv_indices=torch.arange(8, dtype=torch.int32), - kv_last_page_len_buffer=torch.tensor([16, 16], dtype=torch.int32), - decode_wrapper=object(), - ) - runtime = FlashInferDecodeAttState( - backend=backend, - _infer_state_weak_ref=weakref.ref(runtime_owner), - kv_starts=torch.tensor([0, 2, 3], dtype=torch.int32), - kv_indices=torch.arange(7, -1, -1, dtype=torch.int32), - kv_last_page_len_buffer=torch.tensor([1, 5], dtype=torch.int32), - ) - buffer_names = ("kv_starts", "kv_indices", "kv_last_page_len_buffer") - captured_ptrs = {name: getattr(captured, name).data_ptr() for name in buffer_names} - - captured.copy_for_decode_cuda_graph(runtime) - - assert len(plan_calls) == 1 - wrapper, plan = plan_calls[0] - assert wrapper is captured.decode_wrapper - assert plan["indptr"] is captured.kv_starts - assert plan["indices"] is captured.kv_indices - assert plan["last_page_len"] is captured.kv_last_page_len_buffer - torch.testing.assert_close(plan["global_override_indptr_cpu"], torch.tensor([0, 2, 4], dtype=torch.int32)) - for name in buffer_names: - assert getattr(captured, name).data_ptr() == captured_ptrs[name] - torch.testing.assert_close(getattr(captured, name), getattr(runtime, name)) - - -def test_mla_graph_copy_refreshes_plan_with_runtime_lengths_and_preserves_buffers(monkeypatch): - plan_calls = [] - monkeypatch.setattr(mla, "fill_mla_decode_plan_for_cuda_graph", lambda **kwargs: plan_calls.append(kwargs)) - backend = SimpleNamespace(tp_q_head_num=4) - captured_owner = InferStateInfo() - captured_owner.batch_size = 2 - runtime_owner = InferStateInfo() - runtime_owner.b_seq_len = torch.tensor([17, 5], dtype=torch.int32) - runtime_owner.max_kv_seq_len = 17 - captured = mla.MlaFlashInferDecodeAttState( - backend=backend, - _infer_state_weak_ref=weakref.ref(captured_owner), - kv_starts=torch.tensor([0, 4, 8], dtype=torch.int32), - kv_indices=torch.arange(8, dtype=torch.int32), - decode_wrapper=object(), - ) - runtime = mla.MlaFlashInferDecodeAttState( - backend=backend, - _infer_state_weak_ref=weakref.ref(runtime_owner), - kv_starts=torch.tensor([0, 2, 3], dtype=torch.int32), - kv_indices=torch.arange(7, -1, -1, dtype=torch.int32), - ) - buffer_names = ("kv_starts", "kv_indices") - captured_ptrs = {name: getattr(captured, name).data_ptr() for name in buffer_names} - - captured.copy_for_decode_cuda_graph(runtime) - - assert len(plan_calls) == 1 - plan = plan_calls[0] - assert plan["decode_wrapper"] is captured.decode_wrapper - assert plan["kv_page_indptr"] is captured.kv_starts - assert plan["batch_size"] == 2 - assert plan["max_kv_len"] == 17 - torch.testing.assert_close(plan["kv_lens"], torch.tensor([17, 5], dtype=torch.int32)) - for name in buffer_names: - assert getattr(captured, name).data_ptr() == captured_ptrs[name] - torch.testing.assert_close(getattr(captured, name), getattr(runtime, name)) diff --git a/unit_tests/common/basemodel/attention/flashinfer/test_decode_wrapper_init.py b/unit_tests/common/basemodel/attention/flashinfer/test_decode_wrapper_init.py deleted file mode 100644 index 1392d5e688..0000000000 --- a/unit_tests/common/basemodel/attention/flashinfer/test_decode_wrapper_init.py +++ /dev/null @@ -1,132 +0,0 @@ -import sys -import weakref -from types import ModuleType, SimpleNamespace -from unittest.mock import Mock - -import pytest -import torch - -from lightllm.common.basemodel.attention.flashinfer import fp, fp8, mla -from lightllm.common.basemodel.attention.flashinfer.utils import should_init_decode_wrapper -from lightllm.common.basemodel.infer_struct import InferStateInfo - - -def test_decode_wrapper_rejects_conflicting_graph_flags(): - infer_state = InferStateInfo() - infer_state.is_cuda_graph_capture = True - infer_state.is_cuda_graph_replay = True - - with pytest.raises(ValueError, match="capture=True, replay=True"): - should_init_decode_wrapper(infer_state) - - -@pytest.fixture( - params=[fp.FlashInferDecodeAttState, fp8.Fp8FlashInferDecodeAttState, mla.MlaFlashInferDecodeAttState], - ids=["fp", "fp8", "mla"], -) -def decode_init(request, monkeypatch): - state_cls = request.param - module = mla if state_cls is mla.MlaFlashInferDecodeAttState else fp - page_size = 1 if state_cls is fp8.Fp8FlashInferDecodeAttState else 4 - max_page_num = 64 // page_size - wrapper = SimpleNamespace(plan=Mock()) - create_wrapper = Mock(return_value=wrapper) - flashinfer = ModuleType("flashinfer") - flashinfer.decode = SimpleNamespace(BatchDecodeWithPagedKVCacheWrapper=create_wrapper) - flashinfer.mla = SimpleNamespace(BatchMLAPagedAttentionWrapper=create_wrapper) - monkeypatch.setitem(sys.modules, "flashinfer", flashinfer) - - def cpu_arange(*args, **kwargs): - kwargs["device"] = "cpu" - return torch.arange(*args, **kwargs) - - monkeypatch.setattr(module, "torch", SimpleNamespace(int32=torch.int32, arange=cpu_arange, empty=torch.empty)) - - def cumsum_params(q_lengths, kv_lengths): - return tuple( - torch.cat((torch.zeros(1, dtype=torch.int32), lengths.cumsum(0, dtype=torch.int32))) - for lengths in (q_lengths, kv_lengths) - ) - - monkeypatch.setattr(module, "gen_cumsum_pad0_tensor", cumsum_params) - - def repack(**kwargs): - kwargs["out_page_indices"].fill_(7) - - repack_indices = Mock(side_effect=repack) - monkeypatch.setattr(module, "repack_kv_index", repack_indices) - can_run = Mock(return_value=True) - backend = SimpleNamespace( - model=SimpleNamespace( - graph=SimpleNamespace(can_run=can_run), graph_max_batch_size=4, graph_max_len_in_batch=64 - ), - infer_page_size=page_size, - max_page_num=max_page_num, - kv_indices_buffer=[torch.empty(4 * max_page_num, dtype=torch.int32)], - tp_q_head_num=4, - tp_kv_head_num=2, - head_dim=8, - kv_lora_rank=8, - qk_rope_head_dim=2, - softmax_scale=0.25, - q_data_type=torch.bfloat16, - kv_data_type=torch.float8_e4m3fn if page_size == 1 else torch.bfloat16, - workspace_buffer_key="decode_init_test", - workspace_buffer_size=1, - get_gpu_workspace_buffer=Mock(return_value=torch.empty(1, dtype=torch.uint8)), - ) - return SimpleNamespace( - module=module, - state=state_cls(backend=backend), - backend=backend, - wrapper=wrapper, - create_wrapper=create_wrapper, - repack_indices=repack_indices, - can_run=can_run, - ) - - -@pytest.mark.parametrize("mode", ["eager", "capture", "replay"]) -def test_decode_init_uses_explicit_graph_mode_for_wrapper_creation(decode_init, mode): - infer_state = InferStateInfo() - infer_state.input_ids = torch.tensor([1, 2], dtype=torch.int64) - infer_state.batch_size = 2 - infer_state.b_seq_len = torch.tensor([5, 2], dtype=torch.int32) - infer_state.max_kv_seq_len = 5 - infer_state.b_req_idx = torch.tensor([2, 3], dtype=torch.int32) - infer_state.microbatch_index = 0 - infer_state.req_manager = SimpleNamespace(req_to_token_indexs=torch.zeros((4, 64), dtype=torch.int32)) - if mode == "capture": - infer_state.is_cuda_graph_capture = True - elif mode == "replay": - infer_state.is_cuda_graph_replay = True - - decode_init.state._infer_state_weak_ref = weakref.ref(infer_state) - decode_init.state.init_state() - - decode_init.can_run.assert_not_called() - expected_starts = [0, 5, 7] if decode_init.backend.infer_page_size == 1 else [0, 2, 3] - torch.testing.assert_close(decode_init.state.kv_starts, torch.tensor(expected_starts, dtype=torch.int32)) - assert decode_init.state.kv_indices.data_ptr() == decode_init.backend.kv_indices_buffer[0].data_ptr() - assert torch.all(decode_init.state.kv_indices == 7) - decode_init.repack_indices.assert_called_once() - repack_args = decode_init.repack_indices.call_args.kwargs - assert repack_args["b_req_idx"] is infer_state.b_req_idx - assert repack_args["b_token_len"] is infer_state.b_seq_len - if decode_init.module is fp: - expected_last_page_lengths = [1, 1] if decode_init.backend.infer_page_size == 1 else [1, 2] - torch.testing.assert_close( - decode_init.state.kv_last_page_len_buffer, torch.tensor(expected_last_page_lengths, dtype=torch.int32) - ) - else: - torch.testing.assert_close(decode_init.state.q_indptr, torch.tensor([0, 1, 2], dtype=torch.int32)) - - if mode == "replay": - assert decode_init.state.decode_wrapper is None - decode_init.create_wrapper.assert_not_called() - decode_init.backend.get_gpu_workspace_buffer.assert_not_called() - else: - assert decode_init.state.decode_wrapper is decode_init.wrapper - decode_init.create_wrapper.assert_called_once() - decode_init.backend.get_gpu_workspace_buffer.assert_called_once() - decode_init.wrapper.plan.assert_called_once() diff --git a/unit_tests/common/basemodel/attention/flashinfer/test_mla_cuda_graph_plan.py b/unit_tests/common/basemodel/attention/flashinfer/test_mla_cuda_graph_plan.py deleted file mode 100644 index ccf5dfd454..0000000000 --- a/unit_tests/common/basemodel/attention/flashinfer/test_mla_cuda_graph_plan.py +++ /dev/null @@ -1,178 +0,0 @@ -from types import SimpleNamespace -import weakref - -import pytest -import torch - -flashinfer = pytest.importorskip("flashinfer") - -from lightllm.common.basemodel.attention.flashinfer.mla import ( # noqa: E402 - MlaFlashInferDecodeAttState, -) -from lightllm.common.basemodel.infer_struct import InferStateInfo # noqa: E402 - - -def _make_wrapper( - q_indptr, - kv_indptr_buf, - kv_indptr, - kv_indices, - kv_lens_buf, - num_heads, - head_dim_ckv, - head_dim_kpe, - sm_scale, - dtype, - page_size, - init_short, -): - workspace = torch.empty(128 * 1024 * 1024, dtype=torch.int8, device="cuda") - wrapper = flashinfer.mla.BatchMLAPagedAttentionWrapper( - workspace, - use_cuda_graph=True, - qo_indptr=q_indptr, - kv_indices=kv_indices, - kv_indptr=kv_indptr_buf, - kv_len_arr=kv_lens_buf, - ) - if init_short: - batch_size = kv_lens_buf.numel() - init_lens = torch.full((batch_size,), 2, dtype=torch.int32, device="cuda") - init_indptr = torch.arange(batch_size + 1, dtype=torch.int32, device="cuda") * ( - (2 + page_size - 1) // page_size - ) - wrapper.plan( - q_indptr, - init_indptr, - kv_indices, - init_lens, - num_heads, - head_dim_ckv, - head_dim_kpe, - page_size, - False, - sm_scale, - dtype, - dtype, - ) - else: - wrapper.plan( - q_indptr, - kv_indptr, - kv_indices, - kv_lens_buf, - num_heads, - head_dim_ckv, - head_dim_kpe, - page_size, - False, - sm_scale, - dtype, - dtype, - ) - return wrapper - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") -@pytest.mark.parametrize("page_size", [1, 3, 4, 16, 64, 128]) -@pytest.mark.parametrize("lengths,num_heads", [([5, 64, 65], 32), ([1000, 32768], 32), ([1000, 131073], 128)]) -def test_mla_cuda_graph_triton_plan_matches_flashinfer_plan(lengths, num_heads, page_size): - torch.manual_seed(0) - dtype = torch.bfloat16 - head_dim_ckv = 512 - head_dim_kpe = 64 - batch_size = len(lengths) - total_pages = sum((length + page_size - 1) // page_size for length in lengths) - sm_scale = (head_dim_ckv + head_dim_kpe) ** -0.5 - - q_nope = torch.randn((batch_size, num_heads, head_dim_ckv), dtype=dtype, device="cuda") - q_pe = torch.randn((batch_size, num_heads, head_dim_kpe), dtype=dtype, device="cuda") - ckv = torch.randn((total_pages, page_size, head_dim_ckv), dtype=dtype, device="cuda") - kpe = torch.randn((total_pages, page_size, head_dim_kpe), dtype=dtype, device="cuda") - kv_lens = torch.tensor(lengths, dtype=torch.int32, device="cuda") - kv_indptr = torch.empty(batch_size + 1, dtype=torch.int32, device="cuda") - kv_indptr[0] = 0 - kv_indptr[1:] = torch.cumsum((kv_lens + page_size - 1) // page_size, dim=0) - kv_indices = torch.randperm(total_pages, dtype=torch.int32, device="cuda") - q_indptr = torch.arange(batch_size + 1, dtype=torch.int32, device="cuda") - - ref_wrapper = _make_wrapper( - q_indptr, - kv_indptr.clone(), - kv_indptr, - kv_indices, - kv_lens.clone(), - num_heads, - head_dim_ckv, - head_dim_kpe, - sm_scale, - dtype, - page_size, - init_short=False, - ) - graph_wrapper = _make_wrapper( - q_indptr, - kv_indptr.clone(), - kv_indptr, - kv_indices, - kv_lens.clone(), - num_heads, - head_dim_ckv, - head_dim_kpe, - sm_scale, - dtype, - page_size, - init_short=True, - ) - backend = SimpleNamespace(tp_q_head_num=num_heads, infer_page_size=page_size) - captured_owner = InferStateInfo() - captured_owner.batch_size = batch_size - captured_owner.b_seq_len = graph_wrapper._kv_len_arr_buf - graph_state = MlaFlashInferDecodeAttState( - backend=backend, - _infer_state_weak_ref=weakref.ref(captured_owner), - kv_starts=graph_wrapper._kv_indptr_buf, - kv_indices=graph_wrapper._kv_indices_buf, - decode_wrapper=graph_wrapper, - ) - ref_out = torch.empty((batch_size, num_heads, head_dim_ckv), dtype=dtype, device="cuda") - graph_out = torch.empty_like(ref_out) - graph_wrapper.run(q_nope, q_pe, ckv, kpe, out=graph_out, return_lse=False) - torch.cuda.synchronize() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - graph_wrapper.run(q_nope, q_pe, ckv, kpe, out=graph_out, return_lse=False) - - # 复用同一 graph，覆盖页表变化、末页空位及 split/exact 切换后的 plan 更新。 - for replay_lengths in (lengths, [1] * batch_size, lengths): - replay_lens = torch.tensor(replay_lengths, dtype=torch.int32, device="cuda") - replay_indptr = torch.zeros_like(kv_indptr) - replay_indptr[1:] = torch.cumsum((replay_lens + page_size - 1) // page_size, dim=0) - runtime_owner = InferStateInfo() - runtime_owner.batch_size = batch_size - runtime_owner.b_seq_len = replay_lens - runtime_owner.max_kv_seq_len = max(replay_lengths) - new_state = MlaFlashInferDecodeAttState( - backend=backend, - _infer_state_weak_ref=weakref.ref(runtime_owner), - kv_starts=replay_indptr, - kv_indices=kv_indices, - ) - graph_state.copy_for_decode_cuda_graph(new_state) - graph.replay() - ref_wrapper.plan( - q_indptr, - replay_indptr, - kv_indices, - replay_lens, - num_heads, - head_dim_ckv, - head_dim_kpe, - page_size, - False, - sm_scale, - dtype, - dtype, - ) - ref_wrapper.run(q_nope, q_pe, ckv, kpe, out=ref_out, return_lse=False) - torch.testing.assert_close(ref_out, graph_out, atol=1e-2, rtol=1e-2) diff --git a/unit_tests/common/basemodel/attention/flashinfer/test_mla_nope.py b/unit_tests/common/basemodel/attention/flashinfer/test_mla_nope.py deleted file mode 100644 index 4c33d8cf8a..0000000000 --- a/unit_tests/common/basemodel/attention/flashinfer/test_mla_nope.py +++ /dev/null @@ -1,106 +0,0 @@ -"""Check the FlashInfer SM90 path used for unpadded GLM NoPE MLA.""" - -import inspect - -import pytest -import torch - -flashinfer = pytest.importorskip("flashinfer") - -pytestmark = pytest.mark.skipif( - not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 9, - reason="FlashInfer NoPE FA3 MLA requires Hopper", -) - - -@pytest.mark.parametrize( - "num_tokens,num_heads,kv_dtype", - [ - (1, 16, torch.bfloat16), - (8, 16, torch.bfloat16), - (128, 16, torch.bfloat16), - (8, 64, torch.bfloat16), - (8, 16, torch.float8_e4m3fn), - ], -) -def test_nope_sparse_mla_eager_and_graph_replay(num_tokens, num_heads, kv_dtype): - wrapper_cls = flashinfer.mla.BatchMLAPagedAttentionWrapper - assert "ckv_scale_arr" in inspect.signature(wrapper_cls.run).parameters, "FlashInfer >= 0.6.18 is required" - - torch.manual_seed(17) - dim, num_slots, width = 512, 4096, 2176 - sm_scale = 0.0625 - q = torch.randn(num_tokens, num_heads, dim, dtype=torch.bfloat16, device="cuda") - q_pe = q[..., :0] - # BF16 KV shares its token row with the 144-byte index region. The - # attention kernel must respect the resulting non-contiguous token stride. - tail = 72 if kv_dtype == torch.bfloat16 else 0 - packed = torch.full((num_slots, 1, dim + tail), 7.0, dtype=kv_dtype, device="cuda") - ckv = packed[..., :dim] - raw_kv = torch.randn(num_slots, 1, dim, dtype=torch.bfloat16, device="cuda") - scale = 0.125 if kv_dtype == torch.float8_e4m3fn else 1.0 - ckv.copy_((raw_kv.float() / scale).to(kv_dtype)) - kpe = packed[..., dim:dim] - reference_kv = (ckv.float() * scale).to(torch.bfloat16).float() - scale_kwargs = {"ckv_scale": scale, "kpe_scale": 1.0} if scale != 1.0 else {} - - qo_indptr = torch.arange(num_tokens + 1, dtype=torch.int32) - kv_indptr = qo_indptr * width - slots = torch.zeros(num_tokens, width, dtype=torch.int32, device="cuda") - wrapper = wrapper_cls( - torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device="cuda"), - use_cuda_graph=True, - qo_indptr=qo_indptr.to("cuda"), - kv_indptr=kv_indptr.to("cuda"), - kv_indices=slots.flatten(), - kv_len_arr=torch.empty(num_tokens, dtype=torch.int32, device="cuda"), - backend="fa3", - ) - - def plan(step): - choices = [1, 3, 127, 128, 511, 2048, 2049, 2051] - lengths = [choices[(row + step) % len(choices)] for row in range(num_tokens)] - slots.zero_() - for row, length in enumerate(lengths): - slots[row, :length] = (torch.arange(length, device="cuda") * 5 + row * 31 + step * 7) % num_slots - # Lengths are compiled into the host-side plan. Replan before graph - # replay when they change, while keeping all captured buffers stable. - wrapper.plan( - qo_indptr, - kv_indptr, - slots.flatten(), - torch.tensor(lengths, dtype=torch.int32), - num_heads, - dim, - 0, - 1, - False, - sm_scale, - torch.bfloat16, - kv_dtype, - ) - return lengths - - def check(output, lengths): - expected = [] - for row, length in enumerate(lengths): - keys = reference_kv[slots[row, :length].long(), 0] - scores = q[row].float() @ keys.T * sm_scale - expected.append(scores.softmax(-1) @ keys) - torch.testing.assert_close(output.float(), torch.stack(expected), atol=0.012, rtol=0.015) - - lengths = plan(0) - output = torch.empty_like(q) - wrapper.run(q, q_pe, ckv, kpe, out=output, **scale_kwargs) - check(output, lengths) - - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - wrapper.run(q, q_pe, ckv, kpe, out=output, **scale_kwargs) - graph.replay() - check(output, lengths) - - lengths = plan(1) - q.normal_() - graph.replay() - check(output, lengths) diff --git a/unit_tests/common/basemodel/attention/flashinfer/test_workspace.py b/unit_tests/common/basemodel/attention/flashinfer/test_workspace.py deleted file mode 100644 index 924dead570..0000000000 --- a/unit_tests/common/basemodel/attention/flashinfer/test_workspace.py +++ /dev/null @@ -1,56 +0,0 @@ -from lightllm.common.basemodel.attention import base_att -from lightllm.common.basemodel.attention.base_att import BaseAttBackend -from lightllm.common.basemodel.attention.flashinfer.fp import FlashInferAttBackend -from lightllm.common.basemodel.attention.flashinfer.fp8 import Fp8FlashInferAttBackend -from lightllm.common.basemodel.attention.flashinfer.mla import MlaFlashInferAttBackend - - -def test_workspace_is_shared_by_backend_family_and_device(monkeypatch): - allocations = [] - current_device_id = 0 - - def fake_empty(size, dtype, device): - workspace = object() - allocations.append((workspace, size, dtype, device)) - return workspace - - monkeypatch.setattr(base_att, "get_current_device_id", lambda: current_device_id) - monkeypatch.setattr(base_att.torch, "empty", fake_empty) - monkeypatch.setattr(BaseAttBackend, "_workspace_buffers", {}) - - def get_workspace(backend_cls): - return backend_cls.get_gpu_workspace_buffer( - key_name=backend_cls.workspace_buffer_key, - workspace_size=backend_cls.workspace_buffer_size, - ) - - fp_workspace = get_workspace(FlashInferAttBackend) - fp8_workspace = get_workspace(Fp8FlashInferAttBackend) - mla_workspace = get_workspace(MlaFlashInferAttBackend) - - assert fp8_workspace is fp_workspace - assert mla_workspace is not fp_workspace - assert [allocation[1] for allocation in allocations] == [ - 512 * 1024 * 1024, - 256 * 1024 * 1024, - ] - - current_device_id = 1 - other_device_workspace = get_workspace(FlashInferAttBackend) - - assert other_device_workspace is not fp_workspace - assert allocations[-1][3] == 1 - - -def test_workspace_key_includes_size_and_dtype(monkeypatch): - monkeypatch.setattr(base_att, "get_current_device_id", lambda: 0) - monkeypatch.setattr(base_att.torch, "empty", lambda *args, **kwargs: object()) - monkeypatch.setattr(BaseAttBackend, "_workspace_buffers", {}) - - small_workspace = BaseAttBackend.get_gpu_workspace_buffer("shared", 1024) - large_workspace = BaseAttBackend.get_gpu_workspace_buffer("shared", 2048) - fp16_workspace = BaseAttBackend.get_gpu_workspace_buffer("shared", 1024, dtype=base_att.torch.float16) - - assert BaseAttBackend.get_gpu_workspace_buffer("shared", 1024) is small_workspace - assert large_workspace is not small_workspace - assert fp16_workspace is not small_workspace diff --git a/unit_tests/common/basemodel/attention/linear/test_gdn.py b/unit_tests/common/basemodel/attention/linear/test_gdn.py deleted file mode 100644 index 344bfea57f..0000000000 --- a/unit_tests/common/basemodel/attention/linear/test_gdn.py +++ /dev/null @@ -1,438 +0,0 @@ -import weakref -import inspect -import sys -from types import ModuleType, SimpleNamespace - -import pytest -import torch - -from lightllm.common.attention_state import AttentionStateManager -from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState -from lightllm.common.basemodel.infer_struct import InferStateInfo -import lightllm.common.basemodel.attention.linear.gdn as gdn -import lightllm.common.basemodel.attention.create_linear_utils as linear_create_utils -import lightllm.common.basemodel.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.config import Qwen3NextResourceConfig -from lightllm.server.api_cli import make_argument_parser -import lightllm.utils.backend_validator as backend_validator - - -@pytest.mark.parametrize("cache_dtype", [torch.bfloat16, torch.float32]) -def test_prefill_casts_final_state_to_cache_dtype(monkeypatch, cache_dtype): - ssm_states = torch.zeros((1, 1, 2, 2), dtype=cache_dtype) - final_state = torch.ones((1, 1, 2, 2), dtype=torch.float32) - - monkeypatch.setattr(gdn, "fused_gdn_gating", lambda _log, a, b, _bias: (a, b)) - monkeypatch.setattr(gdn, "causal_conv1d_fn", lambda mixed, *args, **kwargs: mixed) - qkv = torch.zeros((1, 3), dtype=cache_dtype) - q = torch.zeros((1, 1, 1, 1), dtype=cache_dtype) - backend = SimpleNamespace( - mtp_step=0, - activation="silu", - ssm_state_dtype=cache_dtype, - _rearrange_mixed_qkv=lambda mixed: (q, q, q), - prefill_kernel=lambda *args, **kwargs: (None, final_state), - ) - infer_state = InferStateInfo() - infer_state.b1_cu_q_seq_len = torch.tensor([0, 1], dtype=torch.int32) - infer_state.b_ready_cache_len = 0 - state = gdn.LinearAttPrefillAttState( - backend=backend, - _infer_state_weak_ref=weakref.ref(infer_state), - ) - state.b_conv_buffer_idx = torch.tensor([0], dtype=torch.int64) - state.b_ssm_buffer_idx = torch.tensor([0], dtype=torch.int64) - layer_weight = SimpleNamespace( - linear_A_log=SimpleNamespace(weight=None), - linear_dt_bias=SimpleNamespace(weight=None), - linear_conv1d=SimpleNamespace(mm_param=SimpleNamespace(weight=None), bias=None), - ) - - state._gdn_prefill_kernel( - qkv, - torch.zeros((1, 3), dtype=cache_dtype), - ssm_states, - torch.zeros((1, 1), dtype=cache_dtype), - torch.zeros((1, 1), dtype=cache_dtype), - state.infer_state, - layer_weight, - ) - - assert ssm_states.dtype == cache_dtype - assert torch.equal(ssm_states, final_state.to(cache_dtype)) - - -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) - - state.init_state() - - assert state.b_conv_buffer_idx is b_req_idx - assert state.b_ssm_buffer_idx is b_req_idx - assert state.b1_mtp_cu_q_seq_len is None - assert state.b_num_accepted_tokens is None - - -@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( - draft_step=2, - dynamic_layout=False, - b_req_idx=b_req_idx, - req_to_mtp_state_index=req_to_mtp_state_index, - role=role, - ) - - state.init_state() - - torch.testing.assert_close(state.b1_mtp_cu_q_seq_len, torch.tensor([0, 3, 6], dtype=torch.int32)) - torch.testing.assert_close(state.b_conv_buffer_idx, torch.tensor([2, 4], dtype=torch.int32)) - torch.testing.assert_close(state.b_num_accepted_tokens, torch.tensor([2, 3], dtype=torch.int32)) - torch.testing.assert_close( - state.b_ssm_buffer_idx, - torch.tensor([[6, 7, 8], [12, 13, 14]], dtype=torch.int32), - ) - - -@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) - expected_conv_buffer_idx = torch.tensor([2, 4, 5, 5, 5], dtype=torch.int32) - expected_num_accepted_tokens = torch.tensor([2, 3, 1, 1, 1], dtype=torch.int32) - build_calls = [] - - def build_params(**kwargs): - build_calls.append(kwargs) - return expected_cu_q_seq_len, expected_conv_buffer_idx, expected_num_accepted_tokens - - monkeypatch.setattr(gdn, "build_dynamic_mtp_linear_att_state_params", build_params) - state, infer_state = _create_decode_state( - draft_step=2, - dynamic_layout=True, - b_req_idx=b_req_idx, - req_to_mtp_state_index=req_to_mtp_state_index, - role=role, - ) - - state.init_state() - - assert state.b1_mtp_cu_q_seq_len is expected_cu_q_seq_len - assert state.b_conv_buffer_idx is expected_conv_buffer_idx - assert state.b_num_accepted_tokens is expected_num_accepted_tokens - torch.testing.assert_close( - state.b_ssm_buffer_idx, - torch.tensor( - [[6, 7, 8], [12, 13, 14], [15, 16, 17], [15, 16, 17], [15, 16, 17]], - dtype=torch.int32, - ), - ) - assert len(build_calls) == 1 - assert build_calls[0]["b_req_idx"] is state.infer_state.b_req_idx - assert build_calls[0]["b_mtp_index"] is state.infer_state.b_mtp_index - assert build_calls[0]["req_to_mtp_state_index"] is req_to_mtp_state_index - assert build_calls[0]["hold_req_id"] == 5 - - -def test_linear_backend_is_abstract_and_triton_supplies_chunk_kernel(): - assert inspect.isabstract(gdn.LinearAttBackend) - backend = object.__new__(TritonLinearAttBackend) - assert backend.get_prefill_kernel() is fla_ops.chunk_gated_delta_rule - - -@pytest.fixture -def auto_linear_backend_args(monkeypatch): - monkeypatch.setattr( - linear_create_utils, - "get_env_start_args", - lambda: SimpleNamespace( - llm_prefill_att_backend=["fa3"], - llm_decode_att_backend=["flashinfer"], - ssm_state_mode="native", - ), - ) - - -def test_missing_linear_backend_arg_uses_auto_selection(monkeypatch, auto_linear_backend_args): - validate_calls = [] - flashqla = ModuleType("flash_qla") - flashqla.chunk_gated_delta_rule = lambda **kwargs: ("flashqla", kwargs) - monkeypatch.setitem(sys.modules, "flash_qla", flashqla) - monkeypatch.setattr( - linear_create_utils, - "validate", - lambda name: validate_calls.append(name) or True, - ) - - backend_class = linear_create_utils.get_gdn_linear_prefill_att_backend_class(index=1) - - assert backend_class is FlashQlaLinearAttBackend - backend = object.__new__(backend_class) - assert backend.get_prefill_kernel()(q="q")[0] == "flashqla" - assert validate_calls == ["flashqla"] - - -def test_replay_mode_selects_standalone_gdn_backends(monkeypatch): - monkeypatch.setattr( - linear_create_utils, - "get_env_start_args", - lambda: SimpleNamespace( - llm_prefill_att_backend=["fa3", "flashqla"], - llm_decode_att_backend=["flashinfer", "triton"], - ssm_state_mode="replay", - ), - ) - from lightllm.common.basemodel.attention.base_att import BaseAttBackend - from lightllm.common.basemodel.attention.linear.gdn import LinearAttBackend - from lightllm.common.basemodel.attention.linear.replay_gdn import ( - ReplayGDNFlashQlaAttBackend, - ReplayGDNTritonAttBackend, - ) - - prefill_backend = linear_create_utils.get_gdn_linear_prefill_att_backend_class(index=1) - decode_backend = linear_create_utils.get_gdn_linear_decode_att_backend_class(index=1) - - assert prefill_backend is ReplayGDNFlashQlaAttBackend - assert decode_backend is ReplayGDNTritonAttBackend - assert prefill_backend.__bases__ == (BaseAttBackend,) - assert decode_backend.__bases__ == (BaseAttBackend,) - assert not issubclass(prefill_backend, LinearAttBackend) - assert not issubclass(decode_backend, LinearAttBackend) - - -def test_kda_entry_selects_native_or_replay(monkeypatch): - from lightllm.common.basemodel.attention.linear.kda import KDALinearAttBackend - from lightllm.common.basemodel.attention.linear.replay_kda import ReplayKDALinearAttBackend - - monkeypatch.setattr( - linear_create_utils, - "get_env_start_args", - lambda: SimpleNamespace(ssm_state_mode="native"), - ) - assert linear_create_utils.get_kda_linear_att_backend_class() is KDALinearAttBackend - - monkeypatch.setattr( - linear_create_utils, - "get_env_start_args", - lambda: SimpleNamespace(ssm_state_mode="replay"), - ) - assert linear_create_utils.get_kda_linear_att_backend_class() is ReplayKDALinearAttBackend - - -def test_flashqla_validation_loads_linear_cache_config(monkeypatch): - linear_config = SimpleNamespace( - 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, - ssm_state_dtype=torch.float32, - ) - load_calls = [] - monkeypatch.setattr( - Qwen3NextResourceConfig, - "load_from_args", - staticmethod(lambda: load_calls.append(True) or linear_config), - ) - - kernel_calls = [] - - def kernel(**kwargs): - kernel_calls.append(kwargs) - return kwargs["q"], kwargs["initial_state"] - - flashqla = ModuleType("flash_qla") - flashqla.chunk_gated_delta_rule = kernel - monkeypatch.setitem(sys.modules, "flash_qla", flashqla) - monkeypatch.setattr(fla_ops, "chunk_gated_delta_rule", kernel) - - original_randn = torch.randn - original_rand = torch.rand - original_tensor = torch.tensor - - def cpu_randn(*args, **kwargs): - kwargs.pop("device", None) - return original_randn(*args, **kwargs) - - def cpu_rand(*args, **kwargs): - kwargs.pop("device", None) - return original_rand(*args, **kwargs) - - def cpu_tensor(*args, **kwargs): - kwargs.pop("device", None) - return original_tensor(*args, **kwargs) - - monkeypatch.setattr(backend_validator.torch, "randn", cpu_randn) - monkeypatch.setattr(backend_validator.torch, "rand", cpu_rand) - monkeypatch.setattr(backend_validator.torch, "tensor", cpu_tensor) - monkeypatch.setattr(backend_validator.torch.cuda, "synchronize", lambda: None) - - success, error = backend_validator._validate_flashqla() - - assert success is True - assert error is None - assert load_calls == [True] - assert len(kernel_calls) == 2 - assert kernel_calls[0]["q"].shape == (1, 64, 2, 8) - assert kernel_calls[0]["v"].shape == (1, 64, 4, 16) - assert kernel_calls[0]["q"].dtype == torch.bfloat16 - assert kernel_calls[0]["initial_state"].shape == (1, 4, 8, 16) - assert kernel_calls[0]["initial_state"].dtype == torch.float32 - - -def test_explicit_linear_backend_arg_skips_auto_selection(monkeypatch): - monkeypatch.setattr( - linear_create_utils, - "get_env_start_args", - lambda: SimpleNamespace( - llm_prefill_att_backend=["fa3", "triton"], - llm_decode_att_backend=["flashinfer", "triton"], - ssm_state_mode="native", - ), - ) - monkeypatch.setattr( - linear_create_utils, - "validate", - lambda *args: pytest.fail("an explicit linear backend must not be auto-validated"), - ) - - prefill_backend_class = linear_create_utils.get_gdn_linear_prefill_att_backend_class(index=1) - decode_backend_class = linear_create_utils.get_gdn_linear_decode_att_backend_class(index=1) - - assert prefill_backend_class is TritonLinearAttBackend - assert decode_backend_class is TritonLinearAttBackend - - -def test_cli_accepts_linear_backend_at_index_one(): - args = make_argument_parser().parse_args( - [ - "--llm_prefill_att_backend", - "fa3", - "flashqla", - "--llm_decode_att_backend", - "flashinfer", - "triton", - ] - ) - - assert args.llm_prefill_att_backend == ["fa3", "flashqla"] - assert args.llm_decode_att_backend == ["flashinfer", "triton"] - - -@pytest.mark.parametrize( - "get_backend_class", - [ - linear_create_utils.get_gdn_linear_prefill_att_backend_class, - linear_create_utils.get_gdn_linear_decode_att_backend_class, - ], -) -def test_linear_backend_rejects_non_linear_index(auto_linear_backend_args, get_backend_class): - with pytest.raises(AssertionError, match="index must be 1"): - get_backend_class(index=0) - - -def test_missing_linear_decode_backend_arg_defaults_to_triton(auto_linear_backend_args): - backend_class = linear_create_utils.get_gdn_linear_decode_att_backend_class(index=1) - - assert backend_class is TritonLinearAttBackend - - -def test_gdn_decode_backend_uses_priority_list(monkeypatch, auto_linear_backend_args): - validate_calls = [] - monkeypatch.setattr( - linear_create_utils, - "gdn_decode_backends", - {"optimized": FlashQlaLinearAttBackend, "triton": TritonLinearAttBackend}, - ) - monkeypatch.setattr( - linear_create_utils, - "validate", - lambda name: validate_calls.append(name) or True, - ) - - backend_class = linear_create_utils.get_gdn_linear_decode_att_backend_class( - index=1, priority_list=("optimized", "triton") - ) - - assert backend_class is FlashQlaLinearAttBackend - assert validate_calls == ["optimized"] - - -def test_gdn_prefill_backend_falls_back_when_flashqla_validation_fails(monkeypatch, auto_linear_backend_args): - monkeypatch.setattr(linear_create_utils, "validate", lambda *args: False) - - backend_class = linear_create_utils.get_gdn_linear_prefill_att_backend_class(index=1) - - assert backend_class is TritonLinearAttBackend - - -def test_gdn_prefill_backend_tries_candidates_in_order(monkeypatch, auto_linear_backend_args): - validate_calls = [] - monkeypatch.setattr( - linear_create_utils, - "gdn_prefill_backends", - {"flashqla2": TritonLinearAttBackend, "flashqla3": FlashQlaLinearAttBackend}, - ) - monkeypatch.setattr( - linear_create_utils, - "validate", - lambda name, *args: validate_calls.append(name) or name == "flashqla3", - ) - - backend_class = linear_create_utils.get_gdn_linear_prefill_att_backend_class( - index=1, priority_list=("flashqla2", "flashqla3") - ) - - 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, role="main"): - 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, - ) - infer_state.att_state_role = role - if req_to_mtp_state_index is not None: - 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, - ) - 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.register_state(model_role, linear) - resources.finish_state_registration() - infer_state.att_state_manager = resources - 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 diff --git a/unit_tests/common/basemodel/attention/test_dsv4_flashmla.py b/unit_tests/common/basemodel/attention/test_dsv4_flashmla.py deleted file mode 100644 index ae848f13de..0000000000 --- a/unit_tests/common/basemodel/attention/test_dsv4_flashmla.py +++ /dev/null @@ -1,377 +0,0 @@ -import gc -import os -from pathlib import Path -import subprocess -import sys -from types import SimpleNamespace - -import pytest -import torch - -pytest.importorskip("flash_mla") - -from lightllm.common.basemodel.attention.nsa import dsv4_fp8_flashmla_sparse as attention -from lightllm.common.basemodel.layer_weights.meta_weights import TpAttSinkWeight - - -@pytest.fixture(autouse=True) -def _tp_env(monkeypatch): - monkeypatch.setenv("LIGHTLLM_DP_WORLD_SIZE", "1") - monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_DP", "0") - monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") - yield - # Collect infer-state cycles here so old graphs are not destroyed during the next capture. - gc.collect() - - -def test_backend_import_without_vllm(): - script = """ -import importlib.abc -import sys - -class BlockVllm(importlib.abc.MetaPathFinder): - def find_spec(self, fullname, path=None, target=None): - if fullname == 'vllm' or fullname.startswith('vllm.'): - raise ModuleNotFoundError('vllm deliberately unavailable') - -sys.meta_path.insert(0, BlockVllm()) -from lightllm.common.basemodel.attention.nsa import dsv4_fp8_flashmla_sparse as attention -assert attention.flashmla.__name__ == 'flash_mla' -assert not any(name == 'vllm' or name.startswith('vllm.') for name in sys.modules) -""" - root = Path(__file__).resolve().parents[4] - env = dict(os.environ, PYTHONPATH=str(root)) - result = subprocess.run([sys.executable, "-c", script], cwd=root, env=env, capture_output=True, text=True) - 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") - rows = torch.randn((pages, page_size, 512), dtype=torch.bfloat16, device="cuda") * 0.1 - nope = rows[..., :448].to(torch.float8_e4m3fn) - data = buffer[:, : page_size * 576].view(pages, page_size, 576) - data[..., :448] = nope.view(torch.uint8) - data[..., 448:] = rows[..., 448:].contiguous().view(torch.uint8) - buffer[:, page_size * 576 : page_size * 584].view(pages, page_size, 8)[..., :7] = 127 - decoded = torch.cat((nope.float(), rows[..., 448:].float()), dim=-1).view(-1, 512) - return buffer, decoded - - -def _case(heads, ratio): - torch.manual_seed(6014) - tokens = 5 - model = type("TestModel", (), {})() - model.config = {"num_attention_heads": heads, "compress_ratios": [0, 4, 128]} - 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 - q[:, heads:] = 0 - pool, decoded = _pool(128, 2) - extra, extra_decoded = _pool(64 if ratio == 4 else 2, 3) if ratio else (None, None) - indices = torch.full((tokens, 1, 128), -1, dtype=torch.int32, device="cuda") - indices[:, 0, :4] = torch.tensor([0, 3, -1, 7], dtype=torch.int32, device="cuda") - lengths = torch.tensor([1, 2, 3, 4, 4], dtype=torch.int32, device="cuda") - extra_indices = indices.clone() if ratio else None - if ratio: - extra_indices[:, 0, :4] = torch.tensor([0, 1, -1, 5], dtype=torch.int32, device="cuda") - attn_sink = TpAttSinkWeight( - all_q_head_num=heads, weight_name="attn_sink", data_type=torch.float32, head_padding_multiple=64 - ) - attn_sink.load_hf_weights({"attn_sink": torch.linspace(-0.5, 0.5, heads, dtype=torch.float32)}) - params = dict( - compress_ratio=ratio, - head_dim_v=512, - softmax_scale=512 ** -0.5, - swa_indices=indices, - swa_lengths=lengths, - attn_sink=attn_sink.weight, - extra_indices=extra_indices, - extra_lengths=lengths.clone() if ratio else None, - ) - # The attention backend receives all cache inputs through k, without a memory manager. - state_info = SimpleNamespace() - return backend, state_info, q, (pool, extra), decoded, extra_decoded, params - - -def _reference(q, decoded, extra_decoded, params, heads): - outputs = [] - for row in range(q.shape[0]): - indexes = params["swa_indices"][row, 0, : params["swa_lengths"][row].item()].long() - selected = [decoded[indexes[indexes >= 0]]] - if extra_decoded is not None: - indexes = params["extra_indices"][row, 0, : params["extra_lengths"][row].item()].long() - selected.append(extra_decoded[indexes[indexes >= 0]]) - kv = torch.cat(selected) - logits = q[row, :heads].float() @ kv.T * params["softmax_scale"] - logits = torch.cat((logits, params["attn_sink"][:heads, None]), dim=-1) - weights = logits.softmax(-1)[:, :-1] - outputs.append((weights @ kv).to(q.dtype)) - return torch.stack(outputs) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("heads", [16, 64, 128]) -@pytest.mark.parametrize("ratio", [0, 4, 128]) -@pytest.mark.parametrize("provide_out", [False, True]) -def test_prefill_matches_reference_and_preserves_out(monkeypatch, heads, ratio, provide_out): - backend, state_info, q, pool, decoded, extra_decoded, params = _case(heads, ratio) - state = backend.create_att_prefill_state(state_info) - state.init_state() - allocations = [] - - def alloc_output(shape, **kwargs): - buffer = torch.empty(shape, **kwargs) - allocations.append(buffer) - return buffer - - kernel = attention.flashmla.flash_mla_with_kvcache - - def flashmla_with_out(**kwargs): - assert kwargs["out"] is not None - return kernel(**kwargs) - - monkeypatch.setattr(attention.flashmla, "flash_mla_with_kvcache", flashmla_with_out) - out = torch.full_like(q, float("nan")) if provide_out else None - actual = state.prefill_att( - q, - pool, - None, - att_control=attention.AttControl(nsa_prefill=True, nsa_prefill_dict=params), - alloc_func=alloc_output, - out=out, - ) - assert params["attn_sink"].shape == (backend.padded_q_head_num,) - assert torch.count_nonzero(params["attn_sink"][heads:]) == 0 - assert len(allocations) == (0 if provide_out else 1) - output_buffer = out if provide_out else allocations[0] - assert actual.data_ptr() == output_buffer.data_ptr() - assert actual.stride() == output_buffer.stride() - assert actual.shape == (q.shape[0], heads, 512) - torch.testing.assert_close(actual, _reference(q, decoded, extra_decoded, params, heads), rtol=0.02, atol=5e-4) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("heads", [16, 64]) -@pytest.mark.parametrize("ratio", [0, 4, 128]) -@pytest.mark.parametrize("overlap", [False, True]) -def test_decode_cuda_graph_reads_updated_query_and_lengths(monkeypatch, heads, ratio, overlap): - from lightllm.common.basemodel import cuda_graph as cuda_graph_module - - args = SimpleNamespace( - enable_decode_microbatch_overlap=overlap, enable_tpsp_mix_mode=False, enable_torch_memory_saver=False - ) - monkeypatch.setattr(cuda_graph_module, "get_env_start_args", lambda: args) - runner = cuda_graph_module.CudaGraph(1, 2, 1, max_batch_size=5, max_len_in_batch=16384) - cases = [] - for _ in range(2 if overlap else 1): - backend, state_info, q, packed_kv, decoded, extra_decoded, params = _case(heads, ratio) - indices = torch.arange(128, dtype=torch.int32, device=q.device) - params["swa_indices"].copy_(indices) - params["swa_lengths"].fill_(1) - if ratio: - page_size = 64 if ratio == 4 else 2 - extra, extra_decoded = _pool(page_size, 128 // page_size) - packed_kv = (packed_kv[0], extra) - params["extra_indices"].copy_(indices) - params["extra_lengths"].fill_(1) - state_info.input_ids = torch.zeros(q.shape[0], dtype=torch.int64, device=q.device) - state_info.q = q - state_info.packed_kv = packed_kv - state_info.att_control = attention.AttControl(nsa_decode=True, nsa_decode_dict=params) - state_info.decode_att_state = backend.create_att_decode_state(state_info) - state_info.decode_att_state1 = None - state_info.decode_att_state.init_state() - cases.append((state_info, decoded, extra_decoded, params)) - - def decode_func(*states): - outputs = tuple( - state.decode_att_state.decode_att(state.q, state.packed_kv, None, att_control=state.att_control) - for state in states - ) - return outputs if overlap else outputs[0] - - # Exercise the framework's warmup + capture directly, with no operator-specific reset. - outputs = runner.capture_decode(decode_func, *(case[0] for case in cases)) - outputs = outputs if overlap else (outputs,) - for state_info, _, _, _ in cases: - state = state_info.decode_att_state - eager_meta = state.sched_meta_by_ratio[ratio] - capture_meta = state.capture_sched_meta_by_ratio[ratio] - assert eager_meta.tile_scheduler_metadata.data_ptr() != capture_meta.tile_scheduler_metadata.data_ptr() - - # Reuse the same graph as lengths grow, shrink, and vary within the batch. - for lengths in ([128] * 5, [64] * 5, [1] * 5, [1, 64, 65, 127, 128]): - for index, (state_info, _, _, params) in enumerate(cases): - state_info.q.mul_(1.25 + index * 0.1) - new_lengths = torch.tensor(lengths, dtype=torch.int32, device=state_info.q.device) - if index: - new_lengths = new_lengths.flip(0) - params["swa_lengths"].copy_(new_lengths) - if ratio: - params["extra_lengths"].copy_(new_lengths) - runner.graph[5][0].replay() - - for actual, (state_info, decoded, extra_decoded, params) in zip(outputs, cases): - assert actual.shape == (state_info.q.shape[0], heads, 512) - assert actual.stride(0) == backend.padded_q_head_num * 512 - torch.testing.assert_close( - actual, _reference(state_info.q, decoded, extra_decoded, params, heads), rtol=0.02, atol=5e-4 - ) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("heads", [16, 64]) -def test_prefill_out_supports_strided_scratch_buffer(heads): - backend, state_info, q, pool, decoded, extra_decoded, params = _case(heads, 4) - storage = torch.full((q.shape[0], backend.padded_q_head_num, 1024), float("nan"), dtype=q.dtype, device=q.device) - out = storage[..., :512] - state = backend.create_att_prefill_state(state_info) - state.init_state() - actual = state.prefill_att( - q, pool, None, out=out, att_control=attention.AttControl(nsa_prefill=True, nsa_prefill_dict=params) - ) - assert actual.data_ptr() == out.data_ptr() - assert actual.stride() == out.stride() - assert storage[..., 512:].isnan().all() - torch.testing.assert_close(actual, _reference(q, decoded, extra_decoded, params, heads), rtol=0.02, atol=5e-4) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("heads", [16, 64]) -@pytest.mark.parametrize("ratio", [0, 4, 128]) -def test_prefill_cuda_graph_reads_new_request_state(monkeypatch, heads, ratio): - from lightllm.common.basemodel import prefill_cuda_graph as cuda_graph_module - from lightllm.common.basemodel.hidden_collector import NoopHiddenCollector - from lightllm.models.deepseek_v4.infer_struct import DeepseekV4InferStateInfo - from lightllm.models.deepseek_v4.layer_infer.transformer_layer_infer import DeepseekV4TransformerLayerInfer - - args = SimpleNamespace( - enable_prefill_microbatch_overlap=False, - prefill_cudagraph_max_handle_token=8, - batch_max_tokens=8, - enable_tpsp_mix_mode=False, - ) - monkeypatch.setattr(cuda_graph_module, "get_env_start_args", lambda: args) - runner = cuda_graph_module.PrefillCudaGraph(decode_cuda_graph=None, tp_world_size=1) - backend, _, q, packed_kv, decoded, extra_decoded, params = _case(heads, ratio) - indices = torch.arange(128, dtype=torch.int32, device=q.device) - params["swa_indices"].copy_(indices) - if ratio: - page_size = 64 if ratio == 4 else 2 - extra, extra_decoded = _pool(page_size, 128 // page_size) - 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.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) - infer_state.prefill_att_state = backend.create_att_prefill_state(infer_state) - infer_state.prefill_att_state.init_state() - return infer_state - - allocations = [] - callback_states = [] - layer = DeepseekV4TransformerLayerInfer.__new__(DeepseekV4TransformerLayerInfer) - layer.tp_q_head_num_ = heads - layer.head_dim_ = 512 - layer.v_head_dim = 512 - layer.compress_ratio = ratio - layer.softmax_scale = params["softmax_scale"] - layer.layer_num_ = 0 - layer_weight = SimpleNamespace(attn_sink_=SimpleNamespace(weight=params["attn_sink"])) - - def alloc_output(shape, **kwargs): - buffer = torch.empty(shape, **kwargs) - allocations.append(buffer) - return buffer - - def compress_and_index(q_lora, x, infer_state, layer_weight): - assert not torch.cuda.is_current_stream_capturing() - callback_states.append(infer_state) - return infer_state.attn_params - - layer.alloc_tensor = alloc_output - layer._compress_and_index = compress_and_index - q_lora = torch.empty((q.shape[0], 4), dtype=q.dtype, device=q.device) - x = torch.empty((q.shape[0], 8), dtype=q.dtype, device=q.device) - - def prefill_func(inputs, infer_state): - graph_q = inputs[0] * 1.0 - att_out = layer._context_attention_wrapper_run(graph_q, inputs[1], inputs[2], infer_state, layer_weight) - assert att_out.shape == (q.shape[0], heads, 512) - assert att_out.data_ptr() == allocations[-1].data_ptr() - assert att_out.stride(0) == backend.padded_q_head_num * 512 - # A graph segment after attention must read the output filled by the CPU callback. - return [att_out + 0.125] - - infer_state = new_request(1) - layer._context_attention_wrapper_run(q, q_lora, x, infer_state, layer_weight) - outputs = runner.capture_prefill(prefill_func, [q, q_lora, x], infer_state) - torch.testing.assert_close( - outputs[0], - _reference(q, decoded, extra_decoded, infer_state.attn_params, heads) + 0.125, - rtol=0.02, - atol=5e-4, - ) - assert len(infer_state.prefill_cuda_graph_exe_list) == 3 - previous_meta = infer_state.prefill_att_state.sched_meta_by_ratio[ratio] - for length in (128, 64, 1, 64): - q.mul_(1.25) - infer_state = new_request(length) - current_meta = infer_state.prefill_att_state.sched_meta_by_ratio[ratio] - assert current_meta.tile_scheduler_metadata is None - actual = runner.replay([q, q_lora, x], infer_state)[0] - assert callback_states[-1] is infer_state - assert current_meta.tile_scheduler_metadata.data_ptr() != previous_meta.tile_scheduler_metadata.data_ptr() - torch.testing.assert_close( - actual, - _reference(q, decoded, extra_decoded, infer_state.attn_params, heads) + 0.125, - rtol=0.02, - atol=5e-4, - ) - previous_meta = current_meta diff --git a/unit_tests/common/basemodel/attention/test_fp8_logical_layer.py b/unit_tests/common/basemodel/attention/test_fp8_logical_layer.py deleted file mode 100644 index 9d02784676..0000000000 --- a/unit_tests/common/basemodel/attention/test_fp8_logical_layer.py +++ /dev/null @@ -1,117 +0,0 @@ -"""FP8 后端按推理层传入的逻辑层号选择标定，覆盖稀疏层号、共享层和 draft。""" - -from dataclasses import replace -from types import SimpleNamespace -import weakref - -import pytest -import torch - -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 -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 - - -@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=template, - config=replace_logical_layer_groups(template.config, ((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 = AttentionStateManager(request_capacity=1) - manager.register_state(role, state) - manager.finish_state_registration() - infer = LlamaInferStateInfo() - 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 = [] - - # 仅替换外部计算内核，保留推理层、公开 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.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, - 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) - - 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/common/basemodel/attention/test_fp8_page_size.py b/unit_tests/common/basemodel/attention/test_fp8_page_size.py deleted file mode 100644 index 107728c26d..0000000000 --- a/unit_tests/common/basemodel/attention/test_fp8_page_size.py +++ /dev/null @@ -1,45 +0,0 @@ -from types import SimpleNamespace -from unittest.mock import patch - -import pytest - -from lightllm.common.basemodel.attention.fa3 import fp8 as fa3_fp8 -from lightllm.common.basemodel.attention.fa3.fp import Fa3AttBackend -from lightllm.common.basemodel.attention.fa3.fp8 import Fp8Fa3AttBackend -from lightllm.common.basemodel.attention.fa3.mla import MlaFa3AttBackend -from lightllm.common.basemodel.attention.flashinfer import fp8 as flashinfer_fp8 -from lightllm.common.basemodel.attention.flashinfer.fp import FlashInferAttBackend -from lightllm.common.basemodel.attention.flashinfer.fp8 import Fp8FlashInferAttBackend -from lightllm.common.basemodel.attention.flashinfer.mla import MlaFlashInferAttBackend - - -@pytest.mark.parametrize( - "backend_class", - [Fa3AttBackend, MlaFa3AttBackend, FlashInferAttBackend, MlaFlashInferAttBackend], -) -def test_default_infer_page_size_matches_model_page_size(backend_class): - backend = object.__new__(backend_class) - backend.model = SimpleNamespace(args=SimpleNamespace(page_size=16)) - - backend._init_infer_page_size() - - assert backend.infer_page_size == 16 - - -@pytest.mark.parametrize( - ("backend_class", "backend_module"), - [ - (Fp8Fa3AttBackend, fa3_fp8), - (Fp8FlashInferAttBackend, flashinfer_fp8), - ], -) -def test_fp8_attention_backend_uses_single_token_infer_pages(backend_class, backend_module): - backend = object.__new__(backend_class) - backend.model = SimpleNamespace(args=SimpleNamespace(page_size=16)) - - with patch.object(backend_module.logger, "warning") as warning: - backend._init_infer_page_size() - - assert backend.infer_page_size == 1 - warning.assert_called_once() - assert "infer_page_size=1" in warning.call_args.args[0] diff --git a/unit_tests/common/basemodel/attention/test_lazy_import.py b/unit_tests/common/basemodel/attention/test_lazy_import.py deleted file mode 100644 index e7423d30f9..0000000000 --- a/unit_tests/common/basemodel/attention/test_lazy_import.py +++ /dev/null @@ -1,248 +0,0 @@ -"""Check import isolation and selection without loading optional GPU dependencies.""" - -import importlib -import importlib.abc -import importlib.util -from pathlib import Path -import sys -from types import ModuleType, SimpleNamespace -import unittest -from unittest.mock import Mock, patch - - -ROOT = Path(__file__).resolve().parents[4] -ATTENTION = "lightllm.common.basemodel.attention" - - -class BackendImports(importlib.abc.MetaPathFinder, importlib.abc.Loader): - """Reject unselected implementations, and supply lightweight selected backends.""" - - def __init__(self, package=ATTENTION): - self.package = package - self.allowed = {} - self.imported = [] - - def allow(self, suffix, class_name, **attributes): - module = ModuleType(f"{self.package}.{suffix}") - setattr(module, class_name, type(class_name, (), {})) - module.__dict__.update(attributes) - self.allowed[module.__name__] = module - return getattr(module, class_name) - - def find_spec(self, fullname, path=None, target=None): - suffix = fullname.removeprefix(f"{self.package}.") - if not fullname.startswith(f"{self.package}.") or "." not in suffix: - return None - if suffix.split(".")[0] not in {"triton", "fa3", "flashinfer", "nsa", "linear", "sdpa", "xformers"}: - return None - self.imported.append(fullname) - if fullname not in self.allowed: - raise AssertionError(f"Unexpected attention backend import: {fullname}") - return importlib.util.spec_from_loader(fullname, self) - - def create_module(self, spec): - return self.allowed[spec.name] - - def exec_module(self, module): - pass - - -class TestLazyAttentionImports(unittest.TestCase): - def setUp(self): - # Isolate attention from basemodel's unrelated GPU initialization. - modules_patch = patch.dict(sys.modules) - modules_patch.start() - self.addCleanup(modules_patch.stop) - for name in list(sys.modules): - if name == ATTENTION or name.startswith(f"{ATTENTION}."): - del sys.modules[name] - for name in ("lightllm", "lightllm.common", "lightllm.common.basemodel", "lightllm.utils"): - module = ModuleType(name) - module.__path__ = [str(ROOT.joinpath(*name.split(".")))] - sys.modules[name] = module - - self.args = SimpleNamespace( - llm_kv_type="None", - llm_prefill_att_backend=["triton"], - llm_decode_att_backend=["triton"], - enable_ep_moe=False, - mtp_step=0, - ssm_state_mode="native", - ) - self.validate = Mock(return_value=True) - dependencies = { - f"{ATTENTION}.base_att": { - name: type(name, (), {}) - for name in ("BaseAttBackend", "BasePrefillAttState", "BaseDecodeAttState", "AttControl") - }, - "lightllm.utils.envs_utils": {"get_env_start_args": lambda: self.args}, - "lightllm.utils.backend_validator": {"validate": self.validate}, - "lightllm.utils.log_utils": {"init_logger": lambda name: Mock()}, - } - for name, attributes in dependencies.items(): - module = ModuleType(name) - module.__dict__.update(attributes) - sys.modules[name] = module - - self.imports = BackendImports() - sys.meta_path.insert(0, self.imports) - self.addCleanup(sys.meta_path.remove, self.imports) - self.attention = importlib.import_module(ATTENTION) - - def test_importing_selection_apis_loads_no_implementations(self): - importlib.import_module(f"{ATTENTION}.linear") - importlib.import_module(f"{ATTENTION}.nsa") - self.assertEqual(self.imports.imported, []) - - def test_explicit_regular_selection_skips_neo(self): - backend = self.imports.allow("fa3.fp", "Fa3AttBackend") - self.args.llm_prefill_att_backend = ["fa3"] - - self.assertIs(self.attention.get_prefill_att_backend_class(), backend) - self.assertEqual(self.imports.imported, [f"{ATTENTION}.fa3.fp"]) - self.validate.assert_not_called() - - def test_auto_selection_does_not_import_failed_candidates(self): - backend = self.imports.allow("flashinfer.fp", "FlashInferAttBackend") - self.args.llm_prefill_att_backend = ["auto"] - self.validate.side_effect = lambda name: name == "flashinfer" - - self.assertIs(self.attention.get_prefill_att_backend_class(), backend) - self.assertEqual([call.args[0] for call in self.validate.call_args_list], ["fa3", "flashinfer"]) - self.assertEqual(self.imports.imported, [f"{ATTENTION}.flashinfer.fp"]) - - def test_mtp_decode_preserves_fa3_priority(self): - backend = self.imports.allow("fa3.fp", "Fa3AttBackend") - self.args.llm_decode_att_backend = ["auto"] - self.args.mtp_step = 2 - - self.assertIs(self.attention.get_decode_att_backend_class(), backend) - self.validate.assert_called_once_with("fa3") - - def test_expert_parallelism_still_excludes_flashinfer(self): - backend = self.imports.allow("triton.fp", "TritonAttBackend") - self.args.llm_decode_att_backend = ["auto"] - self.args.enable_ep_moe = True - self.validate.side_effect = lambda name: name == "triton" - - self.assertIs(self.attention.get_decode_att_backend_class(), backend) - self.assertEqual([call.args[0] for call in self.validate.call_args_list], ["fa3", "triton"]) - - def test_auto_selection_falls_back_to_triton(self): - backend = self.imports.allow("triton.mla", "MlaTritonAttBackend") - self.args.llm_prefill_att_backend = ["auto"] - self.validate.return_value = False - - self.assertIs(self.attention.get_mla_prefill_att_backend_class(), backend) - self.assertEqual(self.imports.imported, [f"{ATTENTION}.triton.mla"]) - - def test_nsa_selection_does_not_import_other_kv_types(self): - backend = self.imports.allow("nsa.fp8_flashmla_sparse", "NsaFlashMlaFp8SparseAttBackend") - self.args.llm_kv_type = "fp8kv_dsa" - self.args.llm_decode_att_backend = ["flashmla_sparse"] - - self.assertIs(self.attention.get_nsa_decode_att_backend_class(), backend) - self.assertEqual(self.imports.imported, [f"{ATTENTION}.nsa.fp8_flashmla_sparse"]) - - def test_custom_backend_map_uses_import_paths(self): - backend = self.imports.allow("nsa.custom", "CustomBackend") - backend_map = {"None": {"custom": "nsa.custom:CustomBackend"}} - self.args.llm_prefill_att_backend = ["custom"] - self.args.llm_decode_att_backend = ["auto"] - - self.assertIs(self.attention.get_nsa_prefill_att_backend_class(backend_map=backend_map), backend) - self.assertIs( - self.attention.get_nsa_decode_att_backend_class(priority_list=["custom"], backend_map=backend_map), - backend, - ) - self.assertEqual(self.imports.imported, [f"{ATTENTION}.nsa.custom"]) - - def test_neo_triton_selection_does_not_probe_fa3(self): - backend = self.imports.allow("triton.neo", "NeoTritonAttBackend") - - self.assertIs(self.attention.get_neo_prefill_att_backend_class(), backend) - self.args.llm_prefill_att_backend = ["auto"] - self.assertIs(self.attention.get_neo_prefill_att_backend_class(priority_list=["triton"]), backend) - self.assertEqual(self.imports.imported, [f"{ATTENTION}.triton.neo"]) - - def test_neo_auto_selection_still_falls_back_without_image_token_interface(self): - self.imports.allow("fa3.neo", "NeoFa3AttBackend", HAS_FLASH_ATTN_INTERFACE=False) - backend = self.imports.allow("triton.neo", "NeoTritonAttBackend") - self.args.llm_prefill_att_backend = ["auto"] - - self.assertIs(self.attention.get_neo_prefill_att_backend_class(), backend) - self.validate.assert_called_once_with("triton") - self.args.llm_prefill_att_backend = ["fa3"] - with self.assertRaises(KeyError): - self.attention.get_neo_prefill_att_backend_class() - - def test_neo_auto_selection_uses_supported_fa3(self): - backend = self.imports.allow("fa3.neo", "NeoFa3AttBackend", HAS_FLASH_ATTN_INTERFACE=True) - self.args.llm_prefill_att_backend = ["auto"] - - self.assertIs(self.attention.get_neo_prefill_att_backend_class(), backend) - self.assertEqual(self.imports.imported, [f"{ATTENTION}.fa3.neo"]) - self.validate.assert_called_once_with("fa3") - - def _allow_linear_backend_classes(self): - flashqla = self.imports.allow("linear.flashqla", "FlashQlaLinearAttBackend") - triton = self.imports.allow("linear.triton", "TritonLinearAttBackend") - self.imports.allow("linear.replay_gdn", "ReplayGDNTritonAttBackend") - replay_gdn = self.imports.allowed[f"{ATTENTION}.linear.replay_gdn"] - replay_gdn.ReplayGDNFlashQlaAttBackend = type("ReplayGDNFlashQlaAttBackend", (), {}) - self.imports.allow("linear.kda", "KDALinearAttBackend") - self.imports.allow("linear.replay_kda", "ReplayKDALinearAttBackend") - return flashqla, triton - - def test_linear_prefill_auto_selects_validated_flashqla(self): - flashqla, _triton = self._allow_linear_backend_classes() - selection = importlib.import_module(f"{ATTENTION}.create_linear_utils") - - self.assertIs(selection.get_gdn_linear_prefill_att_backend_class(), flashqla) - self.validate.assert_called_once_with("flashqla") - - def test_linear_selection_falls_back_to_triton(self): - _flashqla, triton = self._allow_linear_backend_classes() - selection = importlib.import_module(f"{ATTENTION}.create_linear_utils") - self.validate.return_value = False - - self.assertIs(selection.get_gdn_linear_prefill_att_backend_class(), triton) - self.assertIs(selection.get_gdn_linear_decode_att_backend_class(), triton) - self.args.llm_prefill_att_backend.append("triton") - self.assertIs(selection.get_gdn_linear_prefill_att_backend_class(), triton) - - def test_vit_selection_loads_only_requested_implementation(self): - package = "lightllm.common.basemodel.attention_vit" - for name in list(sys.modules): - if name == package or name.startswith(f"{package}."): - del sys.modules[name] - base = ModuleType(f"{package}.base_att") - base.BaseVitAttBackend = type("BaseVitAttBackend", (), {}) - sys.modules[base.__name__] = base - sys.modules["lightllm.utils.backend_validator"]._validate = self.validate - imports = BackendImports(package) - sys.meta_path.insert(0, imports) - self.addCleanup(sys.meta_path.remove, imports) - selection = importlib.import_module(f"{package}.create_utils") - self.assertEqual(imports.imported, []) - - for name, class_name in ( - ("fa3", "Fa3VitAttBackend"), - ("triton", "TritonVitAttBackend"), - ("sdpa", "SdpaVitAttBackend"), - ("xformers", "XformersVitAttBackend"), - ): - backend = imports.allow(f"{name}.fp", class_name) - self.args.vit_att_backend = [name] - self.assertEqual(selection.init_vit_att_backend(), name) - self.assertIs(selection.get_vit_att_backend_class(name), backend) - self.validate.assert_not_called() - - self.args.vit_att_backend = ["auto"] - self.validate.side_effect = lambda name: name == "sdpa" - self.assertEqual(selection.init_vit_att_backend(), "sdpa") - self.assertEqual([call.args[0] for call in self.validate.call_args_list], ["fa3", "xformers", "sdpa"]) - - -if __name__ == "__main__": - unittest.main() diff --git a/unit_tests/common/basemodel/attention/test_state_lifetime.py b/unit_tests/common/basemodel/attention/test_state_lifetime.py deleted file mode 100644 index 9e049d54b4..0000000000 --- a/unit_tests/common/basemodel/attention/test_state_lifetime.py +++ /dev/null @@ -1,174 +0,0 @@ -import gc -import weakref - -import pytest -import torch - -from lightllm.common.basemodel.attention.base_att import BaseDecodeAttState, BasePrefillAttState -from lightllm.common.basemodel.cuda_graph import CudaGraph -from lightllm.common.basemodel.infer_struct import InferStateInfo - - -class _PrefillState(BasePrefillAttState): - def init_state(self): - self.positions = self.infer_state.position_ids - - def prefill_att(self, q, k, v, **kwargs): - return q - - -class _DecodeState(BaseDecodeAttState): - def init_state(self): - self.positions = self.infer_state.position_ids - - def decode_att(self, q, k, v, **kwargs): - return q - - -@pytest.fixture -def no_cycle_gc(): - enabled = gc.isenabled() - gc.collect() - gc.disable() - try: - yield - finally: - gc.collect() - if enabled: - gc.enable() - - -def _create_state(is_prefill, secondary=False): - owner = InferStateInfo() - owner.is_prefill = is_prefill - owner.position_ids = torch.arange(4) - state_cls = _PrefillState if is_prefill else _DecodeState - slot = "prefill_att_state" if is_prefill else "decode_att_state" - setattr(owner, slot, state_cls(_infer_state_weak_ref=weakref.ref(owner))) - if secondary: - setattr(owner, slot + "1", state_cls(_infer_state_weak_ref=weakref.ref(owner))) - owner.init_att_state() - return owner - - -@pytest.mark.parametrize("is_prefill", [False, True]) -@pytest.mark.parametrize("secondary", [False, True]) -def test_forward_states_release_without_cycle_gc(no_cycle_gc, is_prefill, secondary): - owner = _create_state(is_prefill, secondary) - states = [state for state in vars(owner).values() if isinstance(state, (BasePrefillAttState, BaseDecodeAttState))] - owner_ref = weakref.ref(owner) - tensor_ref = weakref.ref(owner.position_ids) - state_refs = [weakref.ref(state) for state in states] - - # Even an external reference to attention metadata must not retain its owner. - del owner - assert owner_ref() is None - with pytest.raises(ReferenceError): - states[0].infer_state - assert tensor_ref() is not None - del states - assert all(ref() is None for ref in state_refs) - assert tensor_ref() is None - - -@pytest.mark.parametrize("state_cls", [_PrefillState, _DecodeState]) -def test_attention_getter_uses_a_weak_reference_and_is_read_only(state_cls): - state = state_cls() - assert state.infer_state is None - - owner = InferStateInfo() - state = state_cls(_infer_state_weak_ref=weakref.ref(owner)) - assert isinstance(state._infer_state_weak_ref, weakref.ReferenceType) - assert state.infer_state is owner - state.check_no_infer_state() - with pytest.raises(AttributeError): - state.infer_state = owner - - -@pytest.mark.parametrize("is_prefill", [False, True]) -@pytest.mark.parametrize("slot_suffix", ["", "1"]) -def test_check_rejects_infer_state_under_any_attribute_name(is_prefill, slot_suffix): - owner = _create_state(is_prefill, secondary=True) - slot = ("prefill_att_state" if is_prefill else "decode_att_state") + slot_suffix - attention = getattr(owner, slot) - - class _OtherInferState(InferStateInfo): - pass - - # Check the type, rather than just the identity of the expected owner. - attention.saved_context = _OtherInferState() - with pytest.raises(RuntimeError, match="saved_context"): - owner.check_att_states() - - -@pytest.mark.parametrize("is_prefill", [False, True]) -@pytest.mark.parametrize("slot_suffix", ["", "1"]) -def test_init_checks_for_accidental_back_reference(is_prefill, slot_suffix): - owner = _create_state(is_prefill, secondary=True) - slot = ("prefill_att_state" if is_prefill else "decode_att_state") + slot_suffix - base = _PrefillState if is_prefill else _DecodeState - - class _RetainingState(base): - def init_state(self): - self.saved_context = self.infer_state - - setattr(owner, slot, _RetainingState(_infer_state_weak_ref=weakref.ref(owner))) - try: - with pytest.raises(RuntimeError, match="saved_context"): - owner.init_att_state() - finally: - del getattr(owner, slot).saved_context - - -def test_decode_graph_replay_preserves_capture_buffers_and_owner_lifetime(no_cycle_gc): - captured = _create_state(is_prefill=False, secondary=True) - captured.input_ids = torch.zeros(4, dtype=torch.int64) - captured_ref = weakref.ref(captured) - capture_positions_ptr = captured.position_ids.data_ptr() - replay_calls = [] - - class _Graph: - def replay(self): - replay_calls.append(True) - - graph = CudaGraph.__new__(CudaGraph) - output = object() - graph.graph = {4: (_Graph(), captured, output)} - del captured - assert captured_ref() is not None - - runtime = _create_state(is_prefill=False, secondary=True) - runtime.input_ids = torch.ones(4, dtype=torch.int64) - runtime.position_ids.add_(7) - runtime_ref = weakref.ref(runtime) - assert graph._replay(runtime) is output - assert replay_calls == [True] - assert captured_ref().position_ids.data_ptr() == capture_positions_ptr - assert captured_ref().decode_att_state.infer_state is captured_ref() - assert captured_ref().decode_att_state1.infer_state is captured_ref() - torch.testing.assert_close(captured_ref().position_ids, runtime.position_ids) - - del runtime - 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) diff --git a/unit_tests/common/basemodel/test_att_sink_weight.py b/unit_tests/common/basemodel/test_att_sink_weight.py deleted file mode 100644 index 996fda7c86..0000000000 --- a/unit_tests/common/basemodel/test_att_sink_weight.py +++ /dev/null @@ -1,37 +0,0 @@ -import pytest -import torch - -from lightllm.common.basemodel.layer_weights.meta_weights import TpAttSinkWeight - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize( - "heads,tp_size,rank,padding_multiple,data_type,start,end,padded_heads", - [ - (128, 8, 0, 64, torch.float32, 0, 16, 64), - (128, 8, 7, 64, torch.float32, 112, 128, 64), - (128, 2, 1, 64, torch.float32, 64, 128, 64), - (256, 2, 1, 64, torch.float32, 128, 256, 128), - (12, 4, 2, None, torch.bfloat16, 6, 9, 3), - (4, 8, 3, None, torch.bfloat16, 1, 2, 1), - ], -) -def test_load_tp_sink_weight(monkeypatch, heads, tp_size, rank, padding_multiple, data_type, start, end, padded_heads): - monkeypatch.setenv("LIGHTLLM_DP_WORLD_SIZE", str(tp_size)) - monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_DP", str(rank)) - monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") - kwargs = {} if padding_multiple is None else {"head_padding_multiple": padding_multiple} - sink = TpAttSinkWeight(all_q_head_num=heads, weight_name="attn_sink", data_type=data_type, **kwargs) - assert sink.weight.shape == (padded_heads,) - assert not sink.verify_load() - storage_ptr = sink.weight.data_ptr() - checkpoint = torch.arange(heads, dtype=torch.float64) + 0.5 - - sink.load_hf_weights({"attn_sink": checkpoint}) - - assert sink.verify_load() - assert sink.weight.data_ptr() == storage_ptr - assert sink.weight.dtype == data_type - assert sink.weight.is_cuda - torch.testing.assert_close(sink.weight[: end - start], checkpoint[start:end].to(device="cuda", dtype=data_type)) - assert torch.count_nonzero(sink.weight[end - start :]) == 0 diff --git a/unit_tests/common/basemodel/test_cuda_graph_autotune.py b/unit_tests/common/basemodel/test_cuda_graph_autotune.py deleted file mode 100644 index 52b1c1cf2d..0000000000 --- a/unit_tests/common/basemodel/test_cuda_graph_autotune.py +++ /dev/null @@ -1,130 +0,0 @@ -from contextlib import contextmanager -from types import SimpleNamespace - -import pytest -import torch - -from lightllm.common.basemodel import cuda_graph as cuda_graph_module -from lightllm.common.basemodel.cuda_graph import CudaGraph -from lightllm.common.triton_utils import autotuner as autotuner_module -from lightllm.common.triton_utils.autotuner import AutotuneKernelType, AutotuneLevel, Autotuner, autotune - - -@pytest.fixture -def capture_env(monkeypatch): - env = SimpleNamespace(capturing=False, capture_count=0, graphs=[]) - env.args = SimpleNamespace( - enable_decode_microbatch_overlap=False, - enable_tpsp_mix_mode=False, - enable_torch_memory_saver=False, - ) - monkeypatch.setattr(cuda_graph_module, "get_env_start_args", lambda: env.args) - monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", None) - monkeypatch.setattr(autotuner_module.dist, "is_initialized", lambda: False) - monkeypatch.setattr(autotuner_module.KernelConfigs, "get_config_file_name", lambda params: "configs.json") - monkeypatch.setattr(torch.cuda, "synchronize", lambda: None) - monkeypatch.setattr(torch.cuda, "graph_pool_handle", lambda: None) - - class FakeGraph: - def __init__(self): - self.replays = 0 - env.graphs.append(self) - - def replay(self): - self.replays += 1 - - @contextmanager - def capture(graph, **kwargs): - assert not Autotuner.is_autotune_warmup() - env.capture_count += 1 - env.capturing = True - try: - yield - finally: - env.capturing = False - - monkeypatch.setattr(torch.cuda, "CUDAGraph", FakeGraph) - monkeypatch.setattr(torch.cuda, "graph", capture) - return env - - -@pytest.mark.parametrize("overlap", [False, True]) -@pytest.mark.parametrize("level", [0, 1, 2, 3]) -def test_decode_attention_tunes_before_capture_only(capture_env, tmp_path, monkeypatch, overlap, level): - monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: level) - capture_env.args.enable_decode_microbatch_overlap = overlap - graph = CudaGraph(1, 2, 1, max_batch_size=2) - benchmarks = [] - - def make_kernel(kernel_type): - @autotune( - kernel_name=kernel_type.value, - kernel_type=kernel_type, - configs_gen_func=lambda: [{"block": 1}, {"block": 2}], - static_key_func=lambda: {}, - run_key_func=lambda size: size, - ) - def kernel(size, run_config=None): - return size - - def bench(size, run_config): - assert kernel_type == AutotuneKernelType.DECODE_ATTENTION - assert Autotuner.is_kernel_autotune_warmup(kernel_type) - assert not capture_env.capturing - benchmarks.append(run_config) - return 1.0 / run_config["block"] - - cache_dir = tmp_path / kernel_type.value - cache_dir.mkdir() - kernel._cache_dir = str(cache_dir) - monkeypatch.setattr(kernel, "_bench", bench) - return kernel, cache_dir / "configs.json" - - general, general_cache = make_kernel(AutotuneKernelType.GENERAL) - attention, attention_cache = make_kernel(AutotuneKernelType.DECODE_ATTENTION) - states = [SimpleNamespace(input_ids=torch.ones(2, dtype=torch.int64)) for _ in range(2 if overlap else 1)] - outputs = [object() for _ in states] - forward_calls = [] - - def decode_func(*infer_states): - forward_calls.append(capture_env.capturing) - assert Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.DECODE_ATTENTION) == (not capture_env.capturing) - for original, state in zip(states, infer_states): - assert (state is original) == capture_env.capturing - assert not hasattr(state, "temporary_buffer") - state.temporary_buffer = object() - general(state.input_ids.shape[0]) - attention(state.input_ids.shape[0]) - return tuple(outputs) if overlap else outputs[0] - - result = graph.capture_decode(decode_func, *states) - assert result == (tuple(outputs) if overlap else outputs[0]) - assert forward_calls == [False, True] - assert not Autotuner.is_autotune_warmup() - assert capture_env.capture_count == 1 - assert capture_env.graphs[0].replays == 1 - expected_benchmarks = 0 - if level in [AutotuneLevel.ADAPTIVE_AUTOTUNE, AutotuneLevel.FORCE_AUTOTUNE]: - expected_benchmarks = 2 - assert len(benchmarks) == expected_benchmarks - assert attention_cache.exists() == (expected_benchmarks > 0) - assert not general_cache.exists() - - -@pytest.mark.parametrize("overlap", [False, True]) -def test_warmup_failure_restores_phase_without_capturing(capture_env, overlap): - capture_env.args.enable_decode_microbatch_overlap = overlap - graph = CudaGraph(1, 2, 1, max_batch_size=2) - states = [SimpleNamespace(input_ids=torch.ones(2, dtype=torch.int64)) for _ in range(2 if overlap else 1)] - - def decode_func(*infer_states): - assert Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.DECODE_ATTENTION) - raise RuntimeError("decode warmup failed") - - with Autotuner.autotune_warmup(AutotuneKernelType.GENERAL): - with pytest.raises(RuntimeError, match="decode warmup failed"): - graph.capture_decode(decode_func, *states) - assert Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.GENERAL) - assert not Autotuner.is_autotune_warmup() - assert capture_env.capture_count == 0 - assert graph.graph == {} diff --git a/unit_tests/common/basemodel/test_cuda_graph_layout.py b/unit_tests/common/basemodel/test_cuda_graph_layout.py deleted file mode 100644 index 30d4c54368..0000000000 --- a/unit_tests/common/basemodel/test_cuda_graph_layout.py +++ /dev/null @@ -1,133 +0,0 @@ -import math -from types import SimpleNamespace - -import pytest - -import lightllm.common.basemodel.basemodel as basemodel_module -import lightllm.common.basemodel.cuda_graph as cuda_graph_module -import lightllm.common.basemodel.mtp_manager as mtp_manager_module -from lightllm.common.basemodel.basemodel import TpPartBaseModel -from lightllm.common.basemodel.cuda_graph import CudaGraph -from lightllm.common.basemodel.mtp_manager import MtpManager - - -@pytest.fixture(autouse=True) -def _graph_args(monkeypatch): - args = SimpleNamespace( - enable_decode_microbatch_overlap=False, - enable_tpsp_mix_mode=False, - enable_torch_memory_saver=False, - ) - monkeypatch.setattr(cuda_graph_module, "get_env_start_args", lambda: args) - return args - - -def _batch_sizes(max_batch_size, batch_stride=1): - physical_max_batch_size = max_batch_size * batch_stride - graph = CudaGraph( - batch_step_size_before_split=batch_stride, - split_batch_size=4 * batch_stride, - batch_step_size_after_split=2 * batch_stride, - max_batch_size=physical_max_batch_size, - ) - return graph.cuda_graph_batch_sizes - - -def test_dynamic_schedule_uses_compacted_physical_rows(_graph_args): - assert _batch_sizes(max_batch_size=128) == [1, 2, 3, 4, *range(6, 129, 2)] - - -def test_public_static_schedule_preserves_original_static_mtp_default(_graph_args): - assert CudaGraph.gen_cuda_graph_batch_sizes( - batch_step_size_before_split=8, - split_batch_size=32, - batch_step_size_after_split=16, - max_batch_size=32, - ) == [ - 8, - 16, - 24, - 32, - ] - - -def test_instance_and_public_static_schedule_match(_graph_args): - graph = CudaGraph( - batch_step_size_before_split=8, - split_batch_size=32, - batch_step_size_after_split=16, - max_batch_size=128, - ) - - assert graph.cuda_graph_batch_sizes == CudaGraph.gen_cuda_graph_batch_sizes( - batch_step_size_before_split=8, - split_batch_size=32, - batch_step_size_after_split=16, - max_batch_size=graph.max_batch_size, - tp_world_size=graph.tp_world_size, - ) - - -def test_batch_step_size_before_split_controls_capture_range(_graph_args): - assert _batch_sizes(max_batch_size=4, batch_stride=8) == [8, 16, 24, 32] - - -def test_batch_step_size_after_split_controls_capture_range(_graph_args): - assert _batch_sizes(max_batch_size=8, batch_stride=7) == [ - 7, - 14, - 21, - 28, - 42, - 56, - ] - - -@pytest.mark.parametrize("tp_size", [1, 2, 4, 8]) -@pytest.mark.parametrize("mtp_step", [1, 2]) -@pytest.mark.parametrize("dynamic,is_draft", [(False, False), (True, False), (False, True)]) -@pytest.mark.parametrize("overlap", [False, True]) -def test_mtp_tpsp_layout(monkeypatch, _graph_args, tp_size, mtp_step, dynamic, is_draft, overlap): - args = _graph_args - args.enable_tpsp_mix_mode = True - args.enable_decode_microbatch_overlap = overlap - args.mtp_mode = "eagle_with_att" - args.mtp_step = mtp_step - args.mtp_dynamic_verify = dynamic - monkeypatch.setattr(mtp_manager_module, "get_env_start_args", lambda: args) - model = TpPartBaseModel.__new__(TpPartBaseModel) - model.is_mtp_draft_model = is_draft - monkeypatch.setattr(basemodel_module, "get_env_start_args", lambda: args) - monkeypatch.setattr(basemodel_module, "get_llm_data_type", lambda: None) - monkeypatch.setattr(basemodel_module, "get_dp_world_size", lambda: tp_size) - monkeypatch.setattr(MtpManager, "_instance", MtpManager()) - - class StopBeforeWeights(Exception): - pass - - def stop_before_weights(): - raise StopBeforeWeights - - # 执行真实构造函数的容量计算，在读取模型配置、分配 GPU 权重前停止。 - monkeypatch.setattr(model, "_init_config", stop_before_weights) - with pytest.raises(StopBeforeWeights): - model.__init__(dict(run_mode="normal", weight_dir="", max_total_token_num=1024, graph_max_batch_size=7)) - - width = 1 if dynamic or is_draft else mtp_step + 1 - logical_max = 7 // 2 if overlap else 7 - physical_max = logical_max * (1 if is_draft else mtp_step + 1) - alignment = math.lcm(width, tp_size) - assert model.graph_max_batch_size % alignment == 0 - assert physical_max <= model.graph_max_batch_size < physical_max + alignment - - sizes = CudaGraph.gen_cuda_graph_batch_sizes( - batch_step_size_before_split=width, - split_batch_size=4 * width, - batch_step_size_after_split=2 * width, - max_batch_size=model.graph_max_batch_size, - tp_world_size=tp_size, - ) - assert sizes[-1] == model.graph_max_batch_size - assert all(size % width == 0 and size % tp_size == 0 for size in sizes) - if tp_size == 8 and width == 3: - assert sizes == [24] diff --git a/unit_tests/common/basemodel/test_hidden_collector.py b/unit_tests/common/basemodel/test_hidden_collector.py deleted file mode 100644 index 22748b7294..0000000000 --- a/unit_tests/common/basemodel/test_hidden_collector.py +++ /dev/null @@ -1,193 +0,0 @@ -from inspect import isabstract -from types import SimpleNamespace - -import pytest -import torch - -import lightllm.common.basemodel.hidden_collector as hidden_collector_module -from lightllm.common.basemodel.hidden_collector import ( - FinalHiddenCollector, - HiddenCollector, - LayerHiddenCollector, - MtpHeadOutputCollector, - NoopHiddenCollector, -) - - -class _IdentityPreInfer: - @staticmethod - def _tpsp_allgather(input, infer_state): - del infer_state - return input - - -def _mock_target_layer_ids(monkeypatch, layer_ids): - monkeypatch.setattr( - hidden_collector_module, - "get_env_start_args", - lambda: SimpleNamespace(mtp_draft_model_dir=["/models/draft"]), - ) - monkeypatch.setattr( - hidden_collector_module.PretrainedConfig, - "get_config_dict", - lambda _: ({"target_layer_ids": layer_ids}, {}), - ) - - -def test_hidden_collector_is_base_class_for_implementations(): - assert isabstract(HiddenCollector) - assert issubclass(NoopHiddenCollector, HiddenCollector) - assert issubclass(FinalHiddenCollector, HiddenCollector) - assert issubclass(LayerHiddenCollector, HiddenCollector) - - -def test_final_hidden_collectors_are_independent_instances(): - prototype = FinalHiddenCollector() - collector0 = prototype.new_instance() - collector1 = prototype.new_instance() - hidden0 = torch.randn(2, 3) - hidden1 = torch.randn(2, 3) - infer_state = SimpleNamespace(need_dp_prefill_balance=False) - - collector0.add_final_hidden(hidden0) - collected = collector0.finish_output(infer_state=infer_state).spec_hidden - assert collected.data_ptr() == hidden0.data_ptr() - assert collector0.final_hidden is None - - collector0.add_final_hidden(hidden0) - collector1.add_final_hidden(hidden1) - collected0 = collector0.finish_output(infer_state=infer_state).spec_hidden - collected1 = collector1.finish_output(infer_state=infer_state).spec_hidden - assert collected0.data_ptr() == hidden0.data_ptr() - assert collected1.data_ptr() == hidden1.data_ptr() - - -def test_layer_hidden_collector_keeps_microbatch_state_separate(monkeypatch): - _mock_target_layer_ids(monkeypatch, [0]) - model = SimpleNamespace(is_mtp_draft_model=False, layers_num=2, pre_infer=_IdentityPreInfer()) - prototype = LayerHiddenCollector(model=model) - collector0 = prototype.new_instance() - collector1 = prototype.new_instance() - hidden0 = torch.full((2, 3), 1.0) - hidden1 = torch.full((2, 3), 2.0) - infer_state = SimpleNamespace(need_dp_prefill_balance=False) - - collector0.add(layer_index=0, hidden=hidden0) - collector1.add(layer_index=0, hidden=hidden1) - - collected0 = collector0.finish_output(infer_state=infer_state).spec_hidden - collected1 = collector1.finish_output(infer_state=infer_state).spec_hidden - - assert torch.equal(collected0, hidden0) - assert torch.equal(collected1, hidden1) - - -def test_layer_hidden_collector_requires_target_layer_ids_from_config(monkeypatch): - _mock_target_layer_ids(monkeypatch, None) - model = SimpleNamespace(layers_num=2, pre_infer=_IdentityPreInfer()) - - with pytest.raises(AssertionError, match="target_layer_ids is required in draft config"): - LayerHiddenCollector(model=model) - - -def test_noop_collector_keeps_normal_forward_output_minimal(): - final_hidden = torch.randn(2, 3) - collector = NoopHiddenCollector() - - collector.add(layer_index=0, hidden=final_hidden) - collector.add_final_hidden(final_hidden) - - assert collector.finish_output(infer_state=None).spec_hidden is None - - -def test_mtp_head_output_collector_returns_and_clears_outputs(): - collector = MtpHeadOutputCollector() - draft_token_ids = torch.arange(6) - confidence_logits = torch.arange(6).view(2, 3) - - collector.add_mtp_outputs( - draft_token_ids=draft_token_ids, - confidence_logits=confidence_logits, - ) - output = collector.finish_output(infer_state=None) - - assert output.spec_hidden is None - assert output.draft_token_ids is draft_token_ids - assert output.confidence_logits is confidence_logits - assert collector.draft_token_ids is None - assert collector.confidence_logits is None - - -def test_final_collector_returns_final_hidden_without_layer_bookkeeping(): - final_hidden = torch.randn(2, 3) - collector = FinalHiddenCollector() - collector.add_final_hidden(final_hidden) - collected = collector.finish_output(infer_state=None).spec_hidden - - assert collected.data_ptr() == final_hidden.data_ptr() - - -def test_layer_collector_preserves_selected_layers_in_model_order(monkeypatch): - _mock_target_layer_ids(monkeypatch, [0, 2]) - layer0 = torch.full((2, 2), 1.0) - layer1 = torch.full((2, 2), 2.0) - layer2 = torch.full((2, 2), 3.0) - model = SimpleNamespace(layers_num=3, pre_infer=_IdentityPreInfer()) - collector = LayerHiddenCollector(model=model) - - collector.add(layer_index=0, hidden=layer0) - collector.add(layer_index=1, hidden=layer1) - collector.add(layer_index=2, hidden=layer2) - layer0.fill_(9.0) - - collected = collector.finish_output(infer_state=SimpleNamespace(need_dp_prefill_balance=False)).spec_hidden - - assert torch.equal(collected, torch.cat([torch.full((2, 2), 1.0), layer2], dim=-1)) - assert not collector.layer_hiddens - - collector.add(layer_index=0, hidden=layer0) - collector.add(layer_index=2, hidden=layer2) - collected = collector.finish_output(infer_state=SimpleNamespace(need_dp_prefill_balance=False)).spec_hidden - - assert torch.equal(collected, torch.cat([layer0, layer2], dim=-1)) - assert not collector.layer_hiddens - - -def test_layer_collector_restores_graph_state_without_sharing_runtime_container(monkeypatch): - _mock_target_layer_ids(monkeypatch, [0]) - model = SimpleNamespace(layers_num=2, pre_infer=_IdentityPreInfer()) - graph_collector = LayerHiddenCollector(model=model) - collector = graph_collector.new_instance() - infer_state = SimpleNamespace(need_dp_prefill_balance=False) - graph_collector.add(layer_index=0, hidden=torch.full((2, 3), 1.0)) - collector.restore_graph_state(graph_collector) - - collected = collector.finish_output(infer_state=infer_state).spec_hidden - - assert torch.equal(collected, torch.full((2, 3), 1.0)) - assert not collector.layer_hiddens - assert len(graph_collector.layer_hiddens) == 1 - - -def test_layer_collector_releases_graph_tensor_ownership_with_statistics(monkeypatch): - _mock_target_layer_ids(monkeypatch, [0, 2]) - model = SimpleNamespace(layers_num=3, pre_infer=_IdentityPreInfer()) - collector = LayerHiddenCollector(model=model) - hidden0 = torch.randn(2, 3) - hidden2 = torch.randn(2, 3) - converted = [] - - def to_no_ref(hidden): - converted.append(hidden) - return hidden - - monkeypatch.setattr(hidden_collector_module, "tensor_to_no_ref_tensor", to_no_ref) - collector.add(layer_index=0, hidden=hidden0) - collector.add(layer_index=2, hidden=hidden2) - - tensor_count, total_nbytes = collector.release_graph_tensor_ownership() - - assert tensor_count == 2 - assert total_nbytes == hidden0.numel() * hidden0.element_size() + hidden2.numel() * hidden2.element_size() - assert all(actual is expected for actual, expected in zip(converted, collector.layer_hiddens)) - assert NoopHiddenCollector().release_graph_tensor_ownership() == (0, 0) diff --git a/unit_tests/common/basemodel/test_model_input.py b/unit_tests/common/basemodel/test_model_input.py deleted file mode 100644 index 6672cccc31..0000000000 --- a/unit_tests/common/basemodel/test_model_input.py +++ /dev/null @@ -1,212 +0,0 @@ -from types import SimpleNamespace - -import pytest -import torch - -from lightllm.common.basemodel.basemodel import TpPartBaseModel -from lightllm.common.basemodel.batch_objs import ModelInput - - -def _create_model_input(*, is_prefill=False): - batch_size = 2 - kwargs = dict( - batch_size=batch_size, - total_token_num=batch_size, - max_q_seq_len=1, - max_kv_seq_len=1, - b_req_idx=torch.arange(batch_size, dtype=torch.int32), - b_mtp_index=torch.zeros(batch_size, dtype=torch.int32), - b_seq_len=torch.ones(batch_size, dtype=torch.int32), - is_prefill=is_prefill, - multimodal_params=[{"images": [], "audios": []} for _ in range(batch_size)], - ) - if is_prefill: - kwargs["max_cache_len"] = 0 - kwargs["input_ids"] = torch.ones(batch_size, dtype=torch.int64) - kwargs["b_ready_cache_len"] = torch.zeros(batch_size, dtype=torch.int32) - kwargs["b_prefill_start_loc"] = torch.arange(batch_size, dtype=torch.int32) - kwargs["b_is_decode_req"] = torch.zeros(batch_size, dtype=torch.bool) - kwargs["b_prefill_has_output_cpu"] = [False] * batch_size - else: - kwargs["b_position_delta"] = torch.zeros(batch_size, dtype=torch.int32) - kwargs["b_shared_seq_len"] = torch.tensor([4, 4], dtype=torch.int32) - kwargs["b_shared_radix_node_id"] = torch.tensor([10, 10], dtype=torch.int64) - return ModelInput(**kwargs) - - -def test_decode_requires_shared_radix_metadata(): - with pytest.raises(AssertionError): - ModelInput( - batch_size=1, - total_token_num=1, - max_q_seq_len=1, - max_kv_seq_len=1, - b_req_idx=torch.zeros(1, dtype=torch.int32), - b_mtp_index=torch.zeros(1, dtype=torch.int32), - b_seq_len=torch.ones(1, dtype=torch.int32), - b_position_delta=torch.zeros(1, dtype=torch.int32), - is_prefill=False, - multimodal_params=[{"images": [], "audios": []}], - ) - - -def test_decode_carries_raw_shared_radix_metadata(): - model_input = _create_model_input() - - assert torch.equal(model_input.b_shared_seq_len, torch.tensor([4, 4], dtype=torch.int32)) - assert torch.equal(model_input.b_shared_radix_node_id, torch.tensor([10, 10], dtype=torch.int64)) - - -def test_prefill_does_not_require_shared_radix_metadata(): - model_input = _create_model_input(is_prefill=True) - - assert model_input.b_shared_seq_len is None - assert model_input.b_shared_radix_node_id is None - - -def test_decode_requires_position_delta(): - model_input = _create_model_input() - model_input.b_position_delta = None - - with pytest.raises(AssertionError): - model_input.check_input() - - -def test_prefill_requires_prefill_metadata(): - model_input = _create_model_input(is_prefill=True) - model_input.b_ready_cache_len = None - - with pytest.raises(AssertionError): - model_input.check_input() - - -def test_prefill_requires_prefill_output_markers(): - model_input = _create_model_input(is_prefill=True) - model_input.b_prefill_has_output_cpu = None - - with pytest.raises(AssertionError, match="prefill must provide b_prefill_has_output_cpu"): - model_input.check_input() - - -def test_prefill_requires_decode_request_markers(): - model_input = _create_model_input(is_prefill=True) - model_input.b_is_decode_req = None - - with pytest.raises(AssertionError): - model_input.check_input() - - -def test_prefill_rejects_position_delta(): - model_input = _create_model_input(is_prefill=True) - model_input.b_position_delta = torch.zeros(model_input.batch_size, dtype=torch.int32) - - with pytest.raises(AssertionError, match="prefill must not provide b_position_delta"): - model_input.check_input() - - -def test_padded_prefill_adds_non_decode_request_marker(): - model_input = ModelInput( - batch_size=1, - total_token_num=2, - max_q_seq_len=2, - max_kv_seq_len=2, - max_cache_len=0, - input_ids=torch.ones(2, dtype=torch.int64), - b_req_idx=torch.zeros(1, dtype=torch.int32), - b_mtp_index=torch.zeros(1, dtype=torch.int32), - b_seq_len=torch.full((1,), 2, dtype=torch.int32), - b_is_decode_req=torch.ones(1, dtype=torch.bool), - 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=[False], - multimodal_params=[{"images": [], "audios": []}], - ) - model = SimpleNamespace( - mem_manager=SimpleNamespace(HOLD_TOKEN_MEMINDEXES=(-1,), page_size=1), - req_manager=SimpleNamespace(HOLD_REQUEST_ID=-1), - ) - - padded_input = TpPartBaseModel._create_padded_prefill_model_input( - model, - model_input=model_input, - new_handle_token_num=4, - ) - - assert padded_input.b_is_decode_req.dtype == torch.bool - assert padded_input.b_is_decode_req.tolist() == [True, False] - - -def test_padded_prefill_builds_internal_request_for_empty_input(): - model_input = ModelInput( - batch_size=0, - total_token_num=0, - max_q_seq_len=0, - max_kv_seq_len=0, - max_cache_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_is_decode_req=torch.empty((0,), dtype=torch.bool), - b_ready_cache_len=torch.empty((0,), dtype=torch.int32), - b_prefill_start_loc=torch.empty((0,), dtype=torch.int32), - is_prefill=True, - b_prefill_has_output_cpu=[], - multimodal_params=[], - mtp_draft_input_hiddens=torch.empty((0, 4), dtype=torch.float32), - ) - model = SimpleNamespace( - mem_manager=SimpleNamespace(HOLD_TOKEN_MEMINDEXES=(99,), page_size=1), - req_manager=SimpleNamespace(HOLD_REQUEST_ID=88), - ) - - padded_input = TpPartBaseModel._create_padded_prefill_model_input( - model, - model_input=model_input, - new_handle_token_num=1, - ) - - assert model_input.batch_size == 0 - assert padded_input.batch_size == 1 - assert padded_input.input_ids.tolist() == [1] - assert padded_input.b_req_idx.tolist() == [88] - assert padded_input.b_seq_len.tolist() == [1] - assert padded_input.b_prefill_has_output_cpu == [False] - assert torch.equal(padded_input.mtp_draft_input_hiddens, torch.zeros((1, 4))) - - -def test_padded_decode_builds_internal_request_from_empty_token_tensor(): - model_input = 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=[], - ) - model = SimpleNamespace( - mem_manager=SimpleNamespace(HOLD_TOKEN_MEMINDEXES=(99,), page_size=1), - req_manager=SimpleNamespace(HOLD_REQUEST_ID=88), - ) - - padded_input = TpPartBaseModel._create_padded_decode_model_input( - model, - model_input=model_input, - new_batch_size=1, - ) - - assert model_input.batch_size == 0 - assert padded_input.batch_size == 1 - assert padded_input.input_ids.tolist() == [1] - assert padded_input.b_req_idx.tolist() == [88] - assert padded_input.b_seq_len.tolist() == [2] - assert padded_input.b_shared_seq_len.tolist() == [0] - assert padded_input.b_shared_radix_node_id.tolist() == [-1] diff --git a/unit_tests/common/basemodel/test_model_output.py b/unit_tests/common/basemodel/test_model_output.py deleted file mode 100644 index 123f603ff2..0000000000 --- a/unit_tests/common/basemodel/test_model_output.py +++ /dev/null @@ -1,344 +0,0 @@ -from types import SimpleNamespace -import weakref - -import torch -import pytest - -from lightllm.common.basemodel import basemodel -from lightllm.common.basemodel.attention.flashinfer.fp import FlashInferDecodeAttState -from lightllm.common.basemodel.attention.flashinfer.utils import should_init_decode_wrapper -from lightllm.common.basemodel.basemodel import TpPartBaseModel -from lightllm.common.basemodel.batch_objs import ModelInput, ModelMtpOutputCollector, ModelOutput -from lightllm.common.basemodel.infer_struct import InferStateInfo - - -def test_decode_unpad_slices_spec_output_with_logits(): - model = TpPartBaseModel.__new__(TpPartBaseModel) - output = ModelOutput( - logits=torch.arange(24).view(6, 4), - mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.arange(18).view(6, 3)), - ) - - unpadded = model._create_unpad_decode_model_output(output, origin_batch_size=4) - - assert unpadded.logits.shape == (4, 4) - assert unpadded.mtp_collector.spec_hidden.shape == (4, 3) - # Unpadding returns a shallow output copy and leaves the graph-owned - # tensors on the original ModelOutput intact. - assert output.logits.shape == (6, 4) - assert output.mtp_collector.spec_hidden.shape == (6, 3) - - -def test_prefill_unpad_uses_token_rows_for_spec_hidden(): - model = TpPartBaseModel.__new__(TpPartBaseModel) - output = ModelOutput( - logits=torch.arange(20).view(5, 4), - mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.arange(24).view(8, 3)), - prompt_logics=torch.arange(32).view(8, 4), - ) - - unpadded = model._create_unpad_prefill_model_output( - output, - origin_handle_token_num=6, - origin_batch_size=3, - ) - - assert unpadded.logits.shape == (3, 4) - assert unpadded.mtp_collector.spec_hidden.shape == (6, 3) - assert unpadded.prompt_logics.shape == (6, 4) - - -def test_decode_unpad_restores_empty_output(): - model = TpPartBaseModel.__new__(TpPartBaseModel) - output = ModelOutput( - logits=torch.arange(4).view(1, 4), - mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.arange(3).view(1, 3)), - ) - - unpadded = model._create_unpad_decode_model_output(output, origin_batch_size=0) - - assert unpadded.logits.shape == (0, 4) - assert unpadded.mtp_collector.spec_hidden.shape == (0, 3) - - -def test_prefill_unpad_restores_empty_output(): - model = TpPartBaseModel.__new__(TpPartBaseModel) - output = ModelOutput( - logits=torch.arange(4).view(1, 4), - mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.arange(6).view(2, 3)), - prompt_logics=torch.arange(8).view(2, 4), - ) - - unpadded = model._create_unpad_prefill_model_output( - output, - origin_handle_token_num=0, - origin_batch_size=0, - ) - - assert unpadded.logits.shape == (0, 4) - assert unpadded.mtp_collector.spec_hidden.shape == (0, 3) - assert unpadded.prompt_logics.shape == (0, 4) - - -@torch.no_grad() -def test_decode_pads_only_once_after_selecting_execution_path(monkeypatch): - execution_configs = ( - # eager 普通模式：空 batch 只补一个 dummy request。 - (None, False, 1, False, 1), - # eager TPSP 模式：dummy request 仍需对齐到 TP world size。 - (None, True, 2, False, 2), - # CUDA Graph replay：使用最终选中的 graph batch size。 - (4, False, 1, False, 4), - # CUDA Graph capture：必须在初始化 attention state 前设置 graph 标记。 - (4, False, 1, True, 4), - ) - for ( - graph_batch_size, - enable_tpsp_mix_mode, - tp_world_size, - need_capture, - expected_batch_size, - ) in execution_configs: - model = TpPartBaseModel.__new__(TpPartBaseModel) - model.args = SimpleNamespace(enable_tpsp_mix_mode=enable_tpsp_mix_mode, page_size=1) - model.tp_world_size_ = tp_world_size - model.is_mtp_draft_model = False - model.mtp_manager = SimpleNamespace(get_decode_batch_alignment=lambda _: 1) - model.req_manager = SimpleNamespace(HOLD_REQUEST_ID=88, req_to_token_indexs=object()) - - pad_batch_sizes = [] - - def pad_once(model_input, new_batch_size): - pad_batch_sizes.append(new_batch_size) - return TpPartBaseModel._create_padded_decode_model_input(model, model_input, new_batch_size) - - model._create_padded_decode_model_input = pad_once - - graph_flags_at_att_init = [] - executions = [] - - def create_infer_state(model_input): - infer_state = SimpleNamespace( - b_req_idx=model_input.b_req_idx, - b_seq_len=model_input.b_seq_len, - mem_index=torch.empty((model_input.batch_size,), dtype=torch.int32), - init_some_extra_state=lambda _: None, - is_cuda_graph_capture=False, - is_cuda_graph_replay=False, - ) - infer_state.init_att_state = lambda: graph_flags_at_att_init.append( - (infer_state.is_cuda_graph_capture, infer_state.is_cuda_graph_replay) - ) - return infer_state - - def execute(mode, infer_state): - executions.append((mode, infer_state.is_cuda_graph_capture, infer_state.is_cuda_graph_replay)) - return ModelOutput(logits=torch.ones((infer_state.b_req_idx.shape[0], 4))) - - model._create_inferstate = create_infer_state - model._token_forward = lambda infer_state: execute("eager", infer_state) - - graph = None - if graph_batch_size is not None: - graph = SimpleNamespace() - graph.can_run = lambda **kwargs: True - graph.find_closest_graph_batch_size = lambda batch_size: graph_batch_size - graph.need_capture = lambda batch_size: need_capture - graph.capture_decode = lambda decode_func, infer_state: execute("capture", infer_state) - graph.replay = lambda infer_state: execute("replay", infer_state) - model.graph = graph - - output = model._decode(_create_empty_decode_input()) - - expected_replay = graph_batch_size is not None and not need_capture - expected_mode = "eager" - if graph_batch_size is not None: - expected_mode = "capture" if need_capture else "replay" - assert pad_batch_sizes == [expected_batch_size] - assert graph_flags_at_att_init == [(need_capture, expected_replay)] - assert executions == [(expected_mode, need_capture, expected_replay)] - assert output.logits.shape == (0, 4) - - -@pytest.mark.parametrize( - "lengths,graph_mode,expected_mode", - [ - pytest.param((1024, 2048), None, "eager", id="graph-disabled"), - pytest.param((1024, 2048), "capture", "capture", id="capture"), - pytest.param((1024, 2048), "replay", "replay", id="replay"), - pytest.param((1024, 16384), "replay", "eager", id="long-second-microbatch"), - pytest.param((16384, 1024), "capture", "eager", id="long-first-microbatch"), - ], -) -def test_overlap_decode_initializes_attention_for_the_selected_execution_mode(lengths, graph_mode, expected_mode): - model = TpPartBaseModel.__new__(TpPartBaseModel) - model.args = SimpleNamespace(enable_tpsp_mix_mode=True, page_size=1) - model.tp_world_size_ = 2 - model.is_mtp_draft_model = False - model.mtp_manager = SimpleNamespace(get_decode_batch_alignment=lambda _: 1) - model.req_manager = SimpleNamespace(HOLD_REQUEST_ID=88) - initializations = [] - executions = [] - - def graph_flags(state): - return state.is_cuda_graph_capture, state.is_cuda_graph_replay - - class DecodeState(FlashInferDecodeAttState): - def init_state(self): - infer_state = self.infer_state - # Replace the external GPU wrapper allocation while preserving the - # production decision and InferStateInfo's initialization entry point. - if should_init_decode_wrapper(infer_state): - self.decode_wrapper = object() - initializations.append( - (infer_state.microbatch_index, graph_flags(infer_state), self.decode_wrapper is not None) - ) - - def create_state(model_input, microbatch_index=0): - state = InferStateInfo() - state.is_prefill = False - state.input_ids = model_input.input_ids - state.batch_size = model_input.batch_size - state.max_kv_seq_len = model_input.max_kv_seq_len - state.b_req_idx = model_input.b_req_idx - state.b_seq_len = model_input.b_seq_len - state.microbatch_index = microbatch_index - state.decode_att_state = DecodeState(_infer_state_weak_ref=weakref.ref(state)) - state.init_some_extra_state = lambda _: None - return state - - def outputs(state, state1): - return tuple(ModelOutput(logits=torch.stack((s.b_req_idx, s.b_seq_len), dim=1)) for s in (state, state1)) - - def forward(state, infer_state1): - # Eager and capture execute attention using these newly prepared states. - assert state.decode_att_state.decode_wrapper is not None - assert infer_state1.decode_att_state.decode_wrapper is not None - executions.append(("forward", graph_flags(state), graph_flags(infer_state1))) - return outputs(state, infer_state1) - - def capture(func, state, infer_state1): - executions.append(("capture", graph_flags(state), graph_flags(infer_state1))) - return func(state, infer_state1) - - def replay(state, infer_state1): - # Replay uses graph-owned wrappers, so temporary states need none. - assert state.decode_att_state.decode_wrapper is None - assert infer_state1.decode_att_state.decode_wrapper is None - executions.append(("replay", graph_flags(state), graph_flags(infer_state1))) - return outputs(state, infer_state1) - - model._create_inferstate = create_state - model._overlap_tpsp_token_forward = forward - model.graph = None - if graph_mode is not None: - model.graph = SimpleNamespace( - can_run=lambda batch_size, max_len_in_batch: batch_size <= 4 and max_len_in_batch <= 8192, - find_closest_graph_batch_size=lambda batch_size: 4, - need_capture=lambda batch_size: graph_mode == "capture", - capture_decode=capture, - replay=replay, - ) - - inputs = [] - for rows, length in zip((1, 3), lengths): - model_input = model._create_padded_decode_model_input(_create_empty_decode_input(), rows) - model_input.b_req_idx = torch.arange(rows, dtype=torch.int32) - model_input.b_seq_len.fill_(length) - model_input.max_kv_seq_len = length - model_input.total_token_num = rows * length - inputs.append(model_input) - - result = model._microbatch_overlap_decode_cuda(*inputs) - - expected_flags = {"eager": (False, False), "capture": (True, False), "replay": (False, True)}[expected_mode] - assert initializations == [(index, expected_flags, expected_mode != "replay") for index in (0, 1)] - expected_execution = {"eager": ["forward"], "capture": ["capture", "forward"], "replay": ["replay"]}[expected_mode] - assert executions == [(mode, expected_flags, expected_flags) for mode in expected_execution] - for output, model_input in zip(result, inputs): - torch.testing.assert_close(output.logits, torch.stack((model_input.b_req_idx, model_input.b_seq_len), dim=1)) - - -@pytest.mark.parametrize("overlap", [False, True]) -@pytest.mark.parametrize("dynamic", [False, True]) -@pytest.mark.parametrize("graph_mode", ["eager", "capture", "replay"]) -@pytest.mark.parametrize("rows", [0, 3, 9, 24, 27]) -def test_mtp_decode_preserves_groups_and_real_rows(overlap, dynamic, graph_mode, rows): - model = TpPartBaseModel.__new__(TpPartBaseModel) - model.args = SimpleNamespace(enable_tpsp_mix_mode=True) - model.tp_world_size_ = 8 - model.is_mtp_draft_model = False - model.mtp_manager = SimpleNamespace(get_decode_batch_alignment=lambda _: 1 if dynamic else 3) - model.req_manager = SimpleNamespace(HOLD_REQUEST_ID=88) - seen_sizes = [] - - def create_state(model_input, microbatch_index=0): - def check_layout(): - if not dynamic: - assert model_input.batch_size % 3 == 0 - assert model_input.batch_size % 8 == 0 - seen_sizes.append(model_input.batch_size) - - return SimpleNamespace( - b_req_idx=model_input.b_req_idx, - b_seq_len=model_input.b_seq_len, - init_some_extra_state=lambda _: None, - init_att_state=check_layout, - ) - - def forward(state): - return ModelOutput(logits=torch.stack((state.b_req_idx, state.b_seq_len), dim=1)) - - model._create_inferstate = create_state - model._token_forward = forward - model._overlap_tpsp_token_forward = lambda state, infer_state1: (forward(state), forward(infer_state1)) - model.graph = None - if graph_mode != "eager": - model.graph = SimpleNamespace( - can_run=lambda batch_size, max_len_in_batch: batch_size <= 24, - find_closest_graph_batch_size=lambda batch_size: ((batch_size + 7) // 8 * 8 if dynamic else 24), - need_capture=lambda batch_size: graph_mode == "capture", - capture_decode=lambda func, state, **kwargs: func(state, **kwargs), - replay=lambda state, infer_state1=None: ( - forward(state) if infer_state1 is None else (forward(state), forward(infer_state1)) - ), - ) - - model_input = model._create_padded_decode_model_input(_create_empty_decode_input(), rows) - model_input.b_req_idx = torch.arange(rows, dtype=torch.int32) // 3 - model_input.b_seq_len = 10 + torch.arange(rows, dtype=torch.int32) % 3 - if dynamic: - # 压缩后每请求分别保留 1、2、3 个 token，保持各请求内部的 token 顺序。 - request_rows = [(req, step) for req in range(rows) for step in range(req % 3 + 1)][:rows] - model_input.b_req_idx = torch.tensor([req for req, _ in request_rows], dtype=torch.int32) - model_input.b_mtp_index = torch.tensor([step for _, step in request_rows], dtype=torch.int32) - model_input.b_seq_len = 10 + model_input.b_mtp_index - expected = torch.stack((model_input.b_req_idx, model_input.b_seq_len), dim=1) - if overlap: - output, empty_output = model._microbatch_overlap_decode_cuda(model_input, _create_empty_decode_input()) - assert empty_output.logits.shape == (0, 2) - else: - output = model._decode(model_input) - - expected_size = (max(1, rows) + 7) // 8 * 8 if dynamic else (48 if rows > 24 else 24) - 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_moe_route_info_manager.py b/unit_tests/common/basemodel/test_moe_route_info_manager.py deleted file mode 100644 index de2bc7a857..0000000000 --- a/unit_tests/common/basemodel/test_moe_route_info_manager.py +++ /dev/null @@ -1,263 +0,0 @@ -from types import SimpleNamespace - -import numpy as np -import pytest -import torch - -from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.common.basemodel.moe_route_info_manager import MoeRouteInfoManager, get_moe_capture_callback - - -def _skip_without_cuda(): - if not torch.cuda.is_available(): - pytest.skip("CUDA is required for moe route info capture.") - - -def test_get_moe_capture_callback_uses_global_manager(monkeypatch): - calls = [] - - class _Manager: - def is_buffer_initialized(self): - return True - - def get_moe_capture_callback(self, layer_index, mem_indexes): - calls.append((layer_index, mem_indexes)) - return "callback" - - mem_indexes = object() - infer_state = SimpleNamespace(mem_index=mem_indexes) - monkeypatch.setattr(MoeRouteInfoManager, "_instance", _Manager()) - - assert get_moe_capture_callback(infer_state, 7) == "callback" - assert calls == [(7, mem_indexes)] - - -class TestMoeRouteInfoManager: - def test_capture_and_extract_basic(self): - _skip_without_cuda() - - manager = MoeRouteInfoManager( - num_moe_layers=4, - topk=8, - dtype_id=1, - ) - manager.init_capture_buffer(kv_cache_size=1024) - mem_indexes = torch.arange(100, 110, device="cuda") - expected = np.zeros((10, 4, 8), dtype=np.uint8) - - for layer_idx in range(4): - topk_ids = torch.randint(0, 64, (10, 8), device="cuda") - manager.get_moe_capture_callback(layer_idx, mem_indexes)(topk_ids) - expected[:, layer_idx, :] = topk_ids.cpu().numpy().astype(np.uint8) - - result = manager.extract(mem_indexes) - assert result.shape == (10, 4, 8) - assert result.dtype == np.uint8 - np.testing.assert_array_equal(result, expected) - - def test_capture_writes_to_correct_kv_positions(self): - _skip_without_cuda() - - manager = MoeRouteInfoManager( - num_moe_layers=2, - topk=4, - dtype_id=1, - ) - manager.init_capture_buffer(kv_cache_size=256) - mem_indexes = torch.tensor([10, 50, 200], device="cuda") - topk_ids = torch.tensor([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]], device="cuda") - topk_ids_layer1 = topk_ids + 20 - - manager.get_moe_capture_callback(0, mem_indexes)(topk_ids) - manager.get_moe_capture_callback(1, mem_indexes)(topk_ids_layer1) - - result = manager.extract(mem_indexes) - np.testing.assert_array_equal(result[:, 0, :], topk_ids.cpu().numpy().astype(np.uint8)) - np.testing.assert_array_equal(result[:, 1, :], topk_ids_layer1.cpu().numpy().astype(np.uint8)) - - def test_capture_maps_transformer_layer_num_to_routing_slot(self): - _skip_without_cuda() - - manager = MoeRouteInfoManager( - num_moe_layers=2, - topk=2, - dtype_id=1, - layer_index_to_moe_index={3: 0, 7: 1}, - ) - manager.init_capture_buffer(kv_cache_size=256) - mem_indexes = torch.tensor([10, 11], device="cuda") - ids_layer3 = torch.tensor([[1, 2], [3, 4]], device="cuda") - ids_layer7 = torch.tensor([[5, 6], [7, 8]], device="cuda") - - manager.get_moe_capture_callback(3, mem_indexes)(ids_layer3) - manager.get_moe_capture_callback(7, mem_indexes)(ids_layer7) - assert manager.get_moe_capture_callback(4, mem_indexes) is None - - result = manager.extract(mem_indexes) - np.testing.assert_array_equal(result[:, 0, :], ids_layer3.cpu().numpy().astype(np.uint8)) - np.testing.assert_array_equal(result[:, 1, :], ids_layer7.cpu().numpy().astype(np.uint8)) - - def test_capture_rejects_unexpected_topk_width(self): - _skip_without_cuda() - - manager = MoeRouteInfoManager( - num_moe_layers=1, - topk=2, - dtype_id=1, - ) - manager.init_capture_buffer(kv_cache_size=256) - mem_indexes = torch.tensor([10, 11], device="cuda") - topk_ids_with_shared_expert = torch.tensor([[1, 2, 32], [3, 4, 32]], device="cuda") - - with pytest.raises(AssertionError): - manager.get_moe_capture_callback(0, mem_indexes)(topk_ids_with_shared_expert) - - def test_cuda_graph_replay_uses_copied_mem_indexes(self): - _skip_without_cuda() - - class _NoopDecodeState: - def copy_for_decode_cuda_graph(self, other): - return - - manager = MoeRouteInfoManager( - num_moe_layers=1, - topk=2, - dtype_id=1, - ) - manager.init_capture_buffer(kv_cache_size=256) - graph_infer_state = InferStateInfo() - graph_infer_state.decode_att_state = _NoopDecodeState() - graph_infer_state.mem_index = torch.tensor([10, 11], device="cuda") - moe_capture_callback = manager.get_moe_capture_callback(0, graph_infer_state.mem_index) - topk_ids = torch.tensor([[1, 2], [3, 4]], device="cuda") - - graph = torch.cuda.CUDAGraph() - torch.cuda.synchronize() - with torch.cuda.graph(graph): - moe_capture_callback(topk_ids) - - graph.replay() - result = manager.extract(graph_infer_state.mem_index) - np.testing.assert_array_equal(result[:, 0, :], topk_ids.cpu().numpy().astype(np.uint8)) - - new_infer_state = InferStateInfo() - new_infer_state.decode_att_state = _NoopDecodeState() - new_infer_state.mem_index = torch.tensor([20, 21], device="cuda") - new_topk_ids = torch.tensor([[5, 6], [7, 8]], device="cuda") - graph_infer_state.copy_for_cuda_graph(new_infer_state) - topk_ids.copy_(new_topk_ids) - - graph.replay() - - result = manager.extract(new_infer_state.mem_index) - np.testing.assert_array_equal(result[:, 0, :], new_topk_ids.cpu().numpy().astype(np.uint8)) - - def test_microbatch_isolation(self): - _skip_without_cuda() - - manager = MoeRouteInfoManager( - num_moe_layers=1, - topk=4, - dtype_id=1, - ) - manager.init_capture_buffer(kv_cache_size=256) - mem0 = torch.tensor([10, 11], device="cuda") - mem1 = torch.tensor([20, 21], device="cuda") - ids_0 = torch.ones((2, 4), dtype=torch.int64, device="cuda") - ids_1 = torch.ones((2, 4), dtype=torch.int64, device="cuda") * 2 - - capture0 = manager.get_moe_capture_callback(0, mem0) - capture1 = manager.get_moe_capture_callback(0, mem1) - capture0(ids_0) - capture1(ids_1) - - result0 = manager.extract(mem0) - result1 = manager.extract(mem1) - assert result0[0, 0, 0] == 1 - assert result1[0, 0, 0] == 2 - - def test_dtype_selection_uint8(self): - manager = MoeRouteInfoManager( - num_moe_layers=1, - topk=2, - dtype_id=1, - ) - assert manager.get_torch_dtype() == torch.uint8 - assert manager.get_np_dtype() == np.uint8 - assert manager.dtype_id == 1 - assert not manager.is_buffer_initialized() - - def test_dtype_selection_int16(self): - manager = MoeRouteInfoManager( - num_moe_layers=1, - topk=2, - dtype_id=2, - ) - assert manager.get_torch_dtype() == torch.int16 - assert manager.get_np_dtype() == np.int16 - assert manager.dtype_id == 2 - assert not manager.is_buffer_initialized() - - def test_extract_preserves_uint8_values(self): - _skip_without_cuda() - - manager = MoeRouteInfoManager( - num_moe_layers=1, - topk=4, - dtype_id=1, - ) - manager.init_capture_buffer(kv_cache_size=64) - mem_indexes = torch.tensor([0, 1, 2], device="cuda") - topk_ids = torch.tensor([[10, 20, 30, 40], [50, 60, 63, 1], [0, 5, 255, 3]], device="cuda") - - manager.get_moe_capture_callback(0, mem_indexes)(topk_ids) - - result = manager.extract(mem_indexes) - np.testing.assert_array_equal(result[:, 0, :], topk_ids.cpu().numpy().astype(np.uint8)) - - def test_routing_buffer_and_pointer_shape(self): - _skip_without_cuda() - - manager = MoeRouteInfoManager( - num_moe_layers=48, - topk=8, - dtype_id=1, - ) - assert not manager.is_buffer_initialized() - manager.init_capture_buffer(kv_cache_size=2048) - assert manager.routing_buffer.shape == (2048, 48, 8) - assert manager.routing_buffer.dtype == torch.uint8 - assert manager.routing_buffer.device.type == "cpu" - assert manager.routing_buffer.is_pinned() - assert manager.routing_buffer_ptr.shape == (1,) - assert manager.routing_buffer_ptr.dtype == torch.uint64 - assert manager.routing_buffer_ptr.device.type == "cuda" - - def test_partial_token_capture(self): - _skip_without_cuda() - - manager = MoeRouteInfoManager( - num_moe_layers=1, - topk=2, - dtype_id=1, - ) - manager.init_capture_buffer(kv_cache_size=128) - mem_indexes = torch.tensor([10, 11, 12, 13, 14], device="cuda") - topk_ids = torch.tensor([[1, 2], [3, 4], [5, 6]], device="cuda") - - manager.get_moe_capture_callback(0, mem_indexes[:3])(topk_ids) - - result_written = manager.extract(mem_indexes[:3]) - np.testing.assert_array_equal(result_written[:, 0, :], topk_ids.cpu().numpy().astype(np.uint8)) - - result_unwritten = manager.extract(mem_indexes[3:]) - np.testing.assert_array_equal(result_unwritten[:, 0, :], np.zeros((2, 2), dtype=np.uint8)) - - def test_phase1_does_not_allocate_capture_buffer(self): - manager = MoeRouteInfoManager( - num_moe_layers=4, - topk=8, - dtype_id=1, - ) - assert not manager.is_buffer_initialized() - assert manager.routing_buffer_ptr is None diff --git a/unit_tests/common/basemodel/test_mtp_manager.py b/unit_tests/common/basemodel/test_mtp_manager.py deleted file mode 100644 index f7e91f8c79..0000000000 --- a/unit_tests/common/basemodel/test_mtp_manager.py +++ /dev/null @@ -1,133 +0,0 @@ -from types import SimpleNamespace - -import pytest - -import lightllm.common.basemodel.hidden_collector as hidden_collector_module -import lightllm.common.basemodel.mtp_manager as mtp_manager_module -from lightllm.common.basemodel.hidden_collector import ( - FinalHiddenCollector, - LayerHiddenCollector, - MtpHeadOutputCollector, - NoopHiddenCollector, -) -from lightllm.common.basemodel.mtp_manager import MtpManager - - -@pytest.fixture(autouse=True) -def _reset_mtp_manager(): - MtpManager._instance = None - yield - MtpManager._instance = None - - -def _decode_tokens_per_request(monkeypatch, spec_mode, *, is_draft_model, mtp_step=7): - args = SimpleNamespace( - mtp_mode=spec_mode, - mtp_step=mtp_step, - mtp_dynamic_verify=False, - ) - monkeypatch.setattr(mtp_manager_module, "get_env_start_args", lambda: args) - return MtpManager.get_instance().get_decode_tokens_per_request(is_draft_model) - - -@pytest.mark.parametrize( - "spec_mode,is_draft_model,expected", - [ - (None, False, 1), - ("eagle3", False, 8), - ("eagle3", True, 1), - ("vanilla_with_att", True, 1), - ("vanilla_no_att", True, 1), - ("eagle_with_att", True, 1), - ("eagle_no_att", True, 1), - ("dspark", True, 7), - ("dflash", True, 7), - ("dflash2", True, 8), - ], -) -def test_decode_tokens_per_request(monkeypatch, spec_mode, is_draft_model, expected): - assert _decode_tokens_per_request(monkeypatch, spec_mode, is_draft_model=is_draft_model) == expected - - -@pytest.mark.parametrize( - "spec_mode,dynamic_verify,is_draft_model,expected", - [ - ("vanilla_with_att", False, False, 8), - ("vanilla_with_att", True, False, 1), - ("vanilla_with_att", False, True, 1), - ("vanilla_with_att", True, True, 1), - ("dspark", True, True, 7), - ("dflash", True, True, 7), - ], -) -def test_decode_batch_alignment(monkeypatch, spec_mode, dynamic_verify, is_draft_model, expected): - args = SimpleNamespace( - mtp_mode=spec_mode, - mtp_step=7, - mtp_dynamic_verify=dynamic_verify, - ) - monkeypatch.setattr(mtp_manager_module, "get_env_start_args", lambda: args) - - assert MtpManager.get_instance().get_decode_batch_alignment(is_draft_model) == expected - - -@pytest.mark.parametrize( - "spec_mode,is_draft_model,expected", - [ - (None, False, 0), - ("eagle3", False, 7), - ("eagle3", True, 0), - ("vanilla_with_att", True, 0), - ("dspark", True, 6), - ("dflash", True, 6), - ("dflash2", True, 7), - ], -) -def test_decode_draft_step(monkeypatch, spec_mode, is_draft_model, expected): - args = SimpleNamespace(mtp_mode=spec_mode, mtp_step=7, mtp_dynamic_verify=False) - monkeypatch.setattr(mtp_manager_module, "get_env_start_args", lambda: args) - - assert MtpManager.get_instance().get_decode_draft_step(is_draft_model) == expected - - -def test_get_instance_returns_singleton(monkeypatch): - args = SimpleNamespace(mtp_mode="eagle3", mtp_step=7, mtp_dynamic_verify=False) - monkeypatch.setattr(mtp_manager_module, "get_env_start_args", lambda: args) - - assert MtpManager.get_instance() is MtpManager.get_instance() - - -@pytest.mark.parametrize( - "spec_mode,is_draft_model,expected_type", - [ - (None, False, NoopHiddenCollector), - ("vanilla_with_att", False, FinalHiddenCollector), - ("eagle3", False, LayerHiddenCollector), - ("dspark", False, LayerHiddenCollector), - ("eagle3", True, FinalHiddenCollector), - ("dspark", True, MtpHeadOutputCollector), - ("dflash2", True, MtpHeadOutputCollector), - ], -) -def test_create_hidden_collector_selects_implementation(monkeypatch, spec_mode, is_draft_model, expected_type): - args = SimpleNamespace( - mtp_mode=spec_mode, - mtp_step=7, - mtp_dynamic_verify=False, - mtp_draft_model_dir=["/models/draft"], - ) - monkeypatch.setattr(mtp_manager_module, "get_env_start_args", lambda: args) - monkeypatch.setattr(hidden_collector_module, "get_env_start_args", lambda: args) - monkeypatch.setattr( - hidden_collector_module.PretrainedConfig, - "get_config_dict", - lambda _: ({"target_layer_ids": [0]}, {}), - ) - model = SimpleNamespace(is_mtp_draft_model=is_draft_model, layers_num=2) - - prototype = MtpManager.get_instance().create_hidden_collector(model=model) - collector = prototype.new_instance() - - assert isinstance(prototype, expected_type) - assert isinstance(collector, expected_type) - assert collector is not prototype diff --git a/unit_tests/common/basemodel/test_overlap_utils.py b/unit_tests/common/basemodel/test_overlap_utils.py deleted file mode 100644 index ef7c18a061..0000000000 --- a/unit_tests/common/basemodel/test_overlap_utils.py +++ /dev/null @@ -1,213 +0,0 @@ -from types import SimpleNamespace - -import pytest -import torch - -from lightllm.common.basemodel import basemodel -from lightllm.common.basemodel.basemodel import TpPartBaseModel -from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput - - -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) - events = [] - original_to_cuda = ModelInput.to_cuda - - def record_to_cuda(self): - events.append("to_cuda") - original_to_cuda(self) - - monkeypatch.setattr(ModelInput, "to_cuda", record_to_cuda) - - def fake_gather(**kwargs): - events.append("gather") - kwargs["input_ids"][-1] = 90 + events.count("gather") - - monkeypatch.setattr(basemodel, "gather_token_prefill_decode_mixed", fake_gather) - - model = TpPartBaseModel.__new__(TpPartBaseModel) - model.args = SimpleNamespace(enable_prefill_decode_mixed=True, page_size=1) - model.req_manager = SimpleNamespace( - req_sampling_params_manager=SimpleNamespace(req_to_next_token_ids=object()), - ) - captured_inputs = [] - - def fake_overlap_forward(input0, input1): - events.append("forward") - captured_inputs.extend((input0, input1)) - return ( - ModelOutput(logits=torch.zeros((input0.batch_size, 1))), - ModelOutput(logits=torch.zeros((input1.batch_size, 1))), - ) - - model._microbatch_overlap_prefill_cuda = fake_overlap_forward - - outputs = model.microbatch_overlap_prefill(model_input0, model_input1) - - assert events == ["to_cuda", "gather", "to_cuda", "gather", "forward"] - assert captured_inputs == [model_input0, model_input1] - assert captured_inputs[0].input_ids.tolist() == [10, 91] - assert captured_inputs[1].input_ids.tolist() == [92] - assert len(outputs) == 2 - - -def test_base_model_decode_accepts_two_inputs_and_skips_empty_gather(monkeypatch): - model_input0 = _make_decode_input() - model_input0.input_ids = None - model_input1 = _make_empty_decode_input() - model_input1.input_ids = None - events = [] - original_to_cuda = ModelInput.to_cuda - - def record_to_cuda(self): - events.append("to_cuda") - original_to_cuda(self) - - monkeypatch.setattr(ModelInput, "to_cuda", record_to_cuda) - - def fake_gather(**kwargs): - events.append("gather") - return torch.arange( - start=40, - end=46, - dtype=torch.int64, - device="cuda", - ) - - monkeypatch.setattr(basemodel, "gather_token", fake_gather) - - model = TpPartBaseModel.__new__(TpPartBaseModel) - model.args = SimpleNamespace(page_size=1) - model.req_manager = SimpleNamespace( - req_sampling_params_manager=SimpleNamespace(req_to_next_token_ids=object()), - ) - captured_inputs = [] - - def fake_overlap_forward(input0, input1): - events.append("forward") - captured_inputs.extend((input0, input1)) - return ( - ModelOutput(logits=torch.zeros((input0.batch_size, 1))), - ModelOutput(logits=torch.zeros((input1.batch_size, 1))), - ) - - model._microbatch_overlap_decode_cuda = fake_overlap_forward - - model.microbatch_overlap_decode(model_input0, model_input1) - - assert events == ["to_cuda", "gather", "to_cuda", "forward"] - assert captured_inputs == [model_input0, model_input1] - assert captured_inputs[0].input_ids.tolist() == [40, 41, 42, 43, 44, 45] - assert captured_inputs[1].input_ids.numel() == 0 - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") -def test_overlap_decode_cuda_pads_empty_side_and_unpads_outputs(monkeypatch): - model_input0 = _make_decode_input() - model_input0.batch_size = 1 - model_input0.total_token_num = 4 - model_input0.max_kv_seq_len = 4 - model_input0.input_ids = model_input0.input_ids[:1] - model_input0.b_req_idx = model_input0.b_req_idx[:1] - model_input0.b_mtp_index = model_input0.b_mtp_index[:1] - model_input0.b_seq_len = model_input0.b_seq_len[:1] - model_input0.b_position_delta = model_input0.b_position_delta[:1] - model_input0.b_shared_seq_len = model_input0.b_shared_seq_len[:1] - model_input0.b_shared_radix_node_id = model_input0.b_shared_radix_node_id[:1] - model_input0.multimodal_params = model_input0.multimodal_params[:1] - model_input0.check_input() - model_input1 = _make_empty_decode_input() - model_input0.to_cuda() - model_input1.to_cuda() - - model = TpPartBaseModel.__new__(TpPartBaseModel) - model.args = SimpleNamespace(enable_tpsp_mix_mode=True, page_size=1) - model.tp_world_size_ = 2 - 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()) - infer_batch_sizes = [] - - def fake_create_inferstate(model_input, microbatch_index): - infer_batch_sizes.append(model_input.batch_size) - return SimpleNamespace( - b_req_idx=model_input.b_req_idx, - b_seq_len=model_input.b_seq_len, - mem_index=torch.empty((model_input.batch_size,), dtype=torch.int32, device="cuda"), - init_some_extra_state=lambda _: None, - init_att_state=lambda: None, - ) - - model._create_inferstate = fake_create_inferstate - model._overlap_tpsp_token_forward = lambda infer_state0, infer_state1: ( - ModelOutput(logits=torch.zeros((infer_state0.b_req_idx.shape[0], 1), device="cuda")), - ModelOutput(logits=torch.zeros((infer_state1.b_req_idx.shape[0], 1), device="cuda")), - ) - output0, output1 = model._microbatch_overlap_decode_cuda(model_input0, model_input1) - - 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/basemodel/test_prefill_cuda_graph_state.py b/unit_tests/common/basemodel/test_prefill_cuda_graph_state.py deleted file mode 100644 index bed4f3967f..0000000000 --- a/unit_tests/common/basemodel/test_prefill_cuda_graph_state.py +++ /dev/null @@ -1,81 +0,0 @@ -from types import SimpleNamespace - -import torch - -import lightllm.common.basemodel.hidden_collector as hidden_collector_module -from lightllm.common.basemodel.hidden_collector import LayerHiddenCollector -from lightllm.common.basemodel.prefill_cuda_graph import PrefillCudaGraph - - -class _GraphInferState: - def __init__(self, hidden_collector): - self.hidden_collector = hidden_collector - self.input_ids = torch.empty(4, dtype=torch.int64) - self.copied_from = None - self.replayed_with = None - - def copy_for_prefill_cuda_graph(self, new_infer_state): - self.copied_from = new_infer_state - - def prefill_replay(self, new_infer_state): - self.replayed_with = new_infer_state - - -class _IdentityPreInfer: - @staticmethod - def _tpsp_allgather(input, infer_state): - del infer_state - return input - - -def _create_layer_collector(monkeypatch): - monkeypatch.setattr( - hidden_collector_module, - "get_env_start_args", - lambda: SimpleNamespace(mtp_draft_model_dir=["/models/draft"]), - ) - monkeypatch.setattr( - hidden_collector_module.PretrainedConfig, - "get_config_dict", - lambda _: ({"target_layer_ids": [0]}, {}), - ) - model = SimpleNamespace(layers_num=2, pre_infer=_IdentityPreInfer()) - return LayerHiddenCollector(model=model) - - -def test_replay_restores_captured_hidden_state_into_request_collector(monkeypatch): - graph_collector = _create_layer_collector(monkeypatch) - graph_collector.add(layer_index=0, hidden=torch.randn(2, 3)) - request_collector = graph_collector.new_instance() - graph_infer_state = _GraphInferState(hidden_collector=graph_collector) - request_infer_state = SimpleNamespace( - input_ids=torch.empty(4, dtype=torch.int64), - hidden_collector=request_collector, - ) - graph_output = torch.randn(2, 3) - prefill_graph = PrefillCudaGraph.__new__(PrefillCudaGraph) - prefill_graph.graph = {4: (graph_infer_state, [], [graph_output], graph_collector)} - - outputs = prefill_graph._replay(input_tensors=[], infer_state=request_infer_state) - - assert len(outputs) == 1 - assert outputs[0] is graph_output - assert request_infer_state.hidden_collector is request_collector - assert request_collector.layer_hiddens is not graph_collector.layer_hiddens - assert request_collector.layer_hiddens[0] is graph_collector.layer_hiddens[0] - assert graph_infer_state.copied_from is request_infer_state - assert graph_infer_state.replayed_with is request_infer_state - - -def test_first_replay_replaces_capture_collector_with_runtime_instance(monkeypatch): - graph_collector = _create_layer_collector(monkeypatch) - graph_collector.add(layer_index=0, hidden=torch.randn(2, 3)) - graph_infer_state = _GraphInferState(hidden_collector=graph_collector) - graph_output = torch.randn(2, 3) - prefill_graph = PrefillCudaGraph.__new__(PrefillCudaGraph) - prefill_graph.graph = {4: (graph_infer_state, [], [graph_output], graph_collector)} - - prefill_graph._replay(input_tensors=[], infer_state=graph_infer_state) - - assert graph_infer_state.hidden_collector is not graph_collector - assert graph_infer_state.hidden_collector.layer_hiddens[0] is graph_collector.layer_hiddens[0] diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_mtp_diverse.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_mtp_diverse.py deleted file mode 100644 index 7c1af6bbc2..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_mtp_diverse.py +++ /dev/null @@ -1,196 +0,0 @@ -""" -MTP Diverse Attention Unit Test - -测试 MTP diverse attention 算子的正确性和性能。 - -Single Token Mode: -- 每个请求只有 1 个 Q token -- 组内第 i 个请求只能看到前 i 个 KV -- b_mark_shared_group 标记： - - 0: 组内非最后一个请求（跳过） - - N>=1: 一个 N 人组的最后一个请求（需要计算） -""" -import pytest -import torch - - -def gqa_attention_reference( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - req_to_tokens: torch.Tensor, - b_req_idx: torch.Tensor, - b_seq_len: torch.Tensor, -) -> torch.Tensor: - """ - Reference GQA attention implementation for single token mode - - q: [batch_size, num_heads, head_dim] - 每个请求 1 个 Q token - k, v: [kv_pool_size, kv_head_num, head_dim] - KV 池 - req_to_tokens: [batch, max_kv_len] - 每个请求的 KV 索引 - b_seq_len: [batch] - 每个请求可见的 KV 数量 - """ - batch_size = b_seq_len.shape[0] - num_heads = q.shape[1] - kv_head_num = k.shape[1] - head_dim = q.shape[2] - gqa_group_size = num_heads // kv_head_num - - # 输出：[batch, num_heads, head_dim] - output = torch.zeros((batch_size, num_heads, head_dim), dtype=q.dtype, device=q.device) - - for b in range(batch_size): - seq_len = b_seq_len[b].item() # 这个请求可见的 KV 数量 - - # 获取这个请求的 KV 索引：[seq_len] - kv_indices = req_to_tokens[b, :seq_len] - - # 获取 K 和 V：[seq_len, kv_head_num, head_dim] - k_batch = k[kv_indices] - v_batch = v[kv_indices] - - # Process each Q head - for h in range(num_heads): - kv_h = h // gqa_group_size - k_h = k_batch[:, kv_h, :] # [seq_len, head_dim] - v_h = v_batch[:, kv_h, :] # [seq_len, head_dim] - - q_h = q[b, h] # [head_dim] - - # Compute attention scores - att = torch.matmul(q_h, k_h.transpose(0, 1)) # [seq_len] - att = att / (head_dim ** 0.5) - att = torch.softmax(att, dim=-1) - out_h = torch.matmul(att, v_h) # [head_dim] - - output[b, h] = out_h - - return output - - -def setup_mtp_test_data(kv_len, group_size, batch_groups, test_dtype=torch.bfloat16, device="cuda", seed=42): - """ - 设置 MTP 测试数据 - Single Token Mode - - Single Token Mode 的特点： - - 每个请求有 1 个 Q token - - 组内第 i 个请求只能看到前 i+1 个 KV - - 组内请求共享相同的 KV 前缀 - - b_mark_shared_group 标记： - - 0: 组内非最后一个请求（跳过计算） - - N>=1: 一个 N 人组的最后一个请求（需要计算） - """ - torch.manual_seed(seed) - - num_heads = 32 - kv_head_num = 4 # gqa_group_size = 8 - head_dim = 128 - - batch_size = batch_groups * group_size - - # KV 池：[batch_groups * kv_len, kv_head_num, head_dim] - # 每个组使用不同的 KV 范围 - kv_pool_size = batch_groups * kv_len - k = torch.randn(size=(kv_pool_size, kv_head_num, head_dim), dtype=test_dtype, device=device) - v = torch.randn(size=(kv_pool_size, kv_head_num, head_dim), dtype=test_dtype, device=device) - - # req_to_tokens: [batch, max_kv_len] - 每个请求的 KV 索引 - max_kv_len = group_size # 每个请求最多有 group_size 个 KV（可见范围） - req_to_tokens = torch.zeros((batch_size, max_kv_len), dtype=torch.int32, device=device) - - b_req_idx = torch.arange(batch_size, dtype=torch.int32, device=device) - b_seq_len = torch.zeros(batch_size, dtype=torch.int32, device=device) # 每个请求可见的 KV 数量 - b_mark_shared_group = torch.zeros(batch_size, dtype=torch.int32, device=device) - - # Q: [batch_size, num_heads, head_dim] - 每个请求 1 个 Q token - q = torch.randn(size=(batch_size, num_heads, head_dim), dtype=test_dtype, device=device) - - for group_idx in range(batch_groups): - group_start = group_idx * group_size - kv_base = group_idx * kv_len - - for member_idx in range(group_size): - batch_idx = group_start + member_idx - # 第 member_idx 个请求可以看到前 member_idx+1 个 KV - b_seq_len[batch_idx] = member_idx + 1 - - # 设置 KV 索引 - 组内请求共享相同的 KV 前缀 - for kv_pos in range(member_idx + 1): - req_to_tokens[batch_idx, kv_pos] = kv_base + kv_pos - - # b_mark_shared_group: - # - 0: 组内非最后一个请求 - # - group_size: 组内最后一个请求 - if member_idx == group_size - 1: - b_mark_shared_group[batch_idx] = group_size - else: - b_mark_shared_group[batch_idx] = 0 - - return q, k, v, req_to_tokens, b_req_idx, b_seq_len, b_mark_shared_group - - -def mtp_diverse_attention( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - req_to_tokens: torch.Tensor, - b_req_idx: torch.Tensor, - b_seq_len: torch.Tensor, - b_mark_shared_group: torch.Tensor, -) -> torch.Tensor: - """ - MTP Diverse Attention 调用入口 - """ - from lightllm.common.basemodel.triton_kernel.att.decode_att.gqa.mtp_diverse import ( - token_decode_attention_mtp_diverse_single_token, - ) - - return token_decode_attention_mtp_diverse_single_token( - q=q, - k=k, - v=v, - Req_to_tokens=req_to_tokens, - B_req_idx=b_req_idx, - b_seq_len=b_seq_len, - b_mark_shared_group=b_mark_shared_group, - max_kv_len=int(b_seq_len.max().item()), - ) - - -@pytest.mark.parametrize("kv_len", [100, 256]) -@pytest.mark.parametrize("group_size", [2, 3]) -@pytest.mark.parametrize("batch_groups", [1, 2, 4]) -def test_mtp_diverse_vs_reference(kv_len, group_size, batch_groups): - """ - 测试 MTP diverse attention 与 reference 实现的正确性对比 - """ - test_dtype = torch.bfloat16 - device = "cuda" - - q, k, v, req_to_tokens, b_req_idx, b_seq_len, b_mark_shared_group = setup_mtp_test_data( - kv_len, group_size, batch_groups, test_dtype, device - ) - - # 运行 reference 实现（计算所有请求） - reference_out = gqa_attention_reference( - q=q, - k=k, - v=v, - req_to_tokens=req_to_tokens, - b_req_idx=b_req_idx, - b_seq_len=b_seq_len, - ) - - # 运行 MTP diverse（只计算组内最后一个请求） - diverse_out = mtp_diverse_attention( - q=q, - k=k, - v=v, - req_to_tokens=req_to_tokens, - b_req_idx=b_req_idx, - b_seq_len=b_seq_len, - b_mark_shared_group=b_mark_shared_group, - ) - - torch.testing.assert_close(diverse_out, reference_out, atol=1e-2, rtol=1e-2) diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_mtp_diverse_autotune.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_mtp_diverse_autotune.py deleted file mode 100644 index 0989fe6d26..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_mtp_diverse_autotune.py +++ /dev/null @@ -1,301 +0,0 @@ -import weakref -import collections -import inspect -import json -import math -from types import SimpleNamespace - -import pytest -import torch -import torch.nn.functional as F -from torch.nn.attention import SDPBackend, sdpa_kernel - -from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.common.basemodel import cuda_graph as graph_module -from lightllm.common.basemodel.attention.triton import fp as state_module -from lightllm.common.basemodel.triton_kernel.att.decode_att.gqa.mtp_diverse import ( - mtp_diverse_attn as decode_module, - stage1_single_token as stage1_module, - stage2_single_token as stage2_module, -) -from lightllm.common.triton_utils import autotuner as autotuner_module -from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType, AutotuneLevel -from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len - - -@pytest.fixture(autouse=True) -def autotune_environment(monkeypatch): - torch.manual_seed(42) - monkeypatch.delenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", raising=False) - monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", None) - monkeypatch.setattr(stage1_module, "get_triton_autotune_level", lambda: AutotuneLevel.ADAPTIVE_AUTOTUNE) - monkeypatch.setattr(stage2_module, "get_triton_autotune_level", lambda: AutotuneLevel.ADAPTIVE_AUTOTUNE) - monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: AutotuneLevel.ADAPTIVE_AUTOTUNE) - get_decode_attn_autotune_seq_len.cache_clear() - yield - get_decode_attn_autotune_seq_len.cache_clear() - - -@pytest.mark.parametrize("level", [AutotuneLevel.ADAPTIVE_AUTOTUNE, AutotuneLevel.FORCE_AUTOTUNE]) -@pytest.mark.parametrize("configured, expected", [(None, 32768), ("8193", 8704)]) -def test_tuning_key_uses_configured_length(monkeypatch, level, configured, expected): - if configured is not None: - monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", configured) - monkeypatch.setattr(stage1_module, "get_triton_autotune_level", lambda: level) - with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - assert stage1_module.get_run_key(torch.empty(7, 8, 64), 2) == 7_000_000_000 + expected - - -@pytest.mark.parametrize( - "phase,level", - [ - (None, AutotuneLevel.ADAPTIVE_AUTOTUNE), - (None, AutotuneLevel.FORCE_AUTOTUNE), - (AutotuneKernelType.GENERAL, AutotuneLevel.ADAPTIVE_AUTOTUNE), - (AutotuneKernelType.DECODE_ATTENTION, AutotuneLevel.USE_AUTOTUNE_HIS_CONFIG), - ], -) -def test_lookup_uses_actual_length(monkeypatch, phase, level): - monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", "invalid") - monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", phase) - monkeypatch.setattr(stage1_module, "get_triton_autotune_level", lambda: level) - for actual, expected in [(2, 512), (511, 512), (512, 512), (513, 1024), (8193, 8704)]: - assert stage1_module.get_run_key(torch.empty(7, 8, 64), actual) == 7_000_000_000 + expected - - -def make_inputs(device="cpu", batch=7, block_batch=3, head_dim=128, token_count=17, table_width=16, block_num=2): - kv = torch.randn(token_count, 4, head_dim, dtype=torch.bfloat16, device=device) - return dict( - q=torch.randn(batch, 8, head_dim, dtype=kv.dtype, device=device), - k=kv[:, :2], - v=kv[:, 2:], - Req_to_tokens=torch.arange(6 * table_width, dtype=torch.int32, device=device).view(6, table_width) - % max(1, token_count), - B_req_idx=torch.full((batch,), 4, dtype=torch.int32, device=device), - b_seq_len=torch.full((batch,), 2, dtype=torch.int32, device=device), - b_mark_shared_group=torch.ones(batch, dtype=torch.int32, device=device), - max_kv_len=2, - mid_out=torch.full((batch, 8, block_num, head_dim), -100, dtype=kv.dtype, device=device), - mid_out_logsumexp=torch.full((batch, 8, block_num), -100, dtype=torch.float32, device=device), - block_batch=block_batch, - ) - - -def rebuild(inputs): - args, kwargs = stage1_module.rebuild_inputs(**inputs) - return inspect.signature(stage1_module.mtp_diverse_stage1_single_token.fn).bind(*args, **kwargs).arguments - - -@pytest.mark.parametrize("token_count", [3, 128]) -@pytest.mark.parametrize( - "batch, block_batch, length, req_ids, lengths, marks", - [ - (7, 3, 17, [0, 0, 0, 1, 1, 1, 2], [15, 16, 17, 15, 16, 17, 17], [0, 0, 3, 0, 0, 3, 1]), - (5, 4, 4, [0, 0, 0, 0, 1], [1, 2, 3, 4, 4], [0, 0, 0, 4, 1]), - (1, 4, 1, [0], [1], [1]), - ], -) -def test_rebuild_preserves_inputs_and_shared_prefixes( - monkeypatch, token_count, batch, block_batch, length, req_ids, lengths, marks -): - monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", str(length)) - inputs = make_inputs(batch=batch, block_batch=block_batch, token_count=token_count) - snapshots = {name: value.clone() for name, value in inputs.items() if isinstance(value, torch.Tensor)} - rebuilt = rebuild(inputs) - assert rebuilt["B_req_idx"].tolist() == req_ids - assert rebuilt["b_seq_len"].tolist() == lengths - assert rebuilt["b_mark_shared_group"].tolist() == marks - assert rebuilt["max_kv_len"] == length - expected = torch.arange((max(req_ids) + 1) * length).view(-1, length) % token_count - torch.testing.assert_close(rebuilt["Req_to_tokens"].long(), expected) - for name in ["q", "k", "v", "mid_out", "mid_out_logsumexp"]: - assert rebuilt[name] is inputs[name] - assert rebuilt["block_batch"] == block_batch - for name, value in snapshots.items(): - torch.testing.assert_close(inputs[name], value) - - -@pytest.mark.parametrize("batch, block_batch, length", [(7, 3, 2), (2, 4, 1)]) -def test_rebuild_rejects_length_shorter_than_group(monkeypatch, batch, block_batch, length): - monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", str(length)) - with pytest.raises(ValueError, match="at least the largest MTP group size"): - rebuild(make_inputs(batch=batch, block_batch=block_batch)) - - -def test_rebuild_rejects_empty_or_mismatched_cache(): - with pytest.raises(ValueError, match="non-empty KV cache"): - rebuild(make_inputs(token_count=0)) - inputs = make_inputs() - inputs["v"] = inputs["v"][:-1] - with pytest.raises(AssertionError, match="same number of tokens"): - rebuild(inputs) - - -def test_decode_passes_saved_length_and_selected_block_n(monkeypatch): - inputs = make_inputs() - infer_state = InferStateInfo() - infer_state.max_kv_seq_len = 8193 - infer_state.b_req_idx = inputs["B_req_idx"] - infer_state.b_seq_len = inputs["b_seq_len"] - infer_state.req_manager = SimpleNamespace(req_to_token_indexs=inputs["Req_to_tokens"]) - model = SimpleNamespace( - is_mtp_draft_model=False, - mtp_manager=SimpleNamespace(get_decode_draft_step=lambda _: 2), - req_manager=SimpleNamespace(HOLD_REQUEST_ID=-1), - ) - monkeypatch.setattr(state_module, "build_mtp_shared_group_markers", lambda *a, **kw: inputs["b_mark_shared_group"]) - state = state_module.TritonDecodeAttState( - backend=SimpleNamespace(model=model), _infer_state_weak_ref=weakref.ref(infer_state) - ) - state.init_state() - infer_state.max_kv_seq_len = 32768 - calls = [] - - def stage1(**kwargs): - calls.append(kwargs) - return 32 - - monkeypatch.setattr(decode_module, "mtp_diverse_stage1_single_token", stage1) - monkeypatch.setattr(decode_module, "mtp_diverse_stage2_single_token", lambda **kwargs: calls.append(kwargs)) - state.decode_att(q=inputs["q"], k=inputs["k"], v=inputs["v"]) - assert calls[0]["max_kv_len"] == 8193 - assert calls[0]["b_mark_shared_group"] is inputs["b_mark_shared_group"] - assert calls[1]["block_n"] == 32 - assert calls[1]["max_kv_len"] == 8193 - - -def reference_attention(inputs): - results = [] - for row, length in enumerate(inputs["b_seq_len"].tolist()): - indices = inputs["Req_to_tokens"][inputs["B_req_idx"][row], :length].long() - with sdpa_kernel(SDPBackend.MATH): - results.append( - F.scaled_dot_product_attention( - inputs["q"][row].float().unsqueeze(1), - inputs["k"][indices].float().transpose(0, 1), - inputs["v"][indices].float().transpose(0, 1), - enable_gqa=True, - ).squeeze(1) - ) - return torch.stack(results) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("kv_len", [17, 8192, 16384]) -@pytest.mark.parametrize("block_batch", [1, 3, 4]) -@pytest.mark.parametrize("head_dim", [64, 128, 256]) -def test_shared_groups_match_fp32_reference(monkeypatch, kv_len, block_batch, head_dim): - monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", str(kv_len)) - inputs = rebuild(make_inputs("cuda", block_batch=block_batch, head_dim=head_dim, token_count=7 * kv_len)) - reference = reference_attention(inputs) - out = torch.empty_like(inputs["q"]) - # 仅两个 grid block，长 KV 必须循环多次；17 token 同时覆盖末块 mask 和无效分块的归约边界。 - for block_n in [16, 32, 64]: - selected_block_n = stage1_module.mtp_diverse_stage1_single_token.fn( - **inputs, - run_config={"BLOCK_N": block_n, "num_warps": 4, "num_stages": 2, "warp_specialize": False}, - ) - stage2_module.mtp_diverse_stage2_single_token.fn( - inputs["mid_out"], inputs["mid_out_logsumexp"], inputs["b_seq_len"], out, selected_block_n, kv_len - ) - torch.testing.assert_close(out.float(), reference, atol=2e-3, rtol=2e-2) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("kv_len", [8192, 16384]) -def test_full_autotune_before_graph_and_cached_reuse(tmp_path, monkeypatch, kv_len): - monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", str(kv_len)) - monkeypatch.setattr(autotuner_module.dist, "is_initialized", lambda: False) - monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: AutotuneLevel.FORCE_AUTOTUNE) - monkeypatch.setattr(stage1_module, "get_triton_autotune_level", lambda: AutotuneLevel.FORCE_AUTOTUNE) - kernel = stage1_module.mtp_diverse_stage1_single_token - monkeypatch.setattr(kernel, "_cache_dir", str(tmp_path), raising=False) - monkeypatch.setattr(kernel, "cached_configs", {}) - monkeypatch.setattr(kernel, "fast_match_configs", collections.defaultdict(dict)) - monkeypatch.setattr(kernel, "warmuped_configs_set", set()) - assert kernel.mutates_args == [] - # 使用生产路径的小 batch 布局，让 stage2 在长请求下归约完整的 128 个分块。 - inputs = make_inputs("cuda", table_width=kv_len, token_count=3 * kv_len, block_num=128) - stage2 = stage2_module.mtp_diverse_stage2_single_token - stage2_cache = tmp_path / "stage2" - stage2_cache.mkdir() - monkeypatch.setattr(stage2, "_cache_dir", str(stage2_cache), raising=False) - monkeypatch.setattr(stage2, "cached_configs", {}) - monkeypatch.setattr(stage2, "fast_match_configs", collections.defaultdict(dict)) - monkeypatch.setattr(stage2, "warmuped_configs_set", set()) - snapshots = {name: value.clone() for name, value in inputs.items() if isinstance(value, torch.Tensor)} - original_rebuild, original_bench = kernel.rebuild_input_func, kernel._bench - rebuild_count, benchmark_count, valid_benchmark_count = 0, 0, 0 - - def checked_rebuild(*args, **kwargs): - nonlocal rebuild_count - assert not torch.cuda.is_current_stream_capturing() - rebuild_count += 1 - return original_rebuild(*args, **kwargs) - - def checked_bench(*args, **kwargs): - nonlocal benchmark_count, valid_benchmark_count - rebuilt = inspect.signature(kernel.fn).bind(*args, **kwargs).arguments - assert rebuilt["max_kv_len"] == kv_len - assert rebuilt["b_seq_len"].tolist() == [kv_len - 2, kv_len - 1, kv_len] * 2 + [kv_len] - assert rebuilt["b_mark_shared_group"].tolist() == [0, 0, 3, 0, 0, 3, 1] - assert rebuilt["Req_to_tokens"].shape == (3, kv_len) - elapsed = original_bench(*args, **kwargs) - # 不可编译或资源不足的候选按 Autotuner 原有规则返回 inf；完整搜索必须选出有效配置。 - if math.isfinite(elapsed): - valid_benchmark_count += 1 - benchmark_count += 1 - return elapsed - - monkeypatch.setattr(kernel, "rebuild_input_func", checked_rebuild) - monkeypatch.setattr(kernel, "_bench", checked_bench) - monkeypatch.setattr( - graph_module, - "get_env_start_args", - lambda: SimpleNamespace( - enable_decode_microbatch_overlap=False, enable_tpsp_mix_mode=False, enable_torch_memory_saver=False - ), - ) - out = torch.empty_like(inputs["q"]) - - def decode(state): - block_n = kernel(**inputs) - stage2(inputs["mid_out"], inputs["mid_out_logsumexp"], inputs["b_seq_len"], out, block_n, inputs["max_kv_len"]) - return out - - graph = graph_module.CudaGraph(1, 7, 1, max_batch_size=7, max_len_in_batch=32768) - graph.capture_decode(decode, SimpleNamespace(input_ids=torch.ones(7, device="cuda"))) - torch.cuda.synchronize() - assert rebuild_count == 1 and benchmark_count == len(kernel.configs_gen_func()) - assert valid_benchmark_count > 0 - # 两个 KV 的短请求沿用原 MTP 测试的 BF16 容差；下面的长请求仍使用更严格的 FP32 对照。 - torch.testing.assert_close(out.float(), reference_attention(inputs), atol=1e-2, rtol=1e-2) - for name, value in snapshots.items(): - if name not in ["mid_out", "mid_out_logsumexp"]: - torch.testing.assert_close(inputs[name], value) - cache_file = next(tmp_path.glob("*.json")) - saved = cache_file.read_bytes() - assert list(json.loads(saved)) == [str(7_000_000_000 + kv_len)] - assert json.loads(saved)[str(7_000_000_000 + kv_len)] is not None - stage2_configs = json.loads(next(stage2_cache.glob("*.json")).read_bytes()) - assert list(stage2_configs) == [str(7 * 8 * 1_000_000_000 + 128)] - assert next(iter(stage2_configs.values())) is not None - # 同一层配置以及从文件重新加载的配置，在 FORCE 模式下都不能再次搜索。 - with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - kernel(**inputs) - kernel.cached_configs.clear() - kernel.fast_match_configs.clear() - kernel.warmuped_configs_set.clear() - with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - kernel(**inputs) - assert rebuild_count == 1 and benchmark_count == len(kernel.configs_gen_func()) - assert cache_file.read_bytes() == saved - # 原缓冲区改为真实长请求和共享组后回放，验证 Graph 没有捕获 benchmark 的重建输入。 - inputs["B_req_idx"].copy_(torch.tensor([1, 1, 1, 3, 3, 3, 2], device="cuda", dtype=torch.int32)) - inputs["b_mark_shared_group"].copy_(torch.tensor([0, 0, 3, 0, 0, 3, 1], device="cuda", dtype=torch.int32)) - inputs["b_seq_len"].copy_( - torch.tensor([kv_len - 2, kv_len - 1, kv_len] * 2 + [kv_len], device="cuda", dtype=torch.int32) - ) - graph.graph[7][0].replay() - torch.cuda.synchronize() - torch.testing.assert_close(out.float(), reference_attention(inputs), atol=2e-3, rtol=2e-2) diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_stage2_autotune.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_stage2_autotune.py deleted file mode 100644 index 48956c4f36..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_stage2_autotune.py +++ /dev/null @@ -1,220 +0,0 @@ -import collections -import inspect -import json -import math -from types import SimpleNamespace - -import pytest -import torch - -from lightllm.common.basemodel import cuda_graph as graph_module -from lightllm.common.basemodel.triton_kernel.att.decode_att.gqa.mtp_diverse import stage2_single_token as stage2_module -from lightllm.common.triton_utils import autotuner as autotuner_module -from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType, AutotuneLevel -from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len - - -@pytest.fixture(autouse=True) -def autotune_environment(monkeypatch): - torch.manual_seed(42) - monkeypatch.delenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", raising=False) - monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", None) - monkeypatch.setattr(stage2_module, "get_triton_autotune_level", lambda: AutotuneLevel.ADAPTIVE_AUTOTUNE) - monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: AutotuneLevel.ADAPTIVE_AUTOTUNE) - get_decode_attn_autotune_seq_len.cache_clear() - yield - get_decode_attn_autotune_seq_len.cache_clear() - - -@pytest.mark.parametrize("level", [AutotuneLevel.ADAPTIVE_AUTOTUNE, AutotuneLevel.FORCE_AUTOTUNE]) -@pytest.mark.parametrize("configured, capacity, expected", [(None, 128, 128), ("8193", 256, 129), ("65", 128, 2)]) -def test_tuning_key_uses_effective_reduction_blocks(monkeypatch, level, configured, capacity, expected): - if configured is not None: - monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", configured) - monkeypatch.setattr(stage2_module, "get_triton_autotune_level", lambda: level) - with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - assert stage2_module.get_run_key(torch.empty(2, 4, capacity, 64), 64, 2) == 8_000_000_000 + expected - - -@pytest.mark.parametrize( - "phase,level", - [ - (None, AutotuneLevel.ADAPTIVE_AUTOTUNE), - (None, AutotuneLevel.FORCE_AUTOTUNE), - (AutotuneKernelType.GENERAL, AutotuneLevel.ADAPTIVE_AUTOTUNE), - (AutotuneKernelType.DECODE_ATTENTION, AutotuneLevel.USE_AUTOTUNE_HIS_CONFIG), - (AutotuneKernelType.DECODE_ATTENTION, AutotuneLevel.CLOSE_AUTOTUNE), - ], -) -def test_lookup_counts_real_blocks_without_reading_gpu_lengths(monkeypatch, phase, level): - monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", "invalid") - monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", phase) - monkeypatch.setattr(stage2_module, "get_triton_autotune_level", lambda: level) - mid = torch.empty(2, 4, 128, 64) - for length, blocks in [(1, 1), (64, 1), (65, 2), (511, 8), (512, 8), (513, 9), (8192, 128), (16384, 128)]: - assert stage2_module.get_run_key(mid, 64, length) == 8_000_000_000 + blocks - - -def test_key_reuses_equal_work_and_distinguishes_capacity_limits(): - key = stage2_module.get_run_key - small, large = torch.empty(2, 4, 64, 64), torch.empty(2, 4, 128, 64) - assert key(small, 64, 513) == key(large, 64, 513) - assert key(small, 64, 8192) != key(large, 64, 8192) - assert key(large, 64, 8192) == key(large, 64, 16384) - - -def make_inputs(device="cpu", capacity=128, block_n=64, dtype=torch.bfloat16, head_dim=128): - # 模拟短请求：stage1 仅写入第一个分块，其余位置用 NaN 标记为未初始化。 - mid = torch.full((2, 4, capacity, head_dim), float("nan"), dtype=dtype, device=device) - lse = torch.full((2, 4, capacity), float("nan"), dtype=torch.float32, device=device) - mid[:, :, 0] = torch.randn(2, 4, head_dim, dtype=dtype, device=device) - lse[:, :, 0] = 0 - return dict( - mid_out=mid, - mid_out_logsumexp=lse, - B_Seqlen=torch.full((2,), 2, dtype=torch.int32, device=device), - out=torch.full((2, 4, head_dim), -100, dtype=dtype, device=device), - block_n=block_n, - max_kv_len=2, - ) - - -@pytest.mark.parametrize("length", [65, 8193, 16384]) -def test_rebuild_reuses_readonly_intermediate_inputs(monkeypatch, length): - monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", str(length)) - inputs = make_inputs() - snapshots = {name: value.clone() for name, value in inputs.items() if isinstance(value, torch.Tensor)} - args, kwargs = stage2_module.rebuild_inputs(**inputs) - rebuilt = inspect.signature(stage2_module.mtp_diverse_stage2_single_token.fn).bind(*args, **kwargs).arguments - assert rebuilt["max_kv_len"] == length - assert rebuilt["B_Seqlen"].tolist() == [length, length] - assert rebuilt["block_n"] == inputs["block_n"] - assert rebuilt["out"] is inputs["out"] - for name in ["mid_out", "mid_out_logsumexp"]: - assert rebuilt[name] is inputs[name] - assert rebuilt[name].shape == inputs[name].shape - assert rebuilt[name].stride() == inputs[name].stride() - assert rebuilt[name].dtype == inputs[name].dtype - assert rebuilt[name].device == inputs[name].device - assert rebuilt["B_Seqlen"] is not inputs["B_Seqlen"] - assert torch.isfinite(rebuilt["B_Seqlen"]).all() - for name, snapshot in snapshots.items(): - torch.testing.assert_close(inputs[name], snapshot, equal_nan=True) - - -def reduction_reference(inputs): - results = [] - for row, length in enumerate(inputs["B_Seqlen"].tolist()): - blocks = min((length + inputs["block_n"] - 1) // inputs["block_n"], inputs["mid_out"].shape[2]) - weights = torch.softmax(inputs["mid_out_logsumexp"][row, :, :blocks].float(), dim=-1) - results.append((weights.unsqueeze(-1) * inputs["mid_out"][row, :, :blocks].float()).sum(dim=1)) - return torch.stack(results) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("block_n", [16, 64]) -@pytest.mark.parametrize("capacity", [2, 128]) -@pytest.mark.parametrize("dtype,head_dim", [(torch.bfloat16, 128), (torch.float16, 64), (torch.bfloat16, 256)]) -def test_all_configs_reduce_only_valid_blocks(block_n, capacity, dtype, head_dim): - lengths = [1, block_n, block_n + 1, capacity * block_n, capacity * block_n + 8192] - mid = torch.randn(len(lengths), 4, capacity, head_dim, dtype=dtype, device="cuda") - lse = torch.randn(len(lengths), 4, capacity, dtype=torch.float32, device="cuda") * 20 - for row, length in enumerate(lengths): - blocks = min((length + block_n - 1) // block_n, capacity) - mid[row, :, blocks:] = float("nan") - lse[row, :, blocks:] = float("nan") - inputs = dict( - mid_out=mid, - mid_out_logsumexp=lse, - B_Seqlen=torch.tensor(lengths, dtype=torch.int32, device="cuda"), - out=torch.empty(len(lengths), 4, head_dim, dtype=dtype, device="cuda"), - block_n=block_n, - max_kv_len=max(lengths), - ) - reference = reduction_reference(inputs) - for config in stage2_module.get_test_configs(): - stage2_module.mtp_diverse_stage2_single_token.fn(**inputs, run_config=config) - torch.testing.assert_close(inputs["out"].float(), reference, atol=2e-3, rtol=1e-2) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("length", [65, 8192, 16384]) -def test_full_tuning_graph_and_cache_reuse(tmp_path, monkeypatch, length): - monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", str(length)) - monkeypatch.setattr(autotuner_module.dist, "is_initialized", lambda: False) - monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: AutotuneLevel.FORCE_AUTOTUNE) - monkeypatch.setattr(stage2_module, "get_triton_autotune_level", lambda: AutotuneLevel.FORCE_AUTOTUNE) - kernel = stage2_module.mtp_diverse_stage2_single_token - monkeypatch.setattr(kernel, "_cache_dir", str(tmp_path), raising=False) - monkeypatch.setattr(kernel, "cached_configs", {}) - monkeypatch.setattr(kernel, "fast_match_configs", collections.defaultdict(dict)) - monkeypatch.setattr(kernel, "warmuped_configs_set", set()) - inputs = make_inputs("cuda") - snapshots = {name: value.clone() for name, value in inputs.items() if isinstance(value, torch.Tensor)} - original_rebuild, original_bench = kernel.rebuild_input_func, kernel._bench - rebuild_count, benchmark_count = 0, 0 - - def checked_rebuild(*args, **kwargs): - nonlocal rebuild_count - assert not torch.cuda.is_current_stream_capturing() - rebuild_count += 1 - return original_rebuild(*args, **kwargs) - - def checked_bench(*args, **kwargs): - nonlocal benchmark_count - rebuilt = inspect.signature(kernel.fn).bind(*args, **kwargs).arguments - assert rebuilt["B_Seqlen"].tolist() == [length, length] - assert rebuilt["max_kv_len"] == length and rebuilt["block_n"] == 64 - for name in ["mid_out", "mid_out_logsumexp"]: - assert rebuilt[name] is inputs[name] - assert rebuilt[name].shape == inputs[name].shape - elapsed = original_bench(*args, **kwargs) - assert math.isfinite(elapsed) - torch.testing.assert_close(inputs["out"], snapshots["out"]) - benchmark_count += 1 - return elapsed - - monkeypatch.setattr(kernel, "rebuild_input_func", checked_rebuild) - monkeypatch.setattr(kernel, "_bench", checked_bench) - monkeypatch.setattr( - graph_module, - "get_env_start_args", - lambda: SimpleNamespace( - enable_decode_microbatch_overlap=False, enable_tpsp_mix_mode=False, enable_torch_memory_saver=False - ), - ) - - def decode(state): - kernel(**inputs) - return inputs["out"] - - graph = graph_module.CudaGraph(1, 2, 1, max_batch_size=2, max_len_in_batch=32768) - graph.capture_decode(decode, SimpleNamespace(input_ids=torch.ones(2, device="cuda"))) - torch.cuda.synchronize() - assert rebuild_count == 1 and benchmark_count == 20 - torch.testing.assert_close(inputs["out"].float(), reduction_reference(inputs), atol=2e-3, rtol=1e-2) - for name in ["mid_out", "mid_out_logsumexp", "B_Seqlen"]: - torch.testing.assert_close(inputs[name], snapshots[name], equal_nan=True) - cache_file = next(tmp_path.glob("*.json")) - saved = cache_file.read_bytes() - expected_key = str(8_000_000_000 + min((length + 63) // 64, 128)) - assert list(json.loads(saved)) == [expected_key] - with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - kernel(**inputs) - kernel.cached_configs.clear() - kernel.fast_match_configs.clear() - kernel.warmuped_configs_set.clear() - # 8K/16K 都归约 128 个块，切换代表性长度后也应复用同一个 key。 - if length >= 8192: - monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", str(24576 - length)) - get_decode_attn_autotune_seq_len.cache_clear() - with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - kernel(**inputs) - assert rebuild_count == 1 and benchmark_count == 20 - assert cache_file.read_bytes() == saved - inputs["mid_out"].normal_() - inputs["mid_out_logsumexp"].normal_() - inputs["B_Seqlen"].fill_(length) - graph.graph[2][0].replay() - torch.cuda.synchronize() - torch.testing.assert_close(inputs["out"].float(), reduction_reference(inputs), atol=2e-3, rtol=1e-2) diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/test_flash_decoding_autotune.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/test_flash_decoding_autotune.py deleted file mode 100644 index b3068db5fb..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/test_flash_decoding_autotune.py +++ /dev/null @@ -1,270 +0,0 @@ -import weakref -import collections -import inspect -import math -from types import SimpleNamespace - -import pytest -import torch -import torch.nn.functional as F -from frozendict import frozendict -from torch.nn.attention import SDPBackend, sdpa_kernel - -from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.common.basemodel.attention.triton.fp import TritonDecodeAttState -from lightllm.common.basemodel.triton_kernel.att.decode_att.gqa.flash_decoding import ( - gqa_flash_decoding as decode_module, - gqa_flash_decoding_stage1 as stage1_module, -) -from lightllm.common.basemodel.triton_kernel.att.decode_att.gqa.flash_decoding.gqa_flash_decoding_stage2 import ( - flash_decode_stage2, -) -from lightllm.common.triton_utils import autotuner as autotuner_module -from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType, AutotuneLevel -from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len - - -@pytest.fixture(autouse=True) -def autotune_environment(monkeypatch): - torch.manual_seed(42) - monkeypatch.delenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", raising=False) - monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", None) - monkeypatch.setattr(stage1_module, "get_triton_autotune_level", lambda: AutotuneLevel.ADAPTIVE_AUTOTUNE) - monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: AutotuneLevel.ADAPTIVE_AUTOTUNE) - get_decode_attn_autotune_seq_len.cache_clear() - yield - get_decode_attn_autotune_seq_len.cache_clear() - - -@pytest.mark.parametrize("level", [AutotuneLevel.ADAPTIVE_AUTOTUNE, AutotuneLevel.FORCE_AUTOTUNE]) -@pytest.mark.parametrize("configured, expected", [(None, 32768), ("8193", 8704)]) -def test_tuning_key_uses_bucketed_configured_length(monkeypatch, level, configured, expected): - if configured is not None: - monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", configured) - monkeypatch.setattr(stage1_module, "get_triton_autotune_level", lambda: level) - q = torch.empty(2, 8, 64) - with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - assert stage1_module.get_run_key(q, 2) == 2_000_000_000 + expected - - -@pytest.mark.parametrize( - "phase,level", - [ - (None, AutotuneLevel.ADAPTIVE_AUTOTUNE), - (None, AutotuneLevel.FORCE_AUTOTUNE), - (AutotuneKernelType.GENERAL, AutotuneLevel.ADAPTIVE_AUTOTUNE), - (AutotuneKernelType.DECODE_ATTENTION, AutotuneLevel.USE_AUTOTUNE_HIS_CONFIG), - ], -) -def test_non_tuning_key_uses_actual_length(monkeypatch, phase, level): - monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", "invalid") - monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", phase) - monkeypatch.setattr(stage1_module, "get_triton_autotune_level", lambda: level) - q = torch.empty(2, 8, 64) - for actual, bucket in [(2, 512), (511, 512), (512, 512), (513, 1024), (8193, 8704)]: - assert stage1_module.get_run_key(q, actual) == 2_000_000_000 + bucket - - -def make_inputs(device="cpu", head_dim=64, token_count=17, table_width=16): - # K/V 是共享存储中的不连续视图；原请求索引也不是从 0 开始连续排列。 - kv = torch.randn(token_count, 4, head_dim, dtype=torch.bfloat16, device=device) - return dict( - q=torch.randn(2, 8, head_dim, dtype=kv.dtype, device=device), - k=kv[:, :2], - v=kv[:, 2:], - Req_to_tokens=torch.arange(5 * table_width, dtype=torch.int32, device=device).view(5, table_width) - % token_count, - B_req_idx=torch.tensor([4, 2], dtype=torch.int32, device=device), - B_Seqlen=torch.tensor([3, 2], dtype=torch.int32, device=device), - max_len_in_batch=3, - mid_out=torch.full((2, 8, 2, head_dim), -100, dtype=kv.dtype, device=device), - mid_out_logsumexp=torch.full((2, 8, 2), -100, dtype=torch.float32, device=device), - block_seq=256, - sliding_window=(-1, -1), - ) - - -def test_static_key_shares_heads_and_blocks_but_separates_windows(): - kernel = stage1_module.flash_decode_stage1 - inputs = make_inputs() - full_key = frozendict(kernel._static_key(**inputs)) - inputs.pop("sliding_window") - assert frozendict(kernel._static_key(**inputs)) == full_key - - inputs["q"] = torch.empty(2, 4, 64, dtype=torch.bfloat16) - inputs["k"] = torch.empty(17, 1, 64, dtype=torch.bfloat16) - inputs["mid_out"] = torch.empty(2, 4, 128, 64, dtype=torch.bfloat16) - assert frozendict(kernel._static_key(**inputs)) == full_key - - window_key = frozendict(kernel._static_key(**inputs, sliding_window=(511, 0))) - assert window_key != full_key - assert frozendict(kernel._static_key(**inputs, sliding_window=[511, 0])) == window_key - assert frozendict(kernel._static_key(**inputs, sliding_window=(1023, 0))) != window_key - - -@pytest.mark.parametrize("num_tokens", [3, 128]) -def test_rebuild_preserves_layout_and_bounds_indices(monkeypatch, num_tokens): - monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", "17") - inputs = make_inputs(token_count=num_tokens) - inputs["sliding_window"] = (7, 0) - snapshots = {name: value.clone() for name, value in inputs.items() if isinstance(value, torch.Tensor)} - - args, kwargs = stage1_module.rebuild_inputs(**inputs) - rebuilt = inspect.signature(stage1_module.flash_decode_stage1.fn).bind(*args, **kwargs).arguments - assert rebuilt["Req_to_tokens"].shape == (2, 17) - assert rebuilt["Req_to_tokens"].min() >= 0 and rebuilt["Req_to_tokens"].max() < num_tokens - assert rebuilt["B_req_idx"].tolist() == [0, 1] - assert rebuilt["B_Seqlen"].tolist() == [17, 17] - assert rebuilt["max_len_in_batch"] == 17 - for name in ["q", "k", "v", "mid_out", "mid_out_logsumexp"]: - assert rebuilt[name] is inputs[name] - assert rebuilt["block_seq"] == 256 - assert rebuilt["sliding_window"] == (7, 0) - for name, value in snapshots.items(): - torch.testing.assert_close(inputs[name], value) - - -def test_state_retains_actual_length_before_graph_capture(monkeypatch): - infer_state = InferStateInfo() - infer_state.max_kv_seq_len = 8193 - state = TritonDecodeAttState( - backend=SimpleNamespace( - model=SimpleNamespace( - is_mtp_draft_model=False, mtp_manager=SimpleNamespace(get_decode_draft_step=lambda _: 0) - ) - ), - _infer_state_weak_ref=weakref.ref(infer_state), - ) - state.init_state() - state.infer_state.max_kv_seq_len = 32768 - calls = [] - monkeypatch.setattr( - decode_module, "gqa_token_decode_attention_flash_decoding", lambda **kwargs: calls.append(kwargs) - ) - state.decode_att( - q=torch.empty(2, 8, 64), - k=torch.empty(16, 2, 64), - v=torch.empty(16, 2, 64), - alloc_func=lambda shape, dtype: torch.empty(shape, dtype=dtype), - ) - assert calls[0]["max_len_in_batch"] == 8193 - - -def reference_attention(inputs): - outputs = [] - for i, length in enumerate(inputs["B_Seqlen"].tolist()): - start = max(0, length - 1 - inputs["sliding_window"][0]) if inputs["sliding_window"][0] >= 0 else 0 - indices = inputs["Req_to_tokens"][inputs["B_req_idx"][i], start:length].long() - with sdpa_kernel(SDPBackend.MATH): - outputs.append( - F.scaled_dot_product_attention( - inputs["q"][i].float().unsqueeze(1), - inputs["k"][indices].float().transpose(0, 1), - inputs["v"][indices].float().transpose(0, 1), - enable_gqa=True, - ).squeeze(1) - ) - return torch.stack(outputs) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("kv_len", [8192, 16384]) -@pytest.mark.parametrize("head_dim", [64, 256]) -@pytest.mark.parametrize("sliding_window", [(-1, -1), (511, 0)]) -def test_stage1_stage2_long_kv_matches_reference(kv_len, head_dim, sliding_window): - inputs = make_inputs("cuda", head_dim=head_dim, token_count=kv_len * 2, table_width=kv_len) - inputs["B_Seqlen"] = torch.tensor([kv_len, kv_len - 1], dtype=torch.int32, device="cuda") - inputs["max_len_in_batch"] = kv_len - inputs["sliding_window"] = sliding_window - reference = reference_attention(inputs) - # 只有两个 grid block，长 KV 需要每个 program 循环多次；stage2 仍归约同一布局。 - output = torch.empty_like(inputs["q"]) - for block_n in [16, 64, 128]: - stage1_module.flash_decode_stage1.fn(**inputs, run_config={"BLOCK_N": block_n, "num_warps": 4, "num_stages": 2}) - flash_decode_stage2( - inputs["mid_out"], - inputs["mid_out_logsumexp"], - inputs["B_Seqlen"], - output, - inputs["block_seq"], - sliding_window=sliding_window, - ) - assert torch.isfinite(output).all() - torch.testing.assert_close(output.float(), reference, atol=2e-3, rtol=2e-2) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("kv_len", [8192, 16384]) -def test_autotune_rebuilds_once_and_graph_uses_original_inputs(tmp_path, monkeypatch, kv_len): - monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", str(kv_len)) - monkeypatch.setattr(autotuner_module.dist, "is_initialized", lambda: False) - kernel = stage1_module.flash_decode_stage1 - monkeypatch.setattr(kernel, "_cache_dir", str(tmp_path), raising=False) - monkeypatch.setattr(kernel, "cached_configs", {}) - monkeypatch.setattr(kernel, "fast_match_configs", collections.defaultdict(dict)) - monkeypatch.setattr(kernel, "warmuped_configs_set", set()) - assert kernel.mutates_args == [] - monkeypatch.setattr( - kernel, - "configs_gen_func", - lambda: [{"BLOCK_N": n, "num_warps": 4, "num_stages": 2} for n in [16, 64, 128]], - ) - inputs = make_inputs("cuda") - snapshots = {name: value.clone() for name, value in inputs.items() if isinstance(value, torch.Tensor)} - reference = reference_attention(inputs) - rebuild = kernel.rebuild_input_func - benchmark = kernel._bench - rebuild_count = 0 - benchmark_count = 0 - - def checked_rebuild(*args, **kwargs): - nonlocal rebuild_count - rebuild_count += 1 - return rebuild(*args, **kwargs) - - def checked_bench(*args, **kwargs): - nonlocal benchmark_count - rebuilt = inspect.signature(kernel.fn).bind(*args, **kwargs).arguments - assert rebuilt["B_Seqlen"].tolist() == [kv_len, kv_len] - assert rebuilt["Req_to_tokens"].shape == (2, kv_len) - assert rebuilt["B_req_idx"].tolist() == [0, 1] - assert rebuilt["max_len_in_batch"] == kv_len - assert rebuilt["mid_out"] is inputs["mid_out"] - elapsed = benchmark(*args, **kwargs) - assert math.isfinite(elapsed) - benchmark_count += 1 - return elapsed - - monkeypatch.setattr(kernel, "rebuild_input_func", checked_rebuild) - monkeypatch.setattr(kernel, "_bench", checked_bench) - with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - kernel(**inputs) - assert rebuild_count == 1 and benchmark_count == 3 - assert all(list(configs) == [str(2_000_000_000 + kv_len)] for configs in kernel.cached_configs.values()) - # FORCE 在 decode 阶段仍复用缓存，不因模型层重复调用而再次调优。 - monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: AutotuneLevel.FORCE_AUTOTUNE) - monkeypatch.setattr(stage1_module, "get_triton_autotune_level", lambda: AutotuneLevel.FORCE_AUTOTUNE) - with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - kernel(**inputs) - assert rebuild_count == 1 and benchmark_count == 3 - - output = torch.empty_like(inputs["q"]) - - def decode(): - kernel(**inputs) - flash_decode_stage2( - inputs["mid_out"], inputs["mid_out_logsumexp"], inputs["B_Seqlen"], output, inputs["block_seq"] - ) - - decode() - torch.cuda.synchronize() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - decode() - graph.replay() - torch.cuda.synchronize() - torch.testing.assert_close(output.float(), reference, atol=2e-3, rtol=2e-2) - assert rebuild_count == 1 and benchmark_count == 3 - for name, value in snapshots.items(): - if name not in ["mid_out", "mid_out_logsumexp"]: - torch.testing.assert_close(inputs[name], value) diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int4kv/test_int4kv_autotune.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int4kv/test_int4kv_autotune.py deleted file mode 100644 index df357539e1..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int4kv/test_int4kv_autotune.py +++ /dev/null @@ -1,295 +0,0 @@ -import weakref -import collections -import inspect -import json -import math -from types import SimpleNamespace - -import pytest -import torch -import torch.nn.functional as F -from torch.nn.attention import SDPBackend, sdpa_kernel - -from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.common.basemodel.attention.triton.int4kv import Int4kvTritonDecodeAttState -from lightllm.common.basemodel.triton_kernel.att.decode_att.int4kv import ( - int4kv_flash_decoding_stage1 as stage1_module, - ppl_int4kv_flash_decoding as decode_module, -) -from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.normal.int8kv_flash_decoding_stage2 import ( - flash_decode_stage2, -) -from lightllm.common.triton_utils import autotuner as autotuner_module -from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType, AutotuneLevel -from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len - - -@pytest.fixture(autouse=True) -def autotune_environment(monkeypatch): - torch.manual_seed(42) - monkeypatch.delenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", raising=False) - monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", None) - monkeypatch.setattr(stage1_module, "get_triton_autotune_level", lambda: AutotuneLevel.ADAPTIVE_AUTOTUNE) - monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: AutotuneLevel.ADAPTIVE_AUTOTUNE) - get_decode_attn_autotune_seq_len.cache_clear() - yield - get_decode_attn_autotune_seq_len.cache_clear() - - -def make_inputs(device="cpu", head_dim=64, group_size=8, dtype=torch.bfloat16, token_count=17, block_num=2): - # 与内存管理器一致：K/V 和 scale 分别来自共享存储中的非连续视图。 - low = torch.randint(0, 15, (token_count, 4, head_dim // 2), device=device) - high = torch.randint(0, 15, low.shape, device=device) - kv = (low | (high << 4)).to(torch.int8) - scale = (torch.rand(token_count, 4, head_dim // group_size, device=device) * 0.1 + 0.02).to(dtype) - return dict( - q=torch.randn(2, 8, head_dim, dtype=dtype, device=device), - k=kv[:, :2], - k_scale=scale[:, :2], - v=kv[:, 2:], - v_scale=scale[:, 2:], - Req_to_tokens=torch.arange(5 * 16, dtype=torch.int32, device=device).view(5, 16) % max(1, token_count), - B_req_idx=torch.tensor([4, 2], dtype=torch.int32, device=device), - B_Seqlen=torch.tensor([3, 2], dtype=torch.int32, device=device), - max_kv_seq_len=3, - mid_out=torch.full((2, 8, block_num, head_dim), -100, dtype=dtype, device=device), - mid_out_logsumexp=torch.full((2, 8, block_num), -100, dtype=dtype, device=device), - block_seq=256, - ) - - -def rebuild(inputs): - args, kwargs = stage1_module.rebuild_inputs(**inputs) - return inspect.signature(stage1_module.int4kv_flash_decode_stage1.fn).bind(*args, **kwargs).arguments - - -@pytest.mark.parametrize("level", [AutotuneLevel.ADAPTIVE_AUTOTUNE, AutotuneLevel.FORCE_AUTOTUNE]) -@pytest.mark.parametrize("configured, expected", [(None, 32768), ("8193", 8704)]) -def test_tuning_key_uses_configured_length(monkeypatch, level, configured, expected): - if configured is not None: - monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", configured) - monkeypatch.setattr(stage1_module, "get_triton_autotune_level", lambda: level) - with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - assert stage1_module.get_run_key(torch.empty(2, 8, 64), 32768) == 2_000_000_000 + expected - - -@pytest.mark.parametrize( - "phase,level", - [ - (None, AutotuneLevel.ADAPTIVE_AUTOTUNE), - (None, AutotuneLevel.FORCE_AUTOTUNE), - (AutotuneKernelType.GENERAL, AutotuneLevel.ADAPTIVE_AUTOTUNE), - (AutotuneKernelType.DECODE_ATTENTION, AutotuneLevel.USE_AUTOTUNE_HIS_CONFIG), - ], -) -def test_lookup_uses_actual_length(monkeypatch, phase, level): - monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", "invalid") - monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", phase) - monkeypatch.setattr(stage1_module, "get_triton_autotune_level", lambda: level) - for actual, expected in [(2, 512), (511, 512), (512, 512), (513, 1024), (8193, 8704)]: - assert stage1_module.get_run_key(torch.empty(2, 8, 64), actual) == 2_000_000_000 + expected - - -@pytest.mark.parametrize("token_count", [3, 128]) -def test_rebuild_preserves_packed_cache_and_scales(monkeypatch, token_count): - monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", "17") - inputs = make_inputs(token_count=token_count) - snapshots = {name: value.clone() for name, value in inputs.items() if isinstance(value, torch.Tensor)} - rebuilt = rebuild(inputs) - torch.testing.assert_close(rebuilt["Req_to_tokens"].long(), torch.arange(34).view(2, 17) % token_count) - assert rebuilt["B_req_idx"].tolist() == [0, 1] - assert rebuilt["B_Seqlen"].tolist() == [17, 17] - assert rebuilt["max_kv_seq_len"] == 17 - assert rebuilt["block_seq"] == inputs["block_seq"] - for name in ["q", "k", "v", "k_scale", "v_scale", "mid_out", "mid_out_logsumexp"]: - assert rebuilt[name] is inputs[name] - for name, value in snapshots.items(): - torch.testing.assert_close(inputs[name], value) - - -@pytest.mark.parametrize("name", ["v", "k_scale", "v_scale"]) -def test_rebuild_rejects_inconsistent_token_capacity(name): - inputs = make_inputs() - inputs[name] = inputs[name][:-1] - with pytest.raises(AssertionError, match="same number of tokens"): - rebuild(inputs) - - -def test_rebuild_rejects_empty_cache(): - with pytest.raises(ValueError, match="non-empty KV cache"): - rebuild(make_inputs(token_count=0)) - - -def test_static_key_separates_quant_groups_without_head_or_block_counts(): - kernel = stage1_module.int4kv_flash_decode_stage1 - inputs = make_inputs() - key = kernel._static_key(**inputs) - inputs["q"] = inputs["q"][:, :4] - for name in ["k", "k_scale", "v", "v_scale"]: - inputs[name] = inputs[name][:, :1] - inputs["mid_out"] = torch.empty(2, 4, 128, 64, dtype=torch.bfloat16) - assert kernel._static_key(**inputs) == key - inputs["k_scale"] = inputs["k_scale"][:, :, :2] - assert kernel._static_key(**inputs) != key - - -def test_decode_passes_saved_length_to_stage1(monkeypatch): - inputs = make_inputs() - infer_state = InferStateInfo() - infer_state.batch_size = 2 - infer_state.max_kv_seq_len = 8193 - infer_state.b_req_idx = inputs["B_req_idx"] - infer_state.b_seq_len = inputs["B_Seqlen"] - infer_state.req_manager = SimpleNamespace(req_to_token_indexs=inputs["Req_to_tokens"]) - state = Int4kvTritonDecodeAttState( - backend=SimpleNamespace(), - _infer_state_weak_ref=weakref.ref(infer_state), - ) - state.init_state() - state.infer_state.max_kv_seq_len = 32768 - calls = [] - monkeypatch.setattr(stage1_module, "int4kv_flash_decode_stage1", lambda **kwargs: calls.append(kwargs)) - stage2 = inspect.getmodule(flash_decode_stage2) - monkeypatch.setattr(stage2, "flash_decode_stage2", lambda *args: None) - state.decode_att( - inputs["q"], - (inputs["k"], inputs["k_scale"]), - (inputs["v"], inputs["v_scale"]), - alloc_func=lambda shape, dtype, device: torch.empty(shape, dtype=dtype), - ) - assert calls[0]["max_kv_seq_len"] == 8193 - assert calls[0]["block_seq"] == 256 - - -def reference_attention(inputs): - def dequant(packed, scale): - packed = packed.to(torch.uint8) - unpacked = torch.stack((packed & 15, packed >> 4), dim=-1).flatten(-2).float() - 7 - group_size = unpacked.shape[-1] // scale.shape[-1] - # 与算子一致，以 scale 的 dtype 产生解量化 K/V，再用 FP32 SDPA 验证 attention 本身。 - return (unpacked * scale.float().repeat_interleave(group_size, dim=-1)).to(scale.dtype).float() - - k, v = dequant(inputs["k"], inputs["k_scale"]), dequant(inputs["v"], inputs["v_scale"]) - outputs = [] - for row, length in enumerate(inputs["B_Seqlen"].tolist()): - indices = inputs["Req_to_tokens"][inputs["B_req_idx"][row], :length].long() - with sdpa_kernel(SDPBackend.MATH): - outputs.append( - F.scaled_dot_product_attention( - inputs["q"][row].float().unsqueeze(1), - k[indices].transpose(0, 1), - v[indices].transpose(0, 1), - enable_gqa=True, - ).squeeze(1) - ) - return torch.stack(outputs) - - -def reduce_output(inputs): - output = torch.empty_like(inputs["q"]) - flash_decode_stage2(inputs["mid_out"], inputs["mid_out_logsumexp"], inputs["B_Seqlen"], output, inputs["block_seq"]) - return output - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("kv_len", [257, 8192, 16384]) -@pytest.mark.parametrize( - "head_dim,group_size,dtype", [(64, 8, torch.bfloat16), (128, 32, torch.float16), (256, 8, torch.bfloat16)] -) -def test_long_kv_matches_dequantized_fp32_reference(monkeypatch, kv_len, head_dim, group_size, dtype): - monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", str(kv_len)) - inputs = rebuild(make_inputs("cuda", head_dim, group_size, dtype, token_count=2 * kv_len)) - inputs["B_Seqlen"][1] -= 1 - reference = reference_attention(inputs) - # 两个中间分块迫使长请求循环处理多个 BLOCK_SEQ，同时验证末尾 mask。 - for block_n in [16, 32, 64, 128]: - stage1_module.int4kv_flash_decode_stage1.fn( - **inputs, run_config={"BLOCK_N": block_n, "num_warps": 4, "num_stages": 2} - ) - output = reduce_output(inputs) - assert torch.isfinite(output).all() - torch.testing.assert_close(output.float(), reference, atol=2e-3, rtol=2e-2) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("kv_len", [8193, 16384]) -def test_full_autotune_and_graph_reuse(tmp_path, monkeypatch, kv_len): - monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", str(kv_len)) - monkeypatch.setattr(autotuner_module.dist, "is_initialized", lambda: False) - kernel = stage1_module.int4kv_flash_decode_stage1 - monkeypatch.setattr(kernel, "_cache_dir", str(tmp_path), raising=False) - monkeypatch.setattr(kernel, "cached_configs", {}) - monkeypatch.setattr(kernel, "fast_match_configs", collections.defaultdict(dict)) - monkeypatch.setattr(kernel, "warmuped_configs_set", set()) - assert kernel.mutates_args == [] - inputs = make_inputs("cuda", token_count=2 * kv_len, block_num=128) - inputs["Req_to_tokens"] = torch.arange(5 * kv_len, dtype=torch.int32, device="cuda").view(5, kv_len) % (2 * kv_len) - snapshots = {name: value.clone() for name, value in inputs.items() if isinstance(value, torch.Tensor)} - reference = reference_attention(inputs) - rebuild_func, benchmark = kernel.rebuild_input_func, kernel._bench - rebuild_count, timings = 0, [] - - def checked_rebuild(*args, **kwargs): - nonlocal rebuild_count - rebuild_count += 1 - return rebuild_func(*args, **kwargs) - - def checked_bench(*args, **kwargs): - rebuilt = inspect.signature(kernel.fn).bind(*args, **kwargs).arguments - assert rebuilt["B_Seqlen"].tolist() == [kv_len, kv_len] - assert rebuilt["Req_to_tokens"].shape == (2, kv_len) - assert rebuilt["max_kv_seq_len"] == kv_len - for name in ["q", "k", "v", "k_scale", "v_scale", "mid_out", "mid_out_logsumexp"]: - assert rebuilt[name] is inputs[name] - elapsed = benchmark(*args, **kwargs) - timings.append(elapsed) - return elapsed - - monkeypatch.setattr(kernel, "rebuild_input_func", checked_rebuild) - monkeypatch.setattr(kernel, "_bench", checked_bench) - with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - kernel(**inputs) - assert rebuild_count == 1 and len(timings) == 48 and any(math.isfinite(t) for t in timings) - run_key = str(2_000_000_000 + (kv_len + 511) // 512 * 512) - assert all(list(configs) == [run_key] for configs in kernel.cached_configs.values()) - assert len(list(tmp_path.glob("*.json"))) == 1 - assert run_key in json.loads(next(tmp_path.glob("*.json")).read_text()) - torch.testing.assert_close(reduce_output(inputs).float(), reference, atol=2e-3, rtol=2e-2) - - # 清空内存配置后从文件重载，FORCE 下仍不能跨层重复搜索。 - monkeypatch.setattr(kernel, "cached_configs", {}) - monkeypatch.setattr(kernel, "fast_match_configs", collections.defaultdict(dict)) - monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: AutotuneLevel.FORCE_AUTOTUNE) - monkeypatch.setattr(stage1_module, "get_triton_autotune_level", lambda: AutotuneLevel.FORCE_AUTOTUNE) - with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - kernel(**inputs) - assert rebuild_count == 1 and len(timings) == 48 - - output = torch.empty_like(inputs["q"]) - - def decode(): - kernel(**inputs) - flash_decode_stage2( - inputs["mid_out"], inputs["mid_out_logsumexp"], inputs["B_Seqlen"], output, inputs["block_seq"] - ) - - decode() - torch.cuda.synchronize() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - decode() - graph.replay() - torch.cuda.synchronize() - torch.testing.assert_close(output.float(), reference, atol=2e-3, rtol=2e-2) - assert rebuild_count == 1 and len(timings) == 48 - for name, snapshot in snapshots.items(): - if name not in ["mid_out", "mid_out_logsumexp"]: - torch.testing.assert_close(inputs[name], snapshot) - - # 捕获后更新为真正的长请求，验证 Graph 读取原始输入缓冲区，stage2 仍按 BLOCK_SEQ 归约。 - inputs["B_Seqlen"].fill_(kv_len) - inputs["max_kv_seq_len"] = kv_len - graph.replay() - torch.cuda.synchronize() - torch.testing.assert_close(output.float(), reference_attention(inputs), atol=2e-3, rtol=2e-2) - assert rebuild_count == 1 and len(timings) == 48 diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse.py deleted file mode 100644 index e1b6bc7758..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse.py +++ /dev/null @@ -1,139 +0,0 @@ -import pytest -from types import SimpleNamespace - -import torch - - -def alloc_tensor_func(shape, dtype, device): - """兼容的 tensor 分配函数""" - return torch.empty(shape, dtype=dtype, device=device) - - -class MockReqManager: - """Mock request manager for testing""" - - def __init__(self, req_to_token_indexs): - self.req_to_token_indexs = req_to_token_indexs - - -class MockInferState: - """Mock infer state for testing""" - - def __init__( - self, - batch_size, - max_kv_seq_len, - req_to_tokens, - b_req_idx, - b_seq_len, - ): - self.batch_size = batch_size - self.max_kv_seq_len = max_kv_seq_len - self.req_manager = MockReqManager(req_to_tokens) - self.b_req_idx = b_req_idx - self.b_seq_len = b_seq_len - - -# @pytest.mark.parametrize("shared_seq_len", [512]) -@pytest.mark.parametrize("shared_seq_len", [0, 77, 256, 311, 512, 550]) -@pytest.mark.parametrize("batch_size", list(range(6, 121, 6))) -def test_token_decode_attention_flash_decoding_diverse_matches_normal_decode(shared_seq_len, batch_size): - """ - diverse 与 normal 均为仓库内 Triton 实现，应数值一致（无外部 CUDA extension）。 - diverse：int8kv_flash_decoding_diverse；对照：int8kv/normal token_decode_attention_flash_decoding。 - """ - - from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.int8kv_flash_decoding_diverse import ( - token_decode_attention_flash_decoding as diverse_attention, - ) - from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.normal import ( - token_decode_attention_flash_decoding as normal_decode, - ) - - num_heads = 32 - kv_head_num = 8 - mark_shared_group_size = 3 - seq_len = 3547 - head_dim = 128 - quant_group_size = 8 - max_len_in_batch = 8192 - test_dtype = torch.bfloat16 - - # 创建测试数据 - kv_shape = (batch_size * max_len_in_batch, kv_head_num, head_dim) - kv_scale_shape = (batch_size * max_len_in_batch, kv_head_num, head_dim // quant_group_size) - - q = torch.randn(size=(batch_size, num_heads, head_dim), dtype=test_dtype, device="cuda") - - # 生成 cache_k 和 cache_v，使得每 mark_shared_group_size 个 batch 共享相同的 cache - - cache_k = torch.randint(low=-100, high=100, size=kv_shape, dtype=torch.int8, device="cuda") - cache_k_scale = torch.ones(size=kv_scale_shape, dtype=test_dtype, device="cuda") / 100.0 - cache_v = torch.randint(low=-100, high=100, size=kv_shape, dtype=torch.int8, device="cuda") - cache_v_scale = torch.ones(size=kv_scale_shape, dtype=test_dtype, device="cuda") / 100.0 - - req_to_tokens = torch.arange(0, max_len_in_batch * batch_size, dtype=torch.int32, device="cuda").view( - batch_size, max_len_in_batch - ) - for i in range(batch_size): - if i % mark_shared_group_size != 0: - req_to_tokens[i, :shared_seq_len] = req_to_tokens[i - 1, :shared_seq_len] - - b_req_idx = torch.arange(batch_size, dtype=torch.int32, device="cuda") - b_seq_len = torch.full((batch_size,), seq_len, dtype=torch.int32, device="cuda") - b_shared_seq_len = torch.full((batch_size,), shared_seq_len, dtype=torch.int32, device="cuda") - b_mark_shared_group = torch.zeros((batch_size,), dtype=torch.int32, device="cuda") - b_mark_shared_group[mark_shared_group_size - 1 :: mark_shared_group_size] = mark_shared_group_size - - # 标准 int8 decode（单路径 Triton） - baseline_infer_state = MockInferState( - batch_size=batch_size, - max_kv_seq_len=max_len_in_batch, - req_to_tokens=req_to_tokens, - b_req_idx=b_req_idx, - b_seq_len=b_seq_len, - ) - - # diverse：多流 + 共享前缀（Triton） - diverse_infer_state = MockInferState( - batch_size=batch_size, - max_kv_seq_len=max_len_in_batch, - req_to_tokens=req_to_tokens, - b_req_idx=b_req_idx, - b_seq_len=b_seq_len, - ) - diverse_infer_state.decode_att_state = SimpleNamespace( - b_shared_seq_len=b_shared_seq_len, - b_mark_shared_group=b_mark_shared_group, - ) - - # 运行 baseline - normal_out = normal_decode( - q=q.clone(), - infer_state=baseline_infer_state, - cache_k=cache_k, - cache_k_scale=cache_k_scale, - cache_v=cache_v, - cache_v_scale=cache_v_scale, - max_len_in_batch=seq_len, - alloc_tensor_func=alloc_tensor_func, - ) - # 运行 diverse 版本 - diverse_out = diverse_attention( - q=q.clone(), - infer_state=diverse_infer_state, - cache_k=cache_k, - cache_k_scale=cache_k_scale, - cache_v=cache_v, - cache_v_scale=cache_v_scale, - alloc_tensor_func=alloc_tensor_func, - ) - - print(f"\nshared_seq_len={shared_seq_len}\nbatch_size={batch_size}") - print(f"normal_out: {normal_out[0, 0, :4]}") - print(f"diverse_out: {diverse_out[0, 0, :4]}") - print(f"max diff: {(normal_out - diverse_out).abs().max()}") - - assert torch.allclose( - normal_out, diverse_out, atol=1e-2, rtol=1e-2 - ), f"diverse vs normal decode mismatch for shared_seq_len={shared_seq_len}" diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage1.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage1.py deleted file mode 100644 index f3cb8de463..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage1.py +++ /dev/null @@ -1,188 +0,0 @@ -import pytest -import torch -from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.int8kv_flash_decoding_diverse_stage1 import ( - flash_decode_stage1, -) - - -def create_tensors( - batch_size=4, - num_heads=4, - kv_head_num=1, - seq_len=256, - max_len_in_batch=8192, - max_batch_group_size=4, - kv_len=None, - req_to_tokens_len=None, -): - head_dim = 128 - block_seq = 256 - quant_group_size = 8 - - test_dtype = torch.bfloat16 - - kv_len = max_len_in_batch if kv_len is None else kv_len - req_to_tokens_len = max_len_in_batch if req_to_tokens_len is None else req_to_tokens_len - - kv_shape = (batch_size * kv_len, kv_head_num, head_dim) - kv_scale_shape = (batch_size * kv_len, kv_head_num, head_dim // quant_group_size) - - q = torch.randn(size=(batch_size, num_heads, head_dim), dtype=test_dtype, device="cuda") - k = torch.randint(low=-100, high=100, size=kv_shape, dtype=torch.int8, device="cuda") - k_scale = torch.ones(size=kv_scale_shape, dtype=test_dtype, device="cuda") - v = torch.randint(low=-100, high=100, size=kv_shape, dtype=torch.int8, device="cuda") - v_scale = torch.ones(size=kv_scale_shape, dtype=test_dtype, device="cuda") - Req_to_tokens = torch.arange(0, req_to_tokens_len * batch_size, dtype=torch.int32, device="cuda").view( - batch_size, req_to_tokens_len - ) - B_req_idx = torch.arange(batch_size, dtype=torch.int32, device="cuda") - b_shared_seq_len = torch.full((batch_size,), seq_len, dtype=torch.int32, device="cuda") - b_mark_shared_group = torch.ones(batch_size, dtype=torch.int32, device="cuda") - mid_out = torch.zeros( - size=(batch_size, num_heads, (max_len_in_batch // block_seq) + 2, head_dim), dtype=q.dtype, device="cuda" - ) - mid_out_logsumexp = torch.zeros( - size=(batch_size, num_heads, (max_len_in_batch // block_seq) + 2), dtype=q.dtype, device="cuda" - ) - - return { - "q": q, - "k": k, - "k_scale": k_scale, - "v": v, - "v_scale": v_scale, - "Req_to_tokens": Req_to_tokens, - "B_req_idx": B_req_idx, - "b_shared_seq_len": b_shared_seq_len, - "b_mark_shared_group": b_mark_shared_group, - "max_len_in_batch": max_len_in_batch, - "mid_out": mid_out, - "mid_out_logsumexp": mid_out_logsumexp, - "block_seq": block_seq, - "max_batch_group_size": max_batch_group_size, - } - - -@pytest.fixture -def setup_tensors(): - return create_tensors() - - -def test_flash_decode_stage1_execution(setup_tensors): - flash_decode_stage1( - q=setup_tensors["q"], - k=setup_tensors["k"], - k_scale=setup_tensors["k_scale"], - v=setup_tensors["v"], - v_scale=setup_tensors["v_scale"], - Req_to_tokens=setup_tensors["Req_to_tokens"], - B_req_idx=setup_tensors["B_req_idx"], - b_shared_seq_len=setup_tensors["b_shared_seq_len"], - b_mark_shared_group=setup_tensors["b_mark_shared_group"], - max_len_in_batch=setup_tensors["max_len_in_batch"], - mid_out=setup_tensors["mid_out"], - mid_out_logsumexp=setup_tensors["mid_out_logsumexp"], - block_seq=setup_tensors["block_seq"], - max_batch_group_size=setup_tensors["max_batch_group_size"], - ) - - q = setup_tensors["q"] - k = setup_tensors["k"] - v = setup_tensors["v"] - true_mid_out = torch.zeros_like(setup_tensors["mid_out"]) - true_mid_out_logsumexp = torch.zeros_like(setup_tensors["mid_out_logsumexp"]) - new_q = q - new_k = k.to(q.dtype) - new_v = v.to(q.dtype) - - from lightllm.common.basemodel.triton_kernel.att.decode_att.gqa.flash_decoding.gqa_flash_decoding_stage1 import ( - flash_decode_stage1 as gqa_flash_decode_stage1, - ) - - gqa_flash_decode_stage1( - q=new_q, - k=new_k, - v=new_v, - Req_to_tokens=setup_tensors["Req_to_tokens"], - B_req_idx=setup_tensors["B_req_idx"], - B_Seqlen=setup_tensors["b_shared_seq_len"], - max_len_in_batch=setup_tensors["max_len_in_batch"], - mid_out=true_mid_out, - mid_out_logsumexp=true_mid_out_logsumexp, - block_seq=setup_tensors["block_seq"], - ) - print(setup_tensors["mid_out"][0:4, 0, 0, 0], true_mid_out[0:4, 0, 0, 0]) - assert torch.allclose( - setup_tensors["mid_out"][0:4, 0, 0, 0], true_mid_out[0:4, 0, 0, 0], atol=1e-2 - ), "Mid output does not match expected values" - assert torch.allclose( - setup_tensors["mid_out_logsumexp"], true_mid_out_logsumexp, atol=1e-2 - ), "LogSumExp output does not match expected values" - - -def autotune_and_benchmark(): - import triton - - batch_sizes = [8, 16, 32, 64] - seq_lens = [1024, 2048, 4096] - - results = [] - for batch in batch_sizes: - for seq in seq_lens: - # Clear GPU cache to reduce CUDA Graph capture failures. - torch.cuda.empty_cache() - - setup_tensors = create_tensors( - batch_size=batch, - num_heads=32, - kv_head_num=8, - seq_len=seq, - max_len_in_batch=8192, - max_batch_group_size=8, - kv_len=seq, - req_to_tokens_len=seq, - ) - - def fn_triton(st=setup_tensors): - return flash_decode_stage1( - q=st["q"], - k=st["k"], - k_scale=st["k_scale"], - v=st["v"], - v_scale=st["v_scale"], - Req_to_tokens=st["Req_to_tokens"], - B_req_idx=st["B_req_idx"], - b_shared_seq_len=st["b_shared_seq_len"], - b_mark_shared_group=st["b_mark_shared_group"], - max_len_in_batch=st["max_len_in_batch"], - mid_out=st["mid_out"], - mid_out_logsumexp=st["mid_out_logsumexp"], - block_seq=st["block_seq"], - max_batch_group_size=st["max_batch_group_size"], - ) - - ms_triton = triton.testing.do_bench_cudagraph(fn_triton, rep=100) - - results.append( - { - "batch_size": batch, - "seq_len": seq, - "triton_ms": ms_triton, - } - ) - print(results[-1]) - - del setup_tensors - - print(f"\n{'='*80}") - print("SUMMARY - Performance Comparison") - print(f"{'='*80}") - print(f"{'batch_size':<8} {'seq_len':<12} {'triton_ms':<12}") - print(f"{'-'*80}") - for r in results: - print(f"{r['batch_size']:<8} {r['seq_len']:<12} {r['triton_ms']:<12.3f}") - print(f"{'='*80}") - - -if __name__ == "__main__": - autotune_and_benchmark() diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage2.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage2.py deleted file mode 100644 index 985a33289f..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage2.py +++ /dev/null @@ -1,192 +0,0 @@ -import pytest -import torch -from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.int8kv_flash_decoding_diverse_stage2 import ( - flash_decode_stage2, -) - - -def create_tensors( - shared_seq_len, - batch_size=4, - seq_len=256, - max_len_in_batch=8192, - max_batch_group_size=4, - kv_len=None, - req_to_tokens_len=None, -): - num_heads = 32 - kv_head_num = 8 - head_dim = 128 - block_seq = 256 - quant_group_size = 8 - - test_dtype = torch.bfloat16 - - kv_len = max_len_in_batch if kv_len is None else kv_len - req_to_tokens_len = max_len_in_batch if req_to_tokens_len is None else req_to_tokens_len - - kv_shape = (batch_size * kv_len, kv_head_num, head_dim) - kv_scale_shape = (batch_size * kv_len, kv_head_num, head_dim // quant_group_size) - - q = torch.randn(size=(batch_size, num_heads, head_dim), dtype=test_dtype, device="cuda") - k = torch.randint(low=-100, high=100, size=kv_shape, dtype=torch.int8, device="cuda") - k_scale = torch.ones(size=kv_scale_shape, dtype=test_dtype, device="cuda") - v = torch.randint(low=-100, high=100, size=kv_shape, dtype=torch.int8, device="cuda") - v_scale = torch.ones(size=kv_scale_shape, dtype=test_dtype, device="cuda") - Req_to_tokens = torch.arange(0, req_to_tokens_len * batch_size, dtype=torch.int32, device="cuda").view( - batch_size, req_to_tokens_len - ) - B_req_idx = torch.arange(batch_size, dtype=torch.int32, device="cuda") - b_seq_len = torch.full((batch_size,), seq_len, dtype=torch.int32, device="cuda") - b_shared_seq_len = torch.full((batch_size,), shared_seq_len, dtype=torch.int32, device="cuda") - b_mark_shared_group = torch.ones(batch_size, dtype=torch.int32, device="cuda") - mid_out = torch.zeros( - size=(batch_size, num_heads, (max_len_in_batch // block_seq) + 2, head_dim), dtype=q.dtype, device="cuda" - ) - mid_out_logsumexp = torch.zeros( - size=(batch_size, num_heads, (max_len_in_batch // block_seq) + 2), dtype=q.dtype, device="cuda" - ) - - return { - "q": q, - "k": k, - "k_scale": k_scale, - "v": v, - "v_scale": v_scale, - "Req_to_tokens": Req_to_tokens, - "B_req_idx": B_req_idx, - "b_seq_len": b_seq_len, - "b_shared_seq_len": b_shared_seq_len, - "b_mark_shared_group": b_mark_shared_group, - "max_len_in_batch": max_len_in_batch, - "mid_out": mid_out, - "mid_out_logsumexp": mid_out_logsumexp, - "block_seq": block_seq, - "max_batch_group_size": max_batch_group_size, - "head_dim": head_dim, - } - - -@pytest.mark.parametrize("shared_seq_len", [0, 47, 77, 128, 200, 255]) -def test_flash_decode_stage2_execution(shared_seq_len): - setup_tensors = create_tensors(shared_seq_len) - - flash_decode_stage2( - q=setup_tensors["q"], - k=setup_tensors["k"], - k_scale=setup_tensors["k_scale"], - v=setup_tensors["v"], - v_scale=setup_tensors["v_scale"], - Req_to_tokens=setup_tensors["Req_to_tokens"], - B_req_idx=setup_tensors["B_req_idx"], - B_Seqlen=setup_tensors["b_seq_len"], - b_shared_seq_len=setup_tensors["b_shared_seq_len"], - max_len_in_batch=setup_tensors["max_len_in_batch"], - mid_out=setup_tensors["mid_out"], - mid_out_logsumexp=setup_tensors["mid_out_logsumexp"], - block_seq=setup_tensors["block_seq"], - ) - seq_block_idx = (setup_tensors["b_shared_seq_len"][0].item() + setup_tensors["block_seq"] - 1) // setup_tensors[ - "block_seq" - ] - mid_out = setup_tensors["mid_out"][:, :, seq_block_idx:, :] - mid_out_logsumexp = setup_tensors["mid_out_logsumexp"][:, :, seq_block_idx:] - - q = setup_tensors["q"] - k = setup_tensors["k"] - v = setup_tensors["v"] - true_mid_out = torch.zeros_like(mid_out) - true_mid_out_logsumexp = torch.zeros_like(mid_out_logsumexp) - new_q = q - new_k = k.to(q.dtype) - new_v = v.to(q.dtype) - - b_seq_len = setup_tensors["b_seq_len"] - setup_tensors["b_shared_seq_len"] - req_to_tokens = setup_tensors["Req_to_tokens"][:, setup_tensors["b_shared_seq_len"][0].item() :] - - from lightllm.common.basemodel.triton_kernel.att.decode_att.gqa.flash_decoding.gqa_flash_decoding_stage1 import ( - flash_decode_stage1 as gqa_flash_decode_stage1, - ) - - gqa_flash_decode_stage1( - q=new_q, - k=new_k, - v=new_v, - Req_to_tokens=req_to_tokens, - B_req_idx=setup_tensors["B_req_idx"], - B_Seqlen=b_seq_len, - max_len_in_batch=setup_tensors["max_len_in_batch"], - mid_out=true_mid_out, - mid_out_logsumexp=true_mid_out_logsumexp, - block_seq=setup_tensors["block_seq"], - ) - print(f"\nshared_seq_len={shared_seq_len}") - print(f"mid_out: {mid_out[0:4, 0, 0, 0]}") - print(f"true_mid_out: {true_mid_out[0:4, 0, 0, 0]}") - abs_diff = (mid_out - true_mid_out).abs() - max_diff = abs_diff.max() - max_diff_idx = abs_diff.argmax() - max_diff_idx_unraveled = torch.unravel_index(max_diff_idx, abs_diff.shape) - mid_out_value = mid_out[max_diff_idx_unraveled] - true_mid_out_value = true_mid_out[max_diff_idx_unraveled] - print(f"max abs diff: {max_diff}, mid_out value: {mid_out_value}, " f"true_mid_out value: {true_mid_out_value}") - - assert torch.allclose( - mid_out[0:4, 0, 0, 0], true_mid_out[0:4, 0, 0, 0], atol=1e-2 - ), f"Mid output does not match expected values for shared_seq_len={shared_seq_len}" - assert torch.allclose( - mid_out_logsumexp, true_mid_out_logsumexp, atol=1e-2 - ), f"LogSumExp output does not match expected values for shared_seq_len={shared_seq_len}" - - -if __name__ == "__main__": - # 可选：对 Triton diverse stage2 做 cudagraph bench（仅本仓库内核，无外部 CUDA 扩展）。 - - import triton - - batch_sizes = [8, 16, 32, 64] - seq_lens = [32, 64, 128, 256] - - results = [] - for batch in batch_sizes: - for seq in seq_lens: - torch.cuda.empty_cache() - - setup_tensors = create_tensors( - shared_seq_len=0, - batch_size=batch, - seq_len=seq, - max_len_in_batch=8192, - kv_len=seq, - req_to_tokens_len=seq, - ) - st = setup_tensors - - def bench_stage2(): - flash_decode_stage2( - q=st["q"], - k=st["k"], - k_scale=st["k_scale"], - v=st["v"], - v_scale=st["v_scale"], - Req_to_tokens=st["Req_to_tokens"], - B_req_idx=st["B_req_idx"], - B_Seqlen=st["b_seq_len"], - b_shared_seq_len=st["b_shared_seq_len"], - max_len_in_batch=st["max_len_in_batch"], - mid_out=st["mid_out"], - mid_out_logsumexp=st["mid_out_logsumexp"], - block_seq=st["block_seq"], - ) - - ms = triton.testing.do_bench_cudagraph(bench_stage2, rep=100) - results.append({"batch_size": batch, "seq_len": seq, "flash_decode_stage2_ms": ms}) - print(results[-1]) - del setup_tensors - - print(f"\n{'='*80}") - print(f"{'batch_size':<10} {'seq_len':<10} {'flash_decode_stage2_ms':<22}") - print(f"{'-'*80}") - for r in results: - print(f"{r['batch_size']:<10} {r['seq_len']:<10} {r['flash_decode_stage2_ms']:<22.4f}") - print(f"{'='*80}") diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage3.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage3.py deleted file mode 100644 index c1a0ca1e58..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage3.py +++ /dev/null @@ -1,35 +0,0 @@ -import pytest -import torch -from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.int8kv_flash_decoding_diverse_stage3 import ( - flash_diverse_decode_stage3, -) - - -@pytest.mark.parametrize( - "batch, head_num, seq_len, shared_seq_len, block_seq, head_dim", - [ - (2, 4, 256, 256, 256, 128), - (1, 8, 256 * 2, 256, 256, 128), - (3, 2, 256 * 4, 256 * 2, 256, 128), - ], -) -def test_flash_diverse_decode_stage3(batch, head_num, seq_len, shared_seq_len, block_seq, head_dim): - # Initialize inputs - mid_out = torch.randn(batch, head_num, seq_len // block_seq + 2, head_dim, dtype=torch.bfloat16, device="cuda") - mid_out_logexpsum = torch.randn(batch, head_num, seq_len // block_seq + 2, dtype=torch.bfloat16, device="cuda") - B_Seqlen = torch.tensor([seq_len] * batch, dtype=torch.int32, device="cuda") - b_shared_seq_len = torch.tensor([shared_seq_len] * batch, dtype=torch.int32, device="cuda") - out = torch.zeros(batch, head_num, head_dim, dtype=torch.float32, device="cuda") - - # Call the function - flash_diverse_decode_stage3(mid_out, mid_out_logexpsum, B_Seqlen, b_shared_seq_len, out, block_seq) - - true_out = torch.zeros_like(out) - - from lightllm.common.basemodel.triton_kernel.att.decode_att.mha.flash_decoding.flash_decoding_stage2 import ( - flash_decode_stage2, - ) - - flash_decode_stage2(mid_out, mid_out_logexpsum, B_Seqlen, true_out, block_seq) - - assert torch.allclose(out, true_out, atol=1e-2) diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_normal_autotune.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_normal_autotune.py deleted file mode 100644 index 7033489f66..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_normal_autotune.py +++ /dev/null @@ -1,300 +0,0 @@ -import weakref -import collections -import inspect -import json -import math -from types import SimpleNamespace - -import pytest -import torch -import torch.nn.functional as F -from torch.nn.attention import SDPBackend, sdpa_kernel - -from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.common.basemodel.attention.triton import int8kv as state_module -from lightllm.common.basemodel.attention.triton.int8kv import Int8kvTritonDecodeAttState -from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.normal import ( - int8kv_flash_decoding_stage1 as stage1_module, - int8kv_flash_decoding as decode_module, -) -from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.normal.int8kv_flash_decoding_stage2 import ( - flash_decode_stage2, -) -from lightllm.common.triton_utils import autotuner as autotuner_module -from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType, AutotuneLevel -from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len - - -@pytest.fixture(autouse=True) -def autotune_environment(monkeypatch): - torch.manual_seed(42) - monkeypatch.setattr(state_module, "enable_diverse_mode_gqa_decode_fast_kernel", lambda: False) - monkeypatch.delenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", raising=False) - monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", None) - monkeypatch.setattr(stage1_module, "get_triton_autotune_level", lambda: AutotuneLevel.ADAPTIVE_AUTOTUNE) - monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: AutotuneLevel.ADAPTIVE_AUTOTUNE) - get_decode_attn_autotune_seq_len.cache_clear() - yield - get_decode_attn_autotune_seq_len.cache_clear() - - -def make_inputs( - device="cpu", head_dim=64, group_size=8, dtype=torch.bfloat16, token_count=17, block_num=2, scale_dtype=None -): - # 与内存管理器一致：K/V 和 scale 分别来自共享存储中的非连续视图。 - kv = torch.randint(-127, 128, (token_count, 4, head_dim), dtype=torch.int8, device=device) - scale = (torch.rand(token_count, 4, head_dim // group_size, device=device) * 0.006 + 0.002).to(scale_dtype or dtype) - return dict( - q=torch.randn(2, 8, head_dim, dtype=dtype, device=device), - k=kv[:, :2], - k_scale=scale[:, :2], - v=kv[:, 2:], - v_scale=scale[:, 2:], - Req_to_tokens=torch.arange(5 * 16, dtype=torch.int32, device=device).view(5, 16) % max(1, token_count), - B_req_idx=torch.tensor([4, 2], dtype=torch.int32, device=device), - B_seq_len=torch.tensor([3, 2], dtype=torch.int32, device=device), - max_len_in_batch=3, - mid_out=torch.full((2, 8, block_num, head_dim), -100, dtype=dtype, device=device), - mid_out_logsumexp=torch.full((2, 8, block_num), -100, dtype=torch.float32, device=device), - block_seq=256, - ) - - -def rebuild(inputs): - args, kwargs = stage1_module.rebuild_inputs(**inputs) - return inspect.signature(stage1_module.flash_decode_stage1.fn).bind(*args, **kwargs).arguments - - -@pytest.mark.parametrize("level", [AutotuneLevel.ADAPTIVE_AUTOTUNE, AutotuneLevel.FORCE_AUTOTUNE]) -@pytest.mark.parametrize("configured, expected", [(None, 32768), ("8193", 8704)]) -def test_tuning_key_uses_configured_length(monkeypatch, level, configured, expected): - if configured is not None: - monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", configured) - monkeypatch.setattr(stage1_module, "get_triton_autotune_level", lambda: level) - with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - assert stage1_module.get_run_key(torch.empty(2, 8, 64), 32768) == 2_000_000_000 + expected - - -@pytest.mark.parametrize( - "phase,level", - [ - (None, AutotuneLevel.ADAPTIVE_AUTOTUNE), - (None, AutotuneLevel.FORCE_AUTOTUNE), - (AutotuneKernelType.GENERAL, AutotuneLevel.ADAPTIVE_AUTOTUNE), - (AutotuneKernelType.DECODE_ATTENTION, AutotuneLevel.USE_AUTOTUNE_HIS_CONFIG), - ], -) -def test_lookup_uses_actual_length(monkeypatch, phase, level): - monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", "invalid") - monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", phase) - monkeypatch.setattr(stage1_module, "get_triton_autotune_level", lambda: level) - for actual, expected in [(2, 512), (511, 512), (512, 512), (513, 1024), (8193, 8704)]: - assert stage1_module.get_run_key(torch.empty(2, 8, 64), actual) == 2_000_000_000 + expected - - -@pytest.mark.parametrize("token_count", [3, 128]) -def test_rebuild_preserves_int8_cache_and_scales(monkeypatch, token_count): - monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", "17") - inputs = make_inputs(token_count=token_count) - snapshots = {name: value.clone() for name, value in inputs.items() if isinstance(value, torch.Tensor)} - rebuilt = rebuild(inputs) - torch.testing.assert_close(rebuilt["Req_to_tokens"].long(), torch.arange(34).view(2, 17) % token_count) - assert rebuilt["B_req_idx"].tolist() == [0, 1] - assert rebuilt["B_seq_len"].tolist() == [17, 17] - assert rebuilt["max_len_in_batch"] == 17 - assert rebuilt["block_seq"] == inputs["block_seq"] - for name in ["q", "k", "v", "k_scale", "v_scale", "mid_out", "mid_out_logsumexp"]: - assert rebuilt[name] is inputs[name] - for name, value in snapshots.items(): - torch.testing.assert_close(inputs[name], value) - - -@pytest.mark.parametrize("name", ["v", "k_scale", "v_scale"]) -def test_rebuild_rejects_inconsistent_token_capacity(name): - inputs = make_inputs() - inputs[name] = inputs[name][:-1] - with pytest.raises(AssertionError, match="same number of tokens"): - rebuild(inputs) - - -def test_rebuild_rejects_empty_cache(): - with pytest.raises(ValueError, match="non-empty KV cache"): - rebuild(make_inputs(token_count=0)) - - -def test_static_key_separates_quant_groups_without_head_or_block_counts(): - kernel = stage1_module.flash_decode_stage1 - inputs = make_inputs() - key = kernel._static_key(**inputs) - inputs["q"] = inputs["q"][:, :4] - for name in ["k", "k_scale", "v", "v_scale"]: - inputs[name] = inputs[name][:, :1] - inputs["mid_out"] = torch.empty(2, 4, 128, 64, dtype=torch.bfloat16) - assert kernel._static_key(**inputs) == key - inputs["k_scale"] = inputs["k_scale"][:, :, :2] - assert kernel._static_key(**inputs) != key - - -def test_decode_passes_saved_length_to_stage1(monkeypatch): - inputs = make_inputs() - infer_state = InferStateInfo() - infer_state.batch_size = 2 - infer_state.max_kv_seq_len = 8193 - infer_state.b_req_idx = inputs["B_req_idx"] - infer_state.b_seq_len = inputs["B_seq_len"] - infer_state.req_manager = SimpleNamespace(req_to_token_indexs=inputs["Req_to_tokens"]) - state = Int8kvTritonDecodeAttState( - backend=SimpleNamespace(), - _infer_state_weak_ref=weakref.ref(infer_state), - ) - state.init_state() - state.infer_state.max_kv_seq_len = 32768 - calls = [] - monkeypatch.setattr(decode_module, "flash_decode_stage1", lambda **kwargs: calls.append(kwargs)) - monkeypatch.setattr(decode_module, "flash_decode_stage2", lambda **kwargs: None) - state.decode_att( - inputs["q"], - (inputs["k"], inputs["k_scale"]), - (inputs["v"], inputs["v_scale"]), - alloc_func=lambda shape, dtype, device: torch.empty(shape, dtype=dtype), - ) - assert calls[0]["max_len_in_batch"] == 8193 - assert calls[0]["block_seq"] == 256 - - -def reference_attention(inputs): - def dequant(quantized, scale): - group_size = quantized.shape[-1] // scale.shape[-1] - # int8 直接乘分组 scale，转换为 Q 的 dtype 后再用 FP32 SDPA 验证 attention。 - return (quantized.float() * scale.float().repeat_interleave(group_size, dim=-1)).to(inputs["q"].dtype).float() - - k, v = dequant(inputs["k"], inputs["k_scale"]), dequant(inputs["v"], inputs["v_scale"]) - outputs = [] - for row, length in enumerate(inputs["B_seq_len"].tolist()): - indices = inputs["Req_to_tokens"][inputs["B_req_idx"][row], :length].long() - with sdpa_kernel(SDPBackend.MATH): - outputs.append( - F.scaled_dot_product_attention( - inputs["q"][row].float().unsqueeze(1), - k[indices].transpose(0, 1), - v[indices].transpose(0, 1), - enable_gqa=True, - ).squeeze(1) - ) - return torch.stack(outputs) - - -def reduce_output(inputs): - output = torch.empty_like(inputs["q"]) - flash_decode_stage2( - inputs["mid_out"], inputs["mid_out_logsumexp"], inputs["B_seq_len"], output, inputs["block_seq"] - ) - return output - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("kv_len", [257, 8192, 16384]) -@pytest.mark.parametrize( - "head_dim,group_size,dtype,scale_dtype", - [ - (64, 8, torch.bfloat16, torch.bfloat16), - (128, 32, torch.float16, torch.float16), - (256, 8, torch.bfloat16, torch.bfloat16), - (128, 8, torch.bfloat16, torch.float32), - ], -) -def test_long_kv_matches_dequantized_fp32_reference(monkeypatch, kv_len, head_dim, group_size, dtype, scale_dtype): - monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", str(kv_len)) - inputs = rebuild(make_inputs("cuda", head_dim, group_size, dtype, token_count=2 * kv_len, scale_dtype=scale_dtype)) - inputs["B_seq_len"][1] -= 1 - reference = reference_attention(inputs) - # 两个中间分块迫使长请求循环处理多个 BLOCK_SEQ，同时验证末尾 mask。 - for block_n in [16, 32, 64, 128]: - stage1_module.flash_decode_stage1.fn(**inputs, run_config={"BLOCK_N": block_n, "num_warps": 4, "num_stages": 2}) - output = reduce_output(inputs) - assert torch.isfinite(output).all() - torch.testing.assert_close(output.float(), reference, atol=2e-3, rtol=2e-2) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("kv_len", [8193, 16384]) -def test_full_autotune_and_graph_reuse(tmp_path, monkeypatch, kv_len): - monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", str(kv_len)) - monkeypatch.setattr(autotuner_module.dist, "is_initialized", lambda: False) - kernel = stage1_module.flash_decode_stage1 - monkeypatch.setattr(kernel, "_cache_dir", str(tmp_path), raising=False) - monkeypatch.setattr(kernel, "cached_configs", {}) - monkeypatch.setattr(kernel, "fast_match_configs", collections.defaultdict(dict)) - monkeypatch.setattr(kernel, "warmuped_configs_set", set()) - assert kernel.mutates_args == [] - inputs = make_inputs("cuda", token_count=2 * kv_len, block_num=128) - inputs["Req_to_tokens"] = torch.arange(5 * kv_len, dtype=torch.int32, device="cuda").view(5, kv_len) % (2 * kv_len) - snapshots = {name: value.clone() for name, value in inputs.items() if isinstance(value, torch.Tensor)} - reference = reference_attention(inputs) - rebuild_func, benchmark = kernel.rebuild_input_func, kernel._bench - rebuild_count, timings = 0, [] - - def checked_rebuild(*args, **kwargs): - nonlocal rebuild_count - rebuild_count += 1 - return rebuild_func(*args, **kwargs) - - def checked_bench(*args, **kwargs): - rebuilt = inspect.signature(kernel.fn).bind(*args, **kwargs).arguments - assert rebuilt["B_seq_len"].tolist() == [kv_len, kv_len] - assert rebuilt["Req_to_tokens"].shape == (2, kv_len) - assert rebuilt["max_len_in_batch"] == kv_len - for name in ["q", "k", "v", "k_scale", "v_scale", "mid_out", "mid_out_logsumexp"]: - assert rebuilt[name] is inputs[name] - elapsed = benchmark(*args, **kwargs) - timings.append(elapsed) - return elapsed - - monkeypatch.setattr(kernel, "rebuild_input_func", checked_rebuild) - monkeypatch.setattr(kernel, "_bench", checked_bench) - with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - kernel(**inputs) - assert rebuild_count == 1 and len(timings) == 48 and any(math.isfinite(t) for t in timings) - run_key = str(2_000_000_000 + (kv_len + 511) // 512 * 512) - assert all(list(configs) == [run_key] for configs in kernel.cached_configs.values()) - assert len(list(tmp_path.glob("*.json"))) == 1 - assert run_key in json.loads(next(tmp_path.glob("*.json")).read_text()) - torch.testing.assert_close(reduce_output(inputs).float(), reference, atol=2e-3, rtol=2e-2) - - # 清空内存配置后从文件重载，FORCE 下仍不能跨层重复搜索。 - monkeypatch.setattr(kernel, "cached_configs", {}) - monkeypatch.setattr(kernel, "fast_match_configs", collections.defaultdict(dict)) - monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: AutotuneLevel.FORCE_AUTOTUNE) - monkeypatch.setattr(stage1_module, "get_triton_autotune_level", lambda: AutotuneLevel.FORCE_AUTOTUNE) - with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - kernel(**inputs) - assert rebuild_count == 1 and len(timings) == 48 - - output = torch.empty_like(inputs["q"]) - - def decode(): - kernel(**inputs) - flash_decode_stage2( - inputs["mid_out"], inputs["mid_out_logsumexp"], inputs["B_seq_len"], output, inputs["block_seq"] - ) - - decode() - torch.cuda.synchronize() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - decode() - graph.replay() - torch.cuda.synchronize() - torch.testing.assert_close(output.float(), reference, atol=2e-3, rtol=2e-2) - assert rebuild_count == 1 and len(timings) == 48 - for name, snapshot in snapshots.items(): - if name not in ["mid_out", "mid_out_logsumexp"]: - torch.testing.assert_close(inputs[name], snapshot) - - # 捕获后更新为真正的长请求，验证 Graph 读取原始输入缓冲区，stage2 仍按 BLOCK_SEQ 归约。 - inputs["B_seq_len"].fill_(kv_len) - inputs["max_len_in_batch"] = kv_len - graph.replay() - torch.cuda.synchronize() - torch.testing.assert_close(output.float(), reference_attention(inputs), atol=2e-3, rtol=2e-2) - assert rebuild_count == 1 and len(timings) == 48 diff --git a/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_attention_fwd_neo.py b/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_attention_fwd_neo.py deleted file mode 100644 index ecab784529..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_attention_fwd_neo.py +++ /dev/null @@ -1,340 +0,0 @@ -"""Unit test for ``context_attention_fwd_neo``. - -Torch reference expresses the *semantics* of the attention, not the kernel's -internal block structure — it has no notion of BLOCK_N / BLOCK_M. For each -batch element we gather K/V for the whole request (prompt + new tokens) via -``req_to_token_indexs`` and apply:: - - allow[m, k] = (k <= q_pos[m]) OR (k < image_end[m]) - -i.e. normal queries are causal, and image-token queries can only see future -tokens from the same image span. If the Triton kernel disagrees with this -reference, the kernel is wrong. - -Run directly for quick debugging: - - python unit_tests/common/basemodel/triton_kernel/att/prefill_att/\ - test_context_attention_fwd_neo.py - -or via pytest: - - pytest unit_tests/common/basemodel/triton_kernel/att/prefill_att/\ - test_context_attention_fwd_neo.py -x -s -""" - -import math -import pytest -import torch - -from lightllm.common.basemodel.triton_kernel.att.prefill_att.context_attention_fwd_neo import context_attention_fwd_neo - - -def torch_reference_context_attention_neo( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - b_req_idx: torch.Tensor, - b_start_loc: torch.Tensor, - b_seq_len: torch.Tensor, - b_prompt_cache_len: torch.Tensor, - req_to_token_indexs: torch.Tensor, - b_image_token_end: torch.Tensor, -) -> torch.Tensor: - device = q.device - dtype = q.dtype - _, Hq, D = q.shape - Hk = k.shape[1] - kv_group = Hq // Hk - scale = 1.0 / math.sqrt(D) - - out = torch.empty_like(q) - - for b in range(b_seq_len.shape[0]): - req_idx = int(b_req_idx[b].item()) - total = int(b_seq_len[b].item()) - prompt = int(b_prompt_cache_len[b].item()) - new = total - prompt - if new <= 0: - continue - - q_start = int(b_start_loc[b].item()) - q_blk = q[q_start : q_start + new] # [M, Hq, D] - q_image_end = b_image_token_end[q_start : q_start + new].to(torch.int64) - - token_locs = req_to_token_indexs[req_idx, :total].to(torch.int64) - k_blk = k[token_locs] # [total, Hk, D] - v_blk = v[token_locs] - - q_pos = torch.arange(prompt, total, device=device, dtype=torch.int64) # [M] - k_pos = torch.arange(0, total, device=device, dtype=torch.int64) # [total] - causal = k_pos[None, :] <= q_pos[:, None] - allow = causal | (k_pos[None, :] < q_image_end[:, None]) - - out_blk = torch.empty_like(q_blk) - for h in range(Hq): - h_k = h // kv_group - q_h = q_blk[:, h, :].to(torch.float32) - k_h = k_blk[:, h_k, :].to(torch.float32) - v_h = v_blk[:, h_k, :].to(torch.float32) - - scores = (q_h @ k_h.transpose(0, 1)) * scale - scores = torch.where(allow, scores, torch.full_like(scores, -1.0e8)) - probs = torch.softmax(scores, dim=-1) - out_h = (probs @ v_h).to(dtype) - out_blk[:, h, :] = out_h - - out[q_start : q_start + new] = out_blk - - return out - - -def _build_inputs( - batch: int, - Hq: int, - Hk: int, - D: int, - dtype: torch.dtype, - device: str, - max_new: int = 256, - max_prompt: int = 512, - image_prob: float = 0.7, - num_image_spans_max: int = 3, - image_span_len_max: int = 24, - kv_pool_slack: int = 4096, - seed: int = 0, -): - g = torch.Generator(device="cpu").manual_seed(seed) - - new_lens = torch.randint(low=1, high=max_new + 1, size=(batch,), generator=g) - prompt_lens = torch.randint(low=0, high=max_prompt + 1, size=(batch,), generator=g) - total_lens = new_lens + prompt_lens - - sum_new = int(new_lens.sum().item()) - sum_total = int(total_lens.sum().item()) - max_total_len = int(total_lens.max().item()) - max_new_len = int(new_lens.max().item()) - - b_start_loc = torch.zeros(batch, dtype=torch.int32) - cur = 0 - for i in range(batch): - b_start_loc[i] = cur - cur += int(new_lens[i].item()) - - # Permute so batch idx != request idx: exercises the Req_to_tokens indexing. - b_req_idx = torch.randperm(batch, generator=g).to(torch.int32) - - # Global KV pool with scattered, non-contiguous slot assignment per request. - base = 1024 - kv_pool_size = base + sum_total + kv_pool_slack - pool = torch.randperm(kv_pool_size - base, generator=g)[:sum_total] + base - - req_to_token_indexs = torch.zeros((batch, max_total_len), dtype=torch.int32) - p = 0 - for r_logical, req_id in enumerate(b_req_idx.tolist()): - L = int(total_lens[r_logical].item()) - req_to_token_indexs[req_id, :L] = pool[p : p + L].to(torch.int32) - p += L - - b_image_token_end = torch.zeros(sum_new, dtype=torch.int32) - for i in range(batch): - M = int(new_lens[i].item()) - P = int(prompt_lens[i].item()) - start_pack = int(b_start_loc[i].item()) - if M < 2: - continue - if torch.rand((), generator=g).item() > image_prob: - continue - n_spans = int(torch.randint(1, num_image_spans_max + 1, (1,), generator=g).item()) - cursor = 0 - for _ in range(n_spans): - remaining = M - cursor - if remaining <= 0: - break - gap = int(torch.randint(0, remaining, (1,), generator=g).item()) - s_rel = cursor + gap - max_span_len = min(image_span_len_max, M - s_rel) - if max_span_len <= 0: - break - span_len = int(torch.randint(1, max_span_len + 1, (1,), generator=g).item()) - e_rel = s_rel + span_len - image_end = P + e_rel - b_image_token_end[start_pack + s_rel : start_pack + e_rel] = image_end - cursor = e_rel - - b_seq_len = total_lens.to(torch.int32) - b_prompt_cache_len = prompt_lens.to(torch.int32) - - q = torch.randn((sum_new, Hq, D), dtype=dtype, device=device) - k = torch.randn((kv_pool_size, Hk, D), dtype=dtype, device=device) - v = torch.randn((kv_pool_size, Hk, D), dtype=dtype, device=device) - o = torch.empty_like(q) - - return dict( - q=q, - k=k, - v=v, - o=o, - b_req_idx=b_req_idx.to(device), - b_start_loc=b_start_loc.to(device), - b_seq_len=b_seq_len.to(device), - b_prompt_cache_len=b_prompt_cache_len.to(device), - max_new_len=max_new_len, - req_to_token_indexs=req_to_token_indexs.to(device), - b_image_token_end=b_image_token_end.to(device), - new_lens=new_lens, - prompt_lens=prompt_lens, - ) - - -def _report_per_batch_error(out_triton, out_ref, new_lens, b_start_loc, image_token_end, tag=""): - print(f"\n[{tag}] per-batch error breakdown (abs / rel / cos):") - for i in range(new_lens.shape[0]): - s = int(b_start_loc[i].item()) - m = int(new_lens[i].item()) - if m == 0: - continue - a = out_triton[s : s + m].float() - b = out_ref[s : s + m].float() - abs_err = (a - b).abs().max().item() - denom = b.abs().max().item() + 1e-6 - rel_err = abs_err / denom - cos = torch.nn.functional.cosine_similarity(a.flatten(), b.flatten(), dim=0).item() - n_img = int((image_token_end[s : s + m] > 0).sum().item()) - print( - f" batch {i:02d} | M={m:4d} | image_tokens={n_img:4d} | " - f"max_abs={abs_err:.4e} | max_rel={rel_err:.4e} | cos={cos:.6f}" - ) - - -def _run_case( - batch: int, - Hq: int, - Hk: int, - D: int, - dtype: torch.dtype, - seed: int, - max_new: int, - max_prompt: int, - atol: float = 5e-2, - rtol: float = 5e-2, - cos_threshold: float = 0.99, - verbose: bool = True, -): - assert Hq % Hk == 0 - device = "cuda" - - inputs = _build_inputs( - batch=batch, - Hq=Hq, - Hk=Hk, - D=D, - dtype=dtype, - device=device, - max_new=max_new, - max_prompt=max_prompt, - seed=seed, - ) - - context_attention_fwd_neo( - inputs["q"], - inputs["k"], - inputs["v"], - inputs["o"], - inputs["b_req_idx"], - inputs["b_start_loc"], - inputs["b_seq_len"], - inputs["b_prompt_cache_len"], - inputs["max_new_len"], - inputs["req_to_token_indexs"], - inputs["b_image_token_end"], - ) - out_triton = inputs["o"] - - out_ref = torch_reference_context_attention_neo( - inputs["q"], - inputs["k"], - inputs["v"], - inputs["b_req_idx"], - inputs["b_start_loc"], - inputs["b_seq_len"], - inputs["b_prompt_cache_len"], - inputs["req_to_token_indexs"], - inputs["b_image_token_end"], - ) - - a = out_triton.float() - b = out_ref.float() - abs_err = (a - b).abs().max().item() - denom = b.abs().max().item() + 1e-6 - rel_err = abs_err / denom - cos = torch.nn.functional.cosine_similarity(a.flatten(), b.flatten(), dim=0).item() - - n_image = int((inputs["b_image_token_end"] > 0).sum().item()) - n_tokens = int(inputs["b_image_token_end"].numel()) - if verbose: - print( - f"\ncase: batch={batch} Hq={Hq} Hk={Hk} D={D} dtype={dtype} " - f"seed={seed} image_tokens={n_image}/{n_tokens}" - ) - print( - f" global: max_abs={abs_err:.4e} max_rel={rel_err:.4e} cos={cos:.6f} " - f"(allclose atol={atol}, rtol={rtol}? " - f"{torch.allclose(a, b, atol=atol, rtol=rtol)})" - ) - _report_per_batch_error( - out_triton, - out_ref, - inputs["new_lens"], - inputs["b_start_loc"], - inputs["b_image_token_end"], - tag=f"seed={seed}", - ) - - return abs_err, rel_err, cos - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="no CUDA") -@pytest.mark.parametrize( - "batch,Hq,Hk,D,dtype,seed,max_new,max_prompt", - [ - (4, 8, 2, 128, torch.bfloat16, 0, 128, 256), - (4, 8, 2, 128, torch.bfloat16, 1, 256, 512), - (8, 16, 4, 128, torch.bfloat16, 2, 256, 512), - (16, 28, 4, 128, torch.bfloat16, 3, 128, 256), - (4, 8, 2, 128, torch.float16, 4, 256, 512), - (4, 8, 8, 64, torch.bfloat16, 5, 128, 256), - (3, 8, 2, 128, torch.bfloat16, 6, 8, 1024), - ], -) -def test_context_attention_fwd_neo(batch, Hq, Hk, D, dtype, seed, max_new, max_prompt): - abs_err, rel_err, cos = _run_case( - batch=batch, - Hq=Hq, - Hk=Hk, - D=D, - dtype=dtype, - seed=seed, - max_new=max_new, - max_prompt=max_prompt, - verbose=True, - ) - assert cos > 0.99, f"cosine similarity too low: {cos}" - assert rel_err < 5e-2, f"max relative error too large: {rel_err}" - - -if __name__ == "__main__": - if not torch.cuda.is_available(): - print("No CUDA available.") - raise SystemExit(0) - - torch.manual_seed(0) - - cases = [ - dict(batch=4, Hq=8, Hk=2, D=128, dtype=torch.bfloat16, seed=0, max_new=128, max_prompt=256), - dict(batch=8, Hq=16, Hk=4, D=128, dtype=torch.bfloat16, seed=1, max_new=256, max_prompt=512), - dict(batch=16, Hq=28, Hk=4, D=128, dtype=torch.bfloat16, seed=2, max_new=128, max_prompt=256), - dict(batch=4, Hq=8, Hk=2, D=128, dtype=torch.float16, seed=3, max_new=256, max_prompt=512), - ] - - for cfg in cases: - _run_case(**cfg, verbose=True) diff --git a/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_flashattention_nopad.py b/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_flashattention_nopad.py deleted file mode 100644 index d1a53f873f..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_flashattention_nopad.py +++ /dev/null @@ -1,104 +0,0 @@ -import torch -import time -import pytest -from lightllm.common.basemodel.triton_kernel.att.prefill_att.context_flashattention_nopad import ( - context_attention_fwd_contiguous_kv, -) -from lightllm.utils.log_utils import init_logger - -logger = init_logger(__name__) - - -def torch_context_attention_fwd2(q, k, v, o, b_start_loc, b_kv_start_loc, b_seq_len, b_prompt_cache_len): - batch = b_start_loc.shape[0] - - for i in range(batch): - start_loc = b_start_loc[i] - kv_start_loc = b_kv_start_loc[i] - seq_len = b_seq_len[i] - prompt_cache_len = b_prompt_cache_len[i] - cur_q = q[start_loc : start_loc + seq_len - prompt_cache_len, :, :] - cur_q = cur_q.clone().to(torch.float32) - cur_k = k[kv_start_loc : (kv_start_loc + seq_len), :, :] - cur_k = cur_k.clone().to(torch.float32) - - cur_v = v[kv_start_loc : (kv_start_loc + seq_len), :, :] - cur_v = cur_v.clone().to(torch.float32) - - dk = cur_q.shape[-1] - cur_q = cur_q.permute(1, 0, 2) - cur_k = cur_k.permute(1, 2, 0) - cur_v = cur_v.permute(1, 0, 2) - dk = cur_q.shape[-1] - - p = torch.matmul(cur_q, cur_k) / torch.sqrt(torch.tensor(dk, dtype=torch.float32)) - - q_index = (torch.arange(cur_q.shape[1]).to(p.device) + prompt_cache_len).view(-1, 1) - k_index = torch.arange(seq_len).to(p.device).view(1, -1) - - p[:, (q_index < k_index)] = float("-inf") - - s = torch.nn.functional.softmax(p, dim=-1) - - o[start_loc : start_loc + seq_len - prompt_cache_len, :, :] = torch.matmul(s, cur_v).transpose(0, 1) - - -@pytest.mark.parametrize( - "B, H, N_CTX, D_HEAD, prompt_cache_len", - [ - (b, H, N_CTX, D_HEAD, prompt_cache_len) - for b in [1, 2, 4] - for H in [1, 8] - for N_CTX in [3, 10, 1024] - for D_HEAD in [64, 128] - for prompt_cache_len in [0, 56, 200] - ], -) -def test_context_attention_fwd_contiguous_kv(B, H, N_CTX, D_HEAD, prompt_cache_len): - dtype = torch.float16 - prompt_cache_len = 0 - if prompt_cache_len >= N_CTX - 1: - return - - q = torch.empty((B * (N_CTX - prompt_cache_len), H, D_HEAD), dtype=dtype, device="cuda").normal_(mean=0.1, std=0.2) - kv = torch.empty((B * N_CTX, 2 * H, D_HEAD), dtype=dtype, device="cuda").normal_(mean=0.4, std=0.2) - k = kv[:, :H, :] - v = kv[:, H:, :] - - o = torch.empty((B * (N_CTX - prompt_cache_len), H, D_HEAD), dtype=dtype, device="cuda") - torch_o = torch.empty((B * (N_CTX - prompt_cache_len), H, D_HEAD), dtype=dtype, device="cuda") - - max_q_input_len = N_CTX - prompt_cache_len - - b_seq_len = torch.ones((B,), dtype=torch.int32, device="cuda") - b_start_loc = torch.zeros((B,), dtype=torch.int32, device="cuda") - b_seq_len = torch.ones((B,), dtype=torch.int32, device="cuda") - b_prompt_cache_len = torch.zeros(B, dtype=torch.int32, device="cuda") - - for i in range(B): - b_seq_len[i] = N_CTX - if i != 0: - b_start_loc[i] = b_start_loc[i - 1] + N_CTX - prompt_cache_len - b_prompt_cache_len[i] = prompt_cache_len - - b_kv_start_loc = torch.cumsum(b_seq_len, dim=0, dtype=torch.int32) - b_seq_len - torch_context_attention_fwd2(q, k, v, torch_o, b_start_loc, b_kv_start_loc, b_seq_len, b_prompt_cache_len) - context_attention_fwd_contiguous_kv( - q=q, - k=k, - v=v, - o=o, - b_start_loc=b_start_loc, - b_kv_start_loc=b_kv_start_loc, - b_seq_len=b_seq_len, - max_q_input_len=max_q_input_len, - b_prompt_cache_len=b_prompt_cache_len, - ) - - assert torch.allclose(torch_o, o, atol=1e-2, rtol=0) - cos = torch.nn.CosineSimilarity(0) - assert cos(o.flatten().float(), torch_o.flatten().float()) > 0.99 - - -if __name__ == "__main__": - pytest.main() diff --git a/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_flashattention_nopad1.py b/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_flashattention_nopad1.py deleted file mode 100644 index 541594306d..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_flashattention_nopad1.py +++ /dev/null @@ -1,196 +0,0 @@ -import torch -import time -import pytest -import numpy as np -import torch.nn.functional as F -import flashinfer -from lightllm.utils.log_utils import init_logger -from lightllm.common.basemodel.triton_kernel.att.prefill_att.context_flashattention_nopad import ( - context_attention_fwd, - context_attention_fwd_no_prompt_cache, -) -from lightllm.models.llama.infer_struct import LlamaInferStateInfo - -logger = init_logger(__name__) - -seed = 42 -torch.manual_seed(seed) - -if torch.cuda.is_available(): - torch.cuda.manual_seed(seed) - torch.cuda.manual_seed_all(seed) - - -@pytest.mark.parametrize( - "batch, seqlen, q_heads, kv_heads, head_dim", - [ - (a, b, c, d, e) - for a in [1, 16, 32, 128, 512] - for b in [16, 32, 512, 1024] - for c in [28] - for d in [4] - for e in [128] - ], -) -def test_context_attention_fwd(batch, seqlen, q_heads, kv_heads, head_dim): - Z, N_CTX, Q_HEADS, KV_HEADS, HEAD_DIM = batch, seqlen, q_heads, kv_heads, head_dim - dtype = torch.bfloat16 - kv = torch.randn((Z * N_CTX, 2 * KV_HEADS, HEAD_DIM), dtype=dtype, device="cuda") - - max_input_len = Z * N_CTX - req_to_token_indexs = torch.randperm(max_input_len, dtype=torch.int32).cuda().view(Z, N_CTX) - b_seq_len = torch.ones((Z,), dtype=torch.int32, device="cuda") * N_CTX - b_ready_cache_len = torch.zeros_like(b_seq_len, dtype=torch.int32, device="cuda") - b_ready_cache_len = torch.randint_like(b_seq_len, high=N_CTX - 1, dtype=torch.int32, device="cuda") - b_req_idx = torch.randperm(Z, dtype=torch.int32).cuda() - q_lens = b_seq_len - b_ready_cache_len - q_start_loc = q_lens.cumsum(0) - q_lens - kv_start_loc = b_seq_len.cumsum(0) - b_seq_len - - q = torch.randn((q_lens.sum(), Q_HEADS, HEAD_DIM), dtype=dtype, device="cuda") - o = torch.zeros((q_lens.sum(), Q_HEADS, HEAD_DIM), dtype=dtype, device="cuda") - o1 = torch.zeros((q_lens.sum(), Q_HEADS, HEAD_DIM), dtype=dtype, device="cuda") - - infer_state = LlamaInferStateInfo() - infer_state.batch_size = Z - infer_state.max_q_seq_len = N_CTX - infer_state.total_token_num = Z * N_CTX - infer_state.req_manager = type("Object", (), {})() - 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 - infer_state.b_ready_cache_len = b_ready_cache_len - infer_state.b_q_start_loc = q_start_loc - - context_attention_fwd( - q, - kv[:, :KV_HEADS, :], - kv[:, KV_HEADS:, :], - o, - infer_state.b_req_idx, - infer_state.b_q_start_loc, - infer_state.b_seq_len, - infer_state.b_ready_cache_len, - infer_state.max_q_seq_len, - infer_state.req_manager.req_to_token_indexs, - ) - - batch_size = Z - head_dim = HEAD_DIM - q_heads = Q_HEADS - kv_heads = KV_HEADS - page_size = 1 - workspace_buffer = torch.empty(256 * 1024 * 1024, dtype=torch.int8).to(0) - q_starts = torch.zeros((Z + 1,)).int().cuda() - q_starts[1:] = torch.cumsum(b_seq_len - b_ready_cache_len, dim=0) - kv_starts = torch.zeros_like(q_starts) - kv_starts[1:] = torch.cumsum(b_seq_len, dim=0) - q_indptr = q_starts.int() - kv_indptr = kv_starts.int() - kv_indices = torch.arange(Z * N_CTX).cuda().int() - for b, sl, start in zip(b_req_idx, b_seq_len, kv_start_loc): - kv_indices[start : start + sl] = req_to_token_indexs[b][:sl] - kv_last_page_len_buffer = torch.empty(batch_size, device="cuda:0", dtype=torch.int32) - wrapper = flashinfer.prefill.BatchPrefillWithPagedKVCacheWrapper( - workspace_buffer, - qo_indptr_buf=q_indptr, - paged_kv_indptr_buf=kv_indptr, - paged_kv_indices_buf=kv_indices, - paged_kv_last_page_len_buf=kv_last_page_len_buffer, - ) - kv_last_page_len = torch.full((batch_size,), page_size, dtype=torch.int32) - wrapper.plan( - q_indptr, - kv_indptr, - kv_indices, - kv_last_page_len, - q_heads, - kv_heads, - head_dim, - page_size, - causal=True, - pos_encoding_mode="NONE", - logits_soft_cap=0.0, - q_data_type=q.dtype, - kv_data_type=kv.dtype, - ) - kv = kv.unsqueeze(1) - wrapper.run(q, (kv[:, :, :KV_HEADS, :], kv[:, :, KV_HEADS:, :]), out=o1, return_lse=False) - - # assert torch.allclose(o, o1, atol=1e-2, rtol=0) - cos_sim1 = F.cosine_similarity(o, o1).mean() - assert cos_sim1 == 1.0 - - -@pytest.mark.parametrize( - "batch, seqlen, q_heads, kv_heads, head_dim", - [ - (a, b, c, d, e) - for a in [ - 1, - 16, - 32, - ] - for b in [16, 32, 512, 1024] - for c in [28] - for d in [4] - for e in [128] - ], -) -def test_context_attention_fwd_no_prompt_cache(batch, seqlen, q_heads, kv_heads, head_dim): - Z, N_CTX, Q_HEADS, KV_HEADS, HEAD_DIM = batch, seqlen, q_heads, kv_heads, head_dim - dtype = torch.bfloat16 - q = torch.randn((Z * N_CTX, Q_HEADS, HEAD_DIM), dtype=dtype, device="cuda") - k = torch.randn((Z * N_CTX, KV_HEADS, HEAD_DIM), dtype=dtype, device="cuda") - v = torch.randn((Z * N_CTX, KV_HEADS, HEAD_DIM), dtype=dtype, device="cuda") - - b_seq_len = torch.ones((Z,), dtype=torch.int32, device="cuda") * N_CTX - b_start_loc = b_seq_len.cumsum(0) - b_seq_len - - o = torch.zeros((Z * N_CTX, Q_HEADS, HEAD_DIM), dtype=dtype, device="cuda") - o1 = torch.zeros((Z * N_CTX, Q_HEADS, HEAD_DIM), dtype=dtype, device="cuda") - - infer_state = LlamaInferStateInfo() - infer_state.batch_size = Z - infer_state.max_q_seq_len = N_CTX - infer_state.b_seq_len = b_seq_len - infer_state.b_q_start_loc = b_start_loc - - context_attention_fwd_no_prompt_cache( - q, - k, - v, - o, - infer_state.b_q_start_loc, - infer_state.b_seq_len, - infer_state.max_q_seq_len, - ) - - head_dim = HEAD_DIM - q_heads = Q_HEADS - kv_heads = KV_HEADS - workspace_buffer = torch.empty(256 * 1024 * 1024, dtype=torch.int8).to(0) - q_starts = torch.zeros((Z + 1,)).int().cuda() - q_starts[1:] = torch.cumsum(b_seq_len, dim=0) - kv_starts = torch.zeros_like(q_starts) - kv_starts[1:] = torch.cumsum(b_seq_len, dim=0) - q_indptr = q_starts.int() - kv_indptr = kv_starts.int() - wrapper = flashinfer.prefill.BatchPrefillWithRaggedKVCacheWrapper( - workspace_buffer, - ) - wrapper.plan( - qo_indptr=q_indptr, - kv_indptr=kv_indptr, - num_qo_heads=q_heads, - num_kv_heads=kv_heads, - head_dim_qk=head_dim, - head_dim_vo=head_dim, - q_data_type=dtype, - causal=True, - ) - wrapper.run(q, k, v, out=o1, return_lse=False) - - # assert torch.allclose(o, o1, atol=1e-2, rtol=0) - cos_sim1 = F.cosine_similarity(o, o1).mean() - assert cos_sim1 == 1.0 diff --git a/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_fa3_neo.py b/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_fa3_neo.py deleted file mode 100644 index eaa19ac711..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_fa3_neo.py +++ /dev/null @@ -1,561 +0,0 @@ -"""Unit test for the FA3-based prefill path with image-token support. - -Neo FA3's ``flash_attn_with_kvcache`` accepts an ``image_token_end`` tensor -of exclusive KV end positions. Image queries can see their entire image span -and preceding tokens, while text queries retain the causal mask. - -Torch reference expresses the *semantics* of the attention, not FA3's internal -tiling — it has no notion of BLOCK_N / BLOCK_M. For each batch element we -gather K/V for the whole request (prompt + new tokens) and apply:: - - allow[m, k] = (k <= q_pos[m]) OR (k < image_end[m]) - -i.e. image queries cannot see subsequent text or later images. - -Run directly for quick debugging: - - python unit_tests/common/basemodel/triton_kernel/att/prefill_att/\ - test_fa3_neo.py - -or via pytest: - - pytest unit_tests/common/basemodel/triton_kernel/att/prefill_att/\ - test_fa3_neo.py -x -s -""" - -import math -import pytest -import torch - -from lightllm.common.basemodel.attention.fa3.neo import flash_attn_with_kvcache_neo as flash_attn_with_kvcache - -try: - import triton - import triton.testing as triton_testing -except ImportError: - triton = None - triton_testing = None - -try: - from lightllm.common.basemodel.triton_kernel.att.prefill_att.context_attention_fwd_neo import ( - context_attention_fwd_neo, - ) -except ImportError: - context_attention_fwd_neo = None - - -def torch_reference_context_attention_neo( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - b_req_idx: torch.Tensor, - b_q_start_loc: torch.Tensor, - b_seq_len: torch.Tensor, - b_prompt_cache_len: torch.Tensor, - req_to_token_indexs: torch.Tensor, - b_image_token_end: torch.Tensor, -) -> torch.Tensor: - device = q.device - dtype = q.dtype - _, Hq, D = q.shape - Hk = k.shape[1] - kv_group = Hq // Hk - scale = 1.0 / math.sqrt(D) - - out = torch.empty_like(q) - - for b in range(b_seq_len.shape[0]): - req_idx = int(b_req_idx[b].item()) - seq_len = int(b_seq_len[b].item()) - prompt_cache_len = int(b_prompt_cache_len[b].item()) - q_seq_len = seq_len - prompt_cache_len - if q_seq_len <= 0: - continue - - q_start = int(b_q_start_loc[b].item()) - q_blk = q[q_start : q_start + q_seq_len] # [M, Hq, D] - image_end = b_image_token_end[q_start : q_start + q_seq_len].to(torch.int64) - - token_locs = req_to_token_indexs[req_idx, :seq_len].to(torch.int64) - k_blk = k[token_locs] # [seq_len, Hk, D] - v_blk = v[token_locs] - - q_pos = torch.arange(prompt_cache_len, seq_len, device=device, dtype=torch.int64) # [M] - k_pos = torch.arange(0, seq_len, device=device, dtype=torch.int64) # [seq_len] - causal = k_pos[None, :] <= q_pos[:, None] - allow = causal | (k_pos[None, :] < image_end[:, None]) - - out_blk = torch.empty_like(q_blk) - for h in range(Hq): - h_k = h // kv_group - q_h = q_blk[:, h, :].to(torch.float32) - k_h = k_blk[:, h_k, :].to(torch.float32) - v_h = v_blk[:, h_k, :].to(torch.float32) - - scores = (q_h @ k_h.transpose(0, 1)) * scale - scores = torch.where(allow, scores, torch.full_like(scores, -1.0e8)) - probs = torch.softmax(scores, dim=-1) - out_h = (probs @ v_h).to(dtype) - out_blk[:, h, :] = out_h - - out[q_start : q_start + q_seq_len] = out_blk - - return out - - -def _build_inputs( - batch: int, - Hq: int, - Hk: int, - D: int, - dtype: torch.dtype, - device: str, - max_q_seq_len: int = 256, - max_prompt_cache_len: int = 512, - image_prob: float = 0.7, - num_image_spans_max: int = 3, - image_span_len_max: int = 24, - kv_pool_slack: int = 4096, - seed: int = 0, -): - """Build one realistic prefill batch. - - Naming matches lightllm's infer_state: - - ``q_seq_len`` = number of new Q tokens in this prefill call - - ``prompt_cache_len`` = length of the already-cached prefix for this req - - ``seq_len`` = prompt_cache_len + q_seq_len (total KV length) - """ - g = torch.Generator(device="cpu").manual_seed(seed) - - q_seq_lens = torch.randint(low=1, high=max_q_seq_len + 1, size=(batch,), generator=g) - prompt_cache_lens = torch.randint(low=0, high=max_prompt_cache_len + 1, size=(batch,), generator=g) - seq_lens = q_seq_lens + prompt_cache_lens - - sum_q = int(q_seq_lens.sum().item()) - sum_total = int(seq_lens.sum().item()) - max_seq_len_in_batch = int(seq_lens.max().item()) - max_q_seq_len_in_batch = int(q_seq_lens.max().item()) - - b_q_start_loc = torch.zeros(batch, dtype=torch.int32) - cur = 0 - for i in range(batch): - b_q_start_loc[i] = cur - cur += int(q_seq_lens[i].item()) - - # Permute so batch idx != request idx: exercises the page_table indexing. - b_req_idx = torch.randperm(batch, generator=g).to(torch.int32) - - # Global KV pool with scattered, non-contiguous slot assignment per request. - base = 1024 - kv_pool_size = base + sum_total + kv_pool_slack - pool = torch.randperm(kv_pool_size - base, generator=g)[:sum_total] + base - - req_to_token_indexs = torch.zeros((batch, max_seq_len_in_batch), dtype=torch.int32) - p = 0 - for r_logical, req_id in enumerate(b_req_idx.tolist()): - L = int(seq_lens[r_logical].item()) - req_to_token_indexs[req_id, :L] = pool[p : p + L].to(torch.int32) - p += L - - # Randomly place contiguous image-token spans inside each batch's new-Q region. - b_image_token_end = torch.zeros(sum_q, dtype=torch.int32) - for i in range(batch): - M = int(q_seq_lens[i].item()) - if M < 2: - continue - if torch.rand((), generator=g).item() > image_prob: - continue - n_spans = int(torch.randint(1, num_image_spans_max + 1, (1,), generator=g).item()) - start_pack = int(b_q_start_loc[i].item()) - for _ in range(n_spans): - span_len = int(torch.randint(1, max(2, image_span_len_max) + 1, (1,), generator=g).item()) - span_len = min(span_len, M) - s_rel = int(torch.randint(0, M - span_len + 1, (1,), generator=g).item()) - image_span = b_image_token_end[start_pack + s_rel : start_pack + s_rel + span_len] - if (image_span > 0).any(): - continue - image_span.fill_(int(prompt_cache_lens[i].item()) + s_rel + span_len) - - b_seq_len = seq_lens.to(torch.int32) - b_prompt_cache_len = prompt_cache_lens.to(torch.int32) - - q = torch.randn((sum_q, Hq, D), dtype=dtype, device=device) - k = torch.randn((kv_pool_size, Hk, D), dtype=dtype, device=device) - v = torch.randn((kv_pool_size, Hk, D), dtype=dtype, device=device) - - return dict( - q=q, - k=k, - v=v, - b_req_idx=b_req_idx.to(device), - b_q_start_loc=b_q_start_loc.to(device), - b_seq_len=b_seq_len.to(device), - b_prompt_cache_len=b_prompt_cache_len.to(device), - max_seq_len_in_batch=max_seq_len_in_batch, - max_q_seq_len_in_batch=max_q_seq_len_in_batch, - req_to_token_indexs=req_to_token_indexs.to(device), - b_image_token_end=b_image_token_end.to(device), - q_seq_lens=q_seq_lens, - prompt_cache_lens=prompt_cache_lens, - ) - - -def _fa3_prefill_with_image_end(inputs: dict) -> torch.Tensor: - """Drive ``flash_attn_with_kvcache`` with the same prefill semantics as - ``Fa3PrefillAttState._nomarl_prefill_att`` plus an optional - ``image_token_end`` kwarg for image-token bidirectional attention. - """ - q = inputs["q"] - k = inputs["k"] - v = inputs["v"] - device = q.device - - # Build page_table[b, p] = req_to_token_indexs[b_req_idx[b], p]. - page_table = inputs["req_to_token_indexs"][inputs["b_req_idx"].long(), : inputs["max_seq_len_in_batch"]].to( - torch.int32 - ) - - q_seq_lens_t = inputs["b_seq_len"].to(torch.int32) - inputs["b_prompt_cache_len"].to(torch.int32) - - cu_seqlens_q = torch.zeros(q_seq_lens_t.shape[0] + 1, dtype=torch.int32, device=device) - cu_seqlens_q[1:] = q_seq_lens_t.cumsum(0).to(torch.int32) - - cu_seqlens_k = torch.zeros(inputs["b_seq_len"].shape[0] + 1, dtype=torch.int32, device=device) - cu_seqlens_k[1:] = inputs["b_seq_len"].cumsum(0).to(torch.int32) - - sm_scale = 1.0 / math.sqrt(q.shape[-1]) - - # page_size = 1 paged KV cache view. - k_cache = k.view(k.shape[0], 1, k.shape[1], k.shape[2]) - v_cache = v.view(v.shape[0], 1, v.shape[1], v.shape[2]) - - o = flash_attn_with_kvcache( - q=q, - k_cache=k_cache, - v_cache=v_cache, - page_table=page_table, - cache_seqlens=inputs["b_seq_len"], - cu_seqlens_q=cu_seqlens_q, - cu_seqlens_k_new=cu_seqlens_k, - max_seqlen_q=inputs["max_q_seq_len_in_batch"], - softmax_scale=sm_scale, - causal=True, - window_size=(-1, -1), - softcap=0.0, - k_descale=None, - v_descale=None, - return_softmax_lse=False, - # Packed like q: int32 exclusive KV end for each image query; - # zero means ordinary causal attention. - image_token_end=inputs["b_image_token_end"], - ) - return o - - -def _report_per_batch_error(out_fa3, out_ref, q_seq_lens, b_q_start_loc, image_end, tag=""): - print(f"\n[{tag}] per-batch error breakdown (abs / rel / cos):") - for i in range(q_seq_lens.shape[0]): - s = int(b_q_start_loc[i].item()) - m = int(q_seq_lens[i].item()) - if m == 0: - continue - a = out_fa3[s : s + m].float() - b = out_ref[s : s + m].float() - abs_err = (a - b).abs().max().item() - denom = b.abs().max().item() + 1e-6 - rel_err = abs_err / denom - cos = torch.nn.functional.cosine_similarity(a.flatten(), b.flatten(), dim=0).item() - n_img = int((image_end[s : s + m] > 0).sum().item()) - print( - f" batch {i:02d} | M={m:4d} | image_tokens={n_img:4d} | " - f"max_abs={abs_err:.4e} | max_rel={rel_err:.4e} | cos={cos:.6f}" - ) - - -def _run_case( - batch: int, - Hq: int, - Hk: int, - D: int, - dtype: torch.dtype, - seed: int, - max_q_seq_len: int, - max_prompt_cache_len: int, - atol: float = 5e-2, - rtol: float = 5e-2, - cos_threshold: float = 0.99, - verbose: bool = True, -): - assert Hq % Hk == 0 - device = "cuda" - - inputs = _build_inputs( - batch=batch, - Hq=Hq, - Hk=Hk, - D=D, - dtype=dtype, - device=device, - max_q_seq_len=max_q_seq_len, - max_prompt_cache_len=max_prompt_cache_len, - seed=seed, - ) - - out_fa3 = _fa3_prefill_with_image_end(inputs) - - out_ref = torch_reference_context_attention_neo( - inputs["q"], - inputs["k"], - inputs["v"], - inputs["b_req_idx"], - inputs["b_q_start_loc"], - inputs["b_seq_len"], - inputs["b_prompt_cache_len"], - inputs["req_to_token_indexs"], - inputs["b_image_token_end"], - ) - - a = out_fa3.float().reshape_as(out_ref.float()) - b = out_ref.float() - abs_err = (a - b).abs().max().item() - denom = b.abs().max().item() + 1e-6 - rel_err = abs_err / denom - cos = torch.nn.functional.cosine_similarity(a.flatten(), b.flatten(), dim=0).item() - - n_image = int((inputs["b_image_token_end"] > 0).sum().item()) - n_tokens = int(inputs["b_image_token_end"].numel()) - if verbose: - print( - f"\ncase: batch={batch} Hq={Hq} Hk={Hk} D={D} dtype={dtype} " - f"seed={seed} image_tokens={n_image}/{n_tokens}" - ) - print( - f" global: max_abs={abs_err:.4e} max_rel={rel_err:.4e} cos={cos:.6f} " - f"(allclose atol={atol}, rtol={rtol}? " - f"{torch.allclose(a, b, atol=atol, rtol=rtol)})" - ) - _report_per_batch_error( - out_fa3, - out_ref, - inputs["q_seq_lens"], - inputs["b_q_start_loc"], - inputs["b_image_token_end"], - tag=f"seed={seed}", - ) - - return abs_err, rel_err, cos - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="no CUDA") -@pytest.mark.skipif(flash_attn_with_kvcache is None, reason="fa3 not available") -@pytest.mark.parametrize( - "batch,Hq,Hk,D,dtype,seed,max_q_seq_len,max_prompt_cache_len", - [ - (4, 8, 2, 128, torch.bfloat16, 0, 128, 256), - (4, 8, 2, 128, torch.bfloat16, 1, 256, 512), - (8, 16, 4, 128, torch.bfloat16, 2, 256, 512), - (16, 28, 4, 128, torch.bfloat16, 3, 128, 256), - (4, 8, 2, 128, torch.float16, 4, 256, 512), - (4, 8, 8, 64, torch.bfloat16, 5, 128, 256), - (3, 8, 2, 128, torch.bfloat16, 6, 8, 1024), - ], -) -def test_fa3_neo_prefill_with_image_end(batch, Hq, Hk, D, dtype, seed, max_q_seq_len, max_prompt_cache_len): - abs_err, rel_err, cos = _run_case( - batch=batch, - Hq=Hq, - Hk=Hk, - D=D, - dtype=dtype, - seed=seed, - max_q_seq_len=max_q_seq_len, - max_prompt_cache_len=max_prompt_cache_len, - verbose=True, - ) - assert cos > 0.99, f"cosine similarity too low: {cos}" - assert rel_err < 5e-2, f"max relative error too large: {rel_err}" - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="no CUDA") -@pytest.mark.parametrize("backend", ["triton", "fa3"]) -def test_image_queries_stop_at_their_own_image_end(backend): - if backend == "fa3" and flash_attn_with_kvcache is None: - pytest.skip("Neo FA3 with image_token_end support is not available") - if backend == "triton" and context_attention_fwd_neo is None: - pytest.skip("Triton attention is not available") - - def indices(values): - return torch.tensor(values, dtype=torch.int32, device="cuda") - - # Two cached text tokens, then text / image / text / image / text. - # Zero Q/K make each output the mean of its visible V values. - inputs = dict( - q=torch.zeros((8, 2, 64), dtype=torch.float16, device="cuda"), - k=torch.zeros((10, 1, 64), dtype=torch.float16, device="cuda"), - v=torch.arange(10, dtype=torch.float16, device="cuda").view(10, 1, 1).expand(10, 1, 64).contiguous(), - b_req_idx=indices([0]), - b_q_start_loc=indices([0]), - b_seq_len=indices([10]), - b_prompt_cache_len=indices([2]), - req_to_token_indexs=indices([list(range(10))]), - b_image_token_end=indices([0, 5, 5, 0, 8, 8, 0, 0]), - max_seq_len_in_batch=10, - max_q_seq_len_in_batch=8, - ) - if backend == "fa3": - output = _fa3_prefill_with_image_end(inputs) - else: - output = torch.empty_like(inputs["q"]) - context_attention_fwd_neo( - inputs["q"], - inputs["k"], - inputs["v"], - output, - inputs["b_req_idx"], - inputs["b_q_start_loc"], - inputs["b_seq_len"], - inputs["b_prompt_cache_len"], - inputs["max_q_seq_len_in_batch"], - inputs["req_to_token_indexs"], - inputs["b_image_token_end"], - ) - expected = torch.tensor([1, 2, 2, 2.5, 3.5, 3.5, 4, 4.5], dtype=output.dtype, device="cuda") - torch.testing.assert_close(output, expected[:, None, None].expand_as(output), rtol=1e-3, atol=1e-3) - - -def _bench_case( - batch: int, - Hq: int, - Hk: int, - D: int, - dtype: torch.dtype, - seed: int, - max_q_seq_len: int, - max_prompt_cache_len: int, - rep_ms: int = 100, - warmup_iters: int = 3, -): - """Compare FA3 (with image_token_end) vs the original Triton - ``context_attention_fwd_neo`` using ``triton.testing.do_bench_cudagraph``. - - Both kernels are captured into a CUDA graph so scheduling/launch overhead - is minimized and the measurement reflects the kernel cost. - """ - assert Hq % Hk == 0 - device = "cuda" - - inputs = _build_inputs( - batch=batch, - Hq=Hq, - Hk=Hk, - D=D, - dtype=dtype, - device=device, - max_q_seq_len=max_q_seq_len, - max_prompt_cache_len=max_prompt_cache_len, - seed=seed, - ) - - # --- fa3 runner: output tensor is allocated inside flash_attn_with_kvcache. - def fa3_run(): - return _fa3_prefill_with_image_end(inputs) - - # --- triton runner: pre-allocate o so the graph captures only the kernel launch. - o_triton = torch.empty_like(inputs["q"]) - - def triton_run(): - context_attention_fwd_neo( - inputs["q"], - inputs["k"], - inputs["v"], - o_triton, - inputs["b_req_idx"], - inputs["b_q_start_loc"], - inputs["b_seq_len"], - inputs["b_prompt_cache_len"], - inputs["max_q_seq_len_in_batch"], - inputs["req_to_token_indexs"], - inputs["b_image_token_end"], - ) - - # Warm up outside the graph capture so lazy allocations / autotune happen. - for _ in range(warmup_iters): - fa3_run() - triton_run() - torch.cuda.synchronize() - - fa3_ms = triton_testing.do_bench_cudagraph(fa3_run, rep=rep_ms) - triton_ms = triton_testing.do_bench_cudagraph(triton_run, rep=rep_ms) - - n_image = int((inputs["b_image_token_end"] > 0).sum().item()) - n_tokens = int(inputs["b_image_token_end"].numel()) - sum_kv = int(inputs["b_seq_len"].sum().item()) - speedup = triton_ms / fa3_ms if fa3_ms > 0 else float("inf") - - print( - f"bench: batch={batch} Hq={Hq} Hk={Hk} D={D} dtype={str(dtype).split('.')[-1]:<8s} " - f"max_q_seq_len={max_q_seq_len:4d} max_prompt_cache_len={max_prompt_cache_len:4d} " - f"image_tokens={n_image:4d}/{n_tokens:5d} sum_kv={sum_kv:6d} | " - f"fa3 {fa3_ms*1000:8.1f} us | triton {triton_ms*1000:8.1f} us | " - f"speedup {speedup:5.2f}x" - ) - - return fa3_ms, triton_ms - - -if __name__ == "__main__": - if not torch.cuda.is_available(): - print("No CUDA available.") - raise SystemExit(0) - if flash_attn_with_kvcache is None: - print("Neo FA3 flash_attn_with_kvcache with image_token_end support is not available.") - raise SystemExit(0) - - torch.manual_seed(0) - - cases = [ - dict(batch=4, Hq=8, Hk=2, D=128, dtype=torch.bfloat16, seed=0, max_q_seq_len=128, max_prompt_cache_len=256), - dict(batch=8, Hq=16, Hk=4, D=128, dtype=torch.bfloat16, seed=1, max_q_seq_len=256, max_prompt_cache_len=512), - dict(batch=16, Hq=28, Hk=4, D=128, dtype=torch.bfloat16, seed=2, max_q_seq_len=128, max_prompt_cache_len=256), - # FP16 case disabled: current build compiled without FP16 support - # dict(batch=4, Hq=8, Hk=2, D=128, dtype=torch.float16, seed=3, max_q_seq_len=256, max_prompt_cache_len=512), - ] - - print("=" * 100) - print("Correctness") - print("=" * 100) - for cfg in cases: - _run_case(**cfg, verbose=True) - - if triton_testing is None or context_attention_fwd_neo is None: - print("\nSkipping benchmark: triton or context_attention_fwd_neo not available.") - raise SystemExit(0) - - print("\n" + "=" * 100) - print("Benchmark (triton.testing.do_bench_cudagraph)") - print("=" * 100) - - # Cold-prefill sweep: max_prompt_cache_len=0 so seq_len == q_seq_len. - # Head shape matches neo_chat_moe / Qwen3 llm_config: - # num_attention_heads = 32, num_key_value_heads = 8, head_dim = 128 - # (GQA ratio 4:1) - bench_batches = [8, 16, 32, 64, 128] - bench_q_seq_lens = [1024, 4096, 8192] - - bench_cases = [ - dict( - batch=b, - Hq=32, - Hk=8, - D=128, - dtype=torch.bfloat16, - seed=0, - max_q_seq_len=s, - max_prompt_cache_len=0, - ) - for b in bench_batches - for s in bench_q_seq_lens - ] - - for cfg in bench_cases: - _bench_case(**cfg) diff --git a/unit_tests/common/basemodel/triton_kernel/kv_copy/test_mla_destindex_copy_kv.py b/unit_tests/common/basemodel/triton_kernel/kv_copy/test_mla_destindex_copy_kv.py deleted file mode 100644 index ed0c6e369f..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/kv_copy/test_mla_destindex_copy_kv.py +++ /dev/null @@ -1,42 +0,0 @@ -import torch -import pytest -from lightllm.common.basemodel.triton_kernel.kv_copy.mla_copy_kv import destindex_copy_kv -from lightllm.utils.log_utils import init_logger -import torch.nn.functional as F - -logger = init_logger(__name__) - -seed = 42 -torch.manual_seed(seed) - -if torch.cuda.is_available(): - torch.cuda.manual_seed(seed) - torch.cuda.manual_seed_all(seed) - - -@pytest.mark.parametrize( - "batch, seqlen, heads, nope_head, rope_head, copy_len", - [ - (a, b, c, d, e, f) - for a in [1, 16, 32, 128, 512] - for b in [1024, 2048] - for c in [1] - for d in [512] - for e in [64] - for f in [10, 20, 100, 1024] - ], -) -def test_destindex_copy_kv(batch, seqlen, heads, nope_head, rope_head, copy_len): - B, N_CTX, H, NOPE_HEAD, ROPE_HEAD, COPY_LEN = batch, seqlen, heads, nope_head, rope_head, copy_len - dtype = torch.bfloat16 - dest_loc = torch.randperm(COPY_LEN).cuda() - kv = torch.randn((len(dest_loc), H, NOPE_HEAD + ROPE_HEAD), dtype=dtype).cuda() - O_nope = torch.zeros((B * N_CTX, H, NOPE_HEAD), dtype=dtype).cuda() - O_rope = torch.zeros((B * N_CTX, H, ROPE_HEAD), dtype=dtype).cuda() - - kv_nope = kv[:, :, :NOPE_HEAD] - kv_rope = kv[:, :, NOPE_HEAD:] - destindex_copy_kv(kv_nope, kv_rope, dest_loc, O_nope, O_rope) - - assert torch.allclose(O_nope[dest_loc], kv_nope, atol=1e-2, rtol=0) - assert torch.allclose(O_rope[dest_loc], kv_rope, atol=1e-2, rtol=0) diff --git a/unit_tests/common/basemodel/triton_kernel/kv_copy/test_ppl_int4kv_copy_kv.py b/unit_tests/common/basemodel/triton_kernel/kv_copy/test_ppl_int4kv_copy_kv.py deleted file mode 100644 index 83537ec708..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/kv_copy/test_ppl_int4kv_copy_kv.py +++ /dev/null @@ -1,62 +0,0 @@ -import torch -import pytest -import numpy as np -from typing import Tuple -from lightllm.common.basemodel.triton_kernel.kv_copy.ppl_int4kv_copy_kv import destindex_copy_int4kv, dequantize_int4kv -from lightllm.utils.log_utils import init_logger - -logger = init_logger(__name__) - - -def test_quanted_and_dequant(): - """Test quantization followed by dequantization.""" - batch_size = 1 - seq_len = 8 - head_num = 4 - k_head_num = 2 - v_head_num = 2 - assert k_head_num + v_head_num == head_num - head_dim = 64 - quant_group_size = 8 - - # Create original data - original_kv = torch.randn(batch_size * seq_len, head_num, head_dim, dtype=torch.float32).clamp_(-1, 1).cuda() - dest_loc = torch.arange(batch_size * seq_len, dtype=torch.int64).cuda() - - # Quantize - group_count = head_dim // quant_group_size - kv_buffer = torch.zeros(batch_size * seq_len, head_num, head_dim // 2, dtype=torch.int8).cuda() - kv_scale_buffer = torch.zeros(batch_size * seq_len, head_num, group_count, dtype=torch.float32).cuda() - destindex_copy_int4kv(original_kv, dest_loc, kv_buffer, kv_scale_buffer, quant_group_size) - - # Dequantize - req_to_token_indexs = torch.arange(seq_len, dtype=torch.int64).view(1, -1).cuda() - b_seq_len = torch.tensor([seq_len], dtype=torch.int32).cuda() - b_req_idx = torch.tensor([0], dtype=torch.int32).cuda() - b_kv_start_loc = torch.tensor([0], dtype=torch.int32).cuda() - - recovered_kv = torch.zeros(batch_size * seq_len, head_num, head_dim, dtype=torch.float32).cuda() - - dequantize_int4kv( - k=kv_buffer[:, 0:k_head_num, :], - k_scale=kv_scale_buffer[:, 0:k_head_num, :], - v=kv_buffer[:, k_head_num:, :], - v_scale=kv_scale_buffer[:, k_head_num:, :], - req_to_token_indexs=req_to_token_indexs, - b_seq_len=b_seq_len, - b_req_idx=b_req_idx, - b_kv_start_loc=b_kv_start_loc, - k_out=recovered_kv[:, :k_head_num, :], - v_out=recovered_kv[:, k_head_num:, :], - max_len_in_batch=seq_len, - quant_group_size=quant_group_size, - ) - - logger.info("Round-trip test completed!") - assert torch.allclose(recovered_kv, original_kv, atol=2 / 14, rtol=0) - cos = torch.nn.CosineSimilarity(0) - assert cos(recovered_kv.flatten().float(), original_kv.flatten().float()) > 0.99 - - -if __name__ == "__main__": - pytest.main() diff --git a/unit_tests/common/basemodel/triton_kernel/kv_copy/test_ppl_int8kv_copy_kv.py b/unit_tests/common/basemodel/triton_kernel/kv_copy/test_ppl_int8kv_copy_kv.py deleted file mode 100644 index 149c9894ae..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/kv_copy/test_ppl_int8kv_copy_kv.py +++ /dev/null @@ -1,86 +0,0 @@ -import torch -import time -import pytest -from lightllm.common.basemodel.triton_kernel.kv_copy.ppl_int8kv_copy_kv import ( - dequantize_int8kv, - destindex_copy_quantize_kv, -) -from lightllm.utils.log_utils import init_logger - -logger = init_logger(__name__) - - -def torch_dequant(kv, kv_scale, b_req_idx, b_seq_len, req_to_token_indexs, odtype, group_quant_size): - batch = b_req_idx.shape[0] - tmp_out = [] - for i in range(batch): - req_idx = b_req_idx[i] - seq_len = b_seq_len[i] - kv_loc = req_to_token_indexs[req_idx, :seq_len] - head_num = kv.shape[1] - cur_kv = kv[kv_loc, :, :].reshape(seq_len, head_num, -1, group_quant_size).to(odtype) - cur_scale = kv_scale[kv_loc, :, :].reshape(seq_len, head_num, -1, 1) - out = cur_kv * cur_scale - tmp_out.append(out.reshape(seq_len, head_num, -1)) - return torch.cat(tmp_out, dim=0) - - -@pytest.mark.parametrize( - "B, H, N_CTX, D_HEAD, group_quant_size", - [ - (b, H, N_CTX, D_HEAD, group_quant_size) - for b in [1, 2, 4] - for H in [1, 8] - for N_CTX in [3, 10, 1024] - for D_HEAD in [64, 128] - for group_quant_size in [8, 16] - ], -) -def test_dequantize_int8kv(B, H, N_CTX, D_HEAD, group_quant_size): - dtype = torch.bfloat16 - kv = torch.empty((B * N_CTX, 2 * H, D_HEAD), dtype=torch.int8, device="cuda").random_(-10, 10) - kv_scale = torch.randn((B * N_CTX, 2 * H, D_HEAD // group_quant_size), dtype=dtype, device="cuda") - out = torch.empty((B * N_CTX, 2 * H, D_HEAD), dtype=dtype, device="cuda") - req_to_token_indexs = torch.empty((B, N_CTX), dtype=torch.int32, device="cuda") - max_input_len = N_CTX - b_seq_len = torch.ones((B,), dtype=torch.int32, device="cuda") - b_seq_len.fill_(N_CTX) - b_req_idx = torch.arange(0, B, dtype=torch.int32, device="cuda") - req_to_token_indexs.view(-1)[:] = torch.arange(0, B * N_CTX, dtype=torch.int32, device="cuda") - b_kv_start_loc = torch.cumsum(b_seq_len, dim=0, dtype=torch.int32) - b_seq_len - - k = kv[:, :H, :] - v = kv[:, H:, :] - k_scale = kv_scale[:, :H, :] - v_scale = kv_scale[:, H:, :] - - ground_out = torch_dequant( - kv=kv, - kv_scale=kv_scale, - b_req_idx=b_req_idx, - b_seq_len=b_seq_len, - req_to_token_indexs=req_to_token_indexs, - odtype=out.dtype, - group_quant_size=group_quant_size, - ) - dequantize_int8kv( - k=k, - k_scale=k_scale, - v=v, - v_scale=v_scale, - req_to_token_indexs=req_to_token_indexs, - b_seq_len=b_seq_len, - b_req_idx=b_req_idx, - b_kv_start_loc=b_kv_start_loc, - k_out=out[:, :H, :], - v_out=out[:, H:, :], - max_len_in_batch=max_input_len, - quant_group_size=group_quant_size, - ) - assert torch.allclose(out, ground_out, atol=1e-2, rtol=0) - cos = torch.nn.CosineSimilarity(0) - assert cos(out.flatten().float(), ground_out.flatten().float()) > 0.99 - - -if __name__ == "__main__": - pytest.main() diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py deleted file mode 100644 index 7fbb63df5d..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py +++ /dev/null @@ -1,272 +0,0 @@ -import pytest -import torch - -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import advance_cursor -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_fold import fold_prefix -from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix -from lightllm.common.basemodel.triton_kernel.mtp_utils import linear_att_mtp_state_index_update as update -from unit_tests.common.basemodel.triton_kernel.linear_att.replay_history import ( - accept_replay_window, - attach_replay_history as ReplaySSMExecutor, - replay_layer, -) - - -def ssm_snapshot(backend, req): - """折进一份 SSM 副本再读出这个请求。不改 canonical SSM，也不动 cursor。""" - - saved = backend.ssm_state - clone = saved.clone() - backend.ssm_state = clone - try: - fold_prefix(backend, torch.tensor([req], dtype=torch.int32, device=saved.device)) - finally: - backend.ssm_state = saved - return clone[:, req] - - -from types import SimpleNamespace -from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState - - -def make(n, width): - hold = n + 3 - ids = list(reversed(range(n))) + [hold] - lengths = [1 + i % width for i in range(n)] + [1] - starts = [] - reqs = [] - mtps = [] - for r, l in zip(ids, lengths): - starts.append(len(reqs)) - reqs.extend([r] * l) - mtps.extend(range(l)) - tensor = lambda x: torch.tensor(x, device="cuda", dtype=torch.int32) - args = [tensor(starts), tensor(reqs), tensor(mtps), tensor([1] * len(reqs)), width] - ssm_backend = ReplayGDNState.__new__(ReplayGDNState) - ssm_backend.cursors = tensor([7] * (hold + 1)) - ssm_backend.hold_req_id = hold - # 这个用例不触发折回，容量要盖住多步累加后的 cursor。 - ssm_backend.capacity = 4096 - ssm_backend.verify_width = width - manager = SimpleNamespace() - manager.mtp_state_index = tensor([-9] * (hold + 1)) - manager._ssm_executor = ssm_backend - return manager, args, (ids, lengths, starts) - - -def _replay_cu(args): - """args[0] 是每个请求的起点，最后一个请求接到本批 token 末尾。""" - - starts = args[0] - return torch.cat((starts, starts.new_full((1,), args[3].numel()))) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_request_counts_share_compiled_kernel(monkeypatch): - from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import _commit_prefix, commit_ssm_prefix - - reqs = torch.arange(512, dtype=torch.int32, device="cuda") - lengths = reqs % 4 + 1 - window = lengths + 4 - indices = torch.zeros(513, dtype=torch.int32, device="cuda") - cursors = torch.zeros_like(indices) - hashes = set() - run = _commit_prefix.run - - def record(*args, **kwargs): - kernel = run(*args, **kwargs) - hashes.add(kernel.hash) - return kernel - - monkeypatch.setattr(_commit_prefix, "run", record) - for count in (1, 2, 7, 16, 23, 24, 25, 31, 32, 33, 255, 256, 257, 511): - indices.fill_(-1) - cursors.zero_() - cursors[:count] = window[:count] - commit_ssm_prefix(indices, reqs[:count], lengths[:count], window[:count], cursors) - assert torch.equal(indices[:count], lengths[:count] - 1) - assert (indices[count:] == -1).all() - assert torch.equal(cursors[:count], lengths[:count]) - assert not cursors[count:].any() - assert len(hashes) == 1 - - -def accept_unfused(m, args): - update(m.mtp_state_index, *args) - reqs = args[1][args[0].long()] - if m._ssm_executor.verify_width > 1: - accept_replay_window(m._ssm_executor, reqs, m.mtp_state_index, _replay_cu(args)) - else: - accept_replay_window(m._ssm_executor, reqs) - - -def accept_fused(m, args): - reqs = args[1][args[0].long()] - ends = torch.cat((args[0][1:], args[0].new_full((1,), args[3].numel()))) - cumulative = torch.cat((args[3].new_zeros(1), args[3].cumsum(0))) - lengths = cumulative[ends.long()] - cumulative[args[0].long()] - executor = m._ssm_executor - if executor.verify_width > 1: - cu = _replay_cu(args) - advance_cursor(executor, reqs, cu) - commit_ssm_prefix(m.mtp_state_index, reqs, lengths, cu[1:] - cu[:-1], executor.cursors) - else: - # 宽度为 1 时正式路径不提交。对照只写下标，并把每个请求的 cursor 加一。 - m.mtp_state_index[reqs] = (lengths - 1).to(m.mtp_state_index.dtype) - advance_cursor(executor, reqs, torch.arange(reqs.numel() + 1, dtype=torch.int32, device=reqs.device)) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_acceptance_fusion_matches_original_and_cpu_under_graph(): - checks = 0 - for width in [1, 2, 4, 8, 16]: - for n in [1, 7, 33, 128]: - m, args, (ids, lengths, starts) = make(n, width) - ref, _, _ = make(n, width) - # Compile both paths before graph capture, then reset before replay. - accept_unfused(ref, args) - accept_fused(m, args) - torch.cuda.synchronize() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - accept_fused(m, args) - expected_index = [-9] * len(m.mtp_state_index) - expected_cursor = [7] * len(expected_index) - m.mtp_state_index.fill_(-9) - ref.mtp_state_index.fill_(-9) - m._ssm_executor.cursors.fill_(7) - ref._ssm_executor.cursors.fill_(7) - for step in range(24): - mask = [] - for i, (r, l) in enumerate(zip(ids, lengths)): - count = (step + i) % (l + 1) - mask.extend([1] * count + [0] * (l - count)) - expected_index[r] = count - 1 - if r != m._ssm_executor.hold_req_id: - expected_cursor[r] += count if width > 1 else 1 - args[3].copy_(torch.tensor(mask, device="cuda", dtype=torch.int32)) - accept_unfused(ref, args) - if step % 2: - graph.replay() - else: - accept_fused(m, args) - assert torch.equal(m.mtp_state_index, ref.mtp_state_index) - assert torch.equal(m._ssm_executor.cursors, ref._ssm_executor.cursors) - assert m.mtp_state_index.tolist() == expected_index - assert m._ssm_executor.cursors.tolist() == expected_cursor - checks += 1 - - print("ACCEPT_EQUIVALENCE_CHECKS=" + str(checks), flush=True) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -@pytest.mark.parametrize("compute_mode", ["inline", "precompute"]) -def test_fused_acceptance_preserves_folds_snapshots_and_slot_reuse(dtype, compute_mode): - torch.manual_seed(93) - initial = torch.randn(2, 8, 4, 32, 32, device="cuda").to(dtype) * 0.01 - ssm_backends = [ - ReplaySSMExecutor( - initial.clone(), - 8, - 4, - torch.bfloat16, - num_key_heads=2, - compute_mode=compute_mode, - ) - for _ in range(2) - ] - managers = [] - for ssm_backend in ssm_backends: - manager = SimpleNamespace() - manager._ssm_executor = ssm_backend - manager.mtp_state_index = torch.zeros(8, dtype=torch.int32, device="cuda") - managers.append(manager) - reqs = torch.full((4,), 7, dtype=torch.int32, device="cuda") - cu = torch.tensor([0, 3, 6, 9, 12], dtype=torch.int32, device="cuda") - token_reqs = torch.full((12,), 7, dtype=torch.int32, device="cuda") - mtp_indexes = torch.arange(12, dtype=torch.int32, device="cuda") % 3 - accepted_mask = torch.ones(12, dtype=torch.int32, device="cuda") - q = torch.randn(12, 2, 32, dtype=torch.bfloat16, device="cuda") - k = torch.randn_like(q) - v = torch.randn(12, 4, 32, dtype=torch.bfloat16, device="cuda") - a = torch.full((12, 4), -3.0, dtype=torch.bfloat16, device="cuda") - b = torch.randn_like(a) - alog = torch.zeros(4, device="cuda") - args = [cu[:-1], token_reqs, mtp_indexes, accepted_mask, 4] - - def cycle(i): - ssm_backend = ssm_backends[i] - outputs = [ - replay_layer( - ssm_backend, - layer, - q, - k, - v, - a, - b, - alog, - alog, - reqs, - cu, - run_config={"BV": 32, "num_warps": 1, "num_stages": 1}, - ) - for layer in range(2) - ] - (accept_unfused if i == 0 else accept_fused)(managers[i], args) - return outputs - - graphs, outputs = [], [] - for i in range(2): - cycle(i) - torch.cuda.synchronize() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - output = cycle(i) - graphs.append(graph) - outputs.append(output) - for iteration in range(40): - ids = [1, 3, 6] if iteration % 2 else [6, 1, 3] - lengths = [4, 4, 3] if iteration % 3 == 0 else [3, 3, 3] - lengths.append(12 - sum(lengths)) - ids.append(7) - starts, tokens, indexes, masks = [0], [], [], [] - for i, (req, length) in enumerate(zip(ids, lengths)): - count = (iteration + i) % (length + 1) if req != 7 else 0 - tokens.extend([req] * length) - indexes.extend(range(length)) - masks.extend([1] * count + [0] * (length - count)) - starts.append(len(tokens)) - for tensor, values in [ - (reqs, ids), - (cu, starts), - (token_reqs, tokens), - (mtp_indexes, indexes), - (accepted_mask, masks), - ]: - tensor.copy_(tensor.new_tensor(values)) - for tensor in (q, k, v, b): - tensor.normal_() - for graph in graphs: - graph.replay() - for left, right in zip(outputs[0], outputs[1]): - assert torch.equal(left, right) - assert torch.equal(ssm_backends[0].ssm_state, ssm_backends[1].ssm_state) - assert torch.equal(ssm_backends[0].cursors, ssm_backends[1].cursors) - assert torch.equal(managers[0].mtp_state_index, managers[1].mtp_state_index) - if iteration % 7 == 0: - assert torch.equal(ssm_snapshot(ssm_backends[0], 3), ssm_snapshot(ssm_backends[1], 3)) - if iteration % 11 == 10: - for ssm_backend in ssm_backends: - ssm_backend.fold_accepted_prefix(reqs[:3]) - assert torch.equal(ssm_backends[0].ssm_state, ssm_backends[1].ssm_state) - if iteration == 19: - for manager in managers: - manager._ssm_executor.cursors[1].zero_() - manager._ssm_executor.ssm_state[:, 1].zero_() - manager.mtp_state_index[1] = 0 - for ssm_backend in ssm_backends: - ssm_backend.fold_accepted_prefix(reqs[:3]) - assert torch.equal(ssm_backends[0].ssm_state, ssm_backends[1].ssm_state) - assert torch.equal(ssm_backends[0].ssm_state[:, [0, 2, 4, 5, 7]], initial[:, [0, 2, 4, 5, 7]]) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_causal_conv1d_mtp.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_causal_conv1d_mtp.py index f7c742db01..1c984a06a4 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_causal_conv1d_mtp.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_causal_conv1d_mtp.py @@ -1,971 +1,107 @@ -from typing import Optional - import pytest import torch +import torch.nn.functional as F from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_mtp import causal_conv1d_update -def causal_conv1d_ref( - x: torch.Tensor, - conv_state: torch.Tensor, - weight: torch.Tensor, - mtp_step: int, - bias: Optional[torch.Tensor] = None, - activation: Optional[str] = None, - conv_state_indices: Optional[torch.Tensor] = None, - num_accepted_tokens: Optional[torch.Tensor] = None, - query_start_loc: Optional[torch.Tensor] = None, - pad_slot_id: int = -1, -) -> torch.Tensor: - """ - Pure-PyTorch reference for causal_conv1d_update. - - The kernel's algorithm per token iteration: - acc = bias - for j in range(width): - if j == 0: acc += col0 * weight[:, 0] - elif j < width-1: acc += col{j} * weight[:, j] - else: acc += x[t] * weight[:, width-1] - shift: col0=col1, col1=col2, ..., col_{w-3}=col_{w-2}, col_{w-2}=x[t] - """ - assert conv_state_indices is not None - batch = conv_state_indices.size(0) - dim = x.size(1) - assert x.size(0) % batch == 0 - seqlen = x.size(0) // batch - _, width = weight.shape - _, _, state_len = conv_state.size() - assert state_len == width - 1 + mtp_step +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") - out = x.clone() - conv_state = conv_state.clone() - for b in range(batch): - slot = conv_state_indices[b].item() - if slot == pad_slot_id: +def reference(x, state, weight, bias, activation, slots, accepted, starts): + output, next_state = x.clone(), state.clone() + width = weight.shape[1] + for i, slot in enumerate(slots.tolist()): + start, end = starts[i : i + 2].tolist() + if slot == -1 or start == end: continue - - if query_start_loc is not None: - start = query_start_loc[b].item() - end = query_start_loc[b + 1].item() - local_seqlen = end - start - else: - start = b * seqlen - local_seqlen = seqlen - - accepted = 1 - if num_accepted_tokens is not None: - accepted = num_accepted_tokens[b].item() - offset = accepted - 1 - - state_3d = conv_state[slot] - - # STEP 1: initial columns from conv_state starting at offset - cols = [state_3d[:, offset + k].clone() for k in range(width - 1)] - - # STEP 2: update conv_state in the reference too - new_state = [] - for k in range(width - 2): - new_state.append(state_3d[:, offset + k + 1].clone()) - x_chunk = x[start : start + local_seqlen, :] - for t_ in range(local_seqlen): - new_state.append(x_chunk[t_, :].clone()) - while len(new_state) < state_len: - new_state.append(torch.zeros(dim, device=x.device, dtype=x.dtype)) - stacked = torch.stack(new_state, dim=1) - write_len = min(state_len, stacked.size(1)) - conv_state[slot, :, :write_len] = stacked[:, :write_len] - - # STEP 3-5: compute output - for t_ in range(local_seqlen): - acc = torch.zeros(dim, device=x.device, dtype=torch.float32) - if bias is not None: - acc += bias.float() - - for k in range(width): - if k == 0: - v = cols[0].float() - elif k == width - 1: - v = x_chunk[t_, :].float() - else: - v = cols[k].float() - w = weight[:, k].float() - acc += v * w - - if activation in ("silu", "swish"): - acc = acc / (1 + torch.exp(-acc)) - - out[start + t_, :] = acc.to(x.dtype) - - # shift cols: col0=col1, col1=col2, ..., col_{w-2}=x[t] - for k in range(width - 2): - cols[k] = cols[k + 1].clone() - if width >= 2: - cols[width - 2] = x_chunk[t_, :].clone() - - return out - - -def make_tensors( - batch: int, - dim: int, - width: int, - mtp_step: int, - has_bias: bool, - num_slots: int = None, - dtype: torch.dtype = torch.float16, - seed: int = 42, - device: str = "cuda", -): - torch.manual_seed(seed) - if num_slots is None: - num_slots = batch - - state_len = width - 1 + mtp_step - x = torch.randn(batch * (mtp_step + 1), dim, device=device, dtype=dtype) - weight = torch.randn(dim, width, device=device, dtype=dtype) - bias = torch.randn(dim, device=device, dtype=dtype) if has_bias else None - conv_state = torch.randn(num_slots, dim, state_len, device=device, dtype=dtype) - conv_state.add_(0.5) - - conv_state_indices = torch.arange(batch, device=device, dtype=torch.int32) % num_slots - num_accepted_tokens = torch.ones(batch, device=device, dtype=torch.int32) - query_start_loc = torch.arange(batch + 1, device=device, dtype=torch.int32) * (mtp_step + 1) - - return x, conv_state, weight, bias, conv_state_indices, num_accepted_tokens, query_start_loc - - -@pytest.mark.parametrize("width", [2, 3, 4, 5, 6]) -@pytest.mark.parametrize("dim", [64, 128]) -@pytest.mark.parametrize("mtp_step", [0, 1, 2]) -@pytest.mark.parametrize("has_bias", [True, False]) -@pytest.mark.parametrize("activation", [None, "silu"]) -def test_single_request(width, dim, mtp_step, has_bias, activation): - """Single request, no pad slots, basic functionality.""" - batch = 1 - x, conv_state, weight, bias, conv_state_indices, num_accepted_tokens, query_start_loc = make_tensors( - batch=batch, - dim=dim, - width=width, - mtp_step=mtp_step, - has_bias=has_bias, - dtype=torch.float16, - ) - - x_orig = x.clone() - conv_state_ref = conv_state.clone() - - out_triton = causal_conv1d_update( - x, - conv_state, - weight, - mtp_step=mtp_step, - bias=bias, - activation=activation, - conv_state_indices=conv_state_indices, - num_accepted_tokens=num_accepted_tokens, - query_start_loc=query_start_loc, - ) - - out_ref = causal_conv1d_ref( - x_orig, - conv_state_ref, - weight, - mtp_step=mtp_step, - bias=bias, - activation=activation, - conv_state_indices=conv_state_indices, - num_accepted_tokens=num_accepted_tokens, - query_start_loc=query_start_loc, - ) - - rtol, atol = 1e-2, 1e-2 - assert torch.allclose(out_triton, out_ref, rtol=rtol, atol=atol), ( - f"Output mismatch: width={width}, dim={dim}, mtp_step={mtp_step}, " - f"bias={has_bias}, activation={activation}\n" - f"max diff={torch.abs(out_triton - out_ref).max().item():.6f}" - ) - - -@pytest.mark.parametrize("batch", [2, 4]) -@pytest.mark.parametrize("width", [3, 4, 5]) -@pytest.mark.parametrize("dim", [64, 256]) -@pytest.mark.parametrize("mtp_step", [0, 1, 2]) -def test_multi_request(batch, width, dim, mtp_step): - """Multiple requests.""" - x, conv_state, weight, bias, conv_state_indices, num_accepted_tokens, query_start_loc = make_tensors( - batch=batch, - dim=dim, - width=width, - mtp_step=mtp_step, - has_bias=True, - dtype=torch.float16, - ) - x_orig = x.clone() - conv_state_ref = conv_state.clone() - - out_triton = causal_conv1d_update( - x, - conv_state, - weight, - mtp_step=mtp_step, - bias=bias, - activation="silu", - conv_state_indices=conv_state_indices, - num_accepted_tokens=num_accepted_tokens, - query_start_loc=query_start_loc, - ) - out_ref = causal_conv1d_ref( - x_orig, - conv_state_ref, - weight, - mtp_step=mtp_step, - bias=bias, - activation="silu", - conv_state_indices=conv_state_indices, - num_accepted_tokens=num_accepted_tokens, - query_start_loc=query_start_loc, - ) - - rtol, atol = 1e-2, 1e-2 - assert torch.allclose(out_triton, out_ref, rtol=rtol, atol=atol), ( - f"Output mismatch: batch={batch}, width={width}, dim={dim}, mtp_step={mtp_step}\n" - f"max diff={torch.abs(out_triton - out_ref).max().item():.6f}" - ) - - -@pytest.mark.parametrize("width", [3, 4]) -@pytest.mark.parametrize("dim", [64, 128]) -@pytest.mark.parametrize("mtp_step", [0, 1]) -def test_pad_slots(width, dim, mtp_step): - """Some slots are padded, should produce same output as reference.""" - batch = 4 - num_slots = 4 - x, conv_state, weight, bias, conv_state_indices, num_accepted_tokens, query_start_loc = make_tensors( - batch=batch, - dim=dim, - width=width, - mtp_step=mtp_step, - has_bias=True, - num_slots=num_slots, - dtype=torch.float16, - ) - - conv_state_indices[2:] = -1 - x_orig = x.clone() - conv_state_ref = conv_state.clone() - - out_triton = causal_conv1d_update( - x, - conv_state, - weight, - mtp_step=mtp_step, - bias=bias, - activation="silu", - conv_state_indices=conv_state_indices, - num_accepted_tokens=num_accepted_tokens, - query_start_loc=query_start_loc, - pad_slot_id=-1, - ) - out_ref = causal_conv1d_ref( - x_orig, - conv_state_ref, - weight, - mtp_step=mtp_step, - bias=bias, - activation="silu", - conv_state_indices=conv_state_indices, - num_accepted_tokens=num_accepted_tokens, - query_start_loc=query_start_loc, - pad_slot_id=-1, - ) - - rtol, atol = 1e-2, 1e-2 - assert torch.allclose(out_triton, out_ref, rtol=rtol, atol=atol), ( - f"Output mismatch with pad slots: width={width}, dim={dim}, mtp_step={mtp_step}\n" - f"max diff={torch.abs(out_triton - out_ref).max().item():.6f}" - ) - - for b in [2, 3]: - if conv_state_indices[b] == -1: - start = query_start_loc[b].item() - end = query_start_loc[b + 1].item() - for t_ in range(start, end): - assert torch.allclose( - x_orig[t_], out_triton[t_], rtol=1e-6, atol=1e-6 - ), f"Padded entry should not be modified at batch={b}, token={t_ - start}" - - -@pytest.mark.parametrize("width", [2, 3, 4, 5]) -@pytest.mark.parametrize("dim", [64, 256]) -@pytest.mark.parametrize("mtp_step", [1, 2]) -def test_num_accepted_tokens_varied(width, dim, mtp_step): - """Varying num_accepted_tokens across batch.""" - batch = 3 - x, conv_state, weight, bias, conv_state_indices, num_accepted_tokens, query_start_loc = make_tensors( - batch=batch, - dim=dim, - width=width, - mtp_step=mtp_step, - has_bias=True, - dtype=torch.float16, - ) - - num_accepted_tokens[0] = 1 - num_accepted_tokens[1] = mtp_step + 1 - num_accepted_tokens[2] = max(1, mtp_step // 2 + 1) - - x_orig = x.clone() - conv_state_ref = conv_state.clone() - - out_triton = causal_conv1d_update( - x, - conv_state, - weight, - mtp_step=mtp_step, - bias=bias, - activation="silu", - conv_state_indices=conv_state_indices, - num_accepted_tokens=num_accepted_tokens, - query_start_loc=query_start_loc, - ) - out_ref = causal_conv1d_ref( - x_orig, - conv_state_ref, - weight, - mtp_step=mtp_step, - bias=bias, - activation="silu", - conv_state_indices=conv_state_indices, - num_accepted_tokens=num_accepted_tokens, - query_start_loc=query_start_loc, - ) - - rtol, atol = 1e-2, 1e-2 - assert torch.allclose(out_triton, out_ref, rtol=rtol, atol=atol), ( - f"Output mismatch with varied accepted tokens: width={width}, dim={dim}, mtp_step={mtp_step}\n" - f"max diff={torch.abs(out_triton - out_ref).max().item():.6f}" - ) - - -@pytest.mark.parametrize("width", [2, 3, 4, 5, 6]) -@pytest.mark.parametrize("dim", [64, 128]) -@pytest.mark.parametrize("mtp_step", [0, 1]) -@pytest.mark.parametrize("activation", [None, "silu"]) -def test_conv_state_update_correctness(width, dim, mtp_step, activation): - """Verify that conv_state is updated correctly after the forward pass.""" - batch = 1 - x, conv_state, weight, bias, conv_state_indices, num_accepted_tokens, query_start_loc = make_tensors( - batch=batch, - dim=dim, - width=width, - mtp_step=mtp_step, - has_bias=True, - dtype=torch.float16, - ) - - x_orig = x.clone() - conv_state_before = conv_state.clone() - - causal_conv1d_update( - x, - conv_state, - weight, - mtp_step=mtp_step, - bias=bias, - activation=activation, - conv_state_indices=conv_state_indices, - num_accepted_tokens=num_accepted_tokens, - query_start_loc=query_start_loc, - ) - - seqlen = mtp_step + 1 - state_len = width - 1 + mtp_step - slot = conv_state_indices[0].item() - - expected = torch.zeros_like(conv_state_before[slot]) - offset = num_accepted_tokens[0].item() - 1 - - for k in range(width - 2): - src = offset + k + 1 - if src < state_len: - expected[:, k] = conv_state_before[slot, :, src] - - for t_ in range(seqlen): - idx = (width - 2) + t_ - if idx < state_len: - expected[:, idx] = x_orig[t_, :] - - actual = conv_state[slot] - rtol, atol = 1e-2, 1e-2 - assert torch.allclose(actual, expected, rtol=rtol, atol=atol), ( - f"Conv state mismatch: width={width}, dim={dim}, mtp_step={mtp_step}\n" - f"max diff={torch.abs(actual - expected).max().item():.6f}" - ) - - -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -def test_dtype_consistency(dtype): - """Test that the function runs correctly with different dtypes.""" - width = 4 - dim = 128 - mtp_step = 1 - - x, conv_state, weight, bias, conv_state_indices, num_accepted_tokens, query_start_loc = make_tensors( - batch=2, - dim=dim, - width=width, - mtp_step=mtp_step, - has_bias=True, - dtype=dtype, - ) - x_orig = x.clone() - conv_state_ref = conv_state.clone() - - out_triton = causal_conv1d_update( - x, - conv_state, - weight, - mtp_step=mtp_step, - bias=bias, - activation="silu", - conv_state_indices=conv_state_indices, - num_accepted_tokens=num_accepted_tokens, - query_start_loc=query_start_loc, - ) - - out_ref = causal_conv1d_ref( - x_orig, - conv_state_ref, - weight, - mtp_step=mtp_step, - bias=bias, - activation="silu", - conv_state_indices=conv_state_indices, - num_accepted_tokens=num_accepted_tokens, - query_start_loc=query_start_loc, - ) - - rtol, atol = (1e-2, 1e-2) if dtype == torch.float16 else (2e-2, 2e-2) - assert torch.allclose(out_triton, out_ref, rtol=rtol, atol=atol), ( - f"Output mismatch for dtype {dtype}: " f"max diff={torch.abs(out_triton - out_ref).max().item():.6f}" - ) + offset = accepted[i].item() - 1 + history = state[slot, :, offset : offset + width - 1] + window = torch.cat((history, x[start:end].T), dim=1) + values = ( + F.conv1d( + window.float().unsqueeze(0), + weight.float().unsqueeze(1), + bias=None if bias is None else bias.float(), + groups=x.shape[1], + ) + .squeeze(0) + .T + ) + output[start:end] = F.silu(values) if activation else values + saved = window[:, 1:] + next_state[slot].zero_() + next_state[slot, :, : saved.shape[1]] = saved + return output, next_state @pytest.mark.parametrize("width", [2, 3, 4, 5, 6]) -@pytest.mark.parametrize("activation", [None, "silu"]) -@pytest.mark.parametrize("has_bias", [True, False]) -def test_known_values(width, activation, has_bias): - """Deterministic known values to verify numerically.""" - batch = 1 - dim = 4 - mtp_step = 1 - device = "cuda" - - torch.manual_seed(999) - x = torch.randn(batch * (mtp_step + 1), dim, device=device, dtype=torch.float16) - weight = torch.randn(dim, width, device=device, dtype=torch.float16) - bias = torch.randn(dim, device=device, dtype=torch.float16) if has_bias else None - state_len = width - 1 + mtp_step - conv_state = torch.randn(1, dim, state_len, device=device, dtype=torch.float16) - conv_state.add_(0.5) - - conv_state_indices = torch.zeros(batch, device=device, dtype=torch.int32) - num_accepted_tokens = torch.ones(batch, device=device, dtype=torch.int32) - qsl = torch.arange(batch + 1, device=device, dtype=torch.int32) * (mtp_step + 1) - - x_orig = x.clone() - conv_state_ref = conv_state.clone() - out_triton = causal_conv1d_update( +@pytest.mark.parametrize("mtp_step", [0, 1, 2, 4]) +@pytest.mark.parametrize( + "dtype,dim,has_bias,activation", + [ + (torch.float16, 64, False, None), + (torch.float16, 128, True, "silu"), + (torch.bfloat16, 65, False, "silu"), + (torch.bfloat16, 256, True, None), + ], +) +def test_output_and_state_match_convolution(width, mtp_step, dtype, dim, has_bias, activation): + torch.manual_seed(42) + lengths = [mtp_step + 1, max(1, mtp_step), 0, mtp_step + 1] + starts = torch.tensor([0, *torch.tensor(lengths).cumsum(0).tolist()], dtype=torch.int32, device="cuda") + slots = torch.tensor([2, 0, 3, -1], dtype=torch.int32, device="cuda") + accepted = torch.tensor([1, mtp_step + 1, 1, 1], dtype=torch.int32, device="cuda") + x = torch.randn(sum(lengths), dim, dtype=dtype, device="cuda") + state = torch.randn(5, dim, width - 1 + mtp_step, dtype=dtype, device="cuda") + weight = torch.randn(dim, width, dtype=dtype, device="cuda") + bias = torch.randn(dim, dtype=dtype, device="cuda") if has_bias else None + expected, next_state = reference(x, state, weight, bias, activation, slots, accepted, starts) + + output = causal_conv1d_update( x, - conv_state, + state, weight, mtp_step=mtp_step, bias=bias, activation=activation, - conv_state_indices=conv_state_indices, - num_accepted_tokens=num_accepted_tokens, - query_start_loc=qsl, - ) - out_ref = causal_conv1d_ref( - x_orig, - conv_state_ref, - weight, - mtp_step=mtp_step, - bias=bias, - activation=activation, - conv_state_indices=conv_state_indices, - num_accepted_tokens=num_accepted_tokens, - query_start_loc=qsl, - ) - - rtol, atol = 1e-2, 1e-2 - assert torch.allclose(out_triton, out_ref, rtol=rtol, atol=atol), ( - f"Known value mismatch: width={width}, activation={activation}, bias={has_bias}\n" - f"max diff={torch.abs(out_triton - out_ref).max().item():.6f}" - ) - - -# ============================================================================= -# Edge-case tests for kernel simplification verification -# ============================================================================= - - -@pytest.mark.parametrize("width", [2, 3, 4, 5, 6]) -@pytest.mark.parametrize("mtp_step", [0, 1, 2, 4]) -def test_single_step_kernel_vs_pytorch_conv(width, mtp_step): - """ - Verify the triton kernel against a direct torch.nn.functional.conv1d - reference (transposed to per-channel depthwise conv). This test ensures - the causal_conv1d kernel computes the mathematically correct causal conv. - """ - batch = 1 - dim = 32 - seqlen = mtp_step + 1 - state_len = width - 1 + mtp_step - device = "cuda" - - torch.manual_seed(123) - x = torch.randn(seqlen, dim, device=device, dtype=torch.float32) - weight = torch.randn(dim, width, device=device, dtype=torch.float32) - bias = torch.randn(dim, device=device, dtype=torch.float32) - conv_state = torch.randn(1, dim, state_len, device=device, dtype=torch.float32) - - x_orig = x.clone() - conv_state_ref = conv_state.clone() - - conv_state_indices = torch.zeros(batch, device=device, dtype=torch.int32) - num_accepted_tokens = torch.ones(batch, device=device, dtype=torch.int32) - qsl = torch.tensor([0, seqlen], device=device, dtype=torch.int32) - - out_triton = causal_conv1d_update( - x.clone().half(), - conv_state.clone().half(), - weight.half(), - mtp_step=mtp_step, - bias=bias.half(), - activation=None, - conv_state_indices=conv_state_indices, - num_accepted_tokens=num_accepted_tokens, - query_start_loc=qsl, - ).float() - - # PyTorch causal conv1d reference: pad input with state on the left, - # then run conv1d and slice output - history = conv_state_ref[0, :, : width - 1] - padded_input = torch.cat( - [ - history.T.contiguous(), - x_orig, - ], - dim=0, - ) - - w_3d = weight.unsqueeze(1).float() - pytorch_out = ( - torch.nn.functional.conv1d( - padded_input.T.unsqueeze(0).float(), - w_3d, - bias=bias.float(), - groups=dim, - padding=0, - stride=1, - ) - .squeeze(0) - .T[:seqlen, :] + conv_state_indices=slots, + num_accepted_tokens=accepted, + query_start_loc=starts, ) - rtol, atol = 1e-2, 1e-2 - assert torch.allclose(out_triton, pytorch_out, rtol=rtol, atol=atol), ( - f"Conv1d mismatch against torch.nn.functional.conv1d: " - f"width={width}, mtp_step={mtp_step}\n" - f"max diff={torch.abs(out_triton - pytorch_out).max().item():.6f}" - ) + assert output.data_ptr() == x.data_ptr() + tolerance = 2e-2 if dtype == torch.bfloat16 else 1e-2 + torch.testing.assert_close(output, expected, rtol=tolerance, atol=tolerance) + torch.testing.assert_close(state, next_state, rtol=0, atol=0) + torch.testing.assert_close(output[starts[-2] :], expected[starts[-2] :], rtol=0, atol=0) @pytest.mark.parametrize("width", [2, 3, 4]) -@pytest.mark.parametrize("num_steps", [2, 3, 5]) -def test_multi_step_decode_sliding_window(width, num_steps): - """ - Simulate multiple consecutive decode steps. Each step produces one output - token and shifts a sliding window of conv_state. The outputs over multi-step - must match a single forward pass that sees the full sequence at once. - """ - batch = 1 - dim = 32 - mtp_step = 0 - state_len = width - 1 + mtp_step - device = "cuda" - +@pytest.mark.parametrize("mtp_step", [0, 1, 2, 3]) +def test_repeated_decode_preserves_accepted_history(width, mtp_step): torch.manual_seed(77) - weight = torch.randn(dim, width, device=device, dtype=torch.float32) - bias = torch.randn(dim, device=device, dtype=torch.float32) - init_state = torch.zeros(1, dim, state_len, device=device, dtype=torch.float32) - - all_tokens = torch.randn(num_steps, dim, device=device, dtype=torch.float32) - - idxs = torch.zeros(batch, device=device, dtype=torch.int32) - nat = torch.ones(batch, device=device, dtype=torch.int32) - qsl1 = torch.tensor([0, 1], device=device, dtype=torch.int32) - - conv_state_step = init_state.clone().half() - step_outputs = [] - - for step in range(num_steps): - x_step = all_tokens[step : step + 1, :].clone().half() - step_outputs.append( - causal_conv1d_update( - x_step, - conv_state_step, - weight.half(), - mtp_step=0, - bias=bias.half(), - activation=None, - conv_state_indices=idxs, - num_accepted_tokens=nat, - query_start_loc=qsl1, - ).float() - ) - - step_outputs = torch.cat(step_outputs, dim=0) - - # Build full forward: preload state, then conv1d over full sequence - full_state = init_state.clone().float() - padded = torch.cat( - [ - full_state[0, :, : width - 1].T.contiguous(), - all_tokens, - ], - dim=0, - ) - - w_3d = weight.unsqueeze(1).float() - expected = ( - torch.nn.functional.conv1d( - padded.T.unsqueeze(0).float(), - w_3d, - bias=bias.float(), - groups=dim, - ) - .squeeze(0) - .T[:num_steps, :] - ) - - rtol, atol = 1e-2, 1e-2 - assert torch.allclose(step_outputs, expected, rtol=rtol, atol=atol), ( - f"Multi-step mismatch: width={width}, num_steps={num_steps}\n" - f"max diff={torch.abs(step_outputs - expected).max().item():.6f}" - ) - - -@pytest.mark.parametrize("width", [2, 3, 4]) -@pytest.mark.parametrize("mtp_step", [1, 2, 3]) -def test_mtp_decode_multi_token_per_step(width, mtp_step): - """ - MTP decode: each step processes (mtp_step+1) tokens. After each step, - some tokens are accepted (num_accepted_tokens varies). The conv_state - sliding window must correctly preserve history across steps. - """ - batch = 1 - dim = 32 - seqlen = mtp_step + 1 - state_len = width - 1 + mtp_step - device = "cuda" - num_steps = 3 - - torch.manual_seed(1234) - weight = torch.randn(dim, width, device=device, dtype=torch.float32) - bias = torch.randn(dim, device=device, dtype=torch.float32) - conv_state = torch.randn(1, dim, state_len, device=device, dtype=torch.float32) - conv_state_ref = conv_state.clone() - - idxs = torch.zeros(batch, device=device, dtype=torch.int32) - qsl = torch.tensor([0, seqlen], device=device, dtype=torch.int32) - - # Phase 1: one-step MTP decode triton - x_full = torch.randn(seqlen, dim, device=device, dtype=torch.float32) - out_triton = causal_conv1d_update( - x_full.clone().half(), - conv_state.clone().half(), - weight.half(), - mtp_step=mtp_step, - bias=bias.half(), - activation="silu", - conv_state_indices=idxs, - num_accepted_tokens=torch.ones(batch, device=device, dtype=torch.int32), - query_start_loc=qsl, - ).float() - - # Reference: torch conv1d with state preload, then silu - history = conv_state_ref[0, :, : width - 1].float() - padded = torch.cat([history.T.contiguous(), x_full], dim=0) - w_3d = weight.unsqueeze(1).float() - ref_out = ( - torch.nn.functional.conv1d( - padded.T.unsqueeze(0).float(), - w_3d, - bias=bias.float(), - groups=dim, - ) - .squeeze(0) - .T[:seqlen, :] - ) - ref_out = ref_out / (1 + torch.exp(-ref_out)) - - rtol, atol = 1e-2, 1e-2 - assert torch.allclose(out_triton, ref_out, rtol=rtol, atol=atol), ( - f"MTP decode mismatch: width={width}, mtp_step={mtp_step}\n" - f"max diff={torch.abs(out_triton - ref_out).max().item():.6f}" - ) - - # Phase 2: multi-step MTP decode with varying acceptance - conv_state_step = conv_state.clone().half() - - for step in range(num_steps): - x_step = torch.randn(seqlen, dim, device=device, dtype=torch.float32) - x_step_half = x_step.clone().half() - - accepted = torch.randint(1, seqlen + 1, (batch,), device=device, dtype=torch.int32) - nat_step = accepted.clone() - qsl_step = torch.tensor([0, seqlen], device=device, dtype=torch.int32) - - x_orig_step = x_step_half.clone() - conv_state_before_step = conv_state_step.clone().half() - - step_out = causal_conv1d_update( - x_step_half, - conv_state_step, - weight.half(), - mtp_step=mtp_step, - bias=bias.half(), - activation="silu", - conv_state_indices=idxs, - num_accepted_tokens=nat_step, - query_start_loc=qsl_step, - ).float() - - ref_step = causal_conv1d_ref( - x_orig_step, - conv_state_before_step, - weight.half(), + dim, length = 32, mtp_step + 1 + weight = torch.randn(dim, width, dtype=torch.float16, device="cuda") + bias = torch.randn(dim, dtype=torch.float16, device="cuda") + state = torch.randn(1, dim, width - 1 + mtp_step, dtype=torch.float16, device="cuda") + expected_state = state.clone() + slots = torch.tensor([0], dtype=torch.int32, device="cuda") + starts = torch.tensor([0, length], dtype=torch.int32, device="cuda") + for step in range(5): + accepted = torch.tensor([1 + step % length], dtype=torch.int32, device="cuda") + x = torch.randn(length, dim, dtype=torch.float16, device="cuda") + expected, expected_state = reference(x, expected_state, weight, bias, "silu", slots, accepted, starts) + output = causal_conv1d_update( + x, + state, + weight, mtp_step=mtp_step, - bias=bias.half(), + bias=bias, activation="silu", - conv_state_indices=idxs, - num_accepted_tokens=nat_step, - query_start_loc=qsl_step, - ).float() - - assert torch.allclose(step_out, ref_step, rtol=rtol, atol=atol), ( - f"Multi-step MTP step={step} mismatch: width={width}, mtp_step={mtp_step}\n" - f"max diff={torch.abs(step_out - ref_step).max().item():.6f}" + conv_state_indices=slots, + num_accepted_tokens=accepted, + query_start_loc=starts, ) - - -@pytest.mark.parametrize("mtp_step", [0, 1, 2, 4]) -def test_kernel_width_2_correctness(mtp_step): - """ - KERNEL_WIDTH=2 is a degenerate case: restore_conv_state_len = 0, - meaning no history is preserved in the state update (the mask is always - False). Verify the kernel still computes correct outputs and the single - history token (col0) is properly read from the state offset. - """ - width = 2 - batch = 1 - dim = 64 - seqlen = mtp_step + 1 - state_len = width - 1 + mtp_step - device = "cuda" - - torch.manual_seed(555) - x = torch.randn(seqlen, dim, device=device, dtype=torch.float32) - weight = torch.randn(dim, width, device=device, dtype=torch.float32) - bias = torch.randn(dim, device=device, dtype=torch.float32) - - idxs = torch.zeros(batch, device=device, dtype=torch.int32) - qsl = torch.tensor([0, seqlen], device=device, dtype=torch.int32) - - # Test with multiple offsets (varying num_accepted_tokens) - for nac_val in range(1, min(seqlen + 1, mtp_step + 2)): - conv_state = torch.randn(1, dim, state_len, device=device, dtype=torch.float32) - x_half = x.clone().half() - nat = torch.full((batch,), nac_val, device=device, dtype=torch.int32) - - out_triton = causal_conv1d_update( - x_half, - conv_state.clone().half(), - weight.half(), - mtp_step=mtp_step, - bias=bias.half(), - activation=None, - conv_state_indices=idxs, - num_accepted_tokens=nat, - query_start_loc=qsl, - ).float() - - # Reference: causal conv with history from state at offset nac_val-1 - history = conv_state[0, :, nac_val - 1 : state_len].float() - history_pad = history[:, : width - 1] - padded = torch.cat([history_pad.T.contiguous(), x], dim=0) - w_3d = weight.unsqueeze(1).float() - ref_out = ( - torch.nn.functional.conv1d( - padded.T.unsqueeze(0).float(), - w_3d, - bias=bias.float(), - groups=dim, - ) - .squeeze(0) - .T[:seqlen, :] - ) - - assert torch.allclose(out_triton, ref_out, rtol=1e-2, atol=1e-2), ( - f"Width=2 mismatch: mtp_step={mtp_step}, nac={nac_val}\n" - f"max diff={torch.abs(out_triton - ref_out).max().item():.6f}" - ) - - -def test_kernel_width_2_no_state_history(): - """ - KERNEL_WIDTH=2 with num_accepted_tokens=1 (offset=0): the restore mask - is always False, so the state update is fully overwritten by x tokens - (no history preserved). Verify this yields correct output. - """ - batch, dim, width, mtp_step = 1, 16, 2, 2 - state_len = width - 1 + mtp_step - seqlen = mtp_step + 1 - device = "cuda" - - torch.manual_seed(42) - x = torch.randn(seqlen, dim, device=device, dtype=torch.float32) - weight = torch.randn(dim, width, device=device, dtype=torch.float32) - bias = torch.randn(dim, device=device, dtype=torch.float32) - conv_state = torch.zeros(1, dim, state_len, device=device, dtype=torch.float32) - conv_state[0, :, 0] = 99.0 # Place a known history value at offset 0 - - x_orig = x.clone() - conv_state_ref = conv_state.clone() - - idxs = torch.zeros(batch, device=device, dtype=torch.int32) - nat = torch.ones(batch, device=device, dtype=torch.int32) - qsl = torch.tensor([0, seqlen], device=device, dtype=torch.int32) - - out_triton = causal_conv1d_update( - x.clone().half(), - conv_state.clone().half(), - weight.half(), - mtp_step=mtp_step, - bias=bias.half(), - activation=None, - conv_state_indices=idxs, - num_accepted_tokens=nat, - query_start_loc=qsl, - ).float() - - # Manual reference: for each token t, acc = bias + col0*w0 + x[t]*w1 - # col0 for token 0 is state[0] = 99.0. Then col0 shifts to x[t]. - ref_out = torch.zeros(seqlen, dim, device=device, dtype=torch.float32) - col0 = conv_state_ref[0, :, 0].clone() - for t in range(seqlen): - acc = bias.clone() - acc += col0 * weight[:, 0] - acc += x_orig[t] * weight[:, 1] - ref_out[t] = acc - col0 = x_orig[t].clone() - - assert torch.allclose( - out_triton, ref_out, rtol=1e-2, atol=1e-2 - ), f"Width=2 no-history: max diff={torch.abs(out_triton - ref_out).max().item():.6f}" - - -@pytest.mark.parametrize("width", [2, 3, 4, 5, 6]) -@pytest.mark.parametrize("mtp_step", [0, 1, 2]) -def test_output_overwrites_x_inplace(width, mtp_step): - """ - Verify the kernel indeed overwrites x in-place (out == x). Also - verify that the output values are different from the original input. - """ - batch = 1 - dim = 64 - x, conv_state, weight, bias, conv_state_indices, num_accepted_tokens, query_start_loc = make_tensors( - batch=batch, - dim=dim, - width=width, - mtp_step=mtp_step, - has_bias=True, - dtype=torch.float16, - ) - x_before = x.clone() - - out = causal_conv1d_update( - x, - conv_state, - weight, - mtp_step=mtp_step, - bias=bias, - activation="silu", - conv_state_indices=conv_state_indices, - num_accepted_tokens=num_accepted_tokens, - query_start_loc=query_start_loc, - ) - - assert x.data_ptr() == out.data_ptr(), "out must be the same tensor as x" - assert torch.allclose(x, out, rtol=0, atol=0), "x must equal out after call" - assert not torch.allclose( - x, x_before, rtol=1e-6, atol=1e-6 - ), "x must be overwritten (different from original input)" - - -@pytest.mark.parametrize("width", [3, 4, 5]) -@pytest.mark.parametrize("mtp_step", [0, 1, 2]) -def test_bias_none_gives_same_shape(width, mtp_step): - """bias=None should produce the same output shape as bias specified.""" - batch = 1 - dim = 32 - x, conv_state, weight, _, conv_state_indices, num_accepted_tokens, query_start_loc = make_tensors( - batch=batch, - dim=dim, - width=width, - mtp_step=mtp_step, - has_bias=False, - dtype=torch.float16, - ) - - x_orig = x.clone() - conv_state_ref = conv_state.clone() - - out_triton = causal_conv1d_update( - x, - conv_state, - weight, - mtp_step=mtp_step, - bias=None, - activation=None, - conv_state_indices=conv_state_indices, - num_accepted_tokens=num_accepted_tokens, - query_start_loc=query_start_loc, - ) - out_ref = causal_conv1d_ref( - x_orig, - conv_state_ref, - weight, - mtp_step=mtp_step, - bias=None, - activation=None, - conv_state_indices=conv_state_indices, - num_accepted_tokens=num_accepted_tokens, - query_start_loc=query_start_loc, - ) - - assert out_triton.shape == out_ref.shape - assert out_triton.shape == x_orig.shape - rtol, atol = 1e-2, 1e-2 - assert torch.allclose(out_triton, out_ref, rtol=rtol, atol=atol), ( - f"No-bias mismatch: width={width}, mtp_step={mtp_step}\n" - f"max diff={torch.abs(out_triton - out_ref).max().item():.6f}" - ) - - -if __name__ == "__main__": - pytest.main([__file__, "-v", "--tb=short"]) + torch.testing.assert_close(output, expected, rtol=1e-2, atol=1e-2) + torch.testing.assert_close(state, expected_state, rtol=0, atol=0) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_fused_recurrent_strided.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_fused_recurrent_strided.py deleted file mode 100644 index 969a3f354d..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_fused_recurrent_strided.py +++ /dev/null @@ -1,71 +0,0 @@ -import pytest -import torch - -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.fused_recurrent import ( - fused_recurrent_gated_delta_rule, -) - -if not torch.cuda.is_available(): - pytest.skip("CUDA required", allow_module_level=True) - - -@pytest.mark.parametrize("batch", [1, 2, 16]) -def test_decode_strided_views_match_contiguous(batch): - """q/k/v/a/b passed as column views of one projection output (the decode - path layout) must produce the same result as contiguous copies.""" - torch.manual_seed(0) - H, HV, K, V = 2, 8, 128, 128 - key_dim, value_dim = H * K, HV * V - qkv_dim = 2 * key_dim + value_dim - total_dim = qkv_dim + value_dim + 2 * HV # qkv + z + b + a - cache_slots = 64 - - mixed = torch.randn(batch, total_dim, device="cuda", dtype=torch.bfloat16) - mixed_qkv = mixed[:, :qkv_dim] - b_raw = mixed[:, qkv_dim + value_dim : qkv_dim + value_dim + HV] - a_raw = mixed[:, qkv_dim + value_dim + HV :] - - query, key, value = torch.split(mixed_qkv, [key_dim, key_dim, value_dim], dim=-1) - q = query.view(batch, 1, H, K) - k = key.view(batch, 1, H, K) - v = value.view(batch, 1, HV, V) - - A_log = torch.randn(HV, device="cuda", dtype=torch.float32) * 0.1 - dt_bias = torch.randn(HV, device="cuda", dtype=torch.float32) * 0.1 - ssm_state = torch.randn(cache_slots, HV, K, V, device="cuda", dtype=torch.bfloat16) - idx = torch.randperm(cache_slots, device="cuda")[:batch].to(torch.int32) - - def run(q_, k_, v_, a_, b_, state): - out, _ = fused_recurrent_gated_delta_rule( - q=q_, - k=k_, - v=v_, - initial_state=state, - inplace_final_state=True, - ssm_state_indices=idx, - use_qk_l2norm_in_kernel=True, - A_log=A_log, - dt_bias=dt_bias, - a_raw=a_, - b_raw=b_, - ) - return out - - state_ref = ssm_state.clone() - out_ref = run(q.contiguous(), k.contiguous(), v.contiguous(), a_raw.contiguous(), b_raw.contiguous(), state_ref) - state_strided = ssm_state.clone() - out_strided = run(q, k, v, a_raw, b_raw, state_strided) - - assert torch.equal(out_ref, out_strided) - assert torch.equal(state_ref, state_strided) - - -# NOTE: the decode-only `cu_seqlens is None` contract from upstream #1349 was -# intentionally lifted on this branch so the Qwen3Next MTP verify path can drive -# the kernel with variable-length verify chunks (cu_seqlens + 2D SSM index -# rows). That varlen path is exercised end-to-end by the MTP GSM8K accuracy -# check rather than a hand-rolled unit test. - - -if __name__ == "__main__": - pytest.main([__file__, "-v"]) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_kda_autotune.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_kda_autotune.py deleted file mode 100644 index 8d92d6cc8d..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_kda_autotune.py +++ /dev/null @@ -1,287 +0,0 @@ -import collections -import functools -import importlib -import inspect -import json -from itertools import accumulate - -import pytest -import torch -import triton - -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops import kda -from lightllm.common.kernel_config import KernelConfigs -from lightllm.common.triton_utils import autotuner as autotuner_module -from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneLevel - - -KERNELS = [ - ("_chunk_kda_scaled_dot_kkt_sub_inter", "chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter"), - ("_chunk_kda_scaled_dot_kkt_sub_intra", "chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra"), - ("recompute_w_u_fwd", "recompute_w_u_fwd_kernel"), - ("chunk_gla_fwd_o_gk", "chunk_gla_fwd_kernel_o"), -] - - -@pytest.fixture(autouse=True) -def autotune_environment(monkeypatch): - torch.manual_seed(42) - monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", None) - monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: AutotuneLevel.CLOSE_AUTOTUNE) - - -def wrapper_inputs(kernel, tokens=65): - q = torch.empty(1, tokens, 2, 128, dtype=torch.bfloat16) - v = torch.empty(1, tokens, 2, 80, dtype=q.dtype) - g = torch.empty_like(q, dtype=torch.float32) - chunk_indices = torch.tensor([[0, 0], *[[1, i] for i in range(triton.cdiv(tokens - 3, 64))]]) - values = dict( - q=q, - k=q, - v=v, - g=g, - gk=g, - beta=torch.empty(1, tokens, 2), - Akk=torch.empty(1, tokens, 2, 64), - Aqk=torch.empty(1, tokens, 2, 64), - R=torch.empty(1, tokens, 2, 64, dtype=q.dtype), - h=torch.empty(1, len(chunk_indices), 2, 128, 80, dtype=q.dtype), - o=torch.empty_like(v), - scale=128 ** -0.5, - cu_seqlens=torch.tensor([0, 3, tokens]), - chunk_indices=chunk_indices, - ) - return {name: values[name] for name in inspect.signature(kernel.fn).parameters if name in values} - - -@pytest.mark.parametrize("wrapper_name,jit_name", KERNELS) -def test_lightllm_tuning_selects_persists_and_reuses_launch_config(monkeypatch, tmp_path, wrapper_name, jit_name): - kernel = getattr(kda, wrapper_name) - assert isinstance(kernel, Autotuner) - inputs = wrapper_inputs(kernel) - candidates = [kernel.configs_gen_func()[0], kernel.configs_gen_func()[-1]] - launches, benchmarks = [], [] - - class LaunchRecorder: - def __getitem__(self, grid): - def launch(**kwargs): - launches.append((grid, kwargs)) - - return launch - - def bench(*args, run_config, **kwargs): - benchmarks.append(run_config) - return 1.0 if run_config == candidates[-1] else 2.0 - - monkeypatch.setattr(kda, jit_name, LaunchRecorder()) - monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: AutotuneLevel.ADAPTIVE_AUTOTUNE) - monkeypatch.setattr(autotuner_module.dist, "is_initialized", lambda: False) - monkeypatch.setattr("lightllm.common.kernel_config.get_current_device_name", lambda: "test-device") - monkeypatch.setattr(kernel, "_cache_dir", str(tmp_path), raising=False) - monkeypatch.setattr(kernel, "cached_configs", {}) - monkeypatch.setattr(kernel, "fast_match_configs", collections.defaultdict(dict)) - monkeypatch.setattr(kernel, "warmuped_configs_set", set()) - monkeypatch.setattr(kernel, "configs_gen_func", lambda: candidates) - monkeypatch.setattr(kernel, "_bench", bench) - - with Autotuner.autotune_warmup(): - kernel(**inputs) - assert benchmarks == candidates - for name, value in candidates[-1].items(): - assert launches[-1][1][name] == value - - cache_file = tmp_path / KernelConfigs.get_config_file_name(kernel._static_key(**inputs)) - assert json.loads(cache_file.read_text()) == {"65": candidates[-1]} - - # A new wrapper invocation must reload the same selected config from disk. - kernel.cached_configs.clear() - kernel.fast_match_configs.clear() - kernel.warmuped_configs_set.clear() - benchmarks.clear() - kernel(**inputs) - assert benchmarks == [] - for name, value in candidates[-1].items(): - assert launches[-1][1][name] == value - - # Packed B=1 still needs different run keys as the token count grows. - assert kernel._run_key(**inputs) == 65 - assert kernel._run_key(**wrapper_inputs(kernel, tokens=129)) == 129 - assert kernel._static_key(**inputs) == kernel._static_key(**wrapper_inputs(kernel, tokens=129)) - # Request boundaries are runtime data; all calls use the packed path. - single_request = dict(inputs, cu_seqlens=torch.tensor([0, 65])) - assert kernel._static_key(**inputs) == kernel._static_key(**single_request) - - -def recurrent_reference(q, k, v, g, beta, initial_state, sequences): - q, k, v, g, beta = [x.float().cpu() for x in (q, k, v, g, beta)] - q = q / (q.square().sum(-1, keepdim=True) + 1e-6).sqrt() / q.shape[-1] ** 0.5 - k = k / (k.square().sum(-1, keepdim=True) + 1e-6).sqrt() - output = torch.empty_like(v) - final = initial_state.float().cpu().clone() - for n, (batch, start, end) in enumerate(sequences): - state = final[n] - for t in range(start, end): - kt = k[batch, t] - state = state * g[batch, t].exp()[..., None] - delta = beta[batch, t, :, None] * (v[batch, t] - torch.einsum("hk,hkv->hv", kt, state)) - state = state + kt[..., None] * delta[:, None, :] - output[batch, t] = torch.einsum("hk,hkv->hv", q[batch, t], state) - final[n] = state - return output, final - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("config_index", [None, 0, -1], ids=["default", "first-candidate", "last-candidate"]) -@pytest.mark.parametrize( - "seq_lens,key_dim,value_dim", - [((65, 65), 64, 96), ((3, 65, 129), 128, 80)], - ids=["equal-length-packed", "ragged-packed"], -) -def test_kda_configs_match_token_recurrence(monkeypatch, config_index, seq_lens, key_dim, value_dim): - triton.set_allocator(lambda size, alignment, stream: torch.empty(size, device="cuda", dtype=torch.int8)) - if config_index is not None: - for wrapper_name, _ in KERNELS: - kernel = getattr(kda, wrapper_name) - config = kernel.configs_gen_func()[config_index] - monkeypatch.setattr(kda, wrapper_name, functools.partial(kernel, run_config=config)) - - # Include partial final chunks, distinct K/V sizes, and V tiles with masked columns. - batch, tokens, heads = 1, sum(seq_lens), 2 - boundaries = [0, *accumulate(seq_lens)] - sequences = [(0, start, end) for start, end in zip(boundaries, boundaries[1:])] - cu_seqlens = torch.tensor(boundaries, device="cuda", dtype=torch.int32) - q, k = [torch.randn(batch, tokens, heads, key_dim, device="cuda", dtype=torch.bfloat16) for _ in range(2)] - v = torch.randn(batch, tokens, heads, value_dim, device="cuda", dtype=torch.bfloat16) - g = -torch.rand(batch, tokens, heads, key_dim, device="cuda") * 0.1 - beta = torch.rand(batch, tokens, heads, device="cuda") - initial = torch.randn(len(sequences), heads, key_dim, value_dim, device="cuda") * 0.1 - initial_before = initial.clone() - expected, expected_final = recurrent_reference(q, k, v, g, beta, initial, sequences) - - actual, final = kda.chunk_kda( - q=q, - k=k, - v=v, - g=g, - beta=beta, - initial_state=initial, - output_final_state=True, - use_qk_l2norm_in_kernel=True, - cu_seqlens=cu_seqlens, - ) - torch.testing.assert_close(actual.float().cpu(), expected, atol=4e-3, rtol=3e-2) - torch.testing.assert_close(final.cpu(), expected_final, atol=8e-3, rtol=3e-2) - torch.testing.assert_close(initial, initial_before, atol=0, rtol=0) - - -@pytest.mark.parametrize("fused_gate", [False, True]) -@pytest.mark.parametrize("invalid_input", ["batch-dimension", "missing-boundaries"]) -def test_kda_requires_packed_inputs(fused_gate, invalid_input): - batch = 2 if invalid_input == "batch-dimension" else 1 - q = torch.empty(batch, 3, 2, 128) - inputs = dict( - q=q, - k=q, - v=q, - beta=torch.empty(batch, 3, 2), - cu_seqlens=None if invalid_input == "missing-boundaries" else torch.tensor([0, 3, 6]), - ) - if fused_gate: - kernel = kda.chunk_kda_with_fused_gate - inputs.update(raw_g=q, A_log=torch.empty(2), g_bias=None) - else: - kernel = kda.chunk_kda - inputs.update(g=q) - expected_message = "packed q" if invalid_input == "batch-dimension" else "cu_seqlens is required" - with pytest.raises(AssertionError, match=expected_message): - kernel(**inputs) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize( - "heads,input_dtype,output_dtype,precision", - [ - (2, torch.float32, torch.bfloat16, "ieee"), - (32, torch.bfloat16, torch.bfloat16, "ieee"), - (64, torch.float16, torch.float16, "ieee"), - (64, torch.float32, torch.float32, "ieee"), - (64, torch.float32, torch.bfloat16, "tf32"), - (64, torch.float32, torch.bfloat16, "tf32x3"), - ], -) -def test_split_diagonal_matches_unsplit(monkeypatch, heads, input_dtype, output_dtype, precision): - if precision != "ieee" and torch.cuda.get_device_capability()[0] < 8: - pytest.skip("TF32 requires SM80+") - solver = importlib.import_module("lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.solve_tril") - monkeypatch.setattr(solver, "FLA_TRIL_PRECISION", precision) - lengths = (1, 16, 1519) - cu_seqlens = torch.tensor([0, *accumulate(lengths)], device="cuda", dtype=torch.int32) - rows = torch.cat([torch.arange(length, device="cuda") % 64 for length in lengths]) - matrix = torch.randn(1, sum(lengths), heads, 64, device="cuda", dtype=input_dtype) * 0.025 - matrix.mul_((rows[:, None] > torch.arange(64, device="cuda"))[None, :, None, :]) - expected = solver.solve_tril(matrix, cu_seqlens, output_dtype) - launches = [] - original = solver._invert_64_diagonal_tiles.run - - def observe(*args, **kwargs): - launches.append(True) - return original(*args, **kwargs) - - monkeypatch.setattr(solver._invert_64_diagonal_tiles, "run", observe) - actual = solver.solve_tril(matrix, cu_seqlens, output_dtype, split_diagonal=True) - assert len(launches) == 1 - assert torch.equal(actual.view(torch.uint8), expected.view(torch.uint8)) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("batch,tokens,heads", [(1, 64, 2), (2, 65, 3), (3, 129, 2)]) -def test_split_diagonal_matches_unsplit_dense(monkeypatch, batch, tokens, heads): - solver = importlib.import_module("lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.solve_tril") - rows = torch.arange(tokens, device="cuda") % 64 - matrix = torch.randn(batch, tokens, heads, 64, device="cuda") * 0.025 - matrix.mul_((rows[:, None] > torch.arange(64, device="cuda"))[None, :, None, :]) - expected = solver.solve_tril(matrix, output_dtype=torch.bfloat16) - launches = [] - original = solver._invert_64_diagonal_tiles.run - - def observe(*args, **kwargs): - launches.append(True) - return original(*args, **kwargs) - - monkeypatch.setattr(solver._invert_64_diagonal_tiles, "run", observe) - actual = solver.solve_tril(matrix, output_dtype=torch.bfloat16, split_diagonal=True) - assert len(launches) == 1 - assert torch.equal(actual.view(torch.uint8), expected.view(torch.uint8)) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize( - "lengths,expect_split", - [ - ((512,) * 16, True), - ((64,) * 128, True), - ((16,) * 16, True), - ((128,) * 128, True), - ((1,) * 256, False), - ((32,), False), - ], - ids=["16-requests", "128-requests", "quarter-full-chunks", "16k-tokens", "mostly-empty-chunks", "tiny-batch"], -) -def test_split_diagonal_uses_chunk_occupancy(monkeypatch, lengths, expect_split): - solver = importlib.import_module("lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.solve_tril") - cu_seqlens = torch.tensor([0, *accumulate(lengths)], device="cuda", dtype=torch.int32) - rows = torch.cat([torch.arange(length, device="cuda") % 64 for length in lengths]) - matrix = torch.randn(1, sum(lengths), 2, 64, device="cuda") * 0.025 - matrix.mul_((rows[:, None] > torch.arange(64, device="cuda"))[None, :, None, :]) - expected = solver.solve_tril(matrix, cu_seqlens, torch.bfloat16) - launches = [] - original = solver._invert_64_diagonal_tiles.run - - def observe(*args, **kwargs): - launches.append(True) - return original(*args, **kwargs) - - monkeypatch.setattr(solver._invert_64_diagonal_tiles, "run", observe) - actual = solver.solve_tril(matrix, cu_seqlens, torch.bfloat16, split_diagonal=True) - assert len(launches) == int(expect_split) - assert torch.equal(actual.view(torch.uint8), expected.view(torch.uint8)) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_autotune.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_autotune.py deleted file mode 100644 index 2f8544cf75..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_autotune.py +++ /dev/null @@ -1,374 +0,0 @@ -import collections -import inspect -import json -import math - -import pytest -import torch -import torch.nn.functional as F - -from lightllm.common.basemodel.triton_kernel.linear_att import ( - mtp_fused_recurrent as kernel_module, -) -from lightllm.common.triton_utils import autotuner as autotuner_module -from lightllm.common.triton_utils.autotuner import ( - Autotuner, - AutotuneKernelType, - AutotuneLevel, -) - - -@pytest.fixture(autouse=True) -def autotune_environment(monkeypatch, tmp_path): - torch.manual_seed(42) - monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", None) - monkeypatch.setattr( - autotuner_module, - "get_triton_autotune_level", - lambda: AutotuneLevel.ADAPTIVE_AUTOTUNE, - ) - kernel = kernel_module.mtp_fused_recurrent_gated_delta_rule - monkeypatch.setattr(kernel, "_cache_dir", str(tmp_path), raising=False) - monkeypatch.setattr(kernel, "cached_configs", {}) - monkeypatch.setattr(kernel, "fast_match_configs", collections.defaultdict(dict)) - monkeypatch.setattr(kernel, "warmuped_configs_set", set()) - - -def make_inputs( - lengths=(3, 2, 0), - mtp_size=3, - device="cpu", - head_k_dim=64, - head_v_dim=64, - heads=2, - value_heads=4, - dtype=torch.bfloat16, - state_dtype=torch.float32, - separate_write=False, - padding=0, -): - num_seqs = len(lengths) - num_tokens = sum(lengths) + padding - # 模拟生产环境从 mixed QKV 和 gate 投影切分的非连续 token 视图。 - mixed = torch.randn( - num_tokens, - heads * head_k_dim * 2 + value_heads * head_v_dim, - dtype=dtype, - device=device, - ) - q, k, v = mixed.split([heads * head_k_dim, heads * head_k_dim, value_heads * head_v_dim], dim=-1) - gates = torch.randn(num_tokens, value_heads * 3, dtype=dtype, device=device) - indices = torch.arange(1, num_seqs * mtp_size + 1, dtype=torch.int32, device=device).view(num_seqs, mtp_size) - return dict( - q=q.view(1, num_tokens, heads, head_k_dim), - k=k.view(1, num_tokens, heads, head_k_dim), - v=v.view(1, num_tokens, value_heads, head_v_dim), - initial_state=torch.randn( - 2 * num_seqs * mtp_size + 5, - value_heads, - head_k_dim, - head_v_dim, - dtype=state_dtype, - device=device, - ) - * 0.25, - cu_seqlens=torch.tensor( - [0, *torch.tensor(lengths).cumsum(0).tolist()], - dtype=torch.int64, - device=device, - ), - ssm_state_indices=indices, - ssm_state_write_indices=(indices + num_seqs * mtp_size if separate_write else indices), - num_accepted_tokens=torch.tensor( - [i % mtp_size + 1 for i in range(num_seqs)], - dtype=torch.int32, - device=device, - ), - A_log=torch.randn(value_heads, dtype=torch.float32, device=device) * 0.1, - dt_bias=torch.randn(value_heads, dtype=torch.float32, device=device) * 0.1, - a_raw=gates[:, :value_heads], - b_raw=gates[:, value_heads : 2 * value_heads], - ) - - -def rebuild(inputs): - args, kwargs = kernel_module.rebuild_inputs(**inputs) - return inspect.signature(kernel_module.mtp_fused_recurrent_gated_delta_rule.fn).bind(*args, **kwargs).arguments - - -def reference(inputs): - state = inputs["initial_state"].clone() - output = torch.zeros_like(inputs["v"], dtype=torch.float32) - q, k, v = [inputs[name][0].float() for name in ["q", "k", "v"]] - group_size = v.shape[1] // q.shape[1] - q = (q / (q.square().sum(-1, keepdim=True) + 1e-6).sqrt()).repeat_interleave(group_size, dim=1) - k = (k / (k.square().sum(-1, keepdim=True) + 1e-6).sqrt()).repeat_interleave(group_size, dim=1) - q *= q.shape[-1] ** -0.5 - cumulative = inputs["cu_seqlens"].tolist() - for row, (start, end) in enumerate(zip(cumulative, cumulative[1:])): - if start == end: - continue - read_index = inputs["ssm_state_indices"][row, inputs["num_accepted_tokens"][row] - 1] - h = inputs["initial_state"][read_index].float().clone() - for token in range(start, end): - g = -inputs["A_log"].float().exp() * F.softplus(inputs["a_raw"][token].float() + inputs["dt_bias"].float()) - h *= g.exp()[:, None, None] - delta = (v[token] - (h * k[token, :, :, None]).sum(1)) * inputs["b_raw"][token].float().sigmoid()[:, None] - h += k[token, :, :, None] * delta[:, None, :] - output[0, token] = (h * q[token, :, :, None]).sum(1) - state[inputs["ssm_state_write_indices"][row, token - start]] = h.to(state.dtype) - return output, state - - -@pytest.mark.parametrize("lengths,mtp_size,padding", [((3, 3), 3, 0), ((0,) * 7, 3, 7), ((2, 1, 0), 4, 2)]) -def test_rebuild_uses_independent_state_slots(lengths, mtp_size, padding): - inputs = make_inputs(lengths=lengths, mtp_size=mtp_size, padding=padding) - snapshots = {name: value.clone() for name, value in inputs.items()} - rebuilt = rebuild(inputs) - tokens = inputs["q"].shape[1] - active_seqs = (tokens + mtp_size - 1) // mtp_size - assert rebuilt["cu_seqlens"].tolist() == [min(i * mtp_size, tokens) for i in range(len(lengths) + 1)] - assert rebuilt["initial_state"] is inputs["initial_state"] - assert rebuilt["num_accepted_tokens"].tolist() == [1] * len(lengths) - assert rebuilt["ssm_state_write_indices"][:active_seqs].flatten().tolist() == list(range(active_seqs * mtp_size)) - kernel = kernel_module.mtp_fused_recurrent_gated_delta_rule - assert kernel._run_key(**inputs) == kernel._run_key(**rebuilt) - for name in ["q", "k", "v", "A_log", "dt_bias", "a_raw", "b_raw"]: - assert rebuilt[name] is inputs[name] - for name, snapshot in snapshots.items(): - torch.testing.assert_close(inputs[name], snapshot) - - -def test_rebuild_checks_active_state_capacity(): - inputs = make_inputs(lengths=(0,) * 7, padding=7) - # 7 个有效 token 只需 7 个槽；空序列及末组未使用的索引取余后也应在池内。 - inputs["initial_state"] = inputs["initial_state"][:7] - rebuilt = rebuild(inputs) - assert rebuilt["initial_state"] is inputs["initial_state"] - for name in ["ssm_state_indices", "ssm_state_write_indices"]: - indices = rebuilt[name] - assert torch.all((indices >= 0) & (indices < 7)) - assert indices.flatten()[:7].tolist() == list(range(7)) - assert indices[2].tolist() == [6, 0, 1] - # 槽数不足时仍拒绝调优，不能靠取余让有效请求的状态槽相互重叠。 - inputs["initial_state"] = inputs["initial_state"][:6] - with pytest.raises(AssertionError, match="Not enough SSM state slots"): - rebuild(inputs) - - -def test_keys_use_mtp_workload_and_state_dtype(monkeypatch): - # 历史上下文长度与递推计算量无关，不应读取 full attention 的调优长度环境变量。 - monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", "invalid") - kernel = kernel_module.mtp_fused_recurrent_gated_delta_rule - inputs = make_inputs() - key = kernel._static_key(**inputs) - with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - assert kernel._run_key(**inputs) == 12 - assert kernel._run_key(**make_inputs(lengths=(3, 3, 0))) == kernel._run_key(**inputs) - assert kernel._run_key(**make_inputs(lengths=(3, 2))) == kernel._run_key(**inputs) - assert kernel._run_key(**make_inputs(mtp_size=4)) == 16 - # 不同序列数共享 token 桶，11/12 个 token 共用一桶，第 13 个 token 进入下一桶。 - assert kernel._run_key(**make_inputs(lengths=(3, 3, 3, 2, 0))) == 12 - assert kernel._run_key(**make_inputs(lengths=(3, 3, 3, 3, 0))) == 12 - assert kernel._run_key(**make_inputs(lengths=(3, 3, 3, 3, 1))) == 24 - # head 数不再编码进 run key，但不同模型/TP 的 head 规模应使用不同静态配置文件。 - more_heads = make_inputs(heads=4, value_heads=8) - assert kernel._run_key(**more_heads) == kernel._run_key(**inputs) - assert kernel._static_key(**more_heads) != key - assert kernel._static_key(**make_inputs(state_dtype=torch.bfloat16)) != key - assert kernel._static_key(**make_inputs(mtp_size=4)) != key - assert kernel._static_key(**make_inputs(padding=1)) == key - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize( - "lengths,mtp_size,head_k_dim,head_v_dim,heads,value_heads,dtype,state_dtype,separate_write", - [ - ((1, 0, 1), 1, 64, 64, 2, 8, torch.bfloat16, torch.bfloat16, False), - ((4, 2, 0), 4, 128, 128, 4, 8, torch.bfloat16, torch.float32, True), - ((2, 3, 1), 3, 64, 80, 2, 2, torch.float16, torch.float32, False), - ((3, 3), 3, 128, 128, 16, 16, torch.bfloat16, torch.float32, False), - ((3, 1), 3, 256, 128, 2, 4, torch.bfloat16, torch.bfloat16, True), - ], -) -def test_all_candidates_match_fp32( - lengths, - mtp_size, - head_k_dim, - head_v_dim, - heads, - value_heads, - dtype, - state_dtype, - separate_write, -): - inputs = make_inputs( - lengths, - mtp_size, - "cuda", - head_k_dim, - head_v_dim, - heads, - value_heads, - dtype, - state_dtype, - separate_write, - ) - output_ref, state_ref = reference(inputs) - state_before = inputs["initial_state"].clone() - kernel = kernel_module.mtp_fused_recurrent_gated_delta_rule - for config in kernel_module.get_test_configs(): - inputs["initial_state"].copy_(state_before) - output, state = kernel(**inputs, run_config=config) - assert state is inputs["initial_state"] - torch.testing.assert_close(output.float(), output_ref, atol=2e-3, rtol=2e-2) - torch.testing.assert_close(state.float(), state_ref.float(), atol=2e-3, rtol=2e-2) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("dynamic", [False, True]) -def test_autotune_reuses_state_cache_and_graph(monkeypatch, tmp_path, dynamic): - inputs = make_inputs( - lengths=(3, 2, 0) if dynamic else (3, 3), - device="cuda", - padding=2 if dynamic else 0, - ) - kernel = kernel_module.mtp_fused_recurrent_gated_delta_rule - state_before = inputs["initial_state"].clone() - output_ref, state_ref = reference(inputs) - benchmark = kernel._bench - timings = [] - num_configs = len(kernel_module.get_test_configs()) - post_tuning_reference = None - - def checked_bench(*args, **kwargs): - nonlocal post_tuning_reference - rebuilt = inspect.signature(kernel.fn).bind(*args, **kwargs).arguments - assert rebuilt["cu_seqlens"].tolist() == ([0, 3, 6, 7] if dynamic else [0, 3, 6]) - assert rebuilt["initial_state"] is inputs["initial_state"] - elapsed = benchmark(*args, **kwargs) - timings.append(elapsed) - # benchmark 原地更新池内前 num_tokens 个槽，未使用的槽应保持不变。 - assert torch.isfinite(rebuilt["initial_state"]).all() - num_tokens = inputs["q"].shape[1] - assert not torch.equal(inputs["initial_state"][:num_tokens], state_before[:num_tokens]) - torch.testing.assert_close(inputs["initial_state"][num_tokens:], state_before[num_tokens:], atol=0, rtol=0) - if len(timings) == num_configs: - # 最后一次 benchmark 后的状态作为正式执行的初值，不再假定调优前后状态隔离。 - post_tuning_reference = reference(inputs) - return elapsed - - monkeypatch.setattr(kernel, "_bench", checked_bench) - # 普通 warmup 不搜索；decode 调优会原地更新状态，随后按原始请求布局正式执行。 - with Autotuner.autotune_warmup(AutotuneKernelType.GENERAL): - kernel(**inputs) - assert not timings - inputs["initial_state"].copy_(state_before) - with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - output, _ = kernel(**inputs) - assert len(timings) == num_configs and all(math.isfinite(t) for t in timings) - tuned_output_ref, tuned_state_ref = post_tuning_reference - valid_tokens = int(inputs["cu_seqlens"][-1]) - torch.testing.assert_close( - output[:, :valid_tokens].float(), - tuned_output_ref[:, :valid_tokens], - atol=2e-3, - rtol=2e-2, - ) - torch.testing.assert_close(inputs["initial_state"], tuned_state_ref, atol=2e-3, rtol=2e-2) - config_path = next(tmp_path.glob("*.json")) - assert str(kernel._run_key(**inputs)) in json.loads(config_path.read_text()) - - # 关闭已有配置预热后，文件重载和 FORCE 复用都只能正式执行一次，即使在 decode warmup 阶段。 - monkeypatch.setattr(kernel, "cached_configs", {}) - monkeypatch.setattr(kernel, "fast_match_configs", collections.defaultdict(dict)) - monkeypatch.setattr( - autotuner_module, - "get_triton_autotune_level", - lambda: AutotuneLevel.FORCE_AUTOTUNE, - ) - inputs["initial_state"].copy_(state_before) - with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - kernel(**inputs) - assert len(timings) == num_configs - torch.testing.assert_close(inputs["initial_state"], state_ref, atol=2e-3, rtol=2e-2) - assert not kernel.warmuped_configs_set - # 预热结束后恢复测试请求的初始状态，正常调用及 Graph 回放每次只推进一次。 - inputs["initial_state"].copy_(state_before) - kernel(**inputs) - torch.testing.assert_close(inputs["initial_state"], state_ref, atol=2e-3, rtol=2e-2) - - torch.cuda.synchronize() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - graph_output, graph_state = kernel(**inputs) - inputs["initial_state"].copy_(state_before) - graph.replay() - torch.cuda.synchronize() - torch.testing.assert_close( - graph_output[:, :valid_tokens].float(), - output_ref[:, :valid_tokens], - atol=2e-3, - rtol=2e-2, - ) - torch.testing.assert_close(graph_state, state_ref, atol=2e-3, rtol=2e-2) - # 同一 Graph 更新变长分组和接受位置，确保回放依赖原始输入而非调优构造的元数据。 - inputs["cu_seqlens"].copy_(torch.tensor([0, 1, 3, 3] if dynamic else [0, 2, 3], device="cuda")) - inputs["num_accepted_tokens"].fill_(2) - inputs["initial_state"].copy_(state_before) - new_output_ref, new_state_ref = reference(inputs) - graph.replay() - torch.cuda.synchronize() - torch.testing.assert_close(graph_output[:, :3].float(), new_output_ref[:, :3], atol=2e-3, rtol=2e-2) - torch.testing.assert_close(graph_state, new_state_ref, atol=2e-3, rtol=2e-2) - assert len(timings) == num_configs - - # 改变序列数和 token 数但仍在同一桶内，decode warmup 应复用已有配置，不再搜索。 - same_bucket_inputs = make_inputs(lengths=(3,), device="cuda") - assert kernel._static_key(**same_bucket_inputs) == kernel._static_key(**inputs) - assert kernel._run_key(**same_bucket_inputs) == kernel._run_key(**inputs) - same_output_ref, same_state_ref = reference(same_bucket_inputs) - with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - same_output, same_state = kernel(**same_bucket_inputs) - assert len(timings) == num_configs - torch.testing.assert_close(same_output.float(), same_output_ref, atol=2e-3, rtol=2e-2) - torch.testing.assert_close(same_state, same_state_ref, atol=2e-3, rtol=2e-2) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("level", [0, 1, 2]) -def test_first_decode_loads_history_without_advancing_state_twice(monkeypatch, tmp_path, level): - monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: level) - inputs = make_inputs(lengths=(3, 3), device="cuda") - # 模拟两个刚完成 prefill 的请求，从各自的 canonical 状态槽开始执行首次 decode。 - inputs["ssm_state_indices"] = torch.arange(6, dtype=torch.int32, device="cuda").view(2, 3) - inputs["ssm_state_write_indices"] = inputs["ssm_state_indices"] - inputs["num_accepted_tokens"].fill_(1) - output_ref, state_ref = reference(inputs) - kernel = kernel_module.mtp_fused_recurrent_gated_delta_rule - config = {"BV": 8, "num_warps": 1, "num_stages": 1} - filename = autotuner_module.KernelConfigs.get_config_file_name(kernel._static_key(**inputs)) - cache_file = tmp_path / filename - cache_file.write_text( - json.dumps( - { - str(kernel._run_key(**inputs)): config, - "24": {"BV": 16, "num_warps": 2, "num_stages": 1}, - } - ) - ) - calls = [] - fn = kernel.fn - - def count_calls(*args, **kwargs): - calls.append(kwargs.get("run_config")) - return fn(*args, **kwargs) - - monkeypatch.setattr(kernel, "fn", count_calls) - assert not Autotuner.is_autotune_warmup() - output, state = kernel(**inputs) - assert calls == [config] - assert state is inputs["initial_state"] - assert not kernel.warmuped_configs_set - torch.testing.assert_close(output.float(), output_ref, atol=2e-3, rtol=2e-2) - torch.testing.assert_close(state, state_ref, atol=2e-3, rtol=2e-2) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_equiv.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_equiv.py deleted file mode 100644 index 11035c83a3..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_equiv.py +++ /dev/null @@ -1,167 +0,0 @@ -import pytest -import torch - -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.fused_recurrent import ( - fused_recurrent_gated_delta_rule, -) -from lightllm.common.basemodel.triton_kernel.linear_att.mtp_fused_recurrent import ( - mtp_fused_recurrent_gated_delta_rule, -) - -if not torch.cuda.is_available(): - pytest.skip("CUDA required", allow_module_level=True) - - -def _run_both( - q, - k, - v, - initial_state, - cu_seqlens, - ssm_state_indices, - ssm_state_write_indices, - num_accepted_tokens, - A_log, - dt_bias, - a_raw, - b_raw, -): - """Run old (via autograd.Function) and new (direct kernel) side-by-side.""" - state_old = initial_state.clone() - state_new = initial_state.clone() - - o_old, fs_old = fused_recurrent_gated_delta_rule( - q=q, - k=k, - v=v, - initial_state=state_old, - inplace_final_state=True, - cu_seqlens=cu_seqlens, - ssm_state_indices=ssm_state_indices, - ssm_state_write_indices=ssm_state_write_indices, - num_accepted_tokens=num_accepted_tokens, - use_qk_l2norm_in_kernel=True, - A_log=A_log, - dt_bias=dt_bias, - a_raw=a_raw, - b_raw=b_raw, - ) - - o_new, fs_new = mtp_fused_recurrent_gated_delta_rule( - q=q, - k=k, - v=v, - initial_state=state_new, - cu_seqlens=cu_seqlens, - ssm_state_indices=ssm_state_indices, - ssm_state_write_indices=ssm_state_write_indices, - num_accepted_tokens=num_accepted_tokens, - A_log=A_log, - dt_bias=dt_bias, - a_raw=a_raw, - b_raw=b_raw, - ) - - return o_old, o_new, fs_old, fs_new - - -@pytest.mark.parametrize("mtp_step", [1, 2, 3]) -def test_mtp_verify_path(mtp_step): - """MTP verify path with cu_seqlens and 2D SSM indices.""" - torch.manual_seed(123) - batch, H, HV, K, V = 2, 2, 8, 64, 64 - seqlen = mtp_step + 1 - num_tokens = batch * seqlen - cache_slots = 64 - - q = torch.randn(1, num_tokens, H, K, device="cuda", dtype=torch.bfloat16) - k = torch.randn(1, num_tokens, H, K, device="cuda", dtype=torch.bfloat16) - v = torch.randn(1, num_tokens, HV, V, device="cuda", dtype=torch.bfloat16) - A_log = torch.randn(HV, device="cuda", dtype=torch.float32) * 0.1 - dt_bias = torch.randn(HV, device="cuda", dtype=torch.float32) * 0.1 - a_raw = torch.randn(num_tokens, HV, device="cuda", dtype=torch.bfloat16) - b_raw = torch.randn(num_tokens, HV, device="cuda", dtype=torch.bfloat16) - ssm_state = torch.randn(cache_slots, HV, K, V, device="cuda", dtype=torch.bfloat16) - ssm_idx = torch.randint(0, cache_slots, (batch, seqlen), device="cuda", dtype=torch.int32) - cu_seqlens = torch.arange(batch + 1, device="cuda", dtype=torch.int32) * seqlen - num_accepted = torch.full((batch,), seqlen, device="cuda", dtype=torch.int32) - - o_old, o_new, fs_old, fs_new = _run_both( - q, - k, - v, - ssm_state, - cu_seqlens.to(torch.long), - ssm_idx, - ssm_idx, - num_accepted, - A_log, - dt_bias, - a_raw, - b_raw, - ) - - assert torch.equal( - o_old, o_new - ), f"output mismatch, max diff={torch.abs(o_old.float() - o_new.float()).max().item():.6f}" - if not torch.equal(fs_old, fs_new): - assert torch.allclose(fs_old.float(), fs_new.float(), rtol=1e-2, atol=5.0), ( - f"final_state mismatch at mtp_step={mtp_step}, " - f"max diff={torch.abs(fs_old.float() - fs_new.float()).max().item():.6f}" - ) - - -@pytest.mark.parametrize("seed", [1, 2, 3, 4]) -def test_variable_length_cu_seqlens(seed): - """Variable-length sequences via cu_seqlens: individual lengths vary, - each ≤ mtp_step+1, with some padded to shorter lengths.""" - torch.manual_seed(seed) - batch, mtp_step, H, HV, K, V = 4, 3, 2, 8, 64, 64 - max_len = mtp_step + 1 # 4 - cache_slots = 64 - - # Random per-sequence lengths in [1, max_len] - lengths = torch.randint(1, max_len + 1, (batch,), device="cpu") - cu_seqlens = torch.zeros(batch + 1, dtype=torch.int32) - torch.cumsum(lengths, dim=0, out=cu_seqlens[1:]) - - num_tokens = int(cu_seqlens[-1].item()) - q = torch.randn(1, num_tokens, H, K, device="cuda", dtype=torch.bfloat16) - k = torch.randn(1, num_tokens, H, K, device="cuda", dtype=torch.bfloat16) - v = torch.randn(1, num_tokens, HV, V, device="cuda", dtype=torch.bfloat16) - A_log = torch.randn(HV, device="cuda", dtype=torch.float32) * 0.1 - dt_bias = torch.randn(HV, device="cuda", dtype=torch.float32) * 0.1 - a_raw = torch.randn(num_tokens, HV, device="cuda", dtype=torch.bfloat16) - b_raw = torch.randn(num_tokens, HV, device="cuda", dtype=torch.bfloat16) - ssm_state = torch.randn(cache_slots, HV, K, V, device="cuda", dtype=torch.bfloat16) - # 2D indices: S+1 = max_len columns, unused columns for short seqs are irrelevant - ssm_idx = torch.randint(0, cache_slots, (batch, max_len), device="cuda", dtype=torch.int32) - num_accepted = lengths.to(device="cuda", dtype=torch.int32) - - o_old, o_new, fs_old, fs_new = _run_both( - q, - k, - v, - ssm_state, - cu_seqlens.to(torch.long).cuda(), - ssm_idx, - ssm_idx, - num_accepted, - A_log, - dt_bias, - a_raw, - b_raw, - ) - - assert torch.equal( - o_old, o_new - ), f"seed={seed}: output mismatch, max diff={torch.abs(o_old.float() - o_new.float()).max().item():.6f}" - if not torch.equal(fs_old, fs_new): - assert torch.allclose(fs_old.float(), fs_new.float(), rtol=1e-2, atol=5.0), ( - f"seed={seed}: final_state mismatch, " - f"max diff={torch.abs(fs_old.float() - fs_new.float()).max().item():.6f}" - ) - - -if __name__ == "__main__": - pytest.main([__file__, "-v"]) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_state_params.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_state_params.py deleted file mode 100644 index 8252c445a1..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_state_params.py +++ /dev/null @@ -1,65 +0,0 @@ -import pytest -import torch - -from lightllm.common.attention_state.bounded_att import GDNState, ConvSSMStateConfig - -from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( - build_dynamic_mtp_linear_att_state_params, -) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("request_capacity,token_bucket", [(4, 3), (4, 16), (128, 512)]) -def test_dynamic_sequences_stay_within_request_capacity_under_graph(request_capacity, token_bucket): - hold = request_capacity - reqs = torch.full((token_bucket,), hold, dtype=torch.int32, device="cuda") - mtp_index = torch.zeros_like(reqs) - config = ConvSSMStateConfig( - logical_layer_ids=(0,), - global_k_heads=1, - global_v_heads=1, - key_dim=2, - value_dim=2, - conv_width=3, - conv_dtype=torch.float32, - ssm_dtype=torch.float32, - ) - manager = GDNState(config, request_capacity + 1, mtp_step=3) - manager.mtp_state_index.copy_(torch.arange(request_capacity + 1, dtype=torch.int32, device="cuda") % 4) - sequences = min(request_capacity, token_bucket) - - def build(): - return build_dynamic_mtp_linear_att_state_params(reqs, mtp_index, manager.mtp_state_index, hold) - - build() - torch.cuda.synchronize() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - cu, ids, counts = build() - assert cu.numel() == token_bucket + 1 - assert ids.numel() == counts.numel() == token_bucket - - # Reuse one graph for full request pools, mixed verify lengths, and HOLD-only padding. - for lengths in ([1] * sequences, [min(4, token_bucket // sequences)] * sequences, [2, 1], []): - lengths = lengths[:request_capacity] - if sum(lengths) > token_bucket: - continue - request_ids = list(reversed(range(len(lengths)))) - host_reqs = [req for req, length in zip(request_ids, lengths) for _ in range(length)] - host_mtp = [j for length in lengths for j in range(length)] - tokens = len(host_reqs) - host_reqs += [hold] * (token_bucket - tokens) - host_mtp += [0] * (token_bucket - tokens) - reqs.copy_(torch.tensor(host_reqs, dtype=torch.int32, device="cuda")) - mtp_index.copy_(torch.tensor(host_mtp, dtype=torch.int32, device="cuda")) - graph.replay() - - expected_cu = [0] - for length in lengths: - expected_cu.append(expected_cu[-1] + length) - expected_cu += [tokens] * (token_bucket + 1 - len(expected_cu)) - expected_ids = request_ids + [hold] * (token_bucket - len(request_ids)) - assert cu.tolist() == expected_cu - assert ids.tolist() == expected_ids - expected_counts = [req % 4 + 1 for req in request_ids] + [1] * (token_bucket - len(request_ids)) - assert counts.tolist() == expected_counts diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py index 786376a362..e6c7d7761c 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py @@ -120,35 +120,6 @@ def test_fp16_correction_history_preserves_finite_recurrence(capacity, bv, state torch.testing.assert_close(ssm_snapshot(ssm_backend, 0)[0].float(), expected_state, atol=128, rtol=0.01) -@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -def test_flush_export_updates_checkpoint_and_clears_history(dtype): - state = torch.randn(1, 2, 1, 32, 64, device="cuda", dtype=dtype) * 0.01 - backend = ReplaySSMExecutor( - state, - capacity=4, - verify_width=1, - activation_dtype=torch.bfloat16, - ) - backend.raw_keys.normal_(0, 0.1) - backend.raw_values.normal_(0, 0.1) - backend.gates.fill_(-0.1) - backend.betas.fill_(0.5) - backend.cursors[0] = 2 - expected = ssm_snapshot(backend, 0).clone() - target = ssm_snapshot(backend, 0) - reqs = torch.tensor([0], dtype=torch.int32, device="cuda") - backend.fold_accepted_prefix(reqs) - torch.cuda.synchronize() - - assert backend.cursors[0].item() == 0 - torch.testing.assert_close(target, expected, atol=0, rtol=0) - torch.testing.assert_close(state[:, 0], expected, atol=0, rtol=0) - - empty_history_target = ssm_snapshot(backend, 0) - torch.cuda.synchronize() - torch.testing.assert_close(empty_history_target, expected, atol=0, rtol=0) - - @pytest.mark.parametrize("capacity", [16, 32, 64]) @pytest.mark.parametrize("width", [1, 3, 5, 16]) @pytest.mark.parametrize("dims", [(32, 64), (128, 128)]) @@ -263,254 +234,6 @@ def step(): torch.testing.assert_close(state.float(), expected.to(dtype).float(), atol=state_atol, rtol=state_rtol) -@pytest.mark.parametrize( - "state_dtype,activation_dtype", - [(torch.float32, torch.float32), (torch.float32, torch.bfloat16), (torch.bfloat16, torch.bfloat16)], -) -@pytest.mark.parametrize("compute_mode", ["inline", "precompute"]) -@pytest.mark.parametrize("dims", [(32, 64), (128, 128)]) -@pytest.mark.parametrize("width", [1, 4]) -@pytest.mark.parametrize("gate_shift", [-3, -10]) -def test_fold_matches_reference_across_precisions( - state_dtype, - activation_dtype, - compute_mode, - dims, - width, - gate_shift, -): - test_replay_acceptance_flush_and_graph( - width, - dims, - 8, - state_dtype, - run_config={"BV": 32, "num_warps": 1, "num_stages": 3}, - compute_mode=compute_mode, - activation_dtype=activation_dtype, - gate_shift=gate_shift, - ) - - -@pytest.mark.parametrize("capacity,width", [(4, 1), (4, 3), (4, 4), (8, 1), (8, 3), (8, 4), (8, 5)]) -@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -def test_short_replay_history_matches_independent_reference(capacity, width, dtype): - test_replay_acceptance_flush_and_graph(width, (128, 128), capacity, dtype) - - -@pytest.mark.parametrize("capacity,width", [(4, 1), (4, 3), (4, 4), (8, 1), (8, 4), (8, 5), (16, 4), (16, 16), (32, 4)]) -@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -def test_precomputed_projections_match_independent_reference(capacity, width, dtype): - test_replay_acceptance_flush_and_graph( - width, - (128, 128), - capacity, - dtype, - run_config={"BV": 32, "num_warps": 1, "num_stages": 3}, - compute_mode="precompute", - ) - - -@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -def test_production_gdn_geometry_matches_independent_reference(dtype): - test_replay_acceptance_flush_and_graph( - 4, - (128, 128), - 8, - dtype, - run_config={"BV": 32, "num_warps": 1, "num_stages": 3}, - compute_mode="precompute", - head_geometry=(4, 12), - ) - - -@pytest.mark.parametrize("warps", [1, 4]) -@pytest.mark.parametrize("decay", [-3, -10]) -@pytest.mark.parametrize( - "state_dtype,dtype", - [(torch.float32, torch.float32), (torch.bfloat16, torch.bfloat16), (torch.float32, torch.bfloat16)], -) -@pytest.mark.parametrize("bv", [8, 32]) -@pytest.mark.parametrize("width", [1, 4]) -def test_precomputed_projections_preserve_precision(warps, decay, state_dtype, dtype, width, bv): - torch.manual_seed(35) - state = (torch.randn((1, 3, 4, 128, 128), device="cuda") * 0.1).to(state_dtype) - config = {"BV": bv, "num_warps": warps, "num_stages": 3} - ssm_backends = [ - ReplaySSMExecutor( - state.clone(), - 16 if width == 1 else 8, - width, - dtype, - num_key_heads=2, - compute_mode=mode, - ) - for mode in ("inline", "precompute") - ] - reqs = torch.tensor([1, 0], dtype=torch.int32, device="cuda") - tokens = 2 if width == 1 else 7 - cu = torch.tensor([0, 1, 2] if width == 1 else [0, 4, 7], dtype=torch.int32, device="cuda") - q = torch.randn((tokens, 2, 128), device="cuda", dtype=dtype) - k = torch.randn_like(q) - v = torch.randn((tokens, 4, 128), device="cuda", dtype=dtype) - a = torch.randn((tokens, 4), device="cuda", dtype=dtype) + decay - b = torch.randn_like(a) - alog = torch.zeros(4, device="cuda") - accepted = torch.zeros(3, dtype=torch.int32, device="cuda") - for step in range(64): - for tensor in (q, k, v, b): - tensor.normal_() - accepted[0] = step % max(1, width - 1) - accepted[1] = step % width - outputs = [] - step_lengths = (cu[1:] - cu[:-1]).tolist() - if any( - ssm_backends[0].cursors[req] + length > ssm_backends[0].capacity - for req, length in zip(reqs.tolist(), step_lengths) - ): - overflow = [ - req - for req, length in zip(reqs.tolist(), step_lengths) - if ssm_backends[0].cursors[req] + length > ssm_backends[0].capacity - ] - overflow_reqs = torch.tensor(overflow, dtype=torch.int32, device="cuda") - for ssm_backend in ssm_backends: - ssm_backend.fold_accepted_prefix(overflow_reqs) - for ssm_backend in ssm_backends: - outputs.append(replay_layer(ssm_backend, 0, q, k, v, a, b, alog, alog, reqs, cu, run_config=config)) - accept_replay_window(ssm_backend, reqs, accepted, cu) - if dtype == torch.float32: - torch.testing.assert_close(outputs[0], outputs[1], atol=2e-6, rtol=2e-5) - else: - torch.testing.assert_close(outputs[0], outputs[1], atol=5e-4, rtol=0.01) - error = torch.linalg.vector_norm(outputs[0].float() - outputs[1].float()) - scale = torch.linalg.vector_norm(outputs[0].float()) - assert error < scale * (torch.finfo(dtype).eps / 2) - for ssm_backend in ssm_backends: - ssm_backend.fold_accepted_prefix(reqs) - torch.testing.assert_close(ssm_backends[0].ssm_state, ssm_backends[1].ssm_state, atol=0, rtol=0) - - -def test_bf16_commit_rounds_before_the_next_token(): - # 前缀写满后先压回 SSM，下一 token 从长度 0 重写。两次 commit 必须得到同一块 BF16 状态。 - state = torch.full((1, 2, 1, 32, 64), 0.123, device="cuda", dtype=torch.bfloat16) - ssm_backends = [ReplaySSMExecutor(state.clone(), 16, 1, torch.float32) for _ in range(2)] - reqs = torch.tensor([0], device="cuda", dtype=torch.int32) - for ssm_backend in ssm_backends: - ssm_backend.raw_keys.zero_() - ssm_backend.raw_values.zero_() - ssm_backend.gates.fill_(-0.01) - ssm_backend.betas.zero_() - ssm_backend.cursors[0] = 16 - ssm_backend.fold_accepted_prefix(reqs) - q = torch.ones((1, 1, 32), device="cuda") - k = torch.zeros_like(q) - v = torch.zeros((1, 1, 64), device="cuda") - a = torch.full((1, 1), -3.0, device="cuda") - b = torch.zeros_like(a) - alog = torch.zeros(1, device="cuda") - bias = torch.zeros_like(alog) - outputs = [] - cu = torch.tensor([0, 1], dtype=torch.int32, device="cuda") - for ssm_backend in ssm_backends: - assert ssm_backend.cursors[0].item() == 0 - outputs.append(replay_layer(ssm_backend, 0, q, k, v, a, b, alog, bias, reqs, cu)) - torch.testing.assert_close(outputs[0], outputs[1], atol=1e-7, rtol=1e-6) - torch.testing.assert_close(ssm_backends[0].ssm_state, ssm_backends[1].ssm_state, atol=0, rtol=0) - - -def test_large_bf16_state_precompute_matches_inline(): - # 输出投影用 tf32x3。重建缓存的 k/d 在 bf16 激活下是 fp16，输出允许一个 fp16 单位。 - # 折回 SSM 仍从原始 k/v 在 fp32 中计算，两条路径的 checkpoint 必须一致。 - heads = 16 - config = {"BV": 32, "num_warps": 1, "num_stages": 1} - reqs = torch.tensor([0], dtype=torch.int32, device="cuda") - alog = torch.zeros(heads, device="cuda") - bias = torch.zeros_like(alog) - - def check(capacity, width, steps): - torch.manual_seed(7) - state = (torch.randn((1, 2, heads, 128, 128), device="cuda") * 10).bfloat16() - backends = [ - ReplaySSMExecutor( - state.clone(), - capacity, - width, - torch.bfloat16, - num_key_heads=heads, - compute_mode=mode, - ) - for mode in ("inline", "precompute") - ] - cu = torch.tensor([0, width], dtype=torch.int32, device="cuda") - accepted = torch.tensor([width - 1], dtype=torch.int32, device="cuda") - for step in range(steps): - torch.manual_seed(100 + step) - q = torch.randn((width, heads, 128), device="cuda", dtype=torch.bfloat16) - k = torch.randn_like(q) - v = torch.randn((width, heads, 128), device="cuda", dtype=torch.bfloat16) - a = torch.randn((width, heads), device="cuda", dtype=torch.bfloat16) - 1 - b = torch.randn_like(a) - outputs = [ - replay_layer(backend, 0, q, k, v, a, b, alog, bias, reqs, cu, run_config=config) for backend in backends - ] - # 重建用的 k/d 在 bf16 激活下是 fp16。两条路径在存入前的舍入可以差一个 fp16 单位。 - torch.testing.assert_close(outputs[0], outputs[1], atol=5e-4, rtol=5e-4) - for backend in backends: - accept_replay_window(backend, reqs, accepted, cu) - overflow = reqs[backend.cursors[reqs.long()] + backend.verify_width > backend.capacity] - if overflow.numel(): - backend.fold_accepted_prefix(overflow) - for backend in backends: - backend.fold_accepted_prefix(reqs) - torch.testing.assert_close(backends[0].ssm_state, backends[1].ssm_state, atol=0, rtol=0) - - check(4, 1, 16) - check(16, 2, 4) - - -@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) -@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -@pytest.mark.parametrize("capacity", [4, 8]) -def test_overflow_folds_inside_replay_for_gdn_and_kda(kda, dtype, capacity): - # 放不下时，GDN 和 KDA 都在 replay kernel 内折回。cursor 等两层都算完才改。 - torch.manual_seed(11 + capacity + int(kda)) - hv, kh, dim = 2, 2, 32 - state = torch.randn(2, 2, hv, dim, dim, device="cuda", dtype=dtype) * 0.05 - config = {"BV": 32, "num_warps": 1, "num_stages": 1} - backend = ReplaySSMExecutor(state, capacity, 1, dtype, num_key_heads=kh, kda=kda) - twin = ReplaySSMExecutor(state.clone(), capacity, 1, dtype, num_key_heads=kh, kda=kda) - assert backend.gates.shape[-1] == (dim if kda else capacity) - reqs = torch.tensor([0], device="cuda", dtype=torch.int32) - cu = torch.tensor([0, 1], dtype=torch.int32, device="cuda") - alog = torch.zeros(hv, device="cuda") - bias = torch.zeros(hv * dim if kda else hv, device="cuda") - for step in range(capacity + 1): - q = torch.randn(1, kh, dim, device="cuda", dtype=dtype) - k = torch.randn_like(q) - v = torch.randn(1, hv, dim, device="cuda", dtype=dtype) - a = (torch.randn(1, hv * dim if kda else hv, device="cuda", dtype=dtype) - 1).contiguous() - b = torch.randn(1, hv, device="cuda", dtype=dtype) - if step == capacity: - fold_prefix(twin, reqs, config) - twin.cursors.index_fill_(0, reqs.long(), 0) - assert int(backend.cursors[0]) == capacity - outs, twin_outs = [], [] - for layer in range(2): - outs.append(replay_layer(backend, layer, q, k, v, a, b, alog, bias, reqs, cu, run_config=config)) - twin_outs.append(replay_layer(twin, layer, q, k, v, a, b, alog, bias, reqs, cu, run_config=config)) - if step == capacity: - assert int(backend.cursors[0]) == capacity - for actual, expected in zip(outs, twin_outs): - torch.testing.assert_close(actual, expected, atol=0, rtol=0) - accept_replay_window(backend, reqs) - accept_replay_window(twin, reqs) - assert int(backend.cursors[0]) == 1 - assert int(twin.cursors[0]) == 1 - # bf16 两条路径都先舍入到同一状态。fp32 的独立 kernel 与层内 kernel 归约顺序可差几个 ulp。 - state_atol = 0 if dtype == torch.bfloat16 else 1e-6 - torch.testing.assert_close(backend.ssm_state, twin.ssm_state, atol=state_atol, rtol=0) - - @pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) @pytest.mark.parametrize("compute_mode", ["inline", "precompute"]) @pytest.mark.parametrize("bv", [32, 128]) @@ -581,173 +304,12 @@ def step(owner): torch.testing.assert_close(backend.ssm_state[:, batch], initial[:, batch], rtol=0, atol=0) -@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) -def test_empty_window_keeps_cursor_for_gdn_and_kda(kda): - state = torch.randn(1, 3, 2, 32, 32, device="cuda") - backend = ReplaySSMExecutor(state, 8, 4, torch.float32, num_key_heads=2, kda=kda) - backend.cursors[0] = 3 - reqs = torch.tensor([0, 2], device="cuda", dtype=torch.int32) - cu = torch.tensor([0, 0, 0], device="cuda", dtype=torch.int32) - q = torch.randn(1, 2, 32, device="cuda") - k = torch.randn_like(q) - v = torch.randn(1, 2, 32, device="cuda") - a = torch.zeros(1, 2 * 32 if kda else 2, device="cuda") - b = torch.zeros(1, 2, device="cuda") - bias = torch.zeros(2 * 32 if kda else 2, device="cuda") - replay_layer( - backend, - 0, - q, - k, - v, - a, - b, - bias, - bias, - reqs, - cu, - run_config={"BV": 32, "num_warps": 1, "num_stages": 3}, - ) - accepted = torch.tensor([-1, -1, -1], device="cuda", dtype=torch.int32) - accept_replay_window(backend, reqs, accepted, cu) - assert backend.cursors.tolist() == [3, 0, 0] - - -@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) -@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -@pytest.mark.parametrize("compute_mode", ["inline", "precompute"]) -@pytest.mark.parametrize("capture_metadata", [False, True]) -def test_dynamic_verify_lengths_fill_history_before_folding(kda, dtype, compute_mode, capture_metadata): - torch.manual_seed(31) - state = (torch.randn(2, 4, 4, 128, 128, device="cuda") * 0.01).to(dtype) - ssm_backend = ReplaySSMExecutor( - state, - 8, - 4, - dtype, - num_key_heads=2, - kda=kda, - compute_mode=compute_mode, - ) - expected = state.float().clone() - pending = [0, 0] - reqs = torch.full((3,), 3, dtype=torch.int32, device="cuda") - cu = torch.tensor([0, 4, 8, 8], dtype=torch.int32, device="cuda") - accepted = torch.zeros(4, dtype=torch.int32, device="cuda") - q = torch.randn(8, 2, 128, dtype=dtype, device="cuda") - k = torch.randn_like(q) - v = torch.randn(8, 4, 128, dtype=dtype, device="cuda") - a = torch.full((8, 4 * 128 if kda else 4), -3.0, dtype=dtype, device="cuda") - b = torch.randn(8, 4, dtype=dtype, device="cuda") - alog = torch.zeros(4, device="cuda") - bias = torch.zeros(4 * 128 if kda else 4, device="cuda") - - def forward(): - return [ - replay_layer( - ssm_backend, - layer, - q, - k, - v, - a, - b, - alog, - bias, - reqs, - cu, - run_config={"BV": 32, "num_warps": 1, "num_stages": 3}, - ) - for layer in range(2) - ] - - def step(): - positions = ssm_backend.cursors.index_select(0, reqs.long()) - out = forward() - accept_replay_window(ssm_backend, reqs, accepted, cu) - return positions, out - - step() - torch.cuda.synchronize() - graph = torch.cuda.CUDAGraph() - positions = torch.zeros_like(reqs) - with torch.cuda.graph(graph): - if capture_metadata: - positions, outputs = step() - else: - outputs = forward() - lengths = [(3, 4), (2, 1), (3, 2), (1, 4), (4, 3), (0, 2), (2, 0), (4, 4)] - for iteration in range(24): - ids = [0, 1] if iteration < 3 or iteration % 2 == 0 else [1, 0] - widths = lengths[iteration % len(lengths)] - counts = ( - list(widths) if iteration < 3 else [min(width, 1 + (iteration + i) % 4) for i, width in enumerate(widths)] - ) - reqs.copy_(reqs.new_tensor(ids + [3])) - cu.copy_(cu.new_tensor([0, widths[0], sum(widths), 8])) - accepted.fill_(-1) - for req, count in zip(ids, counts): - accepted[req] = count - 1 - for tensor in (q, k, v, b): - tensor.normal_() - before = state.clone() - for req, width in zip(ids, widths): - if pending[req] + width > 8: - ssm_backend.fold_accepted_prefix(torch.tensor([req], dtype=torch.int32, device="cuda")) - expected[:, req] = expected[:, req].to(dtype).float() - pending[req] = 0 - old_cursors = [pending[req] for req in ids] + [0] - if not capture_metadata: - positions.copy_(ssm_backend.cursors.index_select(0, reqs.long())) - graph.replay() - if not capture_metadata: - accept_replay_window(ssm_backend, reqs, accepted, cu) - assert positions.tolist() == old_cursors - for seq, (req, width, count) in enumerate(zip(ids, widths, counts)): - for layer in range(2): - current = expected[layer, req].clone() - for j in range(width): - t = (0 if seq == 0 else widths[0]) + j - out, current = reference(q[t], k[t], v[t], a[t], b[t], alog, bias, current, kda=kda) - torch.testing.assert_close(outputs[layer][t].float(), out, atol=0.006, rtol=0.02) - if j + 1 == count: - expected[layer, req].copy_(current) - pending[req] += count - assert ssm_backend.cursors.tolist() == pending + [0, 0] - for output in outputs: - assert torch.count_nonzero(output[sum(widths) :]).item() == 0 - torch.testing.assert_close(state[:, 2:], before[:, 2:], atol=0, rtol=0) - if iteration == 2: - assert pending[0] == 8 - torch.testing.assert_close(state[:, 0], before[:, 0], atol=0, rtol=0) - ssm_backend.fold_accepted_prefix(reqs[:2]) - torch.testing.assert_close( - state.float(), - expected.to(dtype).float(), - atol=2e-5 if dtype == torch.float32 else 0.002, - rtol=2e-4 if dtype == torch.float32 else torch.finfo(dtype).eps, - ) - - -def test_snapshot_without_history_copies_active_state(): - state = torch.randn((2, 3, 2, 32, 64), device="cuda", dtype=torch.bfloat16) - ssm_backend = ReplaySSMExecutor(state, 4, 3) - snapshot = ssm_snapshot(ssm_backend, 1) - torch.testing.assert_close(snapshot, state[:, 1], rtol=0, atol=0) - assert ssm_backend.cursors[1].item() == 0 - - -@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -@pytest.mark.parametrize("run_config", get_configs(), ids=lambda c: f"bv{c['BV']}-w{c['num_warps']}-s{c['num_stages']}") -def test_autotune_candidates_match_recurrence(dtype, run_config): - test_replay_acceptance_flush_and_graph(3, (32, 64), 8, dtype, run_config) - - @pytest.mark.parametrize("save_big", [False, True]) @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) def test_checkpoint_restores_pending_history_and_accepted_conv(runtime_start_args, save_big, dtype): + from lightllm.common.attention_state import AttentionStateManager from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig - from unit_tests.common.test_linear_state import make_state + from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState torch.manual_seed(42) state_rtol = 2e-5 if dtype == torch.float32 else torch.finfo(dtype).eps @@ -761,7 +323,11 @@ def test_checkpoint_restores_pending_history_and_accepted_conv(runtime_start_arg conv_dtype=torch.bfloat16, ssm_dtype=dtype, ) - state, _, layout = make_state(mode="replay", config=config) + state = ReplayGDNState(config, 4, mtp_step=2, device="cuda", ssm_replay_config={"cache_len": 4}) + resources = AttentionStateManager(request_capacity=4) + resources.register_state("main", state) + resources.finish_state_registration() + layout = resources.get_pd_operator().get_pd_bounded_att_page_layouts()[state.state_name] state.initialize_page_pools(big_page_capacity=2, small_page_capacity=2) state.mtp_state_index[1] = 1 state.conv_state.normal_() @@ -810,119 +376,6 @@ def test_checkpoint_restores_pending_history_and_accepted_conv(runtime_start_arg assert not state.conv_state[:, 2].any() -@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -@pytest.mark.parametrize("capacity,width", [(4, 1), (4, 3), (8, 4), (16, 5), (32, 4), (64, 4)]) -@pytest.mark.parametrize("compute_mode", ["inline", "precompute"]) -@pytest.mark.parametrize("gate_lower_bound,gate_shift", [(-5.0, 0), (-2.5, -6)]) -@pytest.mark.parametrize("activation_dtype", [torch.float32, torch.bfloat16]) -def test_kda_replay_reference_fold_and_graph( - dtype, capacity, width, compute_mode, gate_lower_bound, gate_shift, activation_dtype -): - test_replay_acceptance_flush_and_graph( - width, - (128, 128), - capacity, - dtype, - run_config={"BV": 32, "num_warps": 2, "num_stages": 1}, - compute_mode=compute_mode, - gate_shift=gate_shift, - activation_dtype=activation_dtype, - head_geometry=(2, 2), - kda=True, - gate_lower_bound=gate_lower_bound, - ) - - -@pytest.mark.parametrize("config", get_configs()) -@pytest.mark.parametrize("compute_mode", ["inline", "precompute"]) -def test_kda_tuning_candidates_match_reference(config, compute_mode): - test_replay_acceptance_flush_and_graph( - 3, - (128, 128), - 4, - torch.bfloat16, - run_config=config, - compute_mode=compute_mode, - head_geometry=(2, 2), - kda=True, - gate_shift=-3, - ) - - -@pytest.mark.parametrize("kda", [False, True]) -@pytest.mark.parametrize( - "bucket,live,folding,width,separate", - [ - (12, 12, 12, 4, True), - (13, 13, 13, 4, False), - (255, 255, 127, 4, True), - (256, 256, 128, 4, False), - (257, 257, 128, 4, False), - (256, 256, 128, 3, True), - (256, 12, 12, 4, True), - (256, 13, 4, 4, True), - ], -) -def test_fold_scheduling_preserves_recurrence(bucket, live, folding, width, separate, kda): - torch.manual_seed(72) - initial = torch.randn(1, live + 1, 1, 32, 32, device="cuda") * 0.01 - backend = ReplaySSMExecutor(initial.clone(), 8, 4, torch.float32, kda=kda) - reqs = torch.cat( - ( - torch.arange(live, device="cuda", dtype=torch.int32), - torch.full((bucket - live,), live, device="cuda", dtype=torch.int32), - ) - ) - indices = torch.zeros(live + 1, device="cuda", dtype=torch.int32) - expected = initial[0, :live].clone() - alog = torch.zeros(1, device="cuda") - bias = torch.zeros(32 if kda else 1, device="cuda") - for step, query_width in enumerate((4, 4, width)): - if step == 2: - backend.cursors.index_fill_(0, reqs[folding:live].long(), 0) - expected[folding:] = initial[0, folding:live] - cu = (torch.arange(bucket + 1, device="cuda", dtype=torch.int32) * query_width).clamp_max(live * query_width) - q = torch.randn(live, query_width, 1, 32, device="cuda") * 0.2 - k, v = torch.randn_like(q), torch.randn_like(q) - a = torch.full((live, query_width, 32 if kda else 1), -2.0, device="cuda") - b = torch.randn(live, query_width, 1, device="cuda") - if (backend.cursors[:live] + query_width > backend.capacity).any(): - backend.fold_accepted_prefix(reqs[:live]) - actual = replay_layer( - backend, - 0, - q.reshape(-1, 1, 32), - k.reshape(-1, 1, 32), - v.reshape(-1, 1, 32), - a.flatten(0, 1), - b.flatten(0, 1), - alog, - bias, - reqs, - cu, - run_config={"BV": 32, "num_warps": 1, "num_stages": 3}, - ).reshape(live, query_width, 1, 32) - for token in range(query_width): - output, expected = reference( - q[:, token], - k[:, token], - v[:, token], - a[:, token], - b[:, token], - alog, - bias.expand(live, 1, 32) if kda else bias, - expected, - kda=kda, - ) - torch.testing.assert_close(actual[:, token], output, atol=3e-4, rtol=3e-4) - lengths = torch.full((live,), query_width, device="cuda", dtype=torch.int32) - advance_cursor(backend, reqs[:live], cu[: live + 1]) - commit_ssm_prefix(indices, reqs[:live], lengths, lengths, backend.cursors) - backend.fold_accepted_prefix(reqs[:live]) - torch.testing.assert_close(backend.ssm_state[0, :live], expected, atol=3e-4, rtol=3e-4) - assert torch.equal(backend.ssm_state[0, live], initial[0, live]) - - @pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) @pytest.mark.parametrize("compute_mode", ["inline", "precompute"]) def test_mtp_commit_keeps_folded_prefix_length(kda, compute_mode): diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py deleted file mode 100644 index 02797b61b7..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py +++ /dev/null @@ -1,530 +0,0 @@ -import pytest -import torch - -from unit_tests.common.basemodel.triton_kernel.linear_att.replay_history import ( - accept_replay_window, - attach_replay_history as ReplaySSMExecutor, - replay_layer, -) -from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( - build_dynamic_mtp_linear_att_state_params, -) -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import advance_cursor -from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix -from types import SimpleNamespace - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") - - -@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) -def test_decode_compilation_uses_shape_buckets(monkeypatch, kda): - from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import _advance_cursor - from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_decode import _decode_replay - - kernels = {"_decode_replay": _decode_replay, "_advance_cursor": _advance_cursor} - hashes = {name: set() for name in kernels} - launches = [] - for name, kernel in kernels.items(): - - def record(*args, run=kernel.run, name=name, **kwargs): - compiled = run(*args, **kwargs) - hashes[name].add(compiled.hash) - launches.append(name) - return compiled - - monkeypatch.setattr(kernel, "run", record) - - torch.manual_seed(137) - initial = torch.randn(1, 513, 1, 32, 32, device="cuda") * 0.01 - eager, captured = [ReplaySSMExecutor(initial.clone(), 8, 1, torch.float32, kda=kda) for _ in range(2)] - reqs = torch.arange(512, device="cuda", dtype=torch.int32) - q = torch.randn(512, 1, 32, device="cuda") * 0.2 - k, v = torch.randn_like(q), torch.randn_like(q) - a = torch.full((512, 32 if kda else 1), -3.0, device="cuda") - beta = torch.randn(512, 1, device="cuda") - alog, bias = torch.zeros(1, device="cuda"), torch.zeros(a.shape[-1], device="cuda") - - def cycle(backend, ids): - count = ids.numel() - cu = torch.arange(count + 1, dtype=torch.int32, device=ids.device) - out = replay_layer( - backend, - 0, - q[:count], - k[:count], - v[:count], - a[:count], - beta[:count], - alog, - bias, - ids, - cu, - run_config={"BV": 32, "num_warps": 1, "num_stages": 3}, - ) - accept_replay_window(backend, ids, cu_seqlens=cu) - return out - - graphs = {} - for bucket in (32, 512): - ids = reqs[:bucket].clone() - cycle(captured, ids) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - out = cycle(captured, ids) - graphs[bucket] = graph, ids, out - captured.ssm_state.copy_(initial) - captured.cursors.zero_() - counts = (1, 2, 7, 16, 23, 24, 25, 31, 32, 33, 255, 256, 257, 511) - for count in counts: - graph, ids, output = graphs[32 if count <= 32 else 512] - ids.fill_(captured.hold_req_id) - ids[:count].copy_(reqs[:count]) - for _ in range(8): - live = reqs[:count] - if (eager.cursors[live.long()] + 1 > eager.capacity).any(): - eager.fold_accepted_prefix(live) - captured.fold_accepted_prefix(ids[:count]) - expected = cycle(eager, reqs[:count]) - calls = len(launches) - graph.replay() - assert len(launches) == calls - torch.testing.assert_close(output[:count], expected, atol=3e-5, rtol=3e-5) - assert torch.equal(captured.cursors, eager.cursors) - torch.testing.assert_close(captured.ssm_state, eager.ssm_state, atol=3e-5, rtol=3e-5) - assert len(hashes["_decode_replay"]) == 1 - assert len(hashes["_advance_cursor"]) == 1 - assert captured.cursors[-1] == 0 - assert torch.equal(captured.ssm_state[:, -1], initial[:, -1]) - - -@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -@pytest.mark.parametrize("capture_metadata", [False, True]) -@pytest.mark.parametrize("varlen", [False, True]) -@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) -def test_request_manager_accept_with_replay_graph(dtype, capture_metadata, varlen, kda): - torch.manual_seed(827) - config = {"BV": 32, "num_warps": 1, "num_stages": 1} - initial = torch.randn(3, 5, 4, 128, 128, device="cuda", dtype=dtype) * 0.01 - ssm_backends = [ - cls( - initial.clone(), - 8, - 4, - torch.bfloat16, - num_key_heads=2, - kda=kda, - compute_mode="precompute", - ) - for cls in [ReplaySSMExecutor, ReplaySSMExecutor] - ] - manager = SimpleNamespace() - manager._ssm_executor = ssm_backends[1] - manager.mtp_state_index = torch.zeros(5, device="cuda", dtype=torch.int32) - rows = torch.full((12,), 4, device="cuda", dtype=torch.int32) - mtp = torch.zeros_like(rows) - flags = torch.zeros_like(rows) - starts = torch.zeros(3, device="cuda", dtype=torch.int32) - accepted = torch.full((5,), -1, device="cuda", dtype=torch.int32) - q = torch.randn(12, 2, 128, device="cuda", dtype=torch.bfloat16) - k = torch.randn_like(q) - v = torch.randn(12, 4, 128, device="cuda", dtype=torch.bfloat16) - a = torch.full((12, 4 * 128 if kda else 4), -3.0, device="cuda", dtype=torch.bfloat16) - beta = torch.randn(a.shape[0], 4, device="cuda", dtype=a.dtype) - alog = torch.zeros(4, device="cuda") - bias = torch.zeros(4 * 128 if kda else 4, device="cuda") - - def inputs(iteration): - ids = [0, 1, 3] if iteration % 2 == 0 else [3, 0, 1] - lengths = [1 + (iteration + j) % 4 for j in range(3)] if varlen else [4] * 3 - counts = [1 + (iteration // 2 + j) % length for j, length in enumerate(lengths)] - offsets, req_rows, indexes, mask = [], [], [], [] - table = [-1] * 5 - for req, length, count in zip(ids, lengths, counts): - offsets.append(len(req_rows)) - req_rows.extend([req] * length) - indexes.extend(range(length)) - mask.extend([int(j < count) for j in range(length)]) - table[req] = count - 1 - logical = len(req_rows) - req_rows.extend([4] * (12 - logical)) - indexes.extend([0] * (12 - logical)) - mask.extend([0] * (12 - logical)) - rows.copy_(rows.new_tensor(req_rows)) - mtp.copy_(mtp.new_tensor(indexes)) - flags.copy_(flags.new_tensor(mask)) - starts.copy_(starts.new_tensor(offsets)) - accepted.copy_(accepted.new_tensor(table)) - return ids, lengths, counts, logical - - def metadata(ssm_backend): - cu, reqs, _counts = build_dynamic_mtp_linear_att_state_params(rows, mtp, manager.mtp_state_index, 4) - return cu, reqs, ssm_backend.cursors.index_select(0, reqs.long()) - - def forward_accept(index, cu, reqs): - ssm_backend = ssm_backends[index] - out = [ - replay_layer(ssm_backend, layer, q, k, v, a, beta, alog, bias, reqs, cu, run_config=config) - for layer in range(3) - ] - if index == 0: - accept_replay_window(ssm_backend, reqs, accepted, cu) - else: - count = starts.numel() - commit_reqs = rows[starts.long()] - window = (cu[1 : count + 1] - cu[:count]).contiguous() - advance_cursor(ssm_backend, commit_reqs, cu[: count + 1]) - commit_ssm_prefix( - manager.mtp_state_index, - commit_reqs, - accepted[commit_reqs.long()] + 1, - window, - ssm_backend.cursors, - ) - return out - - inputs(0) - graphs, captured, outputs = [], [], [] - for index, ssm_backend in enumerate(ssm_backends): - cu, reqs, pos = metadata(ssm_backend) - forward_accept(index, cu, reqs) - ssm_backend.ssm_state.copy_(initial) - ssm_backend.cursors.zero_() - manager.mtp_state_index.zero_() - g = torch.cuda.CUDAGraph() - with torch.cuda.graph(g): - if capture_metadata: - cu, reqs, pos = metadata(ssm_backend) - out = forward_accept(index, cu, reqs) - graphs.append(g) - captured.append((cu, reqs, pos)) - outputs.append(out) - for ssm_backend in ssm_backends: - ssm_backend.ssm_state.copy_(initial) - ssm_backend.cursors.zero_() - manager.mtp_state_index.zero_() - folds = 0 - for iteration in range(32): - ids, lengths, counts, logical = inputs(iteration) - for tensor in [q, k, v, beta]: - tensor.normal_() - overflow = [ - req for req, length in zip(ids, lengths) if ssm_backends[0].cursors[req] + length > ssm_backends[0].capacity - ] - if overflow: - overflow_reqs = torch.tensor(overflow, dtype=torch.int32, device="cuda") - for ssm_backend in ssm_backends: - ssm_backend.fold_accepted_prefix(overflow_reqs) - folds += len(overflow) - for index, ssm_backend in enumerate(ssm_backends): - if not capture_metadata: - for dst, src in zip(captured[index], metadata(ssm_backend)): - dst.copy_(src) - graphs[index].replay() - torch.cuda.synchronize() - for expected, actual in zip(captured[0], captured[1]): - assert torch.equal(actual, expected) - assert torch.equal(ssm_backends[0].cursors, ssm_backends[1].cursors) - for expected, actual in zip(outputs[0], outputs[1]): - torch.testing.assert_close( - actual[:, :logical].float(), expected[:, :logical].float(), atol=0.006, rtol=0.02 - ) - torch.testing.assert_close( - ssm_backends[1].ssm_state.float(), - ssm_backends[0].ssm_state.float(), - atol=0.002, - rtol=torch.finfo(dtype).eps, - ) - table = manager.mtp_state_index.cpu().tolist() - for req, count in zip(ids, counts): - assert table[req] == count - 1 - assert table[2] == table[4] == 0 - assert ssm_backends[1].cursors[4].item() == 0 - assert folds > 0 - - -@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -@pytest.mark.parametrize("capture_metadata", [False, True]) -@pytest.mark.parametrize("stagger", [False, True]) -@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) -def test_decode_accept_order_preserves_forward_snapshot(dtype, capture_metadata, stagger, kda): - torch.manual_seed(731) - initial = torch.randn(3, 3, 4, 128, 128, device="cuda", dtype=dtype) * 0.01 - ssm_backends = [ - ReplaySSMExecutor( - initial.clone(), - 8, - 1, - dtype, - num_key_heads=2, - kda=kda, - ) - for _ in range(2) - ] - for ssm_backend in ssm_backends: - ssm_backend.fold_programs = 3 - reqs = torch.tensor([0, 1, 2], device="cuda", dtype=torch.int32) - q = torch.randn(3, 2, 128, device="cuda", dtype=dtype) - k = torch.randn_like(q) - v = torch.randn(3, 4, 128, device="cuda", dtype=dtype) - a = torch.full((3, 4 * 128 if kda else 4), -3.0, device="cuda", dtype=dtype) - beta = torch.randn(a.shape[0], 4, device="cuda", dtype=a.dtype) - alog = torch.zeros(4, device="cuda") - bias = torch.zeros(4 * 128 if kda else 4, device="cuda") - - cu = torch.arange(reqs.numel() + 1, dtype=torch.int32, device=reqs.device) - - def forward(ssm_backend): - return [ - replay_layer( - ssm_backend, - layer, - q, - k, - v, - a, - beta, - alog, - bias, - reqs, - cu, - run_config={"BV": 32, "num_warps": 1, "num_stages": 1}, - ) - for layer in range(3) - ] - - def warm(ssm_backend): - out = forward(ssm_backend) - accept_replay_window(ssm_backend, reqs) - return out - - graphs, positions, outputs = [], [], [] - for ssm_backend in ssm_backends: - warm(ssm_backend) - ssm_backend.ssm_state.copy_(initial) - ssm_backend.cursors.zero_() - if stagger: - reqs.copy_(reqs.new_tensor([0, 2, 2])) - for _ in range(3): - warm(ssm_backend) - reqs.copy_(reqs.new_tensor([0, 1, 2])) - pos = torch.zeros_like(reqs) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - if capture_metadata: - pos = ssm_backend.cursors.index_select(0, reqs.long()) - out = forward(ssm_backend) - if capture_metadata: - accept_replay_window(ssm_backend, reqs) - graphs.append(graph) - positions.append(pos) - outputs.append(out) - for iteration in range(24): - reqs.copy_(reqs.new_tensor([0, 1, 2] if iteration % 2 == 0 else [1, 0, 2])) - for tensor in [q, k, v, beta]: - tensor.normal_() - overflow = [int(req) for req in reqs.tolist() if ssm_backends[0].cursors[req] + 1 > ssm_backends[0].capacity] - if overflow: - overflow_reqs = torch.tensor(sorted(set(overflow)), dtype=torch.int32, device="cuda") - for ssm_backend in ssm_backends: - ssm_backend.fold_accepted_prefix(overflow_reqs) - old = ssm_backends[0].cursors[reqs.long()].clone() - for ssm_backend, graph, pos in zip(ssm_backends, graphs, positions): - if not capture_metadata: - pos.copy_(ssm_backend.cursors.index_select(0, reqs.long())) - graph.replay() - if not capture_metadata: - accept_replay_window(ssm_backend, reqs) - torch.cuda.synchronize() - assert torch.equal(positions[0], old) - assert torch.equal(positions[0], positions[1]) - assert torch.equal(ssm_backends[0].cursors, ssm_backends[1].cursors) - for expected, actual in zip(outputs[0], outputs[1]): - torch.testing.assert_close(actual.float(), expected.float(), atol=0.006, rtol=0.02) - torch.testing.assert_close( - ssm_backends[1].ssm_state.float(), - ssm_backends[0].ssm_state.float(), - atol=0.002, - rtol=torch.finfo(dtype).eps, - ) - - -@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) -def test_cross_layer_fold_with_state_offset_above_int32(kda): - if torch.cuda.mem_get_info()[0] < 6 * 1024 ** 3: - pytest.skip("large-stride regression requires 6 GiB of free device memory") - torch.manual_seed(791) - shape = (48, 3, 12, 128, 128) - strides = (50_528_256, 196608, 16384, 128, 1) - assert strides[0] < 2 ** 31 < 47 * strides[0] - dense = torch.randn(shape, device="cuda", dtype=torch.bfloat16) * 0.01 - sparse = torch.empty_strided(shape, strides, device="cuda", dtype=dense.dtype) - sparse.copy_(dense) - config = {"BV": 32, "num_warps": 1, "num_stages": 1} - ssm_backends = [ - ReplaySSMExecutor(state, 8, 4, torch.bfloat16, num_key_heads=4, kda=kda) for state in [dense, sparse] - ] - q = torch.randn(8, 4, 128, device="cuda", dtype=torch.bfloat16) - k = torch.randn_like(q) - v = torch.randn(8, 12, 128, device="cuda", dtype=torch.bfloat16) - a = torch.full((8, 12 * 128 if kda else 12), -3.0, device="cuda", dtype=torch.bfloat16) - beta = torch.randn(8, 12, device="cuda", dtype=torch.bfloat16) - alog = torch.zeros(12, device="cuda") - bias = torch.zeros(12 * 128 if kda else 12, device="cuda") - reqs = torch.tensor([0, 2], device="cuda", dtype=torch.int32) - cu = reqs.new_tensor([0, 4, 8]) - - def step(ssm_backend, accepted): - overflow = [int(req) for req in reqs.tolist() if ssm_backend.cursors[req] + 4 > ssm_backend.capacity] - if overflow: - ssm_backend.fold_accepted_prefix(reqs.new_tensor(overflow)) - pos = ssm_backend.cursors.index_select(0, reqs.long()) - out = [ - replay_layer(ssm_backend, layer, q, k, v, a, beta, alog, bias, reqs, cu, run_config=config) - for layer in range(48) - ] - accept_replay_window(ssm_backend, reqs, accepted, cu) - return pos, out - - for count in [2, 3]: - accepted = reqs.new_tensor([count - 1, -1, -1]) - for ssm_backend in ssm_backends: - step(ssm_backend, accepted) - reqs.copy_(reqs.new_tensor([0, 1])) - expected_pos, expected_out = step(ssm_backends[0], reqs.new_tensor([0, 0, -1])) - actual_pos, actual_out = step(ssm_backends[1], reqs.new_tensor([0, 0, -1])) - torch.cuda.synchronize() - assert torch.equal(actual_pos, expected_pos) - assert torch.equal(ssm_backends[0].cursors, ssm_backends[1].cursors) - for layer in range(48): - torch.testing.assert_close(actual_out[layer].float(), expected_out[layer].float(), atol=0.006, rtol=0.02) - torch.testing.assert_close( - sparse[layer].float(), dense[layer].float(), atol=0.002, rtol=torch.finfo(torch.bfloat16).eps - ) - - -@pytest.mark.parametrize("batch", [1, 8, 12]) -@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) -def test_small_batch_matches_larger_graph_padding(batch, dtype, kda): - torch.manual_seed(1087 + batch) - initial = torch.randn(2, 14, 4, 128, 128, device="cuda", dtype=dtype) * 0.01 - config = {"BV": 32, "num_warps": 1, "num_stages": 1} - ssm_backends = [ - ReplaySSMExecutor(initial.clone(), 8, 4, num_key_heads=2, kda=kda, compute_mode="precompute") for _ in range(2) - ] - reqs = torch.arange(batch, device="cuda", dtype=torch.int32) - padded = torch.cat((reqs, reqs.new_full((13 - batch,), 13))) - cu = torch.arange(batch + 1, device="cuda", dtype=torch.int32) * 4 - padded_cu = torch.cat((cu, cu.new_full((13 - batch,), batch * 4))) - q = torch.randn(batch * 4, 2, 128, device="cuda", dtype=torch.bfloat16) - k = torch.randn_like(q) - v = torch.randn(batch * 4, 4, 128, device="cuda", dtype=torch.bfloat16) - a = torch.full((batch * 4, 4 * 128 if kda else 4), -3.0, device="cuda", dtype=torch.bfloat16) - beta = torch.randn(a.shape[0], 4, device="cuda", dtype=a.dtype) - alog = torch.zeros(4, device="cuda") - bias = torch.zeros(4 * 128 if kda else 4, device="cuda") - accepted = torch.zeros(14, device="cuda", dtype=torch.int32) - graphs, outputs = [], [] - for ssm_backend, ids, offsets in zip(ssm_backends, [reqs, padded], [cu, padded_cu]): - - def step(): - out = [ - replay_layer(ssm_backend, layer, q, k, v, a, beta, alog, bias, ids, offsets, run_config=config) - for layer in range(2) - ] - accept_replay_window(ssm_backend, ids, accepted, offsets) - return out - - step() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - out = step() - graphs.append(graph) - outputs.append(out) - ssm_backend.ssm_state.copy_(initial) - ssm_backend.cursors.zero_() - for iteration in range(32): - accepted[:batch] = (torch.arange(batch, device="cuda") + iteration) % 4 - for value in [q, k, v, beta]: - value.normal_() - overflow = [idx for idx in range(batch) if ssm_backends[0].cursors[idx] + 4 > ssm_backends[0].capacity] - if overflow: - overflow_reqs = torch.tensor(overflow, dtype=torch.int32, device="cuda") - for ssm_backend in ssm_backends: - ssm_backend.fold_accepted_prefix(overflow_reqs) - for graph in graphs: - graph.replay() - for small, padded_out in zip(outputs[0], outputs[1]): - torch.testing.assert_close(small.float(), padded_out.float(), atol=0.006, rtol=0.02) - assert torch.equal(ssm_backends[0].cursors, ssm_backends[1].cursors) - torch.testing.assert_close( - ssm_backends[0].ssm_state.float(), - ssm_backends[1].ssm_state.float(), - atol=0.002, - rtol=torch.finfo(dtype).eps, - ) - - -@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) -def test_full_prefix_does_not_address_past_capacity(kda): - torch.manual_seed(4) - capacity = 4 - backend = ReplaySSMExecutor( - torch.randn(1, 3, 2, 32, 32, device="cuda") * 0.01, - capacity, - 1, - torch.bfloat16, - kda=kda, - ) - reqs = torch.tensor([0, 1], device="cuda", dtype=torch.int32) - cu = torch.arange(reqs.numel() + 1, dtype=torch.int32, device=reqs.device) - q = torch.randn(2, 1, 32, device="cuda") * 0.2 - k, v = torch.randn_like(q), torch.randn_like(q) - a = torch.full((2, 32 if kda else 1), -3.0, device="cuda") - beta = torch.randn(2, 1, device="cuda") - alog, bias = torch.zeros(1, device="cuda"), torch.zeros(a.shape[-1], device="cuda") - backend.ssm_state.zero_() - backend.cursors.zero_() - for name in ("keys", "deltas", "gates", "raw_keys", "raw_values", "betas"): - getattr(backend, name).zero_() - backend.cursors[reqs.long()] = capacity + 2 - replay_layer( - backend, - 0, - q, - k, - v, - a, - beta, - alog, - bias, - reqs, - cu, - run_config={"BV": 32, "num_warps": 1, "num_stages": 1}, - ) - backend.fold_accepted_prefix(reqs) - torch.cuda.synchronize() - assert torch.equal(backend.cursors[reqs.long()], torch.zeros(2, device="cuda", dtype=torch.int32)) - assert torch.isfinite(backend.ssm_state).all() - for _ in range(12): - replay_layer( - backend, - 0, - q, - k, - v, - a, - beta, - alog, - bias, - reqs, - cu, - run_config={"BV": 32, "num_warps": 1, "num_stages": 1}, - ) - accept_replay_window(backend, reqs) - overflow = reqs[backend.cursors[reqs.long()] + backend.verify_width > backend.capacity] - if overflow.numel(): - backend.fold_accepted_prefix(overflow) - torch.cuda.synchronize() - assert int(backend.cursors[reqs.long()].max()) <= capacity - assert torch.isfinite(backend.ssm_state).all() diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune_memory.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune_memory.py deleted file mode 100644 index a94411fbd5..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune_memory.py +++ /dev/null @@ -1,70 +0,0 @@ -import pytest -import torch - -from lightllm.common.basemodel.attention.linear import replay_gdn_decode as tuning -from unit_tests.common.basemodel.triton_kernel.linear_att.replay_history import ( - attach_replay_history as ReplaySSMExecutor, -) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("width", [1, 5]) -def test_rebuild_uses_one_layer_scratch(width): - state = torch.zeros((4, 3, 2, 32, 32), dtype=torch.bfloat16, device="cuda") - ssm_backend = ReplaySSMExecutor(state, 8, width, num_key_heads=2) - tokens = 3 * width - q = torch.ones((tokens, 2, 32), dtype=torch.bfloat16, device="cuda") - gates = torch.zeros((tokens, 2), dtype=torch.bfloat16, device="cuda") - a_log = torch.zeros(2, device="cuda") - cu = torch.arange(4, device="cuda", dtype=torch.int32) * width - reqs = torch.arange(3, device="cuda", dtype=torch.int32) - _, kwargs = tuning.rebuild_inputs(ssm_backend, 0, q, q, q, gates, gates, a_log, a_log, reqs, cu) - scratch = kwargs["state"] - assert scratch is not ssm_backend - assert scratch.ssm_state.shape[0] == 1 - assert scratch.capacity == ssm_backend.capacity - assert torch.count_nonzero(ssm_backend.ssm_state) == 0 - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("width", [1, 5]) -def test_rebuild_failure_leaves_serving_state_unchanged(monkeypatch, tmp_path, width): - from lightllm.common.triton_utils import autotuner as autotuner_module - from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType - - state = torch.zeros((2, 4, 2, 32, 32), dtype=torch.bfloat16, device="cuda") - ssm_backend = ReplaySSMExecutor(state, 8, width, num_key_heads=2) - for tensor in vars(ssm_backend).values(): - if isinstance(tensor, torch.Tensor): - tensor.zero_() - tokens = 3 * width - q = torch.ones((tokens, 2, 32), dtype=torch.bfloat16, device="cuda") - gates = torch.zeros((tokens, 2), dtype=torch.bfloat16, device="cuda") - a_log = torch.zeros(2, device="cuda") - bias = torch.zeros(2, device="cuda") - cu = torch.arange(4, device="cuda", dtype=torch.int32) * width - reqs = torch.arange(3, device="cuda", dtype=torch.int32) - before = {name: x.clone() for name, x in vars(ssm_backend).items() if isinstance(x, torch.Tensor)} - - def oom(*args, **kwargs): - raise torch.cuda.OutOfMemoryError("scratch") - - tuner = tuning.replay - monkeypatch.setattr(tuner, "rebuild_input_func", oom) - monkeypatch.setattr(tuner, "_cache_dir", str(tmp_path), raising=False) - monkeypatch.setattr(tuner, "cached_configs", {}) - monkeypatch.setattr(tuner, "fast_match_configs", {}) - monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: 1) - monkeypatch.setattr( - tuner, "_bench", lambda *args, **kwargs: pytest.fail("Failed rebuild must not reach benchmarks") - ) - with pytest.raises(torch.cuda.OutOfMemoryError, match="scratch"): - with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - tuning.replay(ssm_backend, 0, q, q, q, gates, gates, a_log, bias, reqs, cu) - assert not list(tmp_path.iterdir()) - config = {"BV": 32, "num_warps": 1, "num_stages": 3} - for name, expected in before.items(): - torch.testing.assert_close(getattr(ssm_backend, name), expected, atol=0, rtol=0) - - tuning.replay(ssm_backend, 0, q, q, q, gates, gates, a_log, bias, reqs, cu, run_config=config) - torch.cuda.synchronize() diff --git a/unit_tests/common/basemodel/triton_kernel/mhc/test_autotune.py b/unit_tests/common/basemodel/triton_kernel/mhc/test_autotune.py deleted file mode 100644 index b4e1c1ac12..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/mhc/test_autotune.py +++ /dev/null @@ -1,105 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 - - -import pytest -import torch -import triton - -from lightllm.common.basemodel.triton_kernel.mhc import head, post, pre_norm -from lightllm.common.basemodel.triton_kernel.mhc.configs import post_projection_configs -from lightllm.common.triton_utils import autotuner as autotuner_module -from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneLevel -from unit_tests.common.basemodel.triton_kernel.mhc.test_mhc import _hc_post_reference, _hc_pre_reference -from lightllm.common.basemodel.triton_kernel.norm.rmsnorm import rmsnorm_forward - - -@pytest.fixture(autouse=True) -def tuning_environment(monkeypatch): - torch.manual_seed(1525) - monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", None) - monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: AutotuneLevel.CLOSE_AUTOTUNE) - if torch.cuda.is_available(): - triton.set_allocator(lambda size, alignment, stream: torch.empty(size, device="cuda", dtype=torch.int8)) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("tokens,hidden", [(3, 4096), (16, 6144)]) -def test_all_pre_and_post_projection_configs_match_reference(tokens, hidden): - x = torch.randn(tokens, 4 * hidden, device="cuda", dtype=torch.bfloat16) - fn = torch.randn(24, 4 * hidden, device="cuda") * 0.005 - scale, base = torch.randn(3, device="cuda"), torch.randn(24, device="cuda") - norm = torch.randn(hidden, device="cuda", dtype=x.dtype) - expected = _hc_pre_reference(x, fn, scale, base, 4, 1e-6, 1e-6, 20) - - def check(actual, reference): - torch.testing.assert_close(actual[0], rmsnorm_forward(reference[0], norm, 1e-6), atol=0.04, rtol=0.03) - for a, b in zip(actual[1:], reference[1:]): - torch.testing.assert_close(a, b, atol=2e-4, rtol=2e-3) - - for config in pre_norm._hc_prenorm_projection.configs_gen_func(): - gp, sp = pre_norm._hc_prenorm_projection(x, fn, run_config=config) - check(pre_norm._finish_pre_norm(x, gp, sp, scale, base, norm, 4, 1e-6, 1e-6, 1e-6, 20, 2.0), expected) - gp, sp = pre_norm._hc_prenorm_projection(x, fn) - for config in pre_norm._finish_pre_norm.configs_gen_func(): - check( - pre_norm._finish_pre_norm(x, gp, sp, scale, base, norm, 4, 1e-6, 1e-6, 1e-6, 20, 2.0, run_config=config), - expected, - ) - - layer_output = torch.randn(tokens, hidden, device="cuda", dtype=x.dtype) - residual_mix = torch.randn(tokens, 4, 4, device="cuda").softmax(-1) - post_mix = torch.rand(tokens, 4, device="cuda") * 2 - expected_residual = _hc_post_reference(layer_output, x, residual_mix, post_mix, 4) - expected = _hc_pre_reference(expected_residual, fn, scale, base, 4, 1e-6, 1e-6, 20) - before = x.clone() - for config in post_projection_configs(): - updated, gp, sp = pre_norm._hc_post_prenorm_projection( - layer_output, x, residual_mix, post_mix, fn, run_config=config - ) - torch.testing.assert_close(updated, expected_residual, atol=0.04, rtol=0.02) - check(pre_norm._finish_pre_norm(updated, gp, sp, scale, base, norm, 4, 1e-6, 1e-6, 1e-6, 20, 2.0), expected) - torch.testing.assert_close(x, before, atol=0, rtol=0) - for config in pre_norm._hc_post_pre_norm_decode.configs_gen_func(): - actual = pre_norm._hc_post_pre_norm_decode( - layer_output, - x, - residual_mix, - post_mix, - fn, - scale, - base, - norm, - 4, - 1e-6, - 1e-6, - 1e-6, - 20, - 2.0, - run_config=config, - ) - torch.testing.assert_close(actual[1], expected_residual, atol=0.04, rtol=0.02) - check((actual[0], actual[2], actual[3]), expected) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("hidden", [4096, 6144]) -def test_all_post_and_head_configs_match_reference(hidden): - tokens = 19 - x = torch.randn(tokens, 4 * hidden, device="cuda", dtype=torch.bfloat16) - y = torch.randn(tokens, hidden, device="cuda", dtype=x.dtype) - mix, post_mix = torch.randn(tokens, 4, 4, device="cuda").softmax(-1), torch.rand(tokens, 4, device="cuda") * 2 - expected_post = _hc_post_reference(y, x, mix, post_mix, 4) - for config in post.hc_post.configs_gen_func(): - actual = post.hc_post(y, x, mix, post_mix, 4, run_config=config) - torch.testing.assert_close(actual, expected_post, atol=0.04, rtol=0.02) - fn, scale, base = ( - torch.randn(4, 4 * hidden, device="cuda") * 0.005, - torch.randn(1, device="cuda"), - torch.randn(4, device="cuda"), - ) - inv_rms = torch.rsqrt(x.float().square().mean(-1, keepdim=True) + 1e-6) - gates = torch.sigmoid((x.float() @ fn.T) * inv_rms * scale + base) + 1e-6 - expected_head = (x.view(tokens, 4, hidden).float() * gates.unsqueeze(-1)).sum(1).bfloat16() - for config in head.hc_head.configs_gen_func(): - actual = head.hc_head(x, fn, scale, base, 4, 1e-6, 1e-6, run_config=config) - torch.testing.assert_close(actual, expected_head, atol=0.04, rtol=0.02) diff --git a/unit_tests/common/basemodel/triton_kernel/mhc/test_mhc.py b/unit_tests/common/basemodel/triton_kernel/mhc/test_mhc.py deleted file mode 100644 index dd5cece4af..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/mhc/test_mhc.py +++ /dev/null @@ -1,275 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 - -import sys -from types import ModuleType - -import pytest -import torch -import torch.nn.functional as F -import triton - -from lightllm.common.basemodel.triton_kernel.mhc import ( - hc_contract, - hc_expand, - hc_head, - hc_post, - hc_post_pre_norm, - hc_pre_norm, -) -from lightllm.common.basemodel.triton_kernel.mhc import pre_norm -from lightllm.common.basemodel.triton_kernel.norm.rmsnorm import rmsnorm_forward - - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") - - -def _hc_pre_reference(x, fn, scale, base, streams, rms_eps, hc_eps, sinkhorn_iters, post_multiplier=2.0): - """Unfused FP32 reference; returns the merged input before sublayer RMSNorm.""" - assert x.ndim == 2 and x.shape[-1] % streams == 0 - tokens, flattened_hidden = x.shape - hidden = flattened_hidden // streams - residual = x.view(tokens, streams, hidden) - - x_fp32 = x.float() - inv_rms = torch.rsqrt(x_fp32.square().mean(dim=-1, keepdim=True) + rms_eps) - mixes = F.linear(x_fp32, fn) * inv_rms - - pre_raw = mixes[:, :streams] - post_raw = mixes[:, streams : 2 * streams] - residual_raw = mixes[:, 2 * streams :].view(tokens, streams, streams) - - pre = torch.sigmoid(pre_raw * scale[0] + base[:streams]) + hc_eps - post = post_multiplier * torch.sigmoid(post_raw * scale[1] + base[streams : 2 * streams]) - residual_mix = (residual_raw * scale[2] + base[2 * streams :].view(streams, streams)).softmax(dim=-1) - residual_mix = residual_mix + hc_eps - residual_mix = residual_mix / (residual_mix.sum(dim=-2, keepdim=True) + hc_eps) - for _ in range(sinkhorn_iters - 1): - residual_mix = residual_mix / (residual_mix.sum(dim=-1, keepdim=True) + hc_eps) - residual_mix = residual_mix / (residual_mix.sum(dim=-2, keepdim=True) + hc_eps) - - layer_input = (pre.unsqueeze(-1) * residual.float()).sum(dim=1).to(x.dtype) - return layer_input, residual_mix, post - - -def _hc_post_reference(layer_output, residual, residual_mix, post_mix, streams): - """Unfused FP32 reference; residual_mix maps input streams to output streams.""" - tokens, hidden = layer_output.shape - residual_3d = residual.view(tokens, streams, hidden) - mixed_residual = (residual_mix.unsqueeze(-1) * residual_3d.float().unsqueeze(2)).sum(dim=1) - out = post_mix.unsqueeze(-1) * layer_output.float().unsqueeze(1) + mixed_residual - return out.to(layer_output.dtype).reshape(tokens, streams * hidden) - - -@pytest.fixture(autouse=True) -def setup(): - torch.manual_seed(1525) - triton.set_allocator(lambda size, alignment, stream: torch.empty(size, device="cuda", dtype=torch.int8)) - - -@pytest.mark.parametrize("tokens", [3, 19, 3073]) -@pytest.mark.parametrize("disable_deepgemm", [False, True], ids=["deepgemm", "torch"]) -def test_mhc_matches_reference(tokens, disable_deepgemm, monkeypatch): - monkeypatch.setattr(pre_norm, "LIGHTLLM_DISABLE_DEEPGEMM_MHC", disable_deepgemm) - if disable_deepgemm: - monkeypatch.setitem(sys.modules, "deep_gemm", None) - streams, hidden = 4, 4096 - x = torch.randn(tokens, streams * hidden, device="cuda", dtype=torch.bfloat16) - fn = torch.randn(24, streams * hidden, device="cuda") * 0.005 - scale = torch.randn(3, device="cuda") - base = torch.randn(24, device="cuda") - norm = torch.randn(hidden, device="cuda", dtype=torch.bfloat16) - expected = _hc_pre_reference(x, fn, scale, base, streams, 1e-5, 1e-6, 20) - actual = hc_pre_norm(x, fn, scale, base, norm, streams, 1e-5, 1e-5, 1e-6, 20) - torch.testing.assert_close(actual[0], rmsnorm_forward(expected[0], norm, 1e-5), atol=0.04, rtol=0.03) - for a, b in zip(actual[1:], expected[1:]): - # DeepGEMM's mHC projection uses TF32 inputs and FP32 accumulation. - torch.testing.assert_close(a, b, atol=1e-4, rtol=1e-3) - layer_out = torch.randn(tokens, hidden, device="cuda", dtype=torch.bfloat16) - torch.testing.assert_close( - hc_post(layer_out, x, *actual[1:], streams), - _hc_post_reference(layer_out, x, *actual[1:], streams), - atol=0.04, - rtol=0.02, - ) - - -def test_empty_stream_expansion_and_contraction(): - embeddings = torch.empty(0, 4096, device="cuda", dtype=torch.bfloat16) - expanded = hc_expand(embeddings, 4) - assert expanded.shape == (0, 16384) - torch.testing.assert_close(hc_contract(expanded, 4), embeddings) - - -@pytest.mark.parametrize("missing", ["module", "kernel"]) -def test_mhc_reports_missing_deepgemm(monkeypatch, missing): - monkeypatch.setattr(pre_norm, "LIGHTLLM_DISABLE_DEEPGEMM_MHC", False) - monkeypatch.setitem(sys.modules, "deep_gemm", None if missing == "module" else ModuleType("deep_gemm")) - x = torch.empty(1, 16, dtype=torch.bfloat16) - fn = torch.empty(24, 16) - scale = torch.empty(3) - base = torch.empty(24) - norm = torch.empty(4, dtype=torch.bfloat16) - - with pytest.raises(ImportError, match="deep_gemm"): - hc_pre_norm(x, fn, scale, base, norm, 4, 1e-5, 1e-5, 1e-6, 20) - - -@pytest.mark.parametrize("tokens,hidden", [(1, 4096), (8, 4096), (16, 6144), (17, 4096), (129, 4096), (3073, 4096)]) -@pytest.mark.parametrize("disable_deepgemm", [False, True], ids=["fused", "torch"]) -def test_hc_post_pre_norm_matches_reference(tokens, hidden, disable_deepgemm, monkeypatch): - monkeypatch.setattr(pre_norm, "LIGHTLLM_DISABLE_DEEPGEMM_MHC", disable_deepgemm) - streams = 4 - residual = torch.randn(tokens, streams * hidden, device="cuda", dtype=torch.bfloat16) - layer_output = torch.randn(tokens, hidden, device="cuda", dtype=torch.bfloat16) - # Deliberately asymmetric mixing checks input/output stream orientation. - residual_mix = torch.randn(tokens, streams, streams, device="cuda").softmax(-1) - post_mix = torch.rand(tokens, streams, device="cuda") * 2 - fn = torch.randn(24, streams * hidden, device="cuda") * 0.005 - scale = torch.randn(3, device="cuda") - base = torch.randn(24, device="cuda") - norm = torch.randn(hidden, device="cuda", dtype=torch.bfloat16) - residual_before = residual.clone() - expected_residual = _hc_post_reference(layer_output, residual, residual_mix, post_mix, streams) - expected_pre = _hc_pre_reference(expected_residual, fn, scale, base, streams, 1e-6, 1e-6, 20) - actual = hc_post_pre_norm( - layer_output, residual, residual_mix, post_mix, fn, scale, base, norm, streams, 1e-6, 1e-5, 1e-6, 20 - ) - torch.testing.assert_close(residual, residual_before, atol=0, rtol=0) - torch.testing.assert_close(actual[1], expected_residual, atol=0.04, rtol=0.02) - torch.testing.assert_close(actual[0], rmsnorm_forward(expected_pre[0], norm, 1e-5), atol=0.04, rtol=0.03) - for a, b in zip(actual[2:], expected_pre[1:]): - torch.testing.assert_close(a, b, atol=2e-4, rtol=2e-3) - - -@pytest.mark.parametrize("tokens,hidden", [(1, 4096), (9, 6144), (129, 4096), (3073, 4096)]) -@pytest.mark.parametrize("zeros", [False, True]) -def test_hc_head_learned_gates_and_preallocated_output(tokens, hidden, zeros): - streams = 4 - residual = torch.randn(tokens, streams * hidden, device="cuda", dtype=torch.bfloat16) - if zeros: - residual.zero_() - fn = torch.randn(streams, streams * hidden, device="cuda") * 0.005 - scale = torch.randn(1, device="cuda") - base = torch.randn(streams, device="cuda") - inv_rms = torch.rsqrt(residual.float().square().mean(-1, keepdim=True) + 1e-5) - gates = torch.sigmoid(F.linear(residual.float(), fn) * inv_rms * scale + base) + 1e-6 - expected = (residual.view(tokens, streams, hidden).float() * gates.unsqueeze(-1)).sum(1).bfloat16() - output = torch.full((tokens, hidden), float("nan"), device="cuda", dtype=torch.bfloat16) - actual = hc_head(residual, fn, scale, base, streams, 1e-5, 1e-6, out=output) - assert actual is output - torch.testing.assert_close(actual, expected, atol=0.04, rtol=0.02) - - -@pytest.mark.parametrize("tokens", [1, 16, 17, 129]) -def test_shared_mhc_cuda_graph_replay(tokens): - streams, hidden = 4, 4096 - residual = torch.randn(tokens, streams * hidden, device="cuda", dtype=torch.bfloat16) - fn = torch.randn(24, streams * hidden, device="cuda") * 0.005 - scale = torch.randn(3, device="cuda") - base = torch.randn(24, device="cuda") - norm = torch.randn(hidden, device="cuda", dtype=torch.bfloat16) - layer_output = torch.randn(tokens, hidden, device="cuda", dtype=torch.bfloat16) - head_fn, head_scale, head_base = fn[:streams].contiguous(), scale[:1].contiguous(), base[:streams].contiguous() - output = torch.empty(tokens, hidden, device="cuda", dtype=torch.bfloat16) - - def run(): - x, mix, post = hc_pre_norm(residual, fn, scale, base, norm, streams, 1e-6, 1e-5, 1e-6, 20) - assert x.shape == (tokens, hidden) and post.shape == (tokens, streams) - x, updated, mix, post = hc_post_pre_norm( - layer_output, residual, mix, post, fn, scale, base, norm, streams, 1e-6, 1e-5, 1e-6, 20 - ) - final = hc_post(x, updated, mix, post, streams) - return hc_head(final, head_fn, head_scale, head_base, streams, 1e-6, 1e-6, out=output) - - def reference(): - _, mix, post = _hc_pre_reference(residual, fn, scale, base, streams, 1e-6, 1e-6, 20) - updated = _hc_post_reference(layer_output, residual, mix, post, streams) - x, mix, post = _hc_pre_reference(updated, fn, scale, base, streams, 1e-6, 1e-6, 20) - x = rmsnorm_forward(x, norm, 1e-5) - final = _hc_post_reference(x, updated, mix, post, streams) - inv_rms = torch.rsqrt(final.float().square().mean(-1, keepdim=True) + 1e-6) - gates = torch.sigmoid(F.linear(final.float(), head_fn) * inv_rms * head_scale + head_base) + 1e-6 - return (final.view(tokens, streams, hidden).float() * gates.unsqueeze(-1)).sum(1).bfloat16() - - run() - torch.cuda.synchronize() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - actual = run() - for _ in range(2): - residual.normal_() - layer_output.normal_() - graph.replay() - torch.testing.assert_close(actual, reference(), atol=0.06, rtol=0.04) - - -def test_mhc_reduction_reuses_compilation_within_split_buckets(): - hashes = {} - for tokens, splits in [ - (1, 1), - (16, 1), - (2, 3), - (16, 3), - (16, 8), - (3, 3), - (17, 7), - (31, 14), - (129, 26), - (257, 33), - (513, 64), - ]: - partial = torch.randn((splits, tokens, 24), device="cuda") * 0.05 - squares = torch.rand((splits, tokens), device="cuda") * (16384 / splits) + 1 - scale = torch.randn(3, device="cuda") - base = torch.randn(24, device="cuda") - x = torch.randn((tokens, 16384), device="cuda", dtype=torch.bfloat16) - norm = torch.randn(4096, device="cuda", dtype=torch.bfloat16) - - def run(gemm, sqrsum): - post = torch.empty((tokens, 4), device="cuda") - mix = torch.empty((tokens, 4, 4), device="cuda") - output = torch.empty((tokens, 4096), device="cuda", dtype=torch.bfloat16) - block_splits = triton.next_power_of_2(gemm.shape[0]) - compiled = pre_norm._hc_pre_norm_kernel[(tokens,)]( - x, - gemm, - sqrsum, - scale, - base, - post, - mix, - norm, - output, - hidden=4096, - x_stride_m=x.stride(0), - gemm_stride_s=gemm.stride(0), - gemm_stride_m=gemm.stride(1), - sqrsum_stride_s=sqrsum.stride(0), - sqrsum_stride_m=sqrsum.stride(1), - out_stride_m=output.stride(0), - RMS_EPS=1e-5, - NORM_EPS=1e-5, - STREAMS=4, - HC_EPS=1e-6, - POST_MULTIPLIER=2.0, - SINKHORN_ITERS=20, - N_SPLITS=gemm.shape[0], - BLOCK_SPLITS=block_splits, - BLOCK_H=4096, - num_warps=4, - ) - hashes.setdefault(block_splits, set()).add(compiled.hash) - return output, mix, post - - # The fused kernel reduces splits in parallel, so compare to sequential - # FP32 accumulation with tolerance for the different reduction order. - total = torch.zeros_like(partial[:1]) - total_squares = torch.zeros_like(squares[:1]) - for index in range(splits): - total.add_(partial[index : index + 1]) - total_squares.add_(squares[index : index + 1]) - actual, expected = run(partial, squares), run(total, total_squares) - torch.testing.assert_close(actual[0], expected[0], atol=0.04, rtol=0.03) - for output, reference in zip(actual[1:], expected[1:]): - torch.testing.assert_close(output, reference, atol=1e-6, rtol=1e-5) - assert all(len(bucket_hashes) == 1 for bucket_hashes in hashes.values()) diff --git a/unit_tests/common/basemodel/triton_kernel/mla_att/decode_att/test_gqa_flash_decoding.py b/unit_tests/common/basemodel/triton_kernel/mla_att/decode_att/test_gqa_flash_decoding.py deleted file mode 100644 index a5ac9708d2..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/mla_att/decode_att/test_gqa_flash_decoding.py +++ /dev/null @@ -1,112 +0,0 @@ -import torch -import time -import pytest -import numpy as np -import torch.nn.functional as F -import flashinfer -from lightllm.utils.log_utils import init_logger -from lightllm.common.basemodel.triton_kernel.mla_att.decode_att.gqa_flash_decoding import ( - gqa_token_decode_attention_flash_decoding, -) -from lightllm.models.deepseek2.infer_struct import Deepseek2InferStateInfo - -logger = init_logger(__name__) - -seed = 42 -torch.manual_seed(seed) - -if torch.cuda.is_available(): - torch.cuda.manual_seed(seed) - torch.cuda.manual_seed_all(seed) - - -@pytest.mark.parametrize( - "batch, seqlen, heads, nope_head, rope_head", - [ - (a, b, c, d, e) - for a in [1, 16, 32, 128, 512] - for b in [16, 32, 512, 2048] - for c in [16] - for d in [512] - for e in [64] - ], -) -def test_gqa_flash_decoding(batch, seqlen, heads, nope_head, rope_head): - Z, N_CTX, H, D_HEAD, ROPE_HEAD = batch, seqlen, heads, nope_head, rope_head - dtype = torch.bfloat16 - sm_scale = 1.0 / ((D_HEAD + ROPE_HEAD) ** 0.5) - q_nope = torch.randn((Z, H, D_HEAD), dtype=dtype, device="cuda") - q_rope = torch.randn((Z, H, ROPE_HEAD), dtype=dtype, device="cuda") - - kv = torch.randn((Z * N_CTX, 1, D_HEAD + ROPE_HEAD), dtype=dtype, device="cuda") - - max_input_len = Z * N_CTX - req_to_token_indexs = torch.randperm(max_input_len, dtype=torch.int32).cuda().view(Z, N_CTX) - b_seq_len = torch.ones((Z,), dtype=torch.int32, device="cuda") * N_CTX - b_start_loc = torch.arange(Z).cuda().int() * N_CTX - b_req_idx = torch.randperm(Z, dtype=torch.int32).cuda() - kv_starts = torch.cat([b_start_loc, b_start_loc[-1:] + b_seq_len[-1:]], dim=0) - - o = torch.zeros((Z, H, D_HEAD), dtype=dtype, device="cuda") - o1 = torch.zeros((Z, H, D_HEAD), dtype=dtype, device="cuda") - - infer_state = Deepseek2InferStateInfo() - infer_state.batch_size = Z - infer_state.max_len_in_batch = N_CTX - infer_state.total_token_num = Z * N_CTX - infer_state.req_manager = type("Object", (), {})() - 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 - infer_state.kv_starts = kv_starts - - kv_nope = kv[:, :, :D_HEAD] - kv_rope = kv[:, :, D_HEAD:] - gqa_token_decode_attention_flash_decoding( - q_nope, - q_rope, - kv_nope, - kv_rope, - infer_state, - sm_scale, - o, - ) - - batch_size = Z - head_dim_ckv = D_HEAD - head_dim_kpe = ROPE_HEAD - num_heads = H - page_size = 1 - workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.int8).to(0) - q_indptr = torch.arange(batch_size + 1).to(0).int() - kv_indptr = infer_state.kv_starts - kv_indices = torch.arange(Z * N_CTX).cuda().int() - for b, sl, start in zip(b_req_idx, b_seq_len, b_start_loc): - kv_indices[start : start + sl] = req_to_token_indexs[b][:sl] - kv_lens = b_seq_len - wrapper = flashinfer.mla.BatchMLAPagedAttentionWrapper( - workspace_buffer, - use_cuda_graph=True, - qo_indptr=q_indptr, - kv_indices=kv_indices, - kv_indptr=kv_indptr, - kv_len_arr=kv_lens, - ) - wrapper.plan( - q_indptr, - kv_indptr, - kv_indices, - kv_lens, - num_heads, - head_dim_ckv, - head_dim_kpe, - page_size, - False, # causal - sm_scale, - q_nope.dtype, - kv.dtype, - ) - wrapper.run(q_nope, q_rope, kv_nope, kv_rope, out=o1, return_lse=False) - - cos_sim1 = F.cosine_similarity(o, o1).mean() - assert cos_sim1 == 1.0 diff --git a/unit_tests/common/basemodel/triton_kernel/test_add_in_place.py b/unit_tests/common/basemodel/triton_kernel/test_add_in_place.py deleted file mode 100644 index 7e9a55a775..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/test_add_in_place.py +++ /dev/null @@ -1,31 +0,0 @@ -import torch -import time -import pytest -from lightllm.common.basemodel.triton_kernel.sp_pad_copy import sp_pad_copy -from lightllm.common.basemodel.triton_kernel.add_in_place import add_in_place -from lightllm.utils.log_utils import init_logger - -logger = init_logger(__name__) - - -@pytest.mark.parametrize( - "dim1, dim2, alpha", - [ - (dim1, dim2, alpha) - for dim1 in range(1, 1024, 100) - for dim2 in range(1, 1024, 100) - for alpha in [0.1, 0.3, 0.5, 0.7, 0.1] - ], -) -def test_add_in_place(dim1, dim2, alpha): - input = torch.rand((dim1, dim2), device="cuda") - other = torch.rand((dim1, dim2), device="cuda") - - output = input + other * alpha - add_in_place(input, other, alpha=alpha) - rlt = torch.allclose(input, output, atol=1e-5, rtol=0) - assert rlt - - -if __name__ == "__main__": - pytest.main() diff --git a/unit_tests/common/basemodel/triton_kernel/test_apply_constraint_mask.py b/unit_tests/common/basemodel/triton_kernel/test_apply_constraint_mask.py deleted file mode 100644 index c6768053e4..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/test_apply_constraint_mask.py +++ /dev/null @@ -1,104 +0,0 @@ -import pytest -import torch - -from lightllm.common.basemodel.triton_kernel.post_process.apply_constraint_mask import apply_constraint_mask - - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="Constraint mask kernel requires CUDA") - - -@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) -@pytest.mark.parametrize("vocab_size,padding", [(1, 0), (31, 5), (32, 0), (33, 31), (257, 7), (8195, 43)]) -@pytest.mark.parametrize("mtp", [False, True]) -def test_request_indexed_mask_matches_reference(dtype, vocab_size, padding, mtp): - generator = torch.Generator().manual_seed(2026) - masks = torch.randint( - -(1 << 31), - (1 << 31) - 1, - (6, 4, (vocab_size + 31) // 32), - dtype=torch.int32, - generator=generator, - pin_memory=True, - ) - enabled = torch.tensor([False, True, False, False, True, False], pin_memory=True) - req_ids = torch.tensor([4, 1, 4, 2, 1], dtype=torch.int32) - positions = torch.tensor([2, 0, 1, 3, 3], dtype=torch.int32) if mtp else torch.zeros(5, dtype=torch.int32) - # Exercise a padded row stride as well as unused vocabulary columns. - width = vocab_size + padding - storage = torch.randn(5, width + 11, dtype=dtype, device="cuda") - logits = storage[:, :width] - expected = logits.cpu().clone() - token_ids = torch.arange(vocab_size) - for prediction, (req_id, position) in enumerate(zip(req_ids, positions)): - if enabled[req_id]: - words = masks[req_id, position, token_ids // 32] - allowed = ((words >> (token_ids % 32)) & 1).bool() - expected[prediction, :vocab_size].masked_fill_(~allowed, float("-inf")) - expected[:, vocab_size:] = float("-inf") - untouched_padding = storage[:, width:].clone() - - apply_constraint_mask( - logits=logits, - b_req_idx=req_ids.cuda(), - req_to_bitmask=masks.cuda(non_blocking=True), - req_to_bitmask_enabled=enabled.cuda(non_blocking=True), - vocab_size=vocab_size, - b_mtp_index=positions.cuda() if mtp else None, - ) - - torch.testing.assert_close(logits.cpu(), expected, rtol=0, atol=0) - torch.testing.assert_close(storage[:, width:], untouched_padding, rtol=0, atol=0) - - -def test_pinned_masks_can_be_reused_after_sampling_completion(): - masks = torch.zeros(3, 2, 2, dtype=torch.int32, pin_memory=True) - enabled = torch.tensor([True, True, False], pin_memory=True) - masks[0, 0, 0] = -(1 << 31) - masks[1, 1, 1] = 1 - req_ids = torch.tensor([1, 0], dtype=torch.int32, device="cuda") - positions = torch.tensor([1, 0], dtype=torch.int32, device="cuda") - logits = torch.zeros(2, 33, device="cuda") - stream = torch.cuda.Stream() - stream.wait_stream(torch.cuda.current_stream()) - # This represents the existing sampling/post_handle completion boundary. - done = torch.cuda.Event() - with torch.cuda.stream(stream): - apply_constraint_mask( - logits=logits, - b_req_idx=req_ids, - req_to_bitmask=masks.cuda(non_blocking=True), - req_to_bitmask_enabled=enabled.cuda(non_blocking=True), - vocab_size=33, - b_mtp_index=positions, - ) - done.record() - done.synchronize() - assert torch.isfinite(logits[0]).nonzero().flatten().tolist() == [32] - assert torch.isfinite(logits[1]).nonzero().flatten().tolist() == [31] - - # Reuse one slot for an ordinary request, and update the other's prefix. - enabled[1] = False - masks[0].fill_(-1) - with torch.cuda.stream(stream): - logits.fill_(1.25) - apply_constraint_mask( - logits=logits, - b_req_idx=req_ids, - req_to_bitmask=masks.cuda(non_blocking=True), - req_to_bitmask_enabled=enabled.cuda(non_blocking=True), - vocab_size=33, - b_mtp_index=positions, - ) - done.record() - done.synchronize() - torch.testing.assert_close(logits, torch.full_like(logits, 1.25), rtol=0, atol=0) - - -def test_empty_batch_needs_no_kernel_launch(): - apply_constraint_mask( - logits=torch.empty(0, 33, device="cuda"), - b_req_idx=torch.empty(0, dtype=torch.int32, device="cuda"), - req_to_bitmask=torch.empty(1, 1, 2, dtype=torch.int32, device="cuda"), - req_to_bitmask_enabled=torch.zeros(1, dtype=torch.bool, device="cuda"), - vocab_size=33, - ) diff --git a/unit_tests/common/basemodel/triton_kernel/test_apply_invalid_token.py b/unit_tests/common/basemodel/triton_kernel/test_apply_invalid_token.py deleted file mode 100644 index 3b2f159f62..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/test_apply_invalid_token.py +++ /dev/null @@ -1,50 +0,0 @@ -import pytest -import torch - -from lightllm.common.basemodel.triton_kernel.post_process.apply_invalid_token import ( - apply_invalid_token_ids, -) - - -@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) -def test_apply_invalid_token_ids(dtype): - if not torch.cuda.is_available(): - pytest.skip("CUDA is required for Triton kernels.") - - batch_size = 4 - vocab_size = 32 - logits = torch.randn((batch_size, vocab_size), device="cuda", dtype=dtype) - expected = logits.clone() - - invalid_token_ids_per_batch = [ - [1, 3, 5], - [], - [0, 2, 31], - [7], - ] - - flat_ids = [] - cu_invalid_token_num = [0] - invalid_token_num_start = 0 - for ids in invalid_token_ids_per_batch: - flat_ids.extend(ids) - invalid_token_num_start += len(ids) - cu_invalid_token_num.append(invalid_token_num_start) - - invalid_token_ids = torch.tensor(flat_ids, device="cuda", dtype=torch.int32) - cu_invalid_token_num = torch.tensor(cu_invalid_token_num, device="cuda", dtype=torch.int32) - - for batch_idx, ids in enumerate(invalid_token_ids_per_batch): - if ids: - expected[batch_idx, ids] = float("-inf") - - apply_invalid_token_ids( - Logits=logits, - invalid_token_ids=invalid_token_ids, - cu_invalid_token_num=cu_invalid_token_num, - ) - assert torch.equal(logits, expected) - - -if __name__ == "__main__": - pytest.main([__file__]) diff --git a/unit_tests/common/basemodel/triton_kernel/test_atomic_event.py b/unit_tests/common/basemodel/triton_kernel/test_atomic_event.py deleted file mode 100644 index 0afcd5558a..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/test_atomic_event.py +++ /dev/null @@ -1,28 +0,0 @@ -import torch -import pytest -from lightllm.common.basemodel.triton_kernel.atomic_event import wait_value, add_value -from lightllm.utils.log_utils import init_logger - -logger = init_logger(__name__) - - -def test_add_in_place(): - input = torch.zeros((1,), device="cuda", dtype=torch.int32) - wait_value(input, 0) - add_value(input) - wait_value(input, 1) - add_value(input) - wait_value(input, 2) - add_value(input) - wait_value(input, 3) - assert input.item() == 3, "最终值应为 3" - - -# @pytest.mark.timeout(2) -# def test_wait_timeout(): -# input = torch.zeros((1,), device="cuda", dtype=torch.int32) -# wait_value(input, 4) - - -if __name__ == "__main__": - pytest.main() diff --git a/unit_tests/common/basemodel/triton_kernel/test_diverse_utils.py b/unit_tests/common/basemodel/triton_kernel/test_diverse_utils.py deleted file mode 100644 index 0c1c2df08a..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/test_diverse_utils.py +++ /dev/null @@ -1,48 +0,0 @@ -import weakref -from types import SimpleNamespace - -import pytest -import torch - -if not torch.cuda.is_available(): - pytest.skip("requires CUDA", allow_module_level=True) - -from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.common.basemodel.attention.triton import int8kv as int8kv_module -from lightllm.common.basemodel.attention.triton.int8kv import Int8kvTritonDecodeAttState -from lightllm.common.basemodel.triton_kernel import diverse_utils - - -@pytest.mark.parametrize( - "node_ids,expected_markers", - [ - ([10, 10, 10, 10, 20, 20], [0, 0, 3, 1, 0, 2]), - ([10, 10, 20, 10, 10], [0, 2, 1, 0, 2]), - ([10, 10, -1, -1], [0, 2, 1, 1]), - ([-1, -1, -1, -1], [1, 1, 1, 1]), - ], -) -def test_build_diverse_shared_group_markers(monkeypatch, node_ids, expected_markers): - monkeypatch.setattr(diverse_utils, "get_diverse_max_batch_shared_group_size", lambda: 3) - b_shared_radix_node_id = torch.tensor(node_ids, dtype=torch.int64, device="cuda") - - markers = diverse_utils.build_diverse_shared_group_markers( - b_shared_radix_node_id=b_shared_radix_node_id, - ) - - assert markers.cpu().tolist() == expected_markers - - -def test_int8kv_decode_state_rebuilds_diverse_metadata(monkeypatch): - monkeypatch.setattr(int8kv_module, "enable_diverse_mode_gqa_decode_fast_kernel", lambda: True) - monkeypatch.setattr(diverse_utils, "get_diverse_max_batch_shared_group_size", lambda: 3) - infer_state = InferStateInfo() - infer_state.max_kv_seq_len = 8192 - infer_state.b_shared_seq_len = torch.tensor([8, 8, 5, 0], dtype=torch.int32, device="cuda") - infer_state.b_shared_radix_node_id = torch.tensor([10, 10, 20, -1], dtype=torch.int64, device="cuda") - state = Int8kvTritonDecodeAttState(backend=SimpleNamespace(), _infer_state_weak_ref=weakref.ref(infer_state)) - - state.init_state() - - assert state.b_mark_shared_group.cpu().tolist() == [0, 2, 1, 1] - assert state.b_shared_seq_len.cpu().tolist() == [8, 8, 0, 0] diff --git a/unit_tests/common/basemodel/triton_kernel/test_dynamic_mtp_utils.py b/unit_tests/common/basemodel/triton_kernel/test_dynamic_mtp_utils.py deleted file mode 100644 index cd136ae713..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/test_dynamic_mtp_utils.py +++ /dev/null @@ -1,360 +0,0 @@ -import json - -import torch -import pytest -import triton -import numpy as np - -from lightllm.common.basemodel.batch_objs import ModelInput -from lightllm.common.basemodel.mtp_manager import MtpManager -from lightllm.common.basemodel.triton_kernel import dynamic_mtp_utils -from lightllm.common.basemodel.triton_kernel.dynamic_mtp_utils import ( - _fwd_kernel_cumprod_scores, - sample_dynamic_mtp_row_mask, -) -from lightllm.utils.envs_utils import get_env_start_args - - -@pytest.fixture(autouse=True) -def _reset_mtp_manager(): - MtpManager._instance = None - yield - MtpManager._instance = None - - -def test_compact_dynamic_mtp_model_input(monkeypatch): - monkeypatch.setenv( - "LIGHTLLM_START_ARGS", - json.dumps( - { - "diverse_mode": False, - "llm_kv_type": "fp16", - "mtp_mode": "eagle3", - "mtp_dynamic_verify": True, - "mtp_step": 3, - "llm_decode_att_backend": "triton", - } - ), - ) - monkeypatch.setenv("LIGHTLLM_MAX_BATCH_SHARED_GROUP_SIZE", "4") - get_env_start_args.cache_clear() - - model_input = ModelInput( - batch_size=12, - total_token_num=54, - max_q_seq_len=1, - max_kv_seq_len=6, - input_ids=torch.arange(12, dtype=torch.int64, device="cuda") + 1000, - b_req_idx=torch.tensor([0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2], dtype=torch.int32, device="cuda"), - b_mtp_index=torch.tensor([0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3], dtype=torch.int32, device="cuda"), - b_seq_len=torch.tensor([3, 4, 5, 6, 3, 4, 5, 6, 3, 4, 5, 6], dtype=torch.int32, device="cuda"), - b_position_delta=torch.arange(12, dtype=torch.int32, device="cuda") + 200, - b_shared_seq_len=torch.tensor([0, 0, 0, 0, 7, 7, 7, 7, 9, 9, 9, 9], dtype=torch.int32, device="cuda"), - b_shared_radix_node_id=torch.tensor( - [10, 10, 10, 10, 11, 11, 11, 11, 12, 12, 12, 12], dtype=torch.int64, device="cuda" - ), - is_prefill=False, - multimodal_params=[{"row": i, "images": [], "audios": []} for i in range(12)], - mtp_draft_input_hiddens=(torch.arange(12 * 5, dtype=torch.float32, device="cuda").reshape(12, 5) + 0.5), - ) - req_to_next_token_scores = torch.tensor( - [ - [1.0, 0.95, 0.90, 0.10, 0.0, 0.0], - [1.0, 0.20, 0.80, 0.80, 0.0, 0.0], - [1.0, 0.99, 0.99, 0.99, 0.0, 0.0], - ], - dtype=torch.float32, - device="cuda", - ) - - compacted_input, selected_row_mask = dynamic_mtp_utils.prepare_dynamic_mtp_model_input( - model_input=model_input, - req_num=3, - dynamic_batch_size=8, - req_to_next_token_scores=req_to_next_token_scores, - ) - torch.cuda.synchronize() - - expected_selected_mask = torch.tensor([1, 1, 1, 0, 1, 0, 0, 0, 1, 1, 1, 1], dtype=torch.int32) - expected_selected_rows = torch.where(expected_selected_mask == 1)[0] - - assert torch.equal(selected_row_mask.cpu(), expected_selected_mask) - assert compacted_input.batch_size == 8 - assert compacted_input.max_q_seq_len == 1 - assert compacted_input.multimodal_params == [{"images": [], "audios": []}] * 8 - - assert torch.equal( - compacted_input.input_ids.cpu(), torch.arange(12, dtype=torch.int64)[expected_selected_rows] + 1000 - ) - assert torch.equal(compacted_input.b_req_idx.cpu(), torch.tensor([0, 0, 0, 1, 2, 2, 2, 2], dtype=torch.int32)) - assert torch.equal(compacted_input.b_mtp_index.cpu(), torch.tensor([0, 1, 2, 0, 0, 1, 2, 3], dtype=torch.int32)) - assert torch.equal(compacted_input.b_seq_len.cpu(), torch.tensor([3, 4, 5, 3, 3, 4, 5, 6], dtype=torch.int32)) - assert torch.equal( - compacted_input.b_position_delta.cpu(), - torch.tensor([200, 201, 202, 204, 208, 209, 210, 211], dtype=torch.int32), - ) - assert torch.equal( - compacted_input.b_shared_seq_len.cpu(), torch.tensor([0, 0, 0, 7, 9, 9, 9, 9], dtype=torch.int32) - ) - assert torch.equal( - compacted_input.b_shared_radix_node_id.cpu(), - torch.tensor([10, 10, 10, 11, 12, 12, 12, 12], dtype=torch.int64), - ) - expected_hiddens = (torch.arange(12 * 5, dtype=torch.float32).reshape(12, 5) + 0.5)[expected_selected_rows] - assert torch.equal(compacted_input.mtp_draft_input_hiddens.cpu(), expected_hiddens) - - -def test_compaction_preserves_shared_radix_metadata(): - model_input = ModelInput( - batch_size=5, - total_token_num=36, - max_q_seq_len=1, - max_kv_seq_len=6, - input_ids=None, - b_req_idx=torch.tensor([0, 0, 0, 0, 0], dtype=torch.int32, device="cuda"), - b_mtp_index=torch.tensor([0, 1, 2, 3, 4], dtype=torch.int32, device="cuda"), - b_seq_len=torch.tensor([3, 4, 5, 6, 7], dtype=torch.int32, device="cuda"), - b_position_delta=torch.zeros(5, dtype=torch.int32, device="cuda"), - b_shared_seq_len=torch.full((5,), 7, dtype=torch.int32, device="cuda"), - b_shared_radix_node_id=torch.full((5,), 10, dtype=torch.int64, device="cuda"), - is_prefill=False, - multimodal_params=[{"images": [], "audios": []} for _ in range(5)], - ) - selected_row_mask = torch.ones((5,), dtype=torch.int32, device="cuda") - - compacted_input = dynamic_mtp_utils._compact_decode_model_input( - model_input=model_input, - selected_row_mask=selected_row_mask, - dynamic_batch_size=5, - ) - torch.cuda.synchronize() - - assert torch.equal(compacted_input.b_req_idx.cpu(), torch.tensor([0, 0, 0, 0, 0], dtype=torch.int32)) - assert torch.equal(compacted_input.b_mtp_index.cpu(), torch.tensor([0, 1, 2, 3, 4], dtype=torch.int32)) - assert torch.equal(compacted_input.b_shared_seq_len.cpu(), torch.full((5,), 7, dtype=torch.int32)) - assert torch.equal(compacted_input.b_shared_radix_node_id.cpu(), torch.full((5,), 10, dtype=torch.int64)) - - -def _reference_cumprod_scores(req_to_next_token_scores, b_req_idx, max_draft_step: int) -> torch.Tensor: - scores = req_to_next_token_scores.clone() - req_num = b_req_idx.shape[0] // (max_draft_step + 1) - for req_i in range(req_num): - req_idx = int(b_req_idx[req_i * (max_draft_step + 1)].item()) - row = scores[req_idx, : max_draft_step + 1].clone() - row[0] = 1.0 - row[1:] = torch.clamp(row[1:], min=0.01, max=0.99) - scores[req_idx, : max_draft_step + 1] = torch.cumprod(row, dim=0) - return scores - - -def _flat_cumprod_scores( - b_req_idx: torch.Tensor, - req_to_next_token_scores: torch.Tensor, - max_draft_step: int, -) -> torch.Tensor: - scores = _reference_cumprod_scores(req_to_next_token_scores, b_req_idx, max_draft_step) - req_num = b_req_idx.shape[0] // (max_draft_step + 1) - all_num = req_num * (max_draft_step + 1) - flat_scores = [] - for offset in range(all_num): - req_idx = int(b_req_idx[offset].item()) - mtp_index = offset % (max_draft_step + 1) - flat_scores.append(scores[req_idx, mtp_index]) - return torch.stack(flat_scores) - - -def _assert_topk_mask(select: torch.Tensor, flat_scores: torch.Tensor, dynamic_batch_size: int) -> None: - k = dynamic_batch_size - assert int(select.sum().item()) == k - selected_scores = flat_scores[select.bool()] - unselected_scores = flat_scores[(select == 0).bool()] - if unselected_scores.numel() > 0: - assert selected_scores.min() >= unselected_scores.max() - 1e-5 - - -def _make_batch_scores(req_num: int, max_draft_step: int, rows): - max_req = req_num - scores = torch.zeros((max_req + 1, 16), dtype=torch.float32, device="cuda") - for req_idx, row in enumerate(rows): - scores[req_idx, : max_draft_step + 1] = torch.tensor(row, dtype=torch.float32, device="cuda") - b_req_idx = torch.arange(req_num, dtype=torch.int32, device="cuda").repeat_interleave(max_draft_step + 1) - return scores, b_req_idx - - -@pytest.mark.parametrize("max_draft_step", [1, 3]) -def test_cumprod_scores_kernel(max_draft_step: int): - req_num = 2 - scores, b_req_idx = _make_batch_scores( - req_num, - max_draft_step, - rows=[ - [1.0] + [0.5] * max_draft_step, - [1.0] + [0.2] * max_draft_step, - ], - ) - scores_clone = scores.clone() - _fwd_kernel_cumprod_scores[(req_num,)]( - req_to_next_token_scores=scores_clone, - req_to_next_token_scores_stride=scores_clone.stride(0), - b_req_idx=b_req_idx, - max_draft_step=max_draft_step, - BLOCK_SIZE=triton.next_power_of_2(max_draft_step + 1), - num_warps=1, - num_stages=1, - ) - expected = _reference_cumprod_scores(scores, b_req_idx, max_draft_step) - assert torch.allclose( - scores_clone[:, : max_draft_step + 1], expected[:, : max_draft_step + 1], rtol=1e-5, atol=1e-5 - ) - - -def test_cumprod_scores_clamps_invalid_values(): - max_draft_step = 2 - req_num = 1 - scores, b_req_idx = _make_batch_scores(req_num, max_draft_step, rows=[[1.0, 0.0, 1.5]]) - _fwd_kernel_cumprod_scores[(req_num,)]( - req_to_next_token_scores=scores, - req_to_next_token_scores_stride=scores.stride(0), - b_req_idx=b_req_idx, - max_draft_step=max_draft_step, - BLOCK_SIZE=triton.next_power_of_2(max_draft_step + 1), - num_warps=1, - num_stages=1, - ) - row = scores[0, : max_draft_step + 1] - # Index 0 is the always-accepted target sample; clamping applies only to drafts. - assert row[0].item() == pytest.approx(1.0) - assert row[1].item() == pytest.approx(0.01, rel=1e-4) - assert row[2].item() == pytest.approx(0.01 * 0.99, rel=1e-4) - - -def test_cumprod_scores_clamps_boundary_values(): - max_draft_step = 3 - req_num = 1 - scores, b_req_idx = _make_batch_scores(req_num, max_draft_step, rows=[[1.0, 0.995, 0.005, 0.5]]) - raw_scores = scores.clone() - _fwd_kernel_cumprod_scores[(req_num,)]( - req_to_next_token_scores=scores, - req_to_next_token_scores_stride=scores.stride(0), - b_req_idx=b_req_idx, - max_draft_step=max_draft_step, - BLOCK_SIZE=triton.next_power_of_2(max_draft_step + 1), - num_warps=1, - num_stages=1, - ) - expected = _reference_cumprod_scores(raw_scores, b_req_idx, max_draft_step) - row = scores[0, : max_draft_step + 1] - assert torch.allclose(row, expected[0, : max_draft_step + 1], rtol=1e-5, atol=1e-5) - # Draft scores 0.995 and 0.005 clamp to 0.99 and 0.01. - assert row[0].item() == pytest.approx(1.0) - assert row[1].item() == pytest.approx(0.99, rel=1e-4) - assert row[2].item() == pytest.approx(0.99 * 0.01, rel=1e-4) - assert row[3].item() == pytest.approx(0.99 * 0.01 * 0.5, rel=1e-4) - - -def test_sample_select_count(): - max_draft_step = 3 - req_num = 3 - scores, b_req_idx = _make_batch_scores( - req_num, - max_draft_step, - rows=[ - [1.0, 0.95, 0.90, 0.10], - [1.0, 0.20, 0.80, 0.80], - [1.0, 0.99, 0.99, 0.99], - ], - ) - all_num = req_num * (max_draft_step + 1) - for dynamic_batch_size in [3, 8, all_num]: - select = sample_dynamic_mtp_row_mask( - dynamic_batch_size=dynamic_batch_size, - b_req_idx=b_req_idx, - req_to_next_token_scores=scores.clone(), - max_draft_step=max_draft_step, - ) - assert select.dtype == torch.int32 - assert select.shape[0] == all_num - assert int(select.sum().item()) == dynamic_batch_size - assert torch.all((select == 0) | (select == 1)) - - -def test_sample_accepts_numpy_scalar_dynamic_batch_size(): - max_draft_step = 3 - scores, b_req_idx = _make_batch_scores( - 3, - max_draft_step, - rows=[ - [1.0, 0.95, 0.90, 0.10], - [1.0, 0.20, 0.80, 0.80], - [1.0, 0.99, 0.99, 0.99], - ], - ) - select = sample_dynamic_mtp_row_mask( - dynamic_batch_size=np.int64(8), - b_req_idx=b_req_idx, - req_to_next_token_scores=scores, - max_draft_step=np.int64(max_draft_step), - ) - assert int(select.sum().item()) == 8 - - -def test_sample_topk_by_cumprod_score(): - max_draft_step = 3 - scores, b_req_idx = _make_batch_scores( - 3, - max_draft_step, - rows=[ - [1.0, 0.95, 0.90, 0.10], - [1.0, 0.20, 0.80, 0.80], - [1.0, 0.99, 0.99, 0.99], - ], - ) - flat_scores = _flat_cumprod_scores(b_req_idx, scores, max_draft_step) - for dynamic_batch_size in [1, 4, 8, 12]: - select = sample_dynamic_mtp_row_mask( - dynamic_batch_size=dynamic_batch_size, - b_req_idx=b_req_idx, - req_to_next_token_scores=scores.clone(), - max_draft_step=max_draft_step, - ) - _assert_topk_mask(select, flat_scores, dynamic_batch_size) - - -def test_sample_picks_highest_cumprod_rows(): - max_draft_step = 1 - scores, b_req_idx = _make_batch_scores( - 2, - max_draft_step, - rows=[ - [1.0, 0.9], - [1.0, 0.1], - ], - ) - flat_scores = _flat_cumprod_scores(b_req_idx, scores, max_draft_step) - select = sample_dynamic_mtp_row_mask( - dynamic_batch_size=2, - b_req_idx=b_req_idx, - req_to_next_token_scores=scores.clone(), - max_draft_step=max_draft_step, - ) - _assert_topk_mask(select, flat_scores, 2) - # top-2 scores are both 0.99 at mtp_index==0 (req0 and req1 main rows) - assert select[0].item() == 1 - assert select[2].item() == 1 - - -def test_sample_single_request(): - max_draft_step = 2 - scores, b_req_idx = _make_batch_scores(1, max_draft_step, rows=[[1.0, 0.5, 0.25]]) - flat_scores = _flat_cumprod_scores(b_req_idx, scores, max_draft_step) - select = sample_dynamic_mtp_row_mask( - dynamic_batch_size=2, - b_req_idx=b_req_idx, - req_to_next_token_scores=scores.clone(), - max_draft_step=max_draft_step, - ) - _assert_topk_mask(select, flat_scores, 2) - - -if __name__ == "__main__": - pytest.main([__file__, "-v"]) diff --git a/unit_tests/common/basemodel/triton_kernel/test_fa3_utils.py b/unit_tests/common/basemodel/triton_kernel/test_fa3_utils.py deleted file mode 100644 index c1ef686f8e..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/test_fa3_utils.py +++ /dev/null @@ -1,82 +0,0 @@ -import pytest -import torch - -if not torch.cuda.is_available(): - pytest.skip("requires CUDA", allow_module_level=True) - -from lightllm.common.basemodel.triton_kernel.fa3_utils import build_dynamic_spec_fa3_decode_params - - -def _reference_dynamic_spec_fa3_decode_params(b_req_idx, b_seq_len, b_mark_mtp_shared_group, hold_req_id): - batch_size = b_req_idx.shape[0] - b_req_idx_cpu = b_req_idx.cpu() - b_seq_len_cpu = b_seq_len.cpu() - b_mark_cpu = b_mark_mtp_shared_group.cpu() - valid_idx = torch.nonzero(b_mark_cpu > 0, as_tuple=False).flatten() - valid_size = valid_idx.numel() - - b_q_seq_len = torch.zeros((batch_size,), dtype=torch.int32) - b_kv_seq_len = torch.zeros((batch_size,), dtype=torch.int32) - b_att_req_idx = torch.full((batch_size,), hold_req_id, dtype=torch.int32) - b_att_seq_len = torch.zeros((batch_size,), dtype=torch.int32) - - b_q_seq_len[:valid_size] = b_mark_cpu[valid_idx] - b_kv_seq_len[:valid_size] = b_seq_len_cpu[valid_idx] - b_att_req_idx[:valid_size] = b_req_idx_cpu[valid_idx] - b_att_seq_len[:valid_size] = b_seq_len_cpu[valid_idx] - return b_q_seq_len, b_kv_seq_len, b_att_req_idx, b_att_seq_len - - -@pytest.mark.parametrize("batch_size", [1, 7, 256, 257, 777, 1025, 1537]) -def test_build_dynamic_spec_fa3_decode_params(batch_size): - hold_req_id = -1 - b_req_idx = torch.arange(batch_size, dtype=torch.int32, device="cuda") + 100 - b_seq_len = (torch.arange(batch_size, dtype=torch.int32, device="cuda") % 97) + 1 - b_mark_mtp_shared_group = torch.zeros((batch_size,), dtype=torch.int32, device="cuda") - - mark_positions = [0, 3, 17, 255, 256, 511, batch_size - 1] - for pos in mark_positions: - if 0 <= pos < batch_size: - b_mark_mtp_shared_group[pos] = pos % 5 + 1 - - actual = build_dynamic_spec_fa3_decode_params( - b_req_idx=b_req_idx, - b_seq_len=b_seq_len, - b_mark_mtp_shared_group=b_mark_mtp_shared_group, - att_batch_size=batch_size, - hold_req_id=hold_req_id, - ) - expected = _reference_dynamic_spec_fa3_decode_params( - b_req_idx=b_req_idx, - b_seq_len=b_seq_len, - b_mark_mtp_shared_group=b_mark_mtp_shared_group, - hold_req_id=hold_req_id, - ) - - for actual_tensor, expected_tensor in zip(actual, expected, strict=True): - assert torch.equal(actual_tensor.cpu(), expected_tensor) - - -def test_build_dynamic_spec_fa3_decode_params_all_padding(): - batch_size = 513 - hold_req_id = -1 - b_req_idx = torch.arange(batch_size, dtype=torch.int32, device="cuda") - b_seq_len = torch.arange(batch_size, dtype=torch.int32, device="cuda") + 1 - b_mark_mtp_shared_group = torch.zeros((batch_size,), dtype=torch.int32, device="cuda") - - actual = build_dynamic_spec_fa3_decode_params( - b_req_idx=b_req_idx, - b_seq_len=b_seq_len, - b_mark_mtp_shared_group=b_mark_mtp_shared_group, - att_batch_size=batch_size, - hold_req_id=hold_req_id, - ) - expected = _reference_dynamic_spec_fa3_decode_params( - b_req_idx=b_req_idx, - b_seq_len=b_seq_len, - b_mark_mtp_shared_group=b_mark_mtp_shared_group, - hold_req_id=hold_req_id, - ) - - for actual_tensor, expected_tensor in zip(actual, expected, strict=True): - assert torch.equal(actual_tensor.cpu(), expected_tensor) diff --git a/unit_tests/common/basemodel/triton_kernel/test_gen_decode_params.py b/unit_tests/common/basemodel/triton_kernel/test_gen_decode_params.py deleted file mode 100644 index 41bc217b94..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/test_gen_decode_params.py +++ /dev/null @@ -1,26 +0,0 @@ -import torch -import pytest -from lightllm.common.basemodel.triton_kernel.gen_decode_params import gen_decode_params - - -def test_gen_decode_params_basic(): - b_seq_len = torch.ones((9,), dtype=torch.int64, device="cuda") * 8192 - ( - b_q_seq_len, - b1_cu_q_seq_len, - b_kv_seq_len, - b1_cu_kv_seq_len, - position_ids, - ) = gen_decode_params(b_seq_len) - - true_b_q_seq_len = torch.ones_like(b_seq_len) - - assert torch.equal(b_q_seq_len, true_b_q_seq_len) - assert torch.equal(b1_cu_q_seq_len, torch.nn.functional.pad(torch.cumsum(true_b_q_seq_len, dim=0), (1, 0), value=0)) - assert torch.equal(b_kv_seq_len, b_seq_len) - assert torch.equal(b1_cu_kv_seq_len, torch.nn.functional.pad(torch.cumsum(b_seq_len, dim=0), (1, 0), value=0)) - assert torch.equal(position_ids, b_seq_len - 1) - - -if __name__ == "__main__": - pytest.main() diff --git a/unit_tests/common/basemodel/triton_kernel/test_gen_mtp_prefill_params.py b/unit_tests/common/basemodel/triton_kernel/test_gen_mtp_prefill_params.py deleted file mode 100644 index 2702ca9453..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/test_gen_mtp_prefill_params.py +++ /dev/null @@ -1,44 +0,0 @@ -import torch -import pytest -import numpy as np -from lightllm.utils.log_utils import init_logger -from lightllm.common.basemodel.triton_kernel.gen_mtp_prefill_params import gen_mtp_new_input_ids - - -def test_gen_mtp_new_input_ids_empty_batch(): - input_ids = torch.empty((0,), dtype=torch.int64, device="cuda") - b_next_token_ids = torch.empty((0,), dtype=torch.int64, device="cuda") - b_seq_len = torch.empty((0,), dtype=torch.int32, device="cuda") - - new_input_ids = gen_mtp_new_input_ids( - input_ids=input_ids, - b_next_token_ids=b_next_token_ids, - b_seq_len=b_seq_len, - ) - - assert new_input_ids.is_cuda - assert new_input_ids.dtype == torch.int64 - assert new_input_ids.shape == (0,) - - -def test_gen_mtp_new_input_ids_0(): - input_ids = torch.tensor([1, 2, 3, 4, 5, 6, 7, 8, 9]).long().cuda() - b_next_token_ids = torch.tensor([10, 11, 12]).long().cuda() - b_seq_len = torch.tensor([3, 3, 3]).int().cuda() - expected_output = torch.tensor([2, 3, 10, 5, 6, 11, 8, 9, 12]).long().cuda() - new_input_ids = gen_mtp_new_input_ids(input_ids, b_next_token_ids, b_seq_len) - assert torch.equal(new_input_ids, expected_output) - - -def test_gen_mtp_new_input_ids_1(): - input_ids = torch.tensor([1, 2, 3, 4, 5, 6]).long().cuda() - b_next_token_ids = torch.tensor([10, 11, 12]).long().cuda() - b_seq_len = torch.tensor([3, 3, 3]).int().cuda() - b_ready_cache_len = torch.tensor([1, 1, 1]).int().cuda() - expected_output = torch.tensor([2, 10, 4, 11, 6, 12]).long().cuda() - new_input_ids = gen_mtp_new_input_ids(input_ids, b_next_token_ids, b_seq_len, b_ready_cache_len=b_ready_cache_len) - assert torch.equal(new_input_ids, expected_output) - - -if __name__ == "__main__": - pytest.main() diff --git a/unit_tests/common/basemodel/triton_kernel/test_gen_nsa_ks_ke.py b/unit_tests/common/basemodel/triton_kernel/test_gen_nsa_ks_ke.py deleted file mode 100644 index ba168f4a2f..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/test_gen_nsa_ks_ke.py +++ /dev/null @@ -1,87 +0,0 @@ -import torch -import pytest -from lightllm.common.basemodel.triton_kernel.gen_nsa_ks_ke import gen_nsa_ks_ke, gen_same_req_mark - - -def test_gen_nsa_ks_ke_basic(): - """Test basic functionality of gen_nsa_ks_ke with simple inputs.""" - # Setup test data - b_seq_len = torch.tensor( - [ - 10, - ], - dtype=torch.int32, - device="cuda", - ) - b_q_seq_len = torch.tensor( - [ - 5, - ], - dtype=torch.int32, - device="cuda", - ) - b_req_idx = torch.tensor( - [ - 1, - ], - dtype=torch.int32, - device="cuda", - ) - q_token_num = b_q_seq_len.sum().item() - - req_to_token_index = torch.arange(0, 1000).cuda().view(100, -1) - ragged_mem_index = torch.empty_like(req_to_token_index.view(-1)) - - ks, ke, lengths = gen_nsa_ks_ke( - b_seq_len, b_q_seq_len, b_req_idx, req_to_token_index, q_token_num, ragged_mem_index - ) - - assert torch.equal(ks, torch.tensor([0, 0, 0, 0, 0], dtype=torch.int32, device="cuda")) - assert torch.equal(ke, torch.tensor([6, 7, 8, 9, 10], dtype=torch.int32, device="cuda")) - assert torch.equal(lengths, ke - ks) - assert torch.equal( - ragged_mem_index[0:10], torch.tensor([10, 11, 12, 13, 14, 15, 16, 17, 18, 19], dtype=torch.int32, device="cuda") - ) - return - - -def test_gen_nsa_ks_ke_batch(): - b_seq_len = torch.tensor([10, 11], dtype=torch.int32, device="cuda") - b_q_seq_len = torch.tensor([1, 1], dtype=torch.int32, device="cuda") - b_req_idx = torch.tensor([1, 1], dtype=torch.int32, device="cuda") - q_token_num = b_q_seq_len.sum().item() - - req_to_token_index = torch.arange(0, 1000).cuda().view(10, -1) - ragged_mem_index = torch.empty_like(req_to_token_index.view(-1)) - - ks, ke, lengths = gen_nsa_ks_ke( - b_seq_len, b_q_seq_len, b_req_idx, req_to_token_index, q_token_num, ragged_mem_index - ) - - assert torch.equal( - ks, - torch.tensor( - [ - 0, - 0, - ], - dtype=torch.int32, - device="cuda", - ), - ) - assert torch.equal(ke, torch.tensor([10, 11], dtype=torch.int32, device="cuda")) - assert torch.equal(lengths, ke - ks) - assert torch.equal(ragged_mem_index[0:11], torch.arange(100, 100 + 11).cuda()) - - -def test_gen_same_req_mark(): - b_req_idx = torch.tensor([0, 0, 1, 1, 1, 2], dtype=torch.int32, device="cuda") - expected_same_req_mark = torch.tensor([0, 2, 0, 0, 3, 1], dtype=torch.int32, device="cuda") - - same_req_mark = gen_same_req_mark(b_req_idx) - - assert torch.equal(same_req_mark, expected_same_req_mark) - - -if __name__ == "__main__": - pytest.main() diff --git a/unit_tests/common/basemodel/triton_kernel/test_gen_prefill_params.py b/unit_tests/common/basemodel/triton_kernel/test_gen_prefill_params.py deleted file mode 100644 index e981b0bb73..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/test_gen_prefill_params.py +++ /dev/null @@ -1,43 +0,0 @@ -import torch -import pytest -import numpy as np -from lightllm.utils.log_utils import init_logger -from lightllm.common.basemodel.triton_kernel.gen_prefill_params import gen_prefill_params - - -def test_gen_prefill_params_basic(): - - b_ready_cache_len = torch.ones((9,), dtype=torch.int64, device="cuda") - b_seq_len = torch.ones((9,), dtype=torch.int64, device="cuda") * 8192 - b_seq_len[0:5] = 4567 - - input_token_num = (b_seq_len - b_ready_cache_len).sum().item() - true_b_q_seq_len = b_seq_len - b_ready_cache_len - - ( - b_q_seq_len, - b1_cu_q_seq_len, - b_kv_seq_len, - b1_cu_kv_seq_len, - position_ids, - ) = gen_prefill_params(input_token_num, b_ready_cache_len, b_seq_len) - - assert torch.equal(b_q_seq_len, true_b_q_seq_len) - assert torch.equal(b1_cu_q_seq_len, torch.nn.functional.pad(torch.cumsum(true_b_q_seq_len, dim=0), (1, 0), value=0)) - assert torch.equal(b_kv_seq_len, b_seq_len) - assert torch.equal(b1_cu_kv_seq_len, torch.nn.functional.pad(torch.cumsum(b_seq_len, dim=0), (1, 0), value=0)) - - b_ready_cache_len_numpy = b_ready_cache_len.cpu().numpy() - b_seq_len_numpy = b_seq_len.cpu().numpy() - true_position_ids = torch.from_numpy( - np.concatenate( - [np.arange(b_ready_cache_len_numpy[i], b_seq_len_numpy[i]) for i in range(len(b_seq_len_numpy))], - axis=0, - ) - ).cuda() - - assert torch.equal(position_ids, true_position_ids) - - -if __name__ == "__main__": - pytest.main() diff --git a/unit_tests/common/basemodel/triton_kernel/test_gen_sampling_params.py b/unit_tests/common/basemodel/triton_kernel/test_gen_sampling_params.py deleted file mode 100644 index e9d0193279..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/test_gen_sampling_params.py +++ /dev/null @@ -1,86 +0,0 @@ -import torch -import time -import pytest -from lightllm.common.basemodel.triton_kernel.gen_sampling_params import token_id_counter, gen_sampling_params -from lightllm.common.basemodel.triton_kernel.gen_sampling_params import update_req_to_token_id_counter -from lightllm.utils.log_utils import init_logger - -logger = init_logger(__name__) - - -def test_token_id_counter(): - test_prompt_ids = torch.randint(0, 100, size=(7853,), device="cuda") - test_token_id_counter = torch.zeros((100,), dtype=torch.int32, device="cuda") - token_id_counter(prompt_ids=test_prompt_ids, out_token_id_counter=test_token_id_counter) - - true_out = torch.bincount(test_prompt_ids, minlength=test_token_id_counter.shape[0]) - assert torch.equal(test_token_id_counter, true_out) - - test_prompt_ids = torch.randint(0, 50 * 10000, size=(128 * 1024,), device="cuda") - test_token_id_counter = torch.zeros((50 * 10000,), device="cuda", dtype=torch.int32) - - start_event = torch.cuda.Event(enable_timing=True) - end_event = torch.cuda.Event(enable_timing=True) - start_event.record() - for _ in range(100): - token_id_counter(prompt_ids=test_prompt_ids, out_token_id_counter=test_token_id_counter) - end_event.record() - end_event.synchronize() - logger.info(f"test_token_id_count cost time: {start_event.elapsed_time(end_event)} ms") - - -def test_gen_sampling_params(): - - # Mocking ReqSamplingParamsManager for testing - class MockReqSamplingParamsManager: - def __init__(self, batch_size, vocab_size): - self.req_to_presence_penalty = torch.rand((batch_size,), dtype=torch.float32, device="cuda") - self.req_to_frequency_penalty = torch.rand((batch_size,), dtype=torch.float32, device="cuda") - self.req_to_repetition_penalty = torch.rand((batch_size,), dtype=torch.float32, device="cuda") - self.req_to_temperature = torch.rand((batch_size,), dtype=torch.float32, device="cuda") - self.req_to_exponential_decay_length_penalty = torch.rand((batch_size,), dtype=torch.float32, device="cuda") - - batch_size = 1083 - vocab_size = 100 - req_sampling_params_manager = MockReqSamplingParamsManager(batch_size, vocab_size) - b_req_idx = torch.arange(batch_size, device="cuda").flip(dims=[0]) - - ( - b_presence_penalty, - b_frequency_penalty, - b_repetition_penalty, - b_temperature, - b_exponential_decay_length_penalty, - ) = gen_sampling_params(b_req_idx, req_sampling_params_manager) - - assert torch.equal(b_presence_penalty, req_sampling_params_manager.req_to_presence_penalty[b_req_idx]) - assert torch.equal(b_frequency_penalty, req_sampling_params_manager.req_to_frequency_penalty[b_req_idx]) - assert torch.equal(b_repetition_penalty, req_sampling_params_manager.req_to_repetition_penalty[b_req_idx]) - assert torch.equal(b_temperature, req_sampling_params_manager.req_to_temperature[b_req_idx]) - assert torch.equal( - b_exponential_decay_length_penalty, - req_sampling_params_manager.req_to_exponential_decay_length_penalty[b_req_idx], - ) - - -def test_update_req_to_token_id_counter(): - b_req_idx = torch.tensor([0, 1, 3, 2], dtype=torch.int32, device="cuda") - next_token_ids = torch.tensor([0, 1, 1, 0], dtype=torch.int32, device="cuda") - req_to_out_token_id_counter = torch.zeros((4, 4), dtype=torch.int32, device="cuda") - update_req_to_token_id_counter(b_req_idx, next_token_ids, req_to_out_token_id_counter) - expected_output = torch.tensor( - [ - [1, 0, 0, 0], - [0, 1, 0, 0], - [1, 0, 0, 0], - [0, 1, 0, 0], - ], - dtype=torch.int32, - device="cuda", - ) - - assert torch.equal(req_to_out_token_id_counter, expected_output) - - -if __name__ == "__main__": - pytest.main() diff --git a/unit_tests/common/basemodel/triton_kernel/test_mtp_utils.py b/unit_tests/common/basemodel/triton_kernel/test_mtp_utils.py deleted file mode 100644 index 9ecfa008c4..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/test_mtp_utils.py +++ /dev/null @@ -1,107 +0,0 @@ -import pytest -import torch - -if not torch.cuda.is_available(): - pytest.skip("requires CUDA", allow_module_level=True) - -from lightllm.common.basemodel.triton_kernel import mtp_utils - - -@pytest.mark.parametrize( - "max_group_size,req_indexes,expected_markers", - [ - (2, [7, 7, 7, 11, 11, 20], [0, 2, 1, 0, 2, 1]), - (8, [7, 7, 7, 11, 11, 11, 20, 20, 20], [0, 0, 3, 0, 0, 3, 0, 0, 3]), - (3, [7, 7, 7, 7, 7, 7], [0, 0, 3, 0, 0, 3]), - (3, [7, 7, -1, -1, 11, 11], [0, 2, 1, 1, 0, 2]), - ], -) -def test_build_mtp_shared_group_markers(monkeypatch, max_group_size, req_indexes, expected_markers): - monkeypatch.setattr(mtp_utils, "get_diverse_max_batch_shared_group_size", lambda: max_group_size) - b_req_idx = torch.tensor(req_indexes, dtype=torch.int32, device="cuda") - - markers = mtp_utils.build_mtp_shared_group_markers(b_req_idx, hold_req_id=-1) - - assert markers.cpu().tolist() == expected_markers - - -@pytest.mark.parametrize("pinned", [False, True]) -def test_mtp_verify_scatter_and_start_locations(pinned): - req_to_next_token_ids = torch.tensor( - [[1, 2, -2, -1, -1], [1, 2, 0, -1, -1], [1, 3, 4, 4, 5]], - dtype=torch.int64, - device="cpu" if pinned else "cuda", - pin_memory=pinned, - ) - b_req_idx = torch.tensor([0, 0, 2, 2, 2], dtype=torch.int32, device="cuda") - b_mtp_index = torch.tensor([0, 1, 0, 1, 2], dtype=torch.int32, device="cuda") - b_req_mtp_start_loc = mtp_utils.gen_b_req_mtp_start_loc(b_mtp_index, num_reqs=2) - new_next_token_ids = torch.tensor([1, 4, 2, 4, 13], dtype=torch.int64, device="cuda") - draft_token_ids = torch.tensor( - [[2, 3], [8, 9]], - dtype=torch.int64, - device="cuda", - ) - req_to_next_token_scores = torch.full((3, 5), -1.0, dtype=torch.float32, device="cuda") - schedule_scores = torch.tensor( - [[0.9, 0.8], [0.7, 0.6]], - dtype=torch.float32, - device="cuda", - ) - - mtp_accept_len, accepted_index = mtp_utils.mtp_verify( - req_to_next_token_ids, b_req_mtp_start_loc, new_next_token_ids, b_req_idx - ) - assert mtp_accept_len.is_cuda and accepted_index.is_cuda - mtp_utils.mtp_scatter_next_token_ids( - req_to_next_token_ids=req_to_next_token_ids, - b_req_mtp_start_loc=b_req_mtp_start_loc, - target_next_token_ids=new_next_token_ids, - draft_token_ids=draft_token_ids, - b_req_idx=b_req_idx, - mtp_accept_len=mtp_accept_len, - req_to_next_token_scores=req_to_next_token_scores, - schedule_scores=schedule_scores, - ) - torch.cuda.synchronize() - - assert torch.equal(b_req_mtp_start_loc.cpu(), torch.tensor([0, 2], dtype=torch.int64)) - assert torch.equal(mtp_accept_len.cpu(), torch.tensor([1, 1], dtype=torch.int32)) - assert torch.equal(accepted_index.cpu(), torch.tensor([1, 0, 1, 0, 0], dtype=torch.int32)) - assert torch.equal( - req_to_next_token_ids.cpu(), - torch.tensor( - [[1, 2, 3, 1, 1], [1, 2, 0, -1, -1], [2, 8, 9, 1, 1]], - dtype=torch.int64, - ), - ) - torch.testing.assert_close( - req_to_next_token_scores.cpu(), - torch.tensor( - [[1.0, 0.9, 0.8, 0.0, 0.0], [-1.0, -1.0, -1.0, -1.0, -1.0], [1.0, 0.7, 0.6, 0.0, 0.0]], - dtype=torch.float32, - ), - ) - - -@pytest.mark.parametrize("pinned", [False, True]) -def test_mtp_scatter_handles_zero_draft_step(pinned): - req_to_next_token_ids = torch.full( - (1, 4), -1, dtype=torch.int64, device="cpu" if pinned else "cuda", pin_memory=pinned - ) - req_to_next_token_scores = torch.full((1, 4), -1.0, dtype=torch.float32, device="cuda") - - mtp_utils.mtp_scatter_next_token_ids( - req_to_next_token_ids=req_to_next_token_ids, - b_req_mtp_start_loc=torch.tensor([0], dtype=torch.int32, device="cuda"), - target_next_token_ids=torch.tensor([42], dtype=torch.int64, device="cuda"), - draft_token_ids=torch.empty((1, 0), dtype=torch.int64, device="cuda"), - b_req_idx=torch.tensor([0], dtype=torch.int32, device="cuda"), - mtp_accept_len=torch.tensor([1], dtype=torch.int32, device="cuda"), - req_to_next_token_scores=req_to_next_token_scores, - schedule_scores=torch.empty((1, 0), dtype=torch.float32, device="cuda"), - ) - torch.cuda.synchronize() - - assert req_to_next_token_ids.cpu().tolist() == [[42, 1, 1, 1]] - assert req_to_next_token_scores.cpu().tolist() == [[1.0, 0.0, 0.0, 0.0]] diff --git a/unit_tests/common/basemodel/triton_kernel/test_multimodal_emb.py b/unit_tests/common/basemodel/triton_kernel/test_multimodal_emb.py deleted file mode 100755 index ae3bcf4aed..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/test_multimodal_emb.py +++ /dev/null @@ -1,54 +0,0 @@ -import torch -import pytest -from lightllm.common.basemodel.triton_kernel.multimodal_emb import mark_multimodal_obj, multimodal_emb -from lightllm.utils.log_utils import init_logger - -logger = init_logger(__name__) - - -def test_mark_multimodal_obj(): - virtual_token_base = 2 ** 62 - obj_start_ids = torch.tensor( - [virtual_token_base, virtual_token_base + 3, virtual_token_base + 99], device="cuda", dtype=torch.int64 - ) - obj_token_lens = torch.tensor([1, 3, 2], device="cuda", dtype=torch.int64) - input_ids = torch.tensor([virtual_token_base, virtual_token_base + 6, 9, 333], device="cuda", dtype=torch.int64) - - mark_obj = mark_multimodal_obj( - obj_start_token_ids=obj_start_ids, obj_token_lens=obj_token_lens, input_ids=input_ids - ) - - assert torch.equal(mark_obj, torch.tensor([1, 0, 0], device="cuda")) - - -def test_multimodal_emb(): - S, D = 1024 * 10, 128 * 64 - vob_size = 3200 - image_size = 10 - image_token_size = 512 - virtual_token_base = 2 ** 62 - - text_weight = torch.randn((vob_size, D), device="cuda", dtype=torch.float16) - img_weight = torch.randn((image_size * image_token_size, 1, D), device="cuda", dtype=torch.float16) - img_token_lens = torch.full((image_size,), image_token_size, device="cuda", dtype=torch.long) - img_start_token_ids = ( - torch.arange(0, image_size * image_token_size, image_token_size, dtype=torch.int64, device="cuda") - + virtual_token_base - ) - img_start_locs = torch.arange(0, image_size * image_token_size, image_token_size).cuda().long() - - prompt_ids = torch.arange(0, S, 1).cuda().long() - prompt_ids[0 : image_size * image_token_size] = virtual_token_base + torch.arange( - 0, image_size * image_token_size, 1, dtype=torch.int64, device="cuda" - ) - - out = torch.zeros((S, D), dtype=torch.float16, device="cuda") - multimodal_emb( - out, prompt_ids, text_weight, img_weight, img_token_lens, img_start_token_ids, img_start_locs, 0, vob_size, 1 - ) - assert torch.allclose(out[0 : image_size * image_token_size, :].view(-1), img_weight.view(-1), atol=0.0001) - return - - -if __name__ == "__main__": - pytest.main() diff --git a/unit_tests/common/basemodel/triton_kernel/test_qk_rmsnorm_fused.py b/unit_tests/common/basemodel/triton_kernel/test_qk_rmsnorm_fused.py deleted file mode 100644 index 1e7d5080ca..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/test_qk_rmsnorm_fused.py +++ /dev/null @@ -1,64 +0,0 @@ -import pytest -import torch - -from lightllm.common.basemodel.triton_kernel.norm.qk_norm import ( - qk_rmsnorm_fused_forward, - qk_rmsnorm_forward, -) - - -def torch_qk_rmsnorm(q, w_q, eps=1e-6): - input_dtype = q.dtype - head_dim = w_q.shape[0] - q_fp32 = q.to(torch.float32) - q_fp32 = q_fp32.view(-1, head_dim) - variance = q_fp32.pow(2).mean(dim=-1, keepdim=True) - q_fp32 = q_fp32 * torch.rsqrt(variance + eps) - return (q_fp32 * w_q.to(input_dtype)).view_as(q) - - -def test_qk_rmsnorm_fused_matches_reference(): - """Compare fused QK RMSNorm with separate reference RMSNorm kernels.""" - if not torch.cuda.is_available(): - pytest.skip("CUDA is required for qk_rmsnorm_fused test") - - torch.manual_seed(0) - - # 模拟配置: Batch=2, Seq=128, Head_Dim=128 - # Q: 16 Heads, K: 4 Heads (GQA 场景) - B, S, D = 2, 128, 128 - H_Q = 16 - H_K = 4 - - q = torch.randn((B * S, H_Q * D), device="cuda", dtype=torch.bfloat16) - k = torch.randn((B * S, H_K * D), device="cuda", dtype=torch.bfloat16) - - w_q = torch.ones((D,), device="cuda", dtype=torch.bfloat16) - w_k = torch.ones((D,), device="cuda", dtype=torch.bfloat16) - - # 复制一份做对比（reference 会在新 tensor 上计算） - q_ref = q.clone() - k_ref = k.clone() - - # fused kernel in-place 计算 - q_out, k_out = qk_rmsnorm_fused_forward(q, k, w_q, w_k, eps=1e-6) - - # reference: 分别对 Q / K 做 RMSNorm - q_ref_out = qk_rmsnorm_forward(q_ref, w_q, eps=1e-6) - k_ref_out = qk_rmsnorm_forward(k_ref, w_k, eps=1e-6) - # q_ref_out = torch_qk_rmsnorm(q_ref, w_q, eps=1e-6) - # k_ref_out = torch_qk_rmsnorm(k_ref, w_k, eps=1e-6) - - # fused 是 in-place 的，返回的 q_out/k_out 应该与 q/k 引用一致 - assert q_out.data_ptr() == q.data_ptr() - assert k_out.data_ptr() == k.data_ptr() - - # 误差容忍度: 由于 bfloat16 计算，设定一个合理的 atol - q_max_diff = (q_out - q_ref_out).abs().max().item() - k_max_diff = (k_out - k_ref_out).abs().max().item() - - print(f"Q max diff: {q_max_diff}") - print(f"K max diff: {k_max_diff}") - - assert q_max_diff < 1e-5, f"Q max diff too large: {q_max_diff}" - assert k_max_diff < 1e-5, f"K max diff too large: {k_max_diff}" diff --git a/unit_tests/common/basemodel/triton_kernel/test_redundancy_topk_ids_repair.py b/unit_tests/common/basemodel/triton_kernel/test_redundancy_topk_ids_repair.py deleted file mode 100644 index 16131ef935..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/test_redundancy_topk_ids_repair.py +++ /dev/null @@ -1,151 +0,0 @@ -import torch -import pytest -from lightllm.common.basemodel.triton_kernel.redundancy_topk_ids_repair import redundancy_topk_ids_repair -from lightllm.common.basemodel.triton_kernel.redundancy_topk_ids_repair import expert_id_counter -from lightllm.utils.log_utils import init_logger - -logger = init_logger(__name__) - - -def test_redundancy_topk_ids_repair(): - ep_expert_num = 4 - global_rank = 0 - redundancy_expert_num = 1 - topk_ids = torch.tensor( - [ - [0, 1, 2, 3], - [7, 9, 10, 11], - [1, 3, 5, 7], - ], - dtype=torch.int64, - device="cuda", - ) - - redundancy_expert_ids = torch.tensor( - [ - 0, - ], - dtype=torch.int64, - device="cuda", - ) - - expert_id_counter = torch.zeros(12, dtype=torch.int64, device="cuda") - - redundancy_topk_ids_repair( - topk_ids=topk_ids, - redundancy_expert_ids=redundancy_expert_ids, - ep_expert_num=ep_expert_num, - global_rank=global_rank, - expert_counter=expert_id_counter, - enable_counter=True, - ) - - ans_topk_ids = torch.tensor( - [ - [0, 1, 2, 3], - [7, 9, 10, 11], - [1, 3, 5, 7], - ], - dtype=torch.int64, - device="cuda", - ) - ans_topk_ids = (ans_topk_ids // ep_expert_num) * redundancy_expert_num + ans_topk_ids - new_redundancy_expert_ids = (redundancy_expert_ids // ep_expert_num) * redundancy_expert_num + redundancy_expert_ids - ans_topk_ids[ans_topk_ids == new_redundancy_expert_ids[0]] = ( - (ep_expert_num + redundancy_expert_num) * global_rank + ep_expert_num + 0 - ) - - assert torch.equal(topk_ids, ans_topk_ids) - assert torch.equal( - expert_id_counter, torch.tensor([1, 2, 1, 2, 0, 1, 0, 2, 0, 1, 1, 1], dtype=torch.int64, device="cuda") - ) - - ep_expert_num = 4 - global_rank = 1 - redundancy_expert_num = 1 - topk_ids = torch.tensor( - [ - [0, 1, 2, 3], - [7, 9, 10, 11], - [1, 3, 5, 7], - ], - dtype=torch.int64, - device="cuda", - ) - - redundancy_expert_ids = torch.tensor( - [ - 5, - ], - dtype=torch.int64, - device="cuda", - ) - redundancy_topk_ids_repair( - topk_ids=topk_ids, - redundancy_expert_ids=redundancy_expert_ids, - ep_expert_num=ep_expert_num, - global_rank=global_rank, - ) - - ans_topk_ids = torch.tensor( - [ - [0, 1, 2, 3], - [7, 9, 10, 11], - [1, 3, 5, 7], - ], - dtype=torch.int64, - device="cuda", - ) - ans_topk_ids = (ans_topk_ids // ep_expert_num) * redundancy_expert_num + ans_topk_ids - new_redundancy_expert_ids = (redundancy_expert_ids // ep_expert_num) * redundancy_expert_num + redundancy_expert_ids - ans_topk_ids[ans_topk_ids == new_redundancy_expert_ids[0]] = ( - (ep_expert_num + redundancy_expert_num) * global_rank + ep_expert_num + 0 - ) - - assert torch.equal(topk_ids, ans_topk_ids) - - -def test_expert_id_counter(): - token_num = 256 - tok_ids = torch.randint( - low=0, - high=12, - size=(token_num, 8), - dtype=torch.int64, - device="cuda", - ) - expert_counter = torch.zeros(12, dtype=torch.int64, device="cuda") - expert_id_counter(topk_ids=tok_ids, expert_counter=expert_counter) - - ans_expert_counter = torch.zeros(12, dtype=torch.int64, device="cuda") - ids, counts = torch.unique(tok_ids.view(-1), return_counts=True) - ans_expert_counter[ids] = counts - - assert torch.equal(expert_counter, ans_expert_counter) - - # test speed - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - for _ in range(100): - tok_ids = torch.randint( - low=0, - high=12, - size=(token_num, 8), - dtype=torch.int64, - device="cuda", - ) - expert_counter = torch.zeros(12, dtype=torch.int64, device="cuda") - expert_id_counter(topk_ids=tok_ids, expert_counter=expert_counter) - graph.replay() - - start_event = torch.cuda.Event(enable_timing=True) - start_event.record() - graph.replay() - end_event = torch.cuda.Event(enable_timing=True) - end_event.record() - torch.cuda.synchronize() - logger.info(f"expert_id_counter time cost: {start_event.elapsed_time(end_event)} ms") - - -if __name__ == "__main__": - pytest.main() diff --git a/unit_tests/common/basemodel/triton_kernel/test_repack_kv_index.py b/unit_tests/common/basemodel/triton_kernel/test_repack_kv_index.py deleted file mode 100644 index 15a4447466..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/test_repack_kv_index.py +++ /dev/null @@ -1,50 +0,0 @@ -import torch -import pytest -from lightllm.utils.log_utils import init_logger -from lightllm.common.basemodel.triton_kernel.repack_kv_index import repack_kv_index - -logger = init_logger(__name__) - -seed = 42 -torch.manual_seed(seed) - -if torch.cuda.is_available(): - torch.cuda.manual_seed(seed) - torch.cuda.manual_seed_all(seed) - - -@pytest.mark.parametrize( - "batch, max_seq_len", - [(a, b) for a in [1, 16, 32, 128, 512] for b in [16, 32, 512, 2048]], -) -def test_repack_kv_index(batch, max_seq_len): - def repack_kv_ref(req_to_token_indexs, b_req_idx, b_seq_len, b_start_loc, output): - for b, sl, start in zip(b_req_idx, b_seq_len, b_start_loc): - output[start : start + sl] = req_to_token_indexs[b][:sl] - - BATCH, MAX_SEQ_LEN = batch, max_seq_len - rand_idx = torch.randperm(2 * MAX_SEQ_LEN * BATCH).cuda().int() - b_req_idx = torch.randperm(BATCH).cuda().int() - b_seq_len = torch.randint(1, MAX_SEQ_LEN, (BATCH,)).cuda().int() - req_to_token_indexs = torch.zeros((2 * BATCH, 2 * MAX_SEQ_LEN)).cuda().int() - b_start_loc = ( - torch.cat([torch.zeros([1], device=b_seq_len.device, dtype=b_seq_len.dtype), b_seq_len[0:-1].cumsum(0)]) - .cuda() - .int() - ) - - output = torch.zeros((b_seq_len.sum(),)).cuda().int() - ref = torch.zeros((b_seq_len.sum(),)).cuda().int() - for b, sl, start in zip(b_req_idx, b_seq_len, b_start_loc): - req_to_token_indexs[b][:sl] = rand_idx[start : start + sl] - - repack_kv_ref(req_to_token_indexs, b_req_idx, b_seq_len, b_start_loc, ref) - repack_kv_index( - req_to_token_indexs=req_to_token_indexs, - b_req_idx=b_req_idx, - b_token_len=b_seq_len, - b_page_start_loc=b_start_loc, - max_token_len=MAX_SEQ_LEN, - out_page_indices=output, - ) - assert torch.allclose(output.float(), ref.float()) diff --git a/unit_tests/common/basemodel/triton_kernel/test_routing_capture.py b/unit_tests/common/basemodel/triton_kernel/test_routing_capture.py deleted file mode 100644 index 80cc75ee33..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/test_routing_capture.py +++ /dev/null @@ -1,137 +0,0 @@ -import pytest -import torch - -from lightllm.common.basemodel.triton_kernel.routing_capture import scatter_routing_topk_to_cpu - - -def _skip_without_cuda(): - if not torch.cuda.is_available(): - pytest.skip("CUDA is required for Triton kernels.") - - -@pytest.mark.parametrize("dtype,dtype_id,max_value", [(torch.uint8, 1, 255), (torch.int16, 2, 1024)]) -def test_scatter_routing_topk_to_pinned_cpu(dtype, dtype_id, max_value): - _skip_without_cuda() - - num_tokens = 5 - num_moe_layers = 3 - topk = 4 - kv_cache_size = 32 - moe_layer_index = 1 - - topk_ids = torch.randint(0, max_value, (num_tokens, topk), dtype=torch.int64, device="cuda") - mem_indexes = torch.tensor([17, 3, 29, 8, 11], dtype=torch.int32, device="cuda") - routing_buffer = torch.zeros( - (kv_cache_size, num_moe_layers, topk), - dtype=dtype, - device="cpu", - pin_memory=True, - ) - routing_buffer_ptr = torch.tensor([routing_buffer.data_ptr()], dtype=torch.uint64, device="cuda") - - scatter_routing_topk_to_cpu( - topk_ids=topk_ids, - mem_indexes=mem_indexes, - routing_buffer_ptr=routing_buffer_ptr, - moe_layer_index=moe_layer_index, - num_moe_layers=num_moe_layers, - kv_cache_size=kv_cache_size, - topk=topk, - dtype_id=dtype_id, - ) - torch.cuda.synchronize() - - expected = torch.zeros_like(routing_buffer) - expected[mem_indexes.cpu().long(), moe_layer_index, :] = topk_ids.cpu().to(dtype) - assert torch.equal(routing_buffer, expected) - - -def test_scatter_routing_topk_respects_layer_index(): - _skip_without_cuda() - - num_tokens = 3 - num_moe_layers = 2 - topk = 2 - kv_cache_size = 16 - - topk_ids = torch.arange(num_tokens * topk, dtype=torch.int64, device="cuda").view(num_tokens, topk) - mem_indexes = torch.tensor([10, 4, 13], dtype=torch.int64, device="cuda") - routing_buffer = torch.zeros( - (kv_cache_size, num_moe_layers, topk), - dtype=torch.int16, - device="cpu", - pin_memory=True, - ) - routing_buffer_ptr = torch.tensor([routing_buffer.data_ptr()], dtype=torch.uint64, device="cuda") - - scatter_routing_topk_to_cpu( - topk_ids=topk_ids, - mem_indexes=mem_indexes, - routing_buffer_ptr=routing_buffer_ptr, - moe_layer_index=1, - num_moe_layers=num_moe_layers, - kv_cache_size=kv_cache_size, - topk=topk, - dtype_id=2, - ) - torch.cuda.synchronize() - - expected = torch.zeros_like(routing_buffer) - expected[mem_indexes.cpu(), 1, :] = topk_ids.cpu().to(torch.int16) - assert torch.equal(routing_buffer, expected) - - -def test_scatter_routing_topk_is_cuda_graph_capturable(): - _skip_without_cuda() - - num_tokens = 4 - num_moe_layers = 2 - topk = 3 - kv_cache_size = 16 - - topk_ids = torch.arange(num_tokens * topk, dtype=torch.int64, device="cuda").view(num_tokens, topk) - mem_indexes = torch.tensor([2, 4, 6, 8], dtype=torch.int32, device="cuda") - routing_buffer = torch.zeros( - (kv_cache_size, num_moe_layers, topk), - dtype=torch.uint8, - device="cpu", - pin_memory=True, - ) - routing_buffer_ptr = torch.tensor([routing_buffer.data_ptr()], dtype=torch.uint64, device="cuda") - - scatter_routing_topk_to_cpu( - topk_ids=topk_ids, - mem_indexes=mem_indexes, - routing_buffer_ptr=routing_buffer_ptr, - moe_layer_index=0, - num_moe_layers=num_moe_layers, - kv_cache_size=kv_cache_size, - topk=topk, - dtype_id=1, - ) - torch.cuda.synchronize() - routing_buffer.zero_() - - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - scatter_routing_topk_to_cpu( - topk_ids=topk_ids, - mem_indexes=mem_indexes, - routing_buffer_ptr=routing_buffer_ptr, - moe_layer_index=0, - num_moe_layers=num_moe_layers, - kv_cache_size=kv_cache_size, - topk=topk, - dtype_id=1, - ) - - graph.replay() - torch.cuda.synchronize() - - expected = torch.zeros_like(routing_buffer) - expected[mem_indexes.cpu(), 0, :] = topk_ids.cpu().to(torch.uint8) - assert torch.equal(routing_buffer, expected) - - -if __name__ == "__main__": - pytest.main([__file__]) diff --git a/unit_tests/common/basemodel/triton_kernel/test_runtime_scalar_specialization.py b/unit_tests/common/basemodel/triton_kernel/test_runtime_scalar_specialization.py deleted file mode 100644 index 713a1c20c4..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/test_runtime_scalar_specialization.py +++ /dev/null @@ -1,146 +0,0 @@ -from collections import Counter -from types import SimpleNamespace - -import pytest -import torch -from triton.runtime.jit import JITFunction - -from lightllm.common.basemodel.triton_kernel import gather_token_id as token_ops -from lightllm.common.basemodel.triton_kernel import gen_sampling_params as sampling_ops -from lightllm.common.basemodel.triton_kernel.post_process.apply_penalty_gpu_cache import ( - apply_penalty_gpu_cache, -) - - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") - - -def _observe(monkeypatch, kernels): - misses = Counter() - hashes = {kernel: set() for kernel in kernels} - original_compile = JITFunction._do_compile - - def compile_kernel(self, *args, **kwargs): - if self in hashes: - misses[self] += 1 - return original_compile(self, *args, **kwargs) - - monkeypatch.setattr(JITFunction, "_do_compile", compile_kernel) - for kernel in kernels: - original_run = kernel.run - - def run(*args, _kernel=kernel, _original=original_run, **kwargs): - compiled = _original(*args, **kwargs) - hashes[_kernel].add(compiled.hash) - return compiled - - monkeypatch.setattr(kernel, "run", run) - return misses, hashes - - -def _assert_reused(misses, hashes): - assert all(len(values) == 1 for values in hashes.values()) - assert all(misses[kernel] <= 1 for kernel in hashes) - - -@pytest.mark.parametrize("has_mask", [False, True]) -@pytest.mark.parametrize("width", [1, 3, 16, 17]) -def test_token_helpers_reuse_runtime_counts_with_fixed_strides(monkeypatch, has_mask, width): - # MTP width and vocabulary are fixed for each service. Preserve their stride - # specializations while varying the request count across integer classes. - kernels = [ - token_ops._fwd_kernel_scatter, - token_ops._fwd_kernel_gather, - token_ops._fwd_kernel_gather_prefill_decode_mixed, - sampling_ops._token_id_counter_update_kernel, - ] - misses, hashes = _observe(monkeypatch, kernels) - table = torch.zeros((128, width), dtype=torch.int64, device="cuda") - counters = torch.zeros((128, width), dtype=torch.int32, device="cuda") - for batch in [1, 3, 16, 17, 33, 65]: - requests = torch.arange(batch, dtype=torch.int32, device="cuda") - columns = requests % width - values = requests.long() + 100 - enabled = requests % 3 != 1 - mask = enabled if has_mask else None - table.zero_() - token_ops.scatter_token(values, table, requests, columns, mask) - actual = token_ops.gather_token(table, requests, columns) - expected = torch.where(enabled, values, 0) if has_mask else values - assert torch.equal(actual, expected) - - input_ids = torch.full((batch * 2,), -1, dtype=torch.int64, device="cuda") - is_decode = requests % 2 == 0 - starts = requests * 2 - expected_input = input_ids.clone() - expected_input[starts[is_decode]] = expected[is_decode] - token_ops.gather_token_prefill_decode_mixed(input_ids, table, requests, columns, is_decode, starts) - assert torch.equal(input_ids, expected_input) - - counters.zero_() - owners = requests % 7 - sampling_ops.update_req_to_token_id_counter(owners, columns, counters, mask) - included = enabled if has_mask else torch.ones_like(enabled) - flat = owners[included].long() * width + columns[included] - expected_counts = torch.bincount(flat, minlength=counters.numel()).reshape_as(counters) - assert torch.equal(counters, expected_counts) - _assert_reused(misses, hashes) - - -@pytest.mark.parametrize("vocab", [1, 3, 16, 17]) -def test_sampling_parameter_and_prompt_counts_reuse_compilation(monkeypatch, vocab): - kernels = [sampling_ops._gen_sampling_params_kernel, sampling_ops._token_id_counter_kernel] - misses, hashes = _observe(monkeypatch, kernels) - fields = [ - "presence_penalty", - "frequency_penalty", - "repetition_penalty", - "temperature", - "exponential_decay_length_penalty", - ] - manager = SimpleNamespace( - **{ - f"req_to_{field}": torch.arange(128, device="cuda", dtype=torch.float32) + index - for index, field in enumerate(fields) - } - ) - for count in [1, 3, 16, 17, 65]: - requests = torch.arange(count, device="cuda", dtype=torch.int32).flip(0) - actual = sampling_ops.gen_sampling_params(requests, manager) - for value, field in zip(actual, fields): - assert torch.equal(value, getattr(manager, f"req_to_{field}")[requests]) - tokens = torch.arange(count, device="cuda") % vocab - counters = torch.zeros(vocab, dtype=torch.int32, device="cuda") - sampling_ops.token_id_counter(tokens, counters) - assert torch.equal(counters, torch.bincount(tokens, minlength=vocab)) - _assert_reused(misses, hashes) - - -def test_penalties_match_reference_across_shapes(): - # Preserve the specialized stride/EOS variants selected by performance - # measurements; their outputs must still agree across shapes. - for batch, vocab, eos_count in [(1, 1, 0), (3, 3, 1), (16, 16, 3), (17, 17, 17), (33, 3, 3)]: - requests = torch.arange(batch, dtype=torch.int32, device="cuda") - counters = (torch.arange(128 * vocab, device="cuda") % 3).int().reshape(128, vocab) - manager = SimpleNamespace( - vocab_size=vocab, - req_to_out_token_id_counter=counters, - req_to_presence_penalty=torch.full((128,), 0.5, device="cuda"), - req_to_frequency_penalty=torch.full((128,), 0.25, device="cuda"), - req_to_repetition_penalty=torch.full((128,), 2.0, device="cuda"), - req_to_exponential_decay_length_penalty=torch.full((128,), 2.0, device="cuda"), - ) - logits = ((torch.arange(batch * vocab, device="cuda") % 9 - 4).float() / 4).reshape(batch, vocab) - counts = counters[requests] - repeated = torch.where(logits > 0, logits / 2, logits * 2) - expected = torch.where(counts > 0, repeated, logits) - counts * 0.25 - expected -= torch.where(counts > 0, 0.5, 0.0) - length = requests % 3 - mask_eos = requests % 4 == 2 - eos_ids = torch.arange(eos_count, device="cuda") % vocab - scales = torch.exp2(length.float()) - 1 - for eos in eos_ids.tolist(): - value = expected[:, eos] - expected[:, eos] = torch.where(mask_eos, -10000000.0, value + value.abs() * scales) - apply_penalty_gpu_cache(logits, requests, length, mask_eos, eos_ids, manager) - torch.testing.assert_close(logits, expected, atol=0, rtol=0) diff --git a/unit_tests/common/basemodel/triton_kernel/test_select_kv_index_from_req.py b/unit_tests/common/basemodel/triton_kernel/test_select_kv_index_from_req.py deleted file mode 100644 index 0dab062934..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/test_select_kv_index_from_req.py +++ /dev/null @@ -1,75 +0,0 @@ -import pytest -import torch -from types import SimpleNamespace - -from lightllm.common.basemodel.basemodel import TpPartBaseModel -from lightllm.common.basemodel.triton_kernel.copy_kv_index_to_req import ( - select_kv_index_from_req, - select_kv_index_from_req_prefill, -) - - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") - - -def test_select_kv_index_from_req_for_decode(): - req_to_token_indexs = torch.tensor( - [ - [10, 11, 12, 13, 14, 15, 16, 17], - [20, 21, 22, 23, 24, 25, 26, 27], - ], - dtype=torch.int32, - device="cuda", - ) - - mem_indexes = select_kv_index_from_req( - req_to_token_indexs=req_to_token_indexs, - b_req_idx=torch.tensor([1, 0], dtype=torch.int32, device="cuda"), - b_seq_len=torch.tensor([3, 6], dtype=torch.int32, device="cuda"), - ) - - assert mem_indexes.cpu().tolist() == [22, 15] - - -def test_select_kv_index_from_req_for_prefill(): - req_to_token_indexs = torch.tensor( - [ - [10, 11, 12, 13, 14, 15, 16, 17], - [20, 21, 22, 23, 24, 25, 26, 27], - ], - dtype=torch.int32, - device="cuda", - ) - - mem_indexes = select_kv_index_from_req_prefill( - req_to_token_indexs=req_to_token_indexs, - b_req_idx=torch.tensor([1, 0], dtype=torch.int32, device="cuda"), - b_seq_len=torch.tensor([4, 7], dtype=torch.int32, device="cuda"), - b_ready_cache_len=torch.tensor([2, 3], dtype=torch.int32, device="cuda"), - b_start_loc=torch.tensor([0, 2], dtype=torch.int32, device="cuda"), - max_q_seq_len=4, - token_num=6, - ) - - assert mem_indexes.cpu().tolist() == [22, 23, 13, 14, 15, 16] - - -def test_model_selects_reserved_indexes_when_execution_starts(): - model = TpPartBaseModel.__new__(TpPartBaseModel) - model.args = SimpleNamespace(page_size=4) - model.req_manager = SimpleNamespace( - req_to_token_indexs=torch.tensor( - [[10, 11, 12, 13], [20, 21, 22, 23]], - dtype=torch.int32, - device="cuda", - ) - ) - model_input = SimpleNamespace( - is_prefill=False, - b_req_idx=torch.tensor([1, 0], dtype=torch.int32, device="cuda"), - b_seq_len=torch.tensor([3, 4], dtype=torch.int32, device="cuda"), - ) - - mem_indexes = model._select_mem_indexes(model_input) - - assert mem_indexes.cpu().tolist() == [22, 13] diff --git a/unit_tests/common/basemodel/triton_kernel/test_sp_pad_kernel.py b/unit_tests/common/basemodel/triton_kernel/test_sp_pad_kernel.py deleted file mode 100644 index 772644c28f..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/test_sp_pad_kernel.py +++ /dev/null @@ -1,36 +0,0 @@ -import torch -import time -import pytest -from lightllm.common.basemodel.triton_kernel.sp_pad_copy import sp_pad_copy -from lightllm.utils.log_utils import init_logger - -logger = init_logger(__name__) - - -@pytest.mark.parametrize( - "token_num, hidden_dim, sp_world_size", - [ - (token_num, hidden_dim, sp_world_size) - for token_num in range(3, 6) - for hidden_dim in [257, 2048] - for sp_world_size in range(2, 5) - ], -) -def test_sp_pad_copy(token_num, hidden_dim, sp_world_size): - in_tensor = torch.randn((token_num, hidden_dim), dtype=torch.float16, device="cuda") - if token_num % sp_world_size: - # TPSP pads the model input before splitting it between ranks. - with pytest.raises(AssertionError, match="in_token_num % sp_world_size != 0"): - sp_pad_copy(in_tensor=in_tensor, sp_rank_id=0, sp_world_size=sp_world_size) - padding = in_tensor[: sp_world_size - token_num % sp_world_size] - in_tensor = torch.cat([in_tensor, padding]) - out_tensors = [ - sp_pad_copy(in_tensor=in_tensor, sp_rank_id=rank_id, sp_world_size=sp_world_size) - for rank_id in range(sp_world_size) - ] - out_tensor = torch.cat(out_tensors, dim=0) - assert torch.equal(in_tensor, out_tensor) - - -if __name__ == "__main__": - pytest.main() diff --git a/unit_tests/common/dsv4_test_utils.py b/unit_tests/common/dsv4_test_utils.py index b39478bd12..ed4175596a 100644 --- a/unit_tests/common/dsv4_test_utils.py +++ b/unit_tests/common/dsv4_test_utils.py @@ -245,10 +245,6 @@ def free(self, indices: torch.Tensor): 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) @@ -282,12 +278,6 @@ def alloc_paged_kv_move_buffer(self, count: int, page_size: int): ) return self.kv_move_buffer - def get_pd_kv_move_buffer_size(self): - 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_operator().get_page_bytes()) - def write_mem_to_page_kv_move_buffer( self, indices: torch.Tensor, diff --git a/unit_tests/common/fused_moe/test_activation_config.py b/unit_tests/common/fused_moe/test_activation_config.py deleted file mode 100644 index c59ab2712d..0000000000 --- a/unit_tests/common/fused_moe/test_activation_config.py +++ /dev/null @@ -1,132 +0,0 @@ -import dataclasses -import json - -import pytest -import torch - -from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.fused_moe_weight import FusedMoeWeight -from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.impl.marlin_impl import FuseMoeMarlin -from lightllm.common.quantization.no_quant import NoQuantization -from lightllm.server.core.objs.start_args_type import StartArgs -from lightllm.utils.envs_utils import get_env_start_args - - -@pytest.fixture(autouse=True) -def runtime(monkeypatch): - monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(dataclasses.asdict(StartArgs()))) - for name, value in { - "GLOBAL_RANK": 0, - "GLOBAL_WORLD_SIZE": 1, - "DP_WORLD_SIZE": 1, - "CURRENT_RANK_IN_DP": 0, - "CURRENT_RANK_IN_NODE": 0, - "CURRENT_DEVICE_ID": 0, - }.items(): - monkeypatch.setenv("LIGHTLLM_" + name, str(value)) - get_env_start_args.cache_clear() - yield - get_env_start_args.cache_clear() - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("activation", ["silu", "clamped_silu", "clamped_silu_add_one", "gelu"]) -def test_call_parameters_reach_expert_activation(monkeypatch, activation): - monkeypatch.setattr( - "lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul.ffn_use_tanh_approximate_gelu", - lambda: activation == "gelu", - ) - dim, count = 128, 4 - config = {"norm_topk_prob": True, "num_experts_per_tok": 2, "scoring_func": "softmax"} - kwargs = { - "clamped_silu": {"alpha": 1.0, "limit": 10.0, "clamp_up_add_one": False}, - "clamped_silu_add_one": {"alpha": 1.702, "limit": 7.0}, - }.get(activation, {}) - weight = FusedMoeWeight( - gate_proj_name="gate", - up_proj_name="up", - down_proj_name="down", - e_score_correction_bias_name="", - weight_prefix="experts", - n_routed_experts=count, - hidden_size=dim, - moe_intermediate_size=dim, - data_type=torch.bfloat16, - quant_method=NoQuantization(), - network_config=config, - ) - eye = torch.eye(dim, device="cuda", dtype=torch.bfloat16) - weights = {} - for i in range(count): - weights[f"experts.{i}.gate.weight"] = eye * (1 + i / 4) - weights[f"experts.{i}.up.weight"] = eye * (2 + i / 8) - weights[f"experts.{i}.down.weight"] = eye - weight.load_hf_weights(weights) - assert weight.verify_load() - x = torch.linspace(-25, 25, 3 * dim, device="cuda", dtype=torch.bfloat16).view(3, dim) - router = torch.tensor([[1, 3, 2, 0], [4, 3, 1, 2], [1, 2, 3, 4]], device="cuda", dtype=torch.float32) - top = router.topk(2, dim=-1) - probs = top.values.softmax(-1) - expected = torch.zeros_like(x, dtype=torch.float32) - for row in range(x.shape[0]): - for choice in range(2): - i = int(top.indices[row, choice]) - gate = (x[row] * (1 + i / 4)).float() - up = (x[row] * (2 + i / 8)).float() - if kwargs: - gate = gate.clamp(max=kwargs["limit"]) - up = up.clamp(-kwargs["limit"], kwargs["limit"]) - gate = gate * torch.sigmoid(kwargs["alpha"] * gate) - if kwargs.get("clamp_up_add_one", True): - up += 1 - elif activation == "gelu": - gate = torch.nn.functional.gelu(gate, approximate="tanh") - else: - gate = torch.nn.functional.silu(gate) - expert_out = (gate.bfloat16() * up.bfloat16()).bfloat16() - expected[row] += (expert_out.float() * probs[row, choice]).bfloat16().float() - if kwargs: - default_output = weight.experts(x.clone(), router, 2, True, False, 0, 0, False) - actual = weight.experts(x.clone(), router, 2, True, False, 0, 0, False, **kwargs) - actual_with_topk = weight.experts_with_topk( - x.clone(), - topk_weights=probs, - topk_ids=top.indices, - is_prefill=False, - **kwargs, - ) - torch.testing.assert_close(actual, expected.bfloat16(), atol=0.125, rtol=0.01) - torch.testing.assert_close(actual_with_topk, expected.bfloat16(), atol=0.125, rtol=0.01) - if kwargs: - # A clamped call must not change subsequent calls on the same weight. - actual_default = weight.experts(x.clone(), router, 2, True, False, 0, 0, False) - torch.testing.assert_close(actual_default, default_output, atol=0, rtol=0) - - -def test_marlin_rejects_clamp_at_call(monkeypatch): - backend = FuseMoeMarlin - monkeypatch.setattr(backend, "create_workspace", lambda self: None) - monkeypatch.setattr(backend, "_select_experts", lambda *args, **kwargs: (None, None)) - impl = backend( - n_routed_experts=4, - num_fused_shared_experts=0, - routed_scaling_factor=1.0, - quant_method=None, - ) - with pytest.raises(NotImplementedError, match="does not support clamped SwiGLU"): - impl( - input_tensor=None, - router_logits=None, - w13=None, - w2=None, - correction_bias=None, - scoring_func="softmax", - top_k=2, - renormalize=True, - use_grouped_topk=False, - topk_group=0, - num_expert_group=0, - is_prefill=False, - alpha=1.0, - limit=10.0, - clamp_up_add_one=False, - ) diff --git a/unit_tests/common/fused_moe/test_append_shared_expert_topk.py b/unit_tests/common/fused_moe/test_append_shared_expert_topk.py deleted file mode 100644 index 4a197e911f..0000000000 --- a/unit_tests/common/fused_moe/test_append_shared_expert_topk.py +++ /dev/null @@ -1,222 +0,0 @@ -from types import SimpleNamespace - -import pytest -import torch - -from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.fused_moe_weight import FusedMoeWeight -from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.impl.triton_impl import FuseMoeTriton -from lightllm.common.basemodel.triton_kernel.fused_moe.append_shared_expert_topk import ( - append_fused_shared_experts, -) - - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for Triton kernels") - - -def test_experts_with_topk_appends_the_fused_shared_expert(monkeypatch): - monkeypatch.setattr( - "lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.fused_moe_weight.get_moe_capture_callback", - lambda *args, **kwargs: None, - ) - seen = {} - - def run(**kwargs): - seen.update(kwargs) - return kwargs["input_tensor"] - - impl = SimpleNamespace( - num_fused_shared_experts=1, - n_routed_experts=4, - _fused_experts=run, - ) - impl._prepare_expert_execution = lambda **kwargs: FuseMoeTriton._prepare_expert_execution(impl, **kwargs) - weight = SimpleNamespace( - layer_num_=0, - w13=torch.ones(1, device="cuda"), - w2=torch.ones(1, device="cuda"), - fuse_moe_impl=impl, - ) - topk_ids = torch.tensor([[0, 2], [1, 3]], dtype=torch.int64, device="cuda") - topk_weights = torch.tensor([[0.2, 0.8], [0.4, 0.6]], device="cuda") - FusedMoeWeight.experts_with_topk( - weight, - torch.ones(2, 4, device="cuda"), - topk_weights, - topk_ids, - is_prefill=False, - ) - assert torch.equal(seen["topk_ids"], torch.tensor([[0, 2, 4], [1, 3, 4]], device="cuda")) - assert torch.allclose(seen["topk_weights"], torch.tensor([[0.2, 0.8, 1.0], [0.4, 0.6, 1.0]], device="cuda")) - - -def test_capture_observes_logical_ids_before_execution_mapping(): - logical_ids = torch.tensor([[0, 2]], dtype=torch.int64, device="cuda") - logical_weights = torch.tensor([[0.2, 0.8]], device="cuda") - captured = {} - executed = {} - impl = FuseMoeTriton.__new__(FuseMoeTriton) - impl.n_routed_experts = 4 - impl.num_fused_shared_experts = 1 - impl._select_experts = lambda **kwargs: (logical_weights, logical_ids) - - def fused(**kwargs): - executed.update(kwargs) - return kwargs["input_tensor"] - - impl._fused_experts = fused - hidden = torch.ones(1, 4, device="cuda") - impl( - input_tensor=hidden, - router_logits=hidden, - w13=None, - w2=None, - correction_bias=None, - scoring_func="softmax", - top_k=2, - renormalize=False, - use_grouped_topk=False, - topk_group=0, - num_expert_group=0, - is_prefill=False, - moe_capture_callback=lambda ids: captured.setdefault("ids", ids), - ) - - assert captured["ids"] is logical_ids - assert torch.equal(logical_ids, torch.tensor([[0, 2]], device="cuda")) - assert torch.equal(executed["topk_ids"], torch.tensor([[0, 2, 4]], device="cuda")) - - -def test_redundancy_repair_does_not_overwrite_logical_ids(monkeypatch): - from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.impl.deepgemm_impl import ( - FuseMoeDeepGEMM, - ) - - logical_ids = torch.tensor([[1, 3]], dtype=torch.int64, device="cuda") - seen = {} - - def repair(topk_ids, **kwargs): - seen["ptr"] = topk_ids.data_ptr() - topk_ids.fill_(-1) - - monkeypatch.setattr( - "lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.impl.deepgemm_impl.redundancy_topk_ids_repair", - repair, - ) - impl = FuseMoeDeepGEMM.__new__(FuseMoeDeepGEMM) - impl.redundancy_expert_num = 1 - impl.redundancy_expert_ids_tensor = None - impl.ep_n_routed_experts = 2 - impl.global_rank_ = 0 - impl.routed_expert_counter_tensor = None - impl.auto_update_redundancy_expert = False - impl.num_fused_shared_experts = 0 - impl.n_routed_experts = 4 - _, execution_ids = FuseMoeDeepGEMM._prepare_expert_execution( - impl, - torch.ones(1, 2, device="cuda"), - logical_ids, - is_prefill=False, - ) - assert seen["ptr"] != logical_ids.data_ptr() - assert torch.equal(logical_ids, torch.tensor([[1, 3]], device="cuda")) - assert torch.equal(execution_ids, torch.full_like(logical_ids, -1)) - - -def test_append_fused_shared_experts_without_gate(): - topk_ids = torch.tensor([[0, 2], [1, 3], [2, 0]], dtype=torch.int32, device="cuda") - topk_weights = torch.tensor([[0.2, 0.8], [0.4, 0.6], [0.7, 0.3]], dtype=torch.float32, device="cuda") - - out_weights, out_ids = append_fused_shared_experts( - topk_weights=topk_weights, - topk_ids=topk_ids, - shared_expert_start_id=4, - num_fused_shared_experts=1, - shared_expert_gate=None, - ) - - expect_ids = torch.tensor([[0, 2, 4], [1, 3, 4], [2, 0, 4]], dtype=torch.int32, device="cuda") - expect_weights = torch.tensor( - [[0.2, 0.8, 1.0], [0.4, 0.6, 1.0], [0.7, 0.3, 1.0]], dtype=torch.float32, device="cuda" - ) - assert torch.equal(out_ids, expect_ids) - assert torch.allclose(out_weights, expect_weights) - - -def test_append_fused_shared_experts_with_gate(): - topk_ids = torch.tensor([[0, 2], [1, 3], [2, 0]], dtype=torch.int32, device="cuda") - topk_weights = torch.tensor([[0.2, 0.8], [0.4, 0.6], [0.7, 0.3]], dtype=torch.float32, device="cuda") - shared_expert_gate = torch.tensor([[0.0], [2.0], [-2.0]], dtype=torch.float32, device="cuda") - - out_weights, out_ids = append_fused_shared_experts( - topk_weights=topk_weights, - topk_ids=topk_ids, - shared_expert_start_id=4, - num_fused_shared_experts=1, - shared_expert_gate=shared_expert_gate, - ) - - expect_ids = torch.tensor([[0, 2, 4], [1, 3, 4], [2, 0, 4]], dtype=torch.int32, device="cuda") - expect_weights = torch.cat([topk_weights, torch.sigmoid(shared_expert_gate)], dim=1) - assert torch.equal(out_ids, expect_ids) - assert torch.allclose(out_weights, expect_weights) - - -def test_append_fused_shared_experts_multiple_tokens_per_grid(): - token_num = 4097 - topk_ids = torch.stack( - [ - torch.arange(token_num, dtype=torch.int32, device="cuda") % 4, - (torch.arange(token_num, dtype=torch.int32, device="cuda") + 1) % 4, - ], - dim=1, - ) - topk_weights = torch.rand((token_num, 2), dtype=torch.float32, device="cuda") - shared_expert_gate = torch.randn((token_num, 1), dtype=torch.float32, device="cuda") - - out_weights, out_ids = append_fused_shared_experts( - topk_weights=topk_weights, - topk_ids=topk_ids, - shared_expert_start_id=4, - num_fused_shared_experts=1, - shared_expert_gate=shared_expert_gate, - ) - - expect_ids = torch.cat( - [topk_ids, torch.full((token_num, 1), 4, dtype=torch.int32, device="cuda")], - dim=1, - ) - expect_weights = torch.cat([topk_weights, torch.sigmoid(shared_expert_gate)], dim=1) - assert torch.equal(out_ids, expect_ids) - assert torch.allclose(out_weights, expect_weights) - - -def test_append_fused_shared_experts_multi_shared_gate(): - token_num = 7 - topk_ids = torch.stack( - [ - torch.arange(token_num, dtype=torch.int32, device="cuda") % 4, - (torch.arange(token_num, dtype=torch.int32, device="cuda") + 1) % 4, - ], - dim=1, - ) - topk_weights = torch.rand((token_num, 2), dtype=torch.float32, device="cuda") - shared_expert_gate = torch.randn((token_num, 2), dtype=torch.float32, device="cuda") - - out_weights, out_ids = append_fused_shared_experts( - topk_weights=topk_weights, - topk_ids=topk_ids, - shared_expert_start_id=4, - num_fused_shared_experts=2, - shared_expert_gate=shared_expert_gate, - ) - - expect_ids = torch.cat( - [ - topk_ids, - torch.tensor([[4, 5]], dtype=torch.int32, device="cuda").repeat(token_num, 1), - ], - dim=1, - ) - expect_weights = torch.cat([topk_weights, torch.sigmoid(shared_expert_gate)], dim=1) - assert torch.equal(out_ids, expect_ids) - assert torch.allclose(out_weights, expect_weights) diff --git a/unit_tests/common/fused_moe/test_deepep.py b/unit_tests/common/fused_moe/test_deepep.py deleted file mode 100644 index ecc92af1c0..0000000000 --- a/unit_tests/common/fused_moe/test_deepep.py +++ /dev/null @@ -1,311 +0,0 @@ -import pytest - -pytest.skip(reason="need special env, install deep_ep and deep_gemm", allow_module_level=True) - -import os -import torch -import torch.distributed as dist -import deep_ep -import numpy as np -from lightllm.common.fused_moe.grouped_fused_moe_ep import fused_experts_impl -from typing import Tuple -from lightllm.utils.log_utils import init_logger - -logger = init_logger(__name__) - -seed = 42 -torch.manual_seed(seed) - -if torch.cuda.is_available(): - torch.cuda.manual_seed(seed) - torch.cuda.manual_seed_all(seed) - - -def per_block_cast_to_fp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: - assert x.dim() == 2 - m, n = x.shape - from deep_gemm import ceil_div - - x_padded = torch.zeros((ceil_div(m, 128) * 128, ceil_div(n, 128) * 128), dtype=x.dtype, device=x.device) - x_padded[:m, :n] = x - x_view = x_padded.view(-1, 128, x_padded.size(1) // 128, 128) - x_amax = x_view.abs().float().amax(dim=(1, 3), keepdim=True).clamp(1e-4) - x_scaled = (x_view * (448.0 / x_amax)).to(torch.float8_e4m3fn) - return x_scaled.view_as(x_padded)[:m, :n].contiguous(), (x_amax / 448.0).view(x_view.size(0), x_view.size(2)) - - -def init_dist(local_rank: int, num_local_ranks: int): - # NOTES: you may rewrite this function with your own cluster settings - ip = os.getenv("MASTER_ADDR", "127.0.0.1") - port = int(os.getenv("MASTER_PORT", "8361")) - num_nodes = int(os.getenv("WORLD_SIZE", 1)) - node_rank = int(os.getenv("RANK", 0)) - assert (num_local_ranks < 8 and num_nodes == 1) or num_local_ranks == 8 - - dist.init_process_group( - backend="nccl", - init_method=f"tcp://{ip}:{port}", - world_size=num_nodes * num_local_ranks, - rank=node_rank * num_local_ranks + local_rank, - ) - torch.set_default_dtype(torch.bfloat16) - torch.set_default_device("cuda") - torch.cuda.set_device(local_rank) - - return dist.get_rank(), dist.get_world_size(), dist.new_group(list(range(num_local_ranks * num_nodes))) - - -def fused_experts_impl_ref( - x: torch.Tensor, # [M, K] - w1: torch.Tensor, # [group, N, K] - w2: torch.Tensor, # [group, K, N/2] - topk_weight: torch.Tensor, # [M, topk] - topk_ids: torch.Tensor, # [M, topk] - num_experts: int, -): - N = w1.shape[1] - ep_size = torch.distributed.get_world_size() - experts_per_rank = num_experts // ep_size - - cnts = topk_ids.new_zeros((topk_ids.shape[0], num_experts)) - cnts.scatter_(1, topk_ids, 1) - tokens_per_expert = cnts.sum(dim=0) - idxs = topk_ids.view(-1).argsort() - sorted_tokens = x[idxs // topk_ids.shape[1]] - sorted_tokens_shape = sorted_tokens.shape - - if ep_size > 1: - tokens_per_ep_rank = tokens_per_expert.view(ep_size, -1).sum(dim=1) - tokens_per_expert_group = tokens_per_expert.new_empty(tokens_per_expert.shape[0]) - dist.all_to_all_single(tokens_per_expert_group, tokens_per_expert) - output_splits = tokens_per_expert_group.view(ep_size, -1).sum(1).cpu().numpy().tolist() - gathered_tokens = sorted_tokens.new_empty( - tokens_per_expert_group.sum(dim=0).cpu().item(), sorted_tokens.shape[1] - ) - input_split_sizes = tokens_per_ep_rank.cpu().numpy().tolist() - dist.all_to_all( - list(gathered_tokens.split(output_splits)), - list(sorted_tokens.split(input_split_sizes)), - ) - tokens_per_expert_post_gather = tokens_per_expert_group.view(ep_size, experts_per_rank).sum(dim=0) - gatherd_idxs = np.zeros(shape=(gathered_tokens.shape[0],), dtype=np.int32) - s = 0 - for i, k in enumerate(tokens_per_expert_group.cpu().numpy()): - gatherd_idxs[s : s + k] = i % experts_per_rank - s += k - gatherd_idxs = gatherd_idxs.argsort() - sorted_tokens = gathered_tokens[gatherd_idxs] - tokens_per_expert = tokens_per_expert_post_gather - tokens_per_expert = tokens_per_expert.cpu().numpy() - - outputs = [] - start_idx = 0 - for i, num_tokens in enumerate(tokens_per_expert): - end_idx = start_idx + num_tokens - if num_tokens == 0: - continue - tokens_for_this_expert = sorted_tokens[start_idx:end_idx] - w1_out = torch.matmul(tokens_for_this_expert, w1[i, : N // 2, :].T) - w2_out = torch.matmul(tokens_for_this_expert, w1[i, N // 2 :, :].T) - tmp = torch.nn.functional.silu(w1_out) - tmp = tmp * w2_out - expert_out = torch.matmul(tmp, w2[i].T) - outputs.append(expert_out) - start_idx = end_idx - - outs = torch.cat(outputs, dim=0) if len(outputs) else sorted_tokens.new_empty(0) - - if ep_size > 1: - new_x = torch.empty_like(outs) - new_x[gatherd_idxs] = outs - gathered_tokens = new_x.new_empty(*sorted_tokens_shape) - dist.all_to_all( - list(gathered_tokens.split(input_split_sizes)), - list(new_x.split(output_splits)), - ) - outs = gathered_tokens - - new_x = torch.empty_like(outs) - new_x[idxs] = outs - final_out = ( - new_x.view(*topk_ids.shape, -1) - .type(topk_weight.dtype) - .mul_(topk_weight.unsqueeze(dim=-1)) - .sum(dim=1) - .type(new_x.dtype) - ) - - return final_out - - -def case1(local_rank: int, num_local_ranks: int): - # Init dist - rank, num_ranks, group = init_dist(local_rank, num_local_ranks) - - torch.manual_seed(rank) - - # Construct inputs - seqlen = 16 - hidden_states = torch.randn((seqlen, 7168), device="cuda", dtype=torch.bfloat16) - w1 = torch.randn((256 // num_local_ranks, 4096, 7168), device="cuda", dtype=torch.bfloat16) - w2 = torch.randn((256 // num_local_ranks, 7168, 2048), device="cuda", dtype=torch.bfloat16) - - w1_fp8 = torch.empty_like(w1, dtype=torch.float8_e4m3fn) - w2_fp8 = torch.empty_like(w2, dtype=torch.float8_e4m3fn) - w1_scale = torch.empty((256 // num_local_ranks, 4096 // 128, 7168 // 128), device="cuda", dtype=torch.float) - w2_scale = torch.empty((256 // num_local_ranks, 7168 // 128, 2048 // 128), device="cuda", dtype=torch.float) - - for i in range(256 // num_local_ranks): - w1_fp8[i], w1_scale[i] = per_block_cast_to_fp8(w1[i]) - w2_fp8[i], w2_scale[i] = per_block_cast_to_fp8(w2[i]) - - topk_weights = torch.randn((seqlen, 8), device="cuda", dtype=torch.float32) - topk_weights = torch.softmax(topk_weights, dim=-1) # 对每行进行softmax归一化 - topk_weights = torch.tensor(topk_weights, device="cuda", dtype=torch.float32) - topk_ids = torch.zeros((seqlen, 8), device="cuda", dtype=torch.int64) - for i in range(seqlen): - topk_ids[i] = torch.randperm(254, device="cuda")[:8] + 1 - - # Init buffer - test_ll_compatibility, num_rdma_bytes = True, 0 - num_max_dispatch_tokens_per_rank = 512 - if test_ll_compatibility: - ll_num_tokens, ll_hidden, ll_num_experts, _ = num_max_dispatch_tokens_per_rank, 7168, 256, 8 - num_rdma_bytes = deep_ep.Buffer.get_low_latency_rdma_size_hint( - ll_num_tokens, ll_hidden, num_ranks, ll_num_experts - ) - - buffer = deep_ep.Buffer( - group, - int(1e9), - num_rdma_bytes, - low_latency_mode=test_ll_compatibility, - num_qps_per_rank=(ll_num_experts // num_ranks if test_ll_compatibility else 1), - ) - - # Test normal - ref_output = fused_experts_impl_ref( - x=hidden_states, w1=w1, w2=w2, topk_weight=topk_weights, topk_ids=topk_ids, num_experts=256 - ) - - output = fused_experts_impl( - hidden_states=hidden_states, - w1=w1_fp8, - w2=w2_fp8, - topk_weights=topk_weights, - topk_idx=topk_ids, - num_experts=256, - buffer=buffer, - is_prefill=True, - use_fp8_w8a8=True, - use_fp8_all2all=True, - use_int8_w8a16=False, - w1_scale=w1_scale, - w2_scale=w2_scale, - previous_event=None, - ) - - # Test ll - if test_ll_compatibility: - buffer.clean_low_latency_buffer(ll_num_tokens, ll_hidden, ll_num_experts) - ll_output = fused_experts_impl( - hidden_states=hidden_states, - w1=w1_fp8, - w2=w2_fp8, - topk_weights=topk_weights, - topk_idx=topk_ids, - num_experts=256, - buffer=buffer, - is_prefill=False, - use_fp8_w8a8=True, - use_fp8_all2all=True, - use_int8_w8a16=False, - w1_scale=w1_scale, - w2_scale=w2_scale, - previous_event=None, - ) - - # Check - norm_sim = torch.nn.functional.cosine_similarity(ref_output, output).mean() - ll_sim = torch.nn.functional.cosine_similarity(ref_output, ll_output).mean() - - assert torch.allclose(norm_sim, torch.ones(1), atol=1e-2, rtol=0) - assert torch.allclose(ll_sim, torch.ones(1), atol=1e-2, rtol=0) - logger.info(f"deepep cosine {norm_sim}") - logger.info(f"deepep ll cosine {ll_sim}") - - # Profile - profile = os.getenv("PROFILE", False) - if profile: - with torch.profiler.profile( - activities=[ - torch.profiler.ProfilerActivity.CPU, - torch.profiler.ProfilerActivity.CUDA, - ], - record_shapes=True, - profile_memory=True, - with_stack=True, - ) as prof: - for i in range(10): - output = fused_experts_impl( - hidden_states=hidden_states, - w1=w1_fp8, - w2=w2_fp8, - topk_weights=topk_weights, - topk_idx=topk_ids, - num_experts=256, - buffer=buffer, - is_prefill=True, - use_fp8_w8a8=True, - use_fp8_all2all=True, - use_int8_w8a16=False, - w1_scale=w1_scale, - w2_scale=w2_scale, - previous_event=None, - ) - # prof.step() - if rank == 0: - print(prof.key_averages().table(sort_by="cuda_time_total")) - prof.export_chrome_trace("normal_trace.json") - - if test_ll_compatibility: - with torch.profiler.profile( - activities=[ - torch.profiler.ProfilerActivity.CPU, - torch.profiler.ProfilerActivity.CUDA, - ], - record_shapes=True, - profile_memory=True, - with_stack=True, - ) as prof: - for i in range(10): - ll_output = fused_experts_impl( - hidden_states=hidden_states, - w1=w1_fp8, - w2=w2_fp8, - topk_weights=topk_weights, - topk_idx=topk_ids, - num_experts=256, - buffer=buffer, - is_prefill=False, - use_fp8_w8a8=True, - use_fp8_all2all=True, - use_int8_w8a16=False, - w1_scale=w1_scale, - w2_scale=w2_scale, - previous_event=None, - ) - if rank == 0: - print(prof.key_averages().table(sort_by="cuda_time_total")) - prof.export_chrome_trace("ll_trace.json") - - dist.barrier() - - -def test_end2end(): - num_processes = 8 - torch.multiprocessing.spawn(case1, args=(num_processes,), nprocs=num_processes) - - -if __name__ == "__main__": - pytest.main() diff --git a/unit_tests/common/fused_moe/test_deepep_runtime_specialization.py b/unit_tests/common/fused_moe/test_deepep_runtime_specialization.py deleted file mode 100644 index 8967ded513..0000000000 --- a/unit_tests/common/fused_moe/test_deepep_runtime_specialization.py +++ /dev/null @@ -1,61 +0,0 @@ -import pytest -import torch - -from lightllm.common.basemodel.triton_kernel.fused_moe.deepep_expanded_layout_kernels import ( - _ep_compact_metadata_kernel, - _ep_gather_chunk_kernel, - ep_compact_metadata, - ep_gather_chunk, -) - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") - - -@pytest.mark.parametrize("aligned_chunks", [False, True]) -def test_ep_runtime_counts_reuse_compilation(monkeypatch, aligned_chunks): - torch.manual_seed(51) - # Production chunks are 128-aligned. Generic unaligned calls remain valid, - # but need not share their boundary specialization with other chunks. - chunks = [(0, 128), (128, 128), (32768, 128)] if aligned_chunks else [(0, 1), (1, 3), (3, 16)] - weight_count = max(start + count for start, count in chunks) + 1 - hashes = {"gather": {}, "compact": set()} - bounds = None - for name, kernel in [("gather", _ep_gather_chunk_kernel), ("compact", _ep_compact_metadata_kernel)]: - original = kernel.run - - def observe(*args, _name=name, _original=original, **kwargs): - compiled = _original(*args, **kwargs) - if _name == "gather": - hashes[_name].setdefault(bounds, set()).add(compiled.hash) - else: - hashes[_name].add(compiled.hash) - return compiled - - monkeypatch.setattr(kernel, "run", observe) - for tokens in [1, 3, 16, 17, 33, 129]: - metadata = torch.randint(-1, 32, (tokens, 10), dtype=torch.int32, device="cuda") - weights = torch.randn(weight_count, dtype=torch.float32, device="cuda") - for start, count in chunks: - bounds = (start, count) - metadata[:, 2:] = torch.randint(start, start + count, (tokens, 8), device="cuda", dtype=torch.int32) - metadata[:, 2] = start - 1 - metadata[:, 3] = start - metadata[:, 4] = start + count - 1 - metadata[:, 5] = start + count - chunk = torch.randn((count, 129), dtype=torch.bfloat16, device="cuda") - output = torch.zeros((tokens, 129), dtype=chunk.dtype, device="cuda") - slots = metadata[:, 2:].long() - valid = (slots >= start) & (slots < start + count) - selected = chunk[(slots - start).clamp(0, count - 1)].float() - reference = (selected * weights[slots.clamp(0, weight_count - 1), None] * valid[..., None]).sum(1) - ep_gather_chunk(chunk, start, weights, metadata, output) - torch.testing.assert_close(output.float(), reference, atol=0.04, rtol=0.01) - original_header = metadata[:, :2].clone() - ep_compact_metadata(metadata) - assert torch.equal(metadata[:, :2], original_header) - assert torch.equal(metadata[:, 2], torch.arange(tokens, device="cuda", dtype=torch.int32)) - assert (metadata[:, 3:] == -1).all() - assert len(hashes["compact"]) == 1 - assert all(len(values) == 1 for values in hashes["gather"].values()), hashes - if aligned_chunks: - assert len(set.union(*hashes["gather"].values())) == 1 diff --git a/unit_tests/common/fused_moe/test_deepgemm_activation.py b/unit_tests/common/fused_moe/test_deepgemm_activation.py deleted file mode 100644 index 528d9a1ade..0000000000 --- a/unit_tests/common/fused_moe/test_deepgemm_activation.py +++ /dev/null @@ -1,96 +0,0 @@ -import pytest -import torch - -from lightllm.common.basemodel.triton_kernel.fused_moe import grouped_fused_moe_ep as ep -from lightllm.common.quantization import Quantcfg - - -pytestmark = pytest.mark.skipif( - not torch.cuda.is_available() or not ep.HAS_DEEPGEMM or torch.cuda.get_device_capability()[0] < 9, - reason="SM90+, DeepEP and DeepGEMM required", -) - - -@pytest.mark.parametrize("is_prefill", [False, True]) -@pytest.mark.parametrize("use_ue8m0_scales", [False, True]) -@pytest.mark.parametrize("alpha,limit,add_one", [(None, None, True), (1.0, 10.0, False), (1.702, 7.0, True)]) -def test_grouped_gemm_activation_matches_reference(monkeypatch, is_prefill, use_ue8m0_scales, alpha, limit, add_one): - monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", "0") - monkeypatch.delenv("LIGHTLLM_USE_UE8M0_SCALES") - config = {"quant_method": "fp8", "weight_block_size": [128, 128]} - if use_ue8m0_scales: - config["scale_fmt"] = "ue8m0" - Quantcfg({"n_layer": 1, "quantization_config": config}) - # Identity projections isolate activation/FP8 behavior from GEMM rounding. - dim, padded = 128, 128 - x = torch.linspace(-24, 24, 5 * dim, device="cuda").view(5, dim).to(torch.float8_e4m3fn) - eye = torch.eye(dim, device="cuda", dtype=torch.bfloat16) - w1 = torch.cat([2 * eye, 3 * eye]).unsqueeze(0).repeat(2, 1, 1).to(torch.float8_e4m3fn) - w2 = eye.unsqueeze(0).repeat(2, 1, 1).to(torch.float8_e4m3fn) - w1_scale = torch.ones(2, 2, 1, device="cuda") - w2_scale = torch.ones(2, 1, 1, device="cuda") - counts = torch.tensor([3, 2], device="cuda", dtype=torch.int32) - activation_args = dict(alpha=alpha, limit=limit, clamp_up_add_one=add_one) - - if is_prefill: - recv = torch.zeros(2, padded, dim, device="cuda", dtype=torch.float8_e4m3fn) - recv[0, :3] = x[:3] - recv[1, :2] = x[3:] - metadata = torch.zeros(5, 3, device="cuda", dtype=torch.int32) - metadata[:, 2] = torch.tensor([0, 1, 2, padded, padded + 1], device="cuda") - actual = ep.chunked_expanded_moe_forward( - num_recv_tokens_per_expert_list=[padded, padded], - num_unaligned_recv_tokens_per_expert=counts, - recv_x=(recv.view(2 * padded, dim), torch.ones(1, 2 * padded, device="cuda").T), - recv_topk_weights=torch.ones(2 * padded, device="cuda"), - recv_src_metadata=metadata, - w1=w1, - w1_scale=w1_scale, - w2=w2, - w2_scale=w2_scale, - block_size_k=128, - hidden_dtype=torch.bfloat16, - **activation_args, - ) - else: - alignment = ep.get_mk_alignment_for_contiguous_layout() - starts = [] - psum_values = [] - cursor = 0 - for count in counts.tolist(): - if starts: - cursor = (cursor + alignment - 1) // alignment * alignment - starts.append(cursor) - cursor += count - psum_values.append(cursor) - psum = torch.tensor(psum_values, device="cuda", dtype=torch.int32) - max_rows = (int(psum[-1]) + alignment - 1) // alignment * alignment - recv = torch.zeros(max_rows, dim, device="cuda", dtype=torch.float8_e4m3fn) - recv[starts[0] : starts[0] + 3] = x[:3] - recv[starts[1] : starts[1] + 2] = x[3:] - output = ep.decode_masked_group_gemm( - (recv, torch.ones(max_rows, 1, device="cuda")), - psum, - alignment, - torch.bfloat16, - w1, - w1_scale, - w2, - w2_scale, - expected_m=3, - **activation_args, - ) - actual = torch.cat([output[starts[0] : starts[0] + 3], output[starts[1] : starts[1] + 2]]) - - gate = (x.float() * 2).bfloat16().float() - up = (x.float() * 3).bfloat16().float() - if limit is not None: - gate = gate.clamp(max=limit) - up = up.clamp(-limit, limit) + int(add_one) - gate = (gate * torch.sigmoid(gate * (alpha or 1.0))).bfloat16().float() - activation = (gate * up).bfloat16().float() - scales = activation.abs().amax(-1, keepdim=True).clamp(min=1e-10) / 448 - if use_ue8m0_scales: - scales = torch.exp2(torch.ceil(torch.log2(scales))) - expected = ((activation / scales).to(torch.float8_e4m3fn).float() * scales).bfloat16() - torch.testing.assert_close(actual, expected, atol=0.05, rtol=0.02) diff --git a/unit_tests/common/fused_moe/test_ep_reduce_decode_output.py b/unit_tests/common/fused_moe/test_ep_reduce_decode_output.py deleted file mode 100644 index 8b2805e3a3..0000000000 --- a/unit_tests/common/fused_moe/test_ep_reduce_decode_output.py +++ /dev/null @@ -1,36 +0,0 @@ -import pytest -import torch - -from lightllm.common.basemodel.triton_kernel.fused_moe.deepep_expanded_layout_kernels import ep_reduce_decode_output - - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") - - -def test_ep_reduce_decode_output_weights_local_rows_and_compacts_metadata(): - hidden = 8 - expert_output = torch.tensor( - [[1.0] * hidden, [2.0] * hidden, [4.0] * hidden], - device="cuda", - dtype=torch.bfloat16, - ) - route_weights = torch.tensor([0.5, 0.25, 0.0], device="cuda", dtype=torch.float32) - recv_src_metadata = torch.tensor( - [[7, 9, 0, 1], [3, 4, -1, -1]], - device="cuda", - dtype=torch.int32, - ) - num_valid_recv_tokens = torch.tensor([1], device="cuda", dtype=torch.int32) - - dense_output, compact_metadata = ep_reduce_decode_output( - expert_output=expert_output, - route_weights=route_weights, - recv_src_metadata=recv_src_metadata, - num_valid_recv_tokens=num_valid_recv_tokens, - ) - - expected = expert_output[0].float() * 0.5 + expert_output[1].float() * 0.25 - torch.testing.assert_close(dense_output[0].float(), expected) - assert compact_metadata[0, :2].tolist() == [7, 9] - assert compact_metadata[0, 2].item() == 0 - assert compact_metadata[0, 3].item() == -1 diff --git a/unit_tests/common/fused_moe/test_grouped_fused_moe.py b/unit_tests/common/fused_moe/test_grouped_fused_moe.py deleted file mode 100644 index 4ecb7ed81c..0000000000 --- a/unit_tests/common/fused_moe/test_grouped_fused_moe.py +++ /dev/null @@ -1,331 +0,0 @@ -import torch -import time -import pytest -import triton -from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe import ( - _moe_align_fused_atomic_token, - fused_experts_impl, - moe_align, - moe_align_fused, - moe_align1, - moe_align2, - grouped_matmul, -) -from lightllm.utils.log_utils import init_logger - -logger = init_logger(__name__) - -seed = 42 -torch.manual_seed(seed) - -if torch.cuda.is_available(): - torch.cuda.manual_seed(seed) - torch.cuda.manual_seed_all(seed) - - -def test_moe_align(): - expert_num = 5 - token_num = 3 - topk = 3 - topk_ids = torch.tensor([[0, 1, 2], [0, 3, 1], [3, 1, 4]], dtype=torch.int32, device="cuda") - out = torch.zeros((expert_num, token_num * topk), dtype=torch.int32, device="cuda") - out.fill_(0) - moe_align(topk_ids, out) - true = torch.tensor( - [ - [1, 0, 0, 1, 0, 0, 0, 0, 0], - [0, 1, 0, 0, 0, 1, 0, 1, 0], - [0, 0, 1, 0, 0, 0, 0, 0, 0], - [0, 0, 0, 0, 1, 0, 1, 0, 0], - [0, 0, 0, 0, 0, 0, 0, 0, 1], - ], - dtype=torch.int32, - device="cuda", - ) - assert torch.equal(out, true) - - -def test_moe_align1(): - experts_info = torch.tensor( - [[1, 0, 0, 0], [0, 1, 1, 0], [0, 0, 0, 1], [0, 0, 0, 0]], - dtype=torch.int32, - device="cuda", - ) - topk_weights = torch.tensor([[0.3, 0.7], [0.2, 0.8]], dtype=torch.float32, device="cuda") - experts_token_num = torch.zeros((4,), dtype=torch.int32, device="cuda") - experts_weights = torch.zeros(experts_info.shape, dtype=torch.float32, device="cuda") - - moe_align1(experts_info, topk_weights, experts_weights, experts_token_num, 2) - - true_experts_token_num = torch.tensor([1, 2, 1, 0], device="cuda", dtype=torch.int32) - true_experts_info = torch.tensor( - [[0, 0, 0, 0], [1, 2, 1, 0], [3, 0, 0, 1], [0, 0, 0, 0]], device="cuda:0", dtype=torch.int32 - ) - true_experts_weights = torch.tensor( - [ - [0.3000, 0.0000, 0.0000, 0.0000], - [0.7000, 0.2000, 0.0000, 0.0000], - [0.8000, 0.0000, 0.0000, 0.0000], - [0.0000, 0.0000, 0.0000, 0.0000], - ], - device="cuda", - dtype=torch.float32, - ) - - assert torch.allclose(true_experts_weights, experts_weights) - assert torch.equal(experts_token_num, true_experts_token_num) - assert torch.equal(experts_info, true_experts_info) - - -def _check_moe_align_fused(topk_ids, topk_weights, expert_num, ordered=True): - expert_to_token_index = torch.empty((expert_num, topk_ids.numel()), dtype=torch.int32, device="cuda") - expert_to_weight = torch.empty((expert_num, topk_ids.numel()), dtype=torch.float32, device="cuda") - expert_token_num = torch.empty((expert_num,), dtype=torch.int32, device="cuda") - - moe_align_fused( - expert_to_token_index, - expert_to_weight, - expert_token_num, - topk_ids, - topk_weights, - ) - torch.cuda.synchronize() - - flat_topk_ids = topk_ids.flatten() - flat_topk_weights = topk_weights.flatten() - expected_token_num = torch.bincount(flat_topk_ids, minlength=expert_num).to(torch.int32) - assert torch.equal(expert_token_num, expected_token_num) - - for expert_id, token_num in enumerate(expected_token_num.tolist()): - expected_index = torch.nonzero(flat_topk_ids == expert_id, as_tuple=False).flatten() - expected_weight = flat_topk_weights[expected_index] - expected_index = expected_index.to(torch.int32) - token_index = expert_to_token_index[expert_id, :token_num] - token_weight = expert_to_weight[expert_id, :token_num] - - if not ordered: - order = torch.argsort(token_index) - token_index = token_index[order] - token_weight = token_weight[order] - - assert torch.equal(token_index, expected_index) - assert torch.allclose(token_weight, expected_weight) - - -def test_moe_align_fused_small_token(): - expert_num = 5 - small_topk_ids = torch.tensor([[0, 1, 2], [0, 3, 1], [3, 1, 4]], dtype=torch.int32, device="cuda") - small_topk_weights = torch.tensor( - [[0.3, 0.7, 0.1], [0.2, 0.8, 0.4], [0.5, 0.6, 0.9]], dtype=torch.float32, device="cuda" - ) - _check_moe_align_fused(small_topk_ids, small_topk_weights, expert_num) - - small_many_topk_ids = torch.arange(128 * 17, dtype=torch.int32, device="cuda").reshape(128, 17) % expert_num - small_many_topk_weights = torch.arange(small_many_topk_ids.numel(), dtype=torch.float32, device="cuda").reshape( - 128, 17 - ) - _check_moe_align_fused(small_many_topk_ids, small_many_topk_weights, expert_num) - - -def test_moe_align_fused_large_token(): - expert_num = 5 - - base_topk_ids = torch.tensor([[0, 1, 2], [0, 3, 1], [3, 1, 4], [2, 0, 4]], dtype=torch.int32, device="cuda") - large_topk_ids = base_topk_ids.repeat(33, 1)[:129].contiguous() - large_topk_weights = torch.arange(large_topk_ids.numel(), dtype=torch.float32, device="cuda").reshape(129, 3) - _check_moe_align_fused(large_topk_ids, large_topk_weights, expert_num, ordered=False) - - medium_topk_ids = base_topk_ids.repeat(1024, 1).contiguous() - medium_topk_weights = torch.arange(medium_topk_ids.numel(), dtype=torch.float32, device="cuda").reshape(4096, 3) - _check_moe_align_fused(medium_topk_ids, medium_topk_weights, expert_num, ordered=False) - - shared_expert_num = 257 - shared_routing = torch.arange(512 * 7, dtype=torch.int32, device="cuda").reshape(512, 7) % 256 - shared_last = torch.full((512, 1), 256, dtype=torch.int32, device="cuda") - shared_topk_ids = torch.cat([shared_routing, shared_last], dim=1).contiguous() - shared_topk_weights = torch.arange(shared_topk_ids.numel(), dtype=torch.float32, device="cuda").reshape(512, 8) - _check_moe_align_fused(shared_topk_ids, shared_topk_weights, shared_expert_num, ordered=False) - - large_atomic_topk_ids = base_topk_ids.repeat(1281, 1)[:5121].contiguous() - large_atomic_topk_weights = torch.arange(large_atomic_topk_ids.numel(), dtype=torch.float32, device="cuda").reshape( - 5121, 3 - ) - _check_moe_align_fused(large_atomic_topk_ids, large_atomic_topk_weights, expert_num, ordered=False) - - sparse_expert_num = 257 - sparse_topk_ids = base_topk_ids.repeat(1281, 1)[:5121].contiguous() - sparse_topk_weights = torch.arange(sparse_topk_ids.numel(), dtype=torch.float32, device="cuda").reshape(5121, 3) - _check_moe_align_fused(sparse_topk_ids, sparse_topk_weights, sparse_expert_num, ordered=False) - - -def test_moe_align_fused_large_token_unordered(): - expert_num = 257 - topk_ids = torch.arange(5121 * 8, dtype=torch.int32, device="cuda").reshape(5121, 8) % expert_num - topk_weights = torch.arange(topk_ids.numel(), dtype=torch.float32, device="cuda").reshape(5121, 8) - _check_moe_align_fused(topk_ids, topk_weights, expert_num, ordered=False) - - -def test_moe_align_fused_atomic_token_unordered(): - expert_num = 9 - topk_ids = torch.arange(257 * 4, dtype=torch.int32, device="cuda").reshape(257, 4) % expert_num - topk_weights = torch.arange(topk_ids.numel(), dtype=torch.float32, device="cuda").reshape(257, 4) - expert_to_token_index = torch.empty((expert_num, topk_ids.numel()), dtype=torch.int32, device="cuda") - expert_to_weight = torch.empty((expert_num, topk_ids.numel()), dtype=torch.float32, device="cuda") - expert_token_num = torch.empty((expert_num,), dtype=torch.int32, device="cuda") - - _moe_align_fused_atomic_token( - expert_to_token_index, - expert_to_weight, - expert_token_num, - topk_ids, - topk_weights, - ) - torch.cuda.synchronize() - - flat_topk_ids = topk_ids.flatten() - flat_topk_weights = topk_weights.flatten() - expected_token_num = torch.bincount(flat_topk_ids, minlength=expert_num).to(torch.int32) - assert torch.equal(expert_token_num, expected_token_num) - - for expert_id, token_num in enumerate(expected_token_num.tolist()): - expected_index = torch.nonzero(flat_topk_ids == expert_id, as_tuple=False).flatten().to(torch.int32) - expected_weight = flat_topk_weights[expected_index] - token_index = expert_to_token_index[expert_id, :token_num] - token_weight = expert_to_weight[expert_id, :token_num] - order = torch.argsort(token_index) - assert torch.equal(token_index[order], expected_index) - assert torch.allclose(token_weight[order], expected_weight) - - -def test_fused_experts_atomic_align_path_is_deterministic(): - token_num = 129 - expert_num = 9 - hidden_size = 64 - intermediate_size = 128 - topk = 4 - hidden_states = torch.randn((token_num, hidden_size), dtype=torch.bfloat16, device="cuda") / 10 - w1 = torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.bfloat16, device="cuda") / 10 - w2 = torch.randn((expert_num, hidden_size, intermediate_size // 2), dtype=torch.bfloat16, device="cuda") / 10 - topk_ids = torch.arange(token_num * topk, dtype=torch.int32, device="cuda").reshape(token_num, topk) % expert_num - topk_weights = torch.softmax(torch.randn((token_num, topk), dtype=torch.float32, device="cuda"), dim=-1) - - out_0 = fused_experts_impl(hidden_states, w1, w2, topk_weights, topk_ids) - out_1 = fused_experts_impl(hidden_states, w1, w2, topk_weights, topk_ids) - torch.cuda.synchronize() - - assert torch.equal(out_0, out_1) - - -def test_moe_align2(): - - experts_token_num = torch.zeros((4,), dtype=torch.int32, device="cuda") - experts_token_num[0] = 8 - experts_token_num[1] = 0 - experts_token_num[2] = 60 - experts_token_num[3] = 16 - - token_num_mul_topk_num = 100 - block_m = 16 - mblocks_to_tuple_info = moe_align2(token_num_mul_topk_num, experts_token_num, block_m=block_m) - expected_expert_ids = torch.tensor([0, 2, 2, 2, 2, 3, -1, -1, -1, -1], device="cuda", dtype=torch.int32) - valid_blocks = expected_expert_ids != -1 - - max_num_tokens_padded = token_num_mul_topk_num + experts_token_num.shape[0] * (block_m - 1) - expected_max_num_m_blocks = min(token_num_mul_topk_num, triton.cdiv(max_num_tokens_padded, block_m)) - assert mblocks_to_tuple_info.shape[0] == expected_max_num_m_blocks - assert torch.equal(mblocks_to_tuple_info[:, 0], expected_expert_ids) - assert torch.equal( - mblocks_to_tuple_info[valid_blocks, 1], - torch.tensor([0, 0, 1, 2, 3, 0], device="cuda", dtype=torch.int32), - ) - - # When expert_num is large relative to token_num, the padded cdiv upper bound - # exceeds token_num; allocation must be capped by token_num_mul_topk_num. - token_num_mul_topk_num = 8 - block_m = 16 - experts_token_num = torch.zeros((10,), dtype=torch.int32, device="cuda") - experts_token_num[0] = 8 - max_num_tokens_padded = token_num_mul_topk_num + experts_token_num.shape[0] * (block_m - 1) - padded_cdiv = triton.cdiv(max_num_tokens_padded, block_m) - assert padded_cdiv > token_num_mul_topk_num - - mblocks_to_tuple_info = moe_align2(token_num_mul_topk_num, experts_token_num, block_m=block_m) - assert mblocks_to_tuple_info.shape[0] == token_num_mul_topk_num - assert torch.equal( - mblocks_to_tuple_info[:1, 0], - torch.tensor([0], device="cuda", dtype=torch.int32), - ) - assert torch.all(mblocks_to_tuple_info[1:, 0] == -1) - - -def test_grouped_matmul(): - test_dtype = torch.bfloat16 - token_inputs = torch.randn((10, 512), dtype=test_dtype, device="cuda") / 10 - experts_token_num = torch.tensor([1, 9], dtype=torch.int32, device="cuda") - experts_to_token_index = torch.tensor( - [ - [0, 0, 0, 0, 0, 0, 0, 0, 0, 0], - [1, 2, 3, 4, 5, 6, 7, 8, 9, 0], - ], - dtype=torch.int32, - device="cuda", - ) - experts_to_weights = torch.tensor( - [ - [0.5, 0, 0, 0, 0, 0, 0, 0, 0, 0], - [0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0], - ], - dtype=torch.float32, - device="cuda", - ) - expert_weights = torch.randn((2, 1024, 512), dtype=test_dtype, device="cuda") / 10 - topk_num = 1 - out = torch.empty((10, 1024), dtype=test_dtype, device="cuda") - # warm up - grouped_matmul( - 10 * 1, - token_inputs, - None, - experts_token_num, - experts_to_token_index, - experts_to_weights, - expert_weights, - None, - topk_num, - out, - mul_routed_weight=True, - use_fp8_w8a8=False, - ) - torch.cuda.synchronize() - start = time.time() - for _ in range(100): - grouped_matmul( - 10 * 1, - token_inputs, - None, - experts_token_num, - experts_to_token_index, - experts_to_weights, - expert_weights, - None, - topk_num, - out, - mul_routed_weight=True, - use_fp8_w8a8=False, - ) - torch.cuda.synchronize() - logger.info(f"grouped_matmul test cost time: {time.time() - start} s") - - ans_list = [] - ans_list.append(torch.matmul(token_inputs[0:1, :], expert_weights[0].transpose(0, 1))) - for i in range(9): - t_ans = torch.matmul(token_inputs[(i + 1) : (i + 2), :], expert_weights[1].transpose(0, 1)) - ans_list.append(t_ans) - - true_out = torch.cat(ans_list, dim=0) - - assert torch.allclose(0.5 * true_out, out, atol=1e-2, rtol=0) - - -if __name__ == "__main__": - pytest.main() diff --git a/unit_tests/common/fused_moe/test_grouped_fused_moe_ep_sms.py b/unit_tests/common/fused_moe/test_grouped_fused_moe_ep_sms.py deleted file mode 100644 index 35b6ce7104..0000000000 --- a/unit_tests/common/fused_moe/test_grouped_fused_moe_ep_sms.py +++ /dev/null @@ -1,11 +0,0 @@ -from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe_ep.common import get_ep_num_sms - - -def test_get_ep_num_sms_selects_overlap_quota(monkeypatch): - import lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe_ep.common as common - - monkeypatch.setattr(common.dist_group_manager, "ep_num_sms", 20, raising=False) - monkeypatch.setattr(common.dist_group_manager, "ep_non_overlap_num_sms", 64, raising=False) - - assert get_ep_num_sms(overlap_with_compute=True) == 20 - assert get_ep_num_sms(overlap_with_compute=False) == 64 diff --git a/unit_tests/common/fused_moe/test_grouped_fused_moe_speed.py b/unit_tests/common/fused_moe/test_grouped_fused_moe_speed.py deleted file mode 100644 index 769002517d..0000000000 --- a/unit_tests/common/fused_moe/test_grouped_fused_moe_speed.py +++ /dev/null @@ -1,51 +0,0 @@ -import torch -import time -import pytest -from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe import moe_align, moe_align1, grouped_matmul -from lightllm.utils.log_utils import init_logger - -seed = 42 -torch.manual_seed(seed) - -if torch.cuda.is_available(): - torch.cuda.manual_seed(seed) - torch.cuda.manual_seed_all(seed) - -logger = init_logger(__name__) - - -@pytest.mark.parametrize("token_num", [200, 256, 8 * 1024]) -def test_moe_align1(token_num): - expert_num = 160 - topk_num = 6 - print(token_num) - - def get_one(): - rnd_logics = torch.randn(token_num, expert_num, device="cuda") - topk_values, topk_ids = torch.topk(rnd_logics, topk_num, dim=1) - - experts_info = torch.zeros((expert_num, token_num * topk_num), dtype=torch.int32, device="cuda") - experts_info.fill_(0) - moe_align(topk_ids, experts_info) - - topk_weights = torch.randn((token_num, topk_num), dtype=torch.float32, device="cuda") - experts_token_num = torch.zeros((expert_num,), dtype=torch.int32, device="cuda") - experts_weights = torch.zeros(experts_info.shape, dtype=torch.float32, device="cuda") - return experts_info, topk_weights, experts_weights, experts_token_num - - test_datas = [get_one() for _ in range(100)] - - moe_align1(*test_datas[0], topk_num) - - torch.cuda.synchronize() - start = time.time() - - for i in range(60): - moe_align1(*test_datas[i + 1], topk_num) - torch.cuda.synchronize() - - print(f"token_num: {token_num} cost time: {time.time() - start} s") - - -if __name__ == "__main__": - pytest.main() diff --git a/unit_tests/common/fused_moe/test_grouped_topk.py b/unit_tests/common/fused_moe/test_grouped_topk.py deleted file mode 100755 index 432e133163..0000000000 --- a/unit_tests/common/fused_moe/test_grouped_topk.py +++ /dev/null @@ -1,100 +0,0 @@ -import torch -import time -import pytest -import numpy as np -from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_topk import triton_grouped_topk -from lightllm.common.basemodel.triton_kernel.fused_moe.topk_select import biased_grouped_topk as grouped_topk -from lightllm.utils.log_utils import init_logger - -logger = init_logger(__name__) - -seed = 42 -torch.manual_seed(seed) - -if torch.cuda.is_available(): - torch.cuda.manual_seed(seed) - torch.cuda.manual_seed_all(seed) - - -@pytest.mark.parametrize( - "expert_num, topk_group, group_num, topk_num, scoring_func, token_num", - [ - (*a, b, c) - for a in [(256, 4, 8, 8), (160, 3, 8, 6)] - for b in [ - "sigmoid", - ] - for c in [1, 8, 256, 1024, 2048, 4096, 8192] - ], -) -def test_grouped_topk(expert_num, topk_group, group_num, topk_num, scoring_func, token_num): - print("test", expert_num, topk_group, group_num, topk_num, scoring_func, token_num) - dtype = torch.float32 - hidden_state = torch.randn((token_num, 1), dtype=dtype, device="cuda") - gating_output = torch.randn((token_num, expert_num), dtype=dtype, device="cuda") * 10 - correction_bias = torch.randn((expert_num,), dtype=dtype, device="cuda") - correction_bias[correction_bias <= 0.0] = 0.0 - - old_topk_weights, old_topk_ids = grouped_topk( - hidden_state, - gating_output=gating_output, - correction_bias=correction_bias, - topk=topk_num, - renormalize=True, - num_expert_group=group_num, - topk_group=topk_group, - scoring_func=scoring_func, - ) - - new_topk_weights, new_topk_ids = triton_grouped_topk( - None, - gating_output=gating_output, - correction_bias=correction_bias, - topk=topk_num, - renormalize=True, - num_expert_group=group_num, - topk_group=topk_group, - scoring_func=scoring_func, - ) - - torch.cuda.synchronize() - start = time.time() - for _ in range(60): - old_topk_weights, old_topk_ids = grouped_topk( - hidden_state, - gating_output=gating_output, - correction_bias=correction_bias, - topk=topk_num, - renormalize=True, - num_expert_group=group_num, - topk_group=topk_group, - scoring_func=scoring_func, - ) - torch.cuda.synchronize() - print(f"old cost time {time.time() - start} s") - - torch.cuda.synchronize() - start = time.time() - for _ in range(60): - new_topk_weights, new_topk_ids = triton_grouped_topk( - None, - gating_output=gating_output, - correction_bias=correction_bias, - topk=topk_num, - renormalize=True, - num_expert_group=group_num, - topk_group=topk_group, - scoring_func=scoring_func, - ) - torch.cuda.synchronize() - print(f"new cost time {time.time() - start} s") - - assert torch.equal(torch.sort(old_topk_ids, dim=1)[0], torch.sort(new_topk_ids, dim=1)[0]) - assert torch.allclose( - torch.sort(old_topk_weights, dim=1)[0], torch.sort(new_topk_weights, dim=1)[0], atol=1e-3, rtol=1e-1 - ), f"max delta {torch.max(torch.sort(old_topk_weights, dim=1)[0] - torch.sort(new_topk_weights, dim=1)[0])}" - return - - -if __name__ == "__main__": - pytest.main() diff --git a/unit_tests/common/fused_moe/test_moe_silu_and_mul_mix_quant_ep.py b/unit_tests/common/fused_moe/test_moe_silu_and_mul_mix_quant_ep.py deleted file mode 100644 index b569ec8429..0000000000 --- a/unit_tests/common/fused_moe/test_moe_silu_and_mul_mix_quant_ep.py +++ /dev/null @@ -1,196 +0,0 @@ -import torch -import pytest - - -def is_fp8_native_supported(): - """检查是否为 H100/B200 等原生支持 FP8 的硬件 (SM90+)""" - if not torch.cuda.is_available(): - return False - major, _ = torch.cuda.get_device_capability() - return major >= 9 - - -if not is_fp8_native_supported(): - pytest.skip(reason="not support fp8 test in this gpu card", allow_module_level=True) - -from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul_mix_quant_ep import ( - silu_and_mul_masked_post_quant_fwd, - silu_and_mul_psum_post_quant_fwd, -) -from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import silu_and_mul_fwd -from lightllm.common.basemodel.triton_kernel.quantization.fp8act_quant_kernel import per_token_group_quant_fp8 -from lightllm.utils.log_utils import init_logger - -logger = init_logger(__name__) - - -@pytest.mark.parametrize( - "expert_num, token_num, hidden_dim", - [ - ( - expert_num, - token_num, - hidden_dim, - ) - for expert_num in range(3, 6) - for hidden_dim in [256, 128 * 4, 2048] - for token_num in range(1, 7, 2) - ], -) -def test_silu_and_mul_masked(expert_num, token_num, hidden_dim): - quant_group_size = 128 - in_tensor = torch.randn((expert_num, token_num, hidden_dim), dtype=torch.bfloat16, device="cuda") - out_tensor = torch.empty((expert_num, token_num, hidden_dim // 2), dtype=torch.float8_e4m3fn, device="cuda") - out_scale_tensor = torch.randn( - (expert_num, token_num, hidden_dim // 2 // quant_group_size), dtype=torch.float32, device="cuda" - ) - - true_out_tensor_mid = torch.empty((expert_num, token_num, hidden_dim // 2), dtype=in_tensor.dtype, device="cuda") - - masked_m = torch.full((expert_num,), token_num, dtype=torch.int32, device="cuda") - - silu_and_mul_fwd(in_tensor.view(-1, hidden_dim), true_out_tensor_mid.view(-1, hidden_dim // 2)) - true_out_tensor, true_out_scale_tensor = per_token_group_quant_fp8( - true_out_tensor_mid.view(-1, hidden_dim // 2), - quant_group_size, - alloc_func=torch.empty, - ) - - silu_and_mul_masked_post_quant_fwd(in_tensor, out_tensor, out_scale_tensor, quant_group_size, masked_m) - - true_out_tensor = true_out_tensor.view(out_tensor.shape) - true_out_scale_tensor = true_out_scale_tensor.view(out_scale_tensor.shape) - - hidden_dim_scale_count = hidden_dim // 2 // quant_group_size - for expert_id, expert_token_num in enumerate(masked_m.cpu().numpy()): - true_scale = true_out_scale_tensor[expert_id, :expert_token_num, :hidden_dim_scale_count] - out_scale = out_scale_tensor[expert_id, :expert_token_num, :hidden_dim_scale_count] - assert torch.allclose( - true_scale, - out_scale, - atol=1e-3, - rtol=1e-2, - ) - - true_out = true_out_tensor[expert_id, :expert_token_num, :].to(torch.float32) - out = out_tensor[expert_id, :expert_token_num, :].to(torch.float32) - true_dequant = true_out * true_scale.repeat_interleave(quant_group_size, dim=-1) - out_dequant = out * out_scale.repeat_interleave(quant_group_size, dim=-1) - assert torch.allclose(true_dequant, out_dequant, atol=1e-1, rtol=1e-1) - return - - -@pytest.mark.parametrize("use_ue8m0_scales", [False, True]) -def test_silu_and_mul_masked_skips_padded_tokens(monkeypatch, use_ue8m0_scales): - monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", str(int(use_ue8m0_scales))) - expert_num = 3 - token_num = 4 - hidden_dim = 256 - quant_group_size = 128 - masked_m = torch.tensor([0, 2, token_num], dtype=torch.int32, device="cuda") - - in_tensor = torch.randn((expert_num, token_num, hidden_dim), dtype=torch.bfloat16, device="cuda") - out_tensor = torch.empty((expert_num, token_num, hidden_dim // 2), dtype=torch.float8_e4m3fn, device="cuda") - out_scale_tensor = torch.empty( - (expert_num, token_num, hidden_dim // 2 // quant_group_size), dtype=torch.float32, device="cuda" - ) - out_tensor.fill_(1.0) - out_scale_tensor.fill_(7.0) - - silu_and_mul_masked_post_quant_fwd(in_tensor, out_tensor, out_scale_tensor, quant_group_size, masked_m) - torch.cuda.synchronize() - - for expert_id, expert_token_num in enumerate(masked_m.cpu().tolist()): - assert torch.equal( - out_tensor[expert_id, expert_token_num:, :], - torch.ones_like(out_tensor[expert_id, expert_token_num:, :]), - ) - assert torch.equal( - out_scale_tensor[expert_id, expert_token_num:, :], - torch.full_like(out_scale_tensor[expert_id, expert_token_num:, :], 7.0), - ) - - -@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) -@pytest.mark.parametrize("use_ue8m0_scales", [False, True]) -@pytest.mark.parametrize("alpha,limit,add_one", [(1.0, 10.0, False), (1.702, 7.0, True)]) -def test_clamped_swiglu_quantization_matches_reference(monkeypatch, dtype, use_ue8m0_scales, alpha, limit, add_one): - monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", str(int(use_ue8m0_scales))) - torch.manual_seed(53) - x = (torch.randn(3, 5, 4096, device="cuda") * 16).to(dtype) - x[1, 0].zero_() - masked_m = torch.tensor([0, 2, 5], dtype=torch.int32, device="cuda") - out = torch.full((3, 5, 2048), 1.0, dtype=torch.float8_e4m3fn, device="cuda") - scales = torch.full((3, 5, 16), 7.0, dtype=torch.float32, device="cuda") - - silu_and_mul_masked_post_quant_fwd( - x, - out, - scales, - 128, - masked_m, - alpha=alpha, - limit=limit, - clamp_up_add_one=add_one, - ) - - gate, up = x.float().chunk(2, dim=-1) - gate = gate.clamp(max=limit) - gate = (gate * torch.sigmoid(alpha * gate)).to(dtype).float() - activation = (gate * (up.clamp(-limit, limit) + int(add_one))).to(dtype).float() - groups = activation.reshape(3, 5, 16, 128) - expected_scales = groups.abs().amax(-1).clamp(min=1e-10) / 448 - if use_ue8m0_scales: - expected_scales = torch.exp2(torch.ceil(torch.log2(expected_scales))) - expected_q = (groups / expected_scales.unsqueeze(-1)).clamp(-448, 448).to(torch.float8_e4m3fn) - expected = (expected_q.float() * expected_scales.unsqueeze(-1)).reshape(3, 5, 2048) - actual = out.float() * scales.repeat_interleave(128, dim=-1) - for expert, count in enumerate(masked_m.tolist()): - torch.testing.assert_close(scales[expert, :count], expected_scales[expert, :count], atol=1e-6, rtol=1e-5) - torch.testing.assert_close(actual[expert, :count], expected[expert, :count], atol=0.05, rtol=0.02) - assert torch.all(out[expert, count:].float() == 1) - assert torch.all(scales[expert, count:] == 7) - - -def test_silu_and_mul_psum_skips_alignment_padding(): - hidden = 128 - alignment = 8 - # expert0 使用 3 行，expert1 从对齐后的第 8 行开始再使用 2 行。 - expert_token_psum = torch.tensor([3, 10], device="cuda", dtype=torch.int32) - rows = 16 - gate = torch.randn((rows, hidden), device="cuda", dtype=torch.bfloat16) - up = torch.randn((rows, hidden), device="cuda", dtype=torch.bfloat16) - packed = torch.cat([gate, up], dim=1).contiguous() - output = torch.zeros((rows, hidden), device="cuda", dtype=torch.float8_e4m3fn) - scales = torch.zeros((rows, hidden // 128), device="cuda", dtype=torch.float32) - - silu_and_mul_psum_post_quant_fwd( - input=packed, - output=output, - output_scale=scales, - expert_token_psum=expert_token_psum, - expert_alignment=alignment, - quant_group_size=128, - ) - - def reference(row): - # Match the kernel: SiLU in fp32, then round the gate back to bf16 before multiplying. - activated_gate = torch.nn.functional.silu(gate[row].float()).to(torch.bfloat16) - activated = activated_gate.float() * up[row].float() - scale = activated.abs().amax().clamp(min=1e-10) / torch.finfo(torch.float8_e4m3fn).max - quantized = (activated / scale).clamp(-448, 448).to(torch.float8_e4m3fn) - return quantized, scale - - for row in (0, 1, 2, 8, 9): - expected_q, expected_scale = reference(row) - torch.testing.assert_close(scales[row], expected_scale.view(1), atol=1e-4, rtol=1e-3) - actual = output[row].float() * scales[row] - expected = expected_q.float() * expected_scale - # FP8 E4M3 bins are coarse, so compare the dequantized activation within one bin. - torch.testing.assert_close(actual, expected, atol=0.25, rtol=0.15) - # expert0 与 expert1 之间的对齐 padding 不应被写成有效值。 - assert torch.count_nonzero(output[3:8].float()) == 0 - - -if __name__ == "__main__": - pytest.main() diff --git a/unit_tests/common/fused_moe/test_softmax_topk.py b/unit_tests/common/fused_moe/test_softmax_topk.py index 7c3e483df8..a4050f236b 100755 --- a/unit_tests/common/fused_moe/test_softmax_topk.py +++ b/unit_tests/common/fused_moe/test_softmax_topk.py @@ -1,118 +1,20 @@ -import torch -import time import pytest -import numpy as np -from lightllm.common.basemodel.triton_kernel.fused_moe.softmax_topk import softmax_topk -from lightllm.utils.log_utils import init_logger - -logger = init_logger(__name__) - - -def benchmark(M, N, K, renorm, runs): - try: - import sgl_kernel as sgl_ops - except Exception as e: - pytest.skip(f"no sgl_kernel error: {str(e)}", allow_module_level=True) - - gating = torch.randn(M, N, device="cuda", dtype=torch.float32) - torch.cuda.synchronize() - - # 1. SGL kernel - sgl_vals = torch.empty((M, K), dtype=torch.float32, device="cuda") - sgl_ids = torch.empty((M, K), dtype=torch.int32, device="cuda") - # Warm-up - sgl_ops.topk_softmax(sgl_vals, sgl_ids, gating) - torch.cuda.synchronize() - start = torch.cuda.Event(True) - end = torch.cuda.Event(True) - start.record() - for _ in range(runs): - sgl_ops.topk_softmax(sgl_vals, sgl_ids, gating) - if renorm: - sgl_vals.div_(sgl_vals.sum(-1, keepdim=True).clamp_min(1e-8)) - - end.record() - torch.cuda.synchronize() - t_sgl = start.elapsed_time(end) / runs - - # 2. Triton kernel - t0 = torch.cuda.Event(True) - t1 = torch.cuda.Event(True) - # Warm-up - softmax_topk(gating, K) - torch.cuda.synchronize() - t0.record() - for _ in range(runs): - triton_vals, triton_ids = softmax_topk(gating, K, renorm) - t1.record() - torch.cuda.synchronize() - t_triton = t0.elapsed_time(t1) / runs - - # 3. Native PyTorch - _ = torch.softmax(gating, dim=-1) - _, _ = torch.topk(_, K, dim=-1) - torch.cuda.synchronize() - - start, end = torch.cuda.Event(True), torch.cuda.Event(True) - start.record() - for _ in range(runs): - probs = torch.softmax(gating, dim=-1) - torch_vals, torch_ids = torch.topk(probs, K, dim=-1) - if renorm: - torch_vals.div_(torch_vals.sum(-1, keepdim=True).clamp_min(1e-8)) - end.record() - torch.cuda.synchronize() - t_torch = start.elapsed_time(end) / runs - - # Compare indices and weights - # Count mismatches of ordered indices - diff_sgl_triton_ids = (sgl_ids != triton_ids).sum().item() - diff_torch_triton_ids = (torch_ids != triton_ids).sum().item() - # Max absolute difference of weights aligned by position - max_err_triton_torch = (triton_vals - torch_vals).abs().max().item() - max_err_triton_torch_sgl = (sgl_vals - torch_vals).abs().max().item() - max_err_triton_sgl = (triton_vals - sgl_vals).abs().max().item() - - assert diff_sgl_triton_ids == 0, f"Mismatch SGL vs Triton ids: {diff_sgl_triton_ids}" - assert diff_torch_triton_ids == 0, f"Mismatch Torch vs Triton ids: {diff_torch_triton_ids}" - assert max_err_triton_torch < 1e-3, f"Max err Triton vs Torch: {max_err_triton_torch}" - assert max_err_triton_torch_sgl < 1e-3, f"Max err Triton vs SGL: {max_err_triton_torch_sgl}" - assert max_err_triton_sgl < 1e-3, f"Max err Torch vs SGL: {max_err_triton_sgl}" +import torch - results = { - "time_sgl": t_sgl, - "time_triton": t_triton, - "time_torch": t_torch, - "mismatch_sgl_triton_ids": diff_sgl_triton_ids, - "mismatch_torch_triton_ids": diff_torch_triton_ids, - "max_err_triton_torch": max_err_triton_torch, - "max_err_triton_sgl": max_err_triton_sgl, - "max_err_triton_torch_sgl": max_err_triton_torch_sgl, - "sgl_ids": sgl_ids, - "triton_ids": triton_ids, - "torch_ids": torch_ids, - "sgl_vals": sgl_vals, - "triton_vals": triton_vals, - "torch_vals": torch_vals, - } - return results +from lightllm.common.basemodel.triton_kernel.fused_moe.softmax_topk import softmax_topk -def test_softmax_topk(): - M, N, K = 8192, 1024, 8 - res = benchmark(M, N, K, False, 1000) - print(f"SGL time: {res['time_sgl']:.6f}ms") - print(f"Triton time: {res['time_triton']:.6f}ms") - print(f"PyTorch time: {res['time_torch']:.6f}ms") - print("Mismatch SGL vs Triton ids:", res["mismatch_sgl_triton_ids"]) - print("Mismatch Torch vs Triton ids:", res["mismatch_torch_triton_ids"]) - print("Max err Triton vs Torch :", res["max_err_triton_torch"]) - print("Max err Triton vs SGL :", res["max_err_triton_sgl"]) - print("Max err Torch vs SGL :", res["max_err_triton_torch_sgl"]) - benchmark(M, N, K, True, 10) - benchmark(M, 256, 5, True, 10) - benchmark(M, 127, 5, True, 10) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("experts,topk", [(1024, 8), (256, 5), (127, 5)]) +@pytest.mark.parametrize("renorm", [False, True]) +def test_softmax_topk(experts, topk, renorm): + torch.manual_seed(42) + gating = torch.randn(32, experts, device="cuda") + expected_weights, expected_ids = torch.softmax(gating, dim=-1).topk(topk, dim=-1) + if renorm: + expected_weights /= expected_weights.sum(dim=-1, keepdim=True) + weights, ids = softmax_topk(gating, topk, renorm) -if __name__ == "__main__": - pytest.main() + torch.testing.assert_close(ids.long(), expected_ids, rtol=0, atol=0) + torch.testing.assert_close(weights, expected_weights, rtol=0, atol=1e-3) diff --git a/unit_tests/common/kv_trans_kernel/test_kv_trans_v2.py b/unit_tests/common/kv_trans_kernel/test_kv_trans_v2.py deleted file mode 100644 index dd34402dd2..0000000000 --- a/unit_tests/common/kv_trans_kernel/test_kv_trans_v2.py +++ /dev/null @@ -1,123 +0,0 @@ -import pytest -import torch -import random -from lightllm.common.kv_trans_kernel.kv_trans_v2 import kv_trans_v2_for_p_node, kv_trans_v2_for_d_node, kv_trans_for_dp - - -@pytest.mark.parametrize( - "token_num", - [token_num for token_num in range(5, 10)], -) -def test_kv_trans_v2_for_p_node(token_num): - dp_size_in_node = 8 - head_num = 2 - head_dim = 512 - kv_buffer_token_num = 512 - mems = [] - for _ in range(dp_size_in_node): - mems.append(torch.randn((kv_buffer_token_num, head_num, head_dim), dtype=torch.float16, device="cuda")) - input_mems = torch.tensor([e.data_ptr() for e in mems], dtype=torch.uint64, device="cuda") - input_idx = [random.randint(0, kv_buffer_token_num - 1) for _ in range(token_num)] - input_idx = torch.tensor(input_idx, dtype=torch.int32, device="cuda") - input_dp_idx = [random.randint(0, dp_size_in_node - 1) for _ in range(token_num)] - input_dp_idx = torch.tensor(input_dp_idx, dtype=torch.int32, device="cuda") - - true_output = torch.zeros((token_num, head_num, head_dim), dtype=torch.float16, device="cuda") - test_output = torch.zeros((token_num, head_num, head_dim), dtype=torch.float16, device="cuda") - output_idx = torch.arange(0, token_num, 1, dtype=torch.int32, device="cuda") - - kv_trans_v2_for_p_node(input_mems, input_idx, input_dp_idx, test_output, output_idx, dp_size_in_node) - - for dest_token_index, token_index, dp_index in zip( - list(range(token_num)), input_idx.cpu().numpy(), input_dp_idx.cpu().numpy() - ): - true_output[dest_token_index, :, :] = mems[dp_index][token_index] - - assert torch.equal(true_output, test_output) - return - - -@pytest.mark.parametrize( - "token_num", - [token_num for token_num in range(5, 10)], -) -def test_kv_trans_v2_for_d_node(token_num): - card_num = 8 - dp_size_in_node = 4 - head_num = 2 - head_dim = 512 - kv_buffer_token_num = 512 - mems = [] - for _ in range(card_num): - mems.append(torch.randn((kv_buffer_token_num, head_num, head_dim), dtype=torch.float16, device="cuda")) - output_mems = torch.tensor([e.data_ptr() for e in mems], dtype=torch.uint64, device="cuda") - output_idx = [random.randint(0, kv_buffer_token_num - 1) for _ in range(token_num)] - output_idx = torch.tensor(output_idx, dtype=torch.int32, device="cuda") - output_dp_idx = [random.randint(0, dp_size_in_node - 1) for _ in range(token_num)] - output_dp_idx = torch.tensor(output_dp_idx, dtype=torch.int32, device="cuda") - - test_input = torch.randn((token_num, head_num, head_dim), dtype=torch.float16, device="cuda") - input_idx = torch.arange(0, token_num, 1, dtype=torch.int32, device="cuda") - - kv_trans_v2_for_d_node(output_mems, output_idx, output_dp_idx, test_input, input_idx, dp_size_in_node) - - for dest_token_index, token_index, dest_token_index, dp_index in zip( - list(range(token_num)), - input_idx.cpu().numpy(), - output_idx.cpu().numpy(), - output_dp_idx.cpu().numpy(), - ): - for mem_index in range(dp_index * card_num // dp_size_in_node, (dp_index + 1) * card_num // dp_size_in_node): - torch.equal(mems[mem_index][dest_token_index, :, :], test_input[token_index, :, :]) - - return - - -@pytest.mark.parametrize( - "token_num", - [token_num for token_num in range(5, 10)], -) -def test_kv_trans_for_dp(token_num): - card_num = 8 - dp_size_in_node = 4 - layer_num = 3 - head_num = 2 - head_dim = 512 - kv_buffer_token_num = 512 - rank_in_dp = 1 - - card_num_per_d = card_num // dp_size_in_node - - # 创建多层的 mem，每个 mem 包含所有层的数据 - mems = [] - for _ in range(card_num): - mems.append( - torch.randn((layer_num, kv_buffer_token_num, head_num, head_dim), dtype=torch.float16, device="cuda") - ) - - input_mems = torch.tensor([e.data_ptr() for e in mems], dtype=torch.uint64, device="cuda") - input_idx = [random.randint(0, kv_buffer_token_num - 1) for _ in range(token_num)] - input_idx = torch.tensor(input_idx, dtype=torch.int32, device="cuda") - input_dp_idx = [random.randint(0, dp_size_in_node - 1) for _ in range(token_num)] - input_dp_idx = torch.tensor(input_dp_idx, dtype=torch.int32, device="cuda") - - true_output = torch.zeros((layer_num, kv_buffer_token_num, head_num, head_dim), dtype=torch.float16, device="cuda") - test_output = torch.zeros((layer_num, kv_buffer_token_num, head_num, head_dim), dtype=torch.float16, device="cuda") - output_idx = torch.arange(0, token_num, 1, dtype=torch.int32, device="cuda") - - kv_trans_for_dp(input_mems, input_idx, input_dp_idx, test_output, output_idx, dp_size_in_node, rank_in_dp) - - # 验证结果 - for dest_token_index, src_token_index, dp_index in zip( - list(range(token_num)), input_idx.cpu().numpy(), input_dp_idx.cpu().numpy() - ): - mem_index = rank_in_dp + dp_index * card_num_per_d - # 所有 layer 都从同一个 mem 的对应层读取 - true_output[:, dest_token_index, :, :] = mems[mem_index][:, src_token_index, :, :] - - assert torch.equal(true_output, test_output), "kv_trans_for_dp output mismatch" - return - - -if __name__ == "__main__": - pytest.main() diff --git a/unit_tests/common/kv_trans_kernel/test_nixl_kv_trans.py b/unit_tests/common/kv_trans_kernel/test_nixl_kv_trans.py deleted file mode 100644 index 3ec72e15d2..0000000000 --- a/unit_tests/common/kv_trans_kernel/test_nixl_kv_trans.py +++ /dev/null @@ -1,141 +0,0 @@ -import os -import pytest -import torch - - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for Triton kernels") - - -def _randn_like(shape, dtype=torch.float32, device="cuda"): - torch.manual_seed(1234) - return torch.randn(shape, dtype=dtype, device=device).contiguous() - - -def _gather_along_size(t, idx): - return t.index_select(dim=1, index=idx) - - -@pytest.mark.parametrize("head_dim", [64, 96]) -@pytest.mark.parametrize( - "kv_head_num,tp_world_size", - [ - (4, 1), - (4, 2), - (8, 1), - (8, 2), - (8, 4), - (16, 2), - ], -) -def test_page_io_roundtrip_with_tp(head_dim, kv_head_num, tp_world_size): - from lightllm.common.kv_trans_kernel.nixl_kv_trans import page_io - - device = "cuda" - dtype = torch.bfloat16 - - # Shapes - layer_num = 3 - size = 32 - assert kv_head_num % 2 == 0 - page_size = 32 - page_head_num = kv_head_num * tp_world_size - - tp_indices = list(range(tp_world_size)) - - kv_buffer = _randn_like((layer_num, size, kv_head_num, head_dim), dtype=dtype, device=device) - page_tensor = torch.zeros((page_size, layer_num, page_head_num, head_dim), dtype=dtype, device=device).contiguous() - - # Select a handful of token positions to move (tid count <= page_size) - mem_indexes = torch.tensor([2, 5, 7, 9, 11], dtype=torch.int64, device=device).contiguous() - - # Write: kv_buffer -> page_tensor, done by all tp ranks to fill their partition - for tp_index in tp_indices: - page_io( - mem_indexes=mem_indexes, - page_tensor=page_tensor, - kv_buffer=kv_buffer, - tp_index=tp_index, - tp_world_size=tp_world_size, - mode="write", - ) - - # After-write expectation check - token_num = mem_indexes.numel() - k_head_num = kv_head_num // 2 - page_k_head_num = page_head_num // 2 - page_write_head_num = page_k_head_num // tp_world_size - - expected_page = torch.zeros_like(page_tensor) - for tid in range(token_num): - mem_index = int(mem_indexes[tid].item()) - for layer_index in range(layer_num): - for tp_index in tp_indices: - page_head_start = tp_index * page_write_head_num - for kv_head_id in range(page_write_head_num): - # K half - expected_page[tid, layer_index, page_head_start + kv_head_id, :] = kv_buffer[ - layer_index, mem_index, kv_head_id, : - ] - # V half - expected_page[tid, layer_index, page_k_head_num + page_head_start + kv_head_id, :] = kv_buffer[ - layer_index, mem_index, k_head_num + kv_head_id, : - ] - - assert torch.allclose(page_tensor[:token_num], expected_page[:token_num], atol=1e-3, rtol=1e-3) - - # Read back to a fresh buffer - kv_buffer_rt = torch.zeros_like(kv_buffer) - for tp_index in tp_indices: - page_io( - mem_indexes=mem_indexes, - page_tensor=page_tensor, - kv_buffer=kv_buffer_rt, - tp_index=tp_index, - tp_world_size=tp_world_size, - mode="read", - ) - - # Check equality only at selected positions along size-dim - ref = _gather_along_size(kv_buffer, mem_indexes) - out = _gather_along_size(kv_buffer_rt, mem_indexes) - assert torch.allclose(out, ref, atol=1e-3, rtol=1e-3) - - -@pytest.mark.parametrize("head_dim", [32, 80]) # include non-power-of-two -def test_mla_page_io_roundtrip(head_dim): - from lightllm.common.kv_trans_kernel.nixl_kv_trans import mla_page_io - - device = "cuda" - dtype = torch.bfloat16 - - # Shapes (single-head) - layer_num = 2 - size = 20 - kv_head_num = 1 - page_size = 10 - page_head_num = 1 - - kv_buffer = _randn_like((layer_num, size, kv_head_num, head_dim), dtype=dtype, device=device) - page_tensor = torch.zeros((page_size, layer_num, page_head_num, head_dim), dtype=dtype, device=device).contiguous() - - mem_indexes = torch.tensor([0, 3, 6, 7], dtype=torch.int64, device=device).contiguous() - - # Write kv -> page - mla_page_io(mem_indexes=mem_indexes, page_tensor=page_tensor, kv_buffer=kv_buffer, mode="write") - - # After-write expectation check - token_num = mem_indexes.numel() - expected_page = torch.zeros_like(page_tensor) - for tid in range(token_num): - mem_index = int(mem_indexes[tid].item()) - for layer_index in range(layer_num): - expected_page[tid, layer_index, 0, :] = kv_buffer[layer_index, mem_index, 0, :] - assert torch.allclose(page_tensor[:token_num], expected_page[:token_num], atol=1e-3, rtol=1e-3) - - # Read back page -> kv - kv_buffer_rt = torch.zeros_like(kv_buffer) - mla_page_io(mem_indexes=mem_indexes, page_tensor=page_tensor, kv_buffer=kv_buffer_rt, mode="read") - - ref = kv_buffer.index_select(dim=1, index=mem_indexes) - out = kv_buffer_rt.index_select(dim=1, index=mem_indexes) - assert torch.allclose(out, ref, atol=1e-3, rtol=1e-3) diff --git a/unit_tests/common/quantization/test_no_quant.py b/unit_tests/common/quantization/test_no_quant.py deleted file mode 100644 index 2b66767539..0000000000 --- a/unit_tests/common/quantization/test_no_quant.py +++ /dev/null @@ -1,43 +0,0 @@ -import pytest -import torch - -from lightllm.common.basemodel.layer_weights.meta_weights.mm_weight.mm_weight import MMWeightTpl -from lightllm.common.quantization.no_quant import NoQuantization -from lightllm.common.quantization.quantize_method import WeightPack - - -@pytest.fixture(autouse=True) -def device_env(monkeypatch): - monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") -@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) -@pytest.mark.parametrize("fp32_output", [False, True]) -def test_output_buffer_controls_matmul_precision(dtype, fp32_output): - torch.manual_seed(1593) - x = torch.randn((17, 64), dtype=dtype, device="cuda") - weight = torch.randn((32, 64), dtype=dtype, device="cuda") - out = torch.empty((17, 32), dtype=torch.float32 if fp32_output else dtype, device="cuda") - - projection = MMWeightTpl.__new__(MMWeightTpl) - projection.quant_method = NoQuantization() - projection.mm_param = WeightPack(weight=weight) - projection.bias = None - result = projection.mm(x, out=out) - - assert result is out - expected = torch.mm(x.float(), weight.float().t()).to(out.dtype) - torch.testing.assert_close(result, expected, rtol=1e-5, atol=1e-5) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") -def test_default_output_and_bias_preserve_input_dtype(): - x = torch.randn((3, 64), dtype=torch.bfloat16, device="cuda") - weight = torch.randn((32, 64), dtype=x.dtype, device=x.device) - bias = torch.randn(32, dtype=x.dtype, device=x.device) - - result = NoQuantization().apply(x, WeightPack(weight=weight), bias=bias, use_custom_tensor_mananger=False) - - assert result.dtype == x.dtype - torch.testing.assert_close(result, torch.addmm(bias, x, weight.t()), rtol=0, atol=0) diff --git a/unit_tests/common/quantization/test_ue8m0_fp8_quantization.py b/unit_tests/common/quantization/test_ue8m0_fp8_quantization.py deleted file mode 100644 index f1c9b228a2..0000000000 --- a/unit_tests/common/quantization/test_ue8m0_fp8_quantization.py +++ /dev/null @@ -1,513 +0,0 @@ -import os -import pytest -import torch -import torch.nn.functional as F - -from lightllm.common.basemodel.triton_kernel.quantization import fp8act_quant_kernel as activation -from lightllm.common.basemodel.triton_kernel.quantization import fp8w8a8_block_quant_kernel as weight -from lightllm.common.quantization import Quantcfg -from lightllm.common.quantization import deepgemm - - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") - - -@pytest.fixture(autouse=True) -def preserve_ue8m0_env(monkeypatch): - monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", os.getenv("LIGHTLLM_USE_UE8M0_SCALES", "0")) - - -def _unpack_scales(packed, groups): - shifts = torch.arange(4, device=packed.device) * 8 - exponents = ((packed.to(torch.int64)[..., None] >> shifts) & 255).flatten(1) - return torch.exp2(exponents[:, :groups].float() - 127) - - -@pytest.fixture -def packed_weight_layout(monkeypatch): - calls = [] - transform = deepgemm.deep_gemm.transform_sf_into_required_layout - sm100 = deepgemm.is_sm100_gpu() - - def record_transform(sf, n, k, recipe, num_groups=None, is_sfa=False): - calls.append((n, k, num_groups)) - if sm100: - return transform(sf, n, k, recipe, num_groups=num_groups, is_sfa=is_sfa) - # Exercise DeepGEMM's actual pack kernel on Hopper, emulating only the SM100 layout dispatch. - assert recipe == (1, 128, 128) and not is_sfa - expanded = sf.index_select(-2, torch.arange(n, device=sf.device) // 128) - return deepgemm.deep_gemm.get_mn_major_tma_aligned_packed_ue8m0_tensor(expanded) - - monkeypatch.setattr(deepgemm.deep_gemm, "transform_sf_into_required_layout", record_transform) - return calls - - -@pytest.mark.skipif(not activation.HAS_SGL_KERNEL, reason="requires SGL kernel") -@pytest.mark.parametrize("use_ue8m0_scales,use_packed_ue8m0", [(False, False), (True, False), (True, True)]) -def test_sgl_dispatch_uses_ue8m0_flag(monkeypatch, use_ue8m0_scales, use_packed_ue8m0): - quantize = activation.sgl_ops.sgl_per_token_group_quant_fp8 - calls = [] - - def record_call(*args, **kwargs): - calls.append(kwargs["scale_ue8m0"]) - return quantize(*args, **kwargs) - - monkeypatch.setattr(activation.sgl_ops, "sgl_per_token_group_quant_fp8", record_call) - x = torch.ones((3, 512), device="cuda", dtype=torch.bfloat16) - _, scales = activation.per_token_group_quant_fp8( - x, 128, use_ue8m0_scales=use_ue8m0_scales, use_packed_ue8m0=use_packed_ue8m0 - ) - assert calls == ([use_packed_ue8m0] if not use_ue8m0_scales or use_packed_ue8m0 else []) - assert scales.dtype == (torch.int32 if use_packed_ue8m0 else torch.float32) - if use_ue8m0_scales: - actual = _unpack_scales(scales, 4) if use_packed_ue8m0 else scales - torch.testing.assert_close(actual, torch.full_like(actual, 2.0 ** -8), rtol=0, atol=0) - - -@pytest.mark.parametrize("use_packed_ue8m0", [False, True]) -def test_ue8m0_rounding_at_power_of_two_boundaries(monkeypatch, use_packed_ue8m0): - monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") - monkeypatch.setattr(activation, "HAS_SGL_KERNEL", False) - powers = torch.tensor([2.0 ** exponent for exponent in (-16, -8, 0, 8, 16, 118)], device="cuda") - above = torch.nextafter(powers, torch.full_like(powers, float("inf"))) - below = torch.nextafter(powers, torch.zeros_like(powers)) - magnitudes = torch.cat((powers, above, below, powers.new_zeros(1))) * 448.0 - x = magnitudes[:, None].expand(-1, 128).contiguous() - expected_act = torch.exp2(torch.ceil(torch.log2(magnitudes.clamp_min(1e-10) / 448.0))) - expected_weight = torch.exp2(torch.ceil(torch.log2(magnitudes.clamp_min(1e-4) / 448.0))) - - _, act_scales = activation.per_token_group_quant_fp8( - x, 128, use_ue8m0_scales=True, use_packed_ue8m0=use_packed_ue8m0 - ) - _, weight_scales = weight.weight_quant(x.repeat_interleave(128, dim=0), use_ue8m0_scales=True) - - actual = _unpack_scales(act_scales, 1) if use_packed_ue8m0 else act_scales - torch.testing.assert_close(actual[:, 0], expected_act, rtol=0, atol=0) - torch.testing.assert_close(weight_scales[:, 0], expected_weight, rtol=0, atol=0) - - -@pytest.mark.parametrize("layout", ["row", "column", "tma"]) -@pytest.mark.parametrize("group_size", [64, 128]) -@pytest.mark.parametrize("use_ue8m0_scales,use_packed_ue8m0", [(False, False), (True, False), (True, True)]) -def test_activation_scales_and_quantized_values(monkeypatch, layout, group_size, use_ue8m0_scales, use_packed_ue8m0): - monkeypatch.setattr(activation, "HAS_SGL_KERNEL", False) - torch.manual_seed(20261008) - rows, groups = 17, 4 - x = torch.randn(rows, groups * group_size, device="cuda", dtype=torch.bfloat16) - x[0].zero_() - x[1].fill_(1e-12) - q, scales = activation.per_token_group_quant_fp8( - x, - group_size, - column_major_scales=layout != "row", - scale_tma_aligned=layout == "tma", - use_ue8m0_scales=use_ue8m0_scales, - use_packed_ue8m0=use_packed_ue8m0, - ) - - amax = x.float().reshape(rows, groups, group_size).abs().amax(dim=-1) - reference_scales = amax.clamp_min(1e-10) / 448.0 - if use_ue8m0_scales: - reference_scales = torch.exp2(torch.ceil(torch.log2(reference_scales.double()))).float() - reference_q = (x.float().reshape(rows, groups, group_size) / reference_scales[..., None]).to(q.dtype) - actual_scales = _unpack_scales(scales, groups) if use_packed_ue8m0 else scales - torch.testing.assert_close(actual_scales, reference_scales, rtol=0, atol=0) - torch.testing.assert_close(q.float(), reference_q.reshape_as(x).float(), rtol=0, atol=0) - assert scales.dtype == (torch.int32 if use_packed_ue8m0 else torch.float32) - # UE8M0 always follows SGL's packed TMA layout, including with default layout flags. - if use_packed_ue8m0 or layout == "tma": - assert scales.stride() == (1, 20) - elif layout == "row": - assert scales.stride() == (groups, 1) - - -@pytest.mark.parametrize("shape", [(2, 3, 256), (2, 3, 4, 256)]) -@pytest.mark.parametrize("group_size", [64, 128]) -@pytest.mark.parametrize("use_ue8m0_scales", [False, True]) -def test_nd_activation_scales_preserve_row_major_output(monkeypatch, shape, group_size, use_ue8m0_scales): - monkeypatch.setattr(activation, "HAS_SGL_KERNEL", False) - x = torch.randn(shape, device="cuda", dtype=torch.bfloat16) - x[0, 0].zero_() - q, scales = activation.per_token_group_quant_fp8(x, group_size, use_ue8m0_scales=use_ue8m0_scales) - groups = shape[-1] // group_size - reference_scales = x.float().view(*shape[:-1], groups, group_size).abs().amax(-1).clamp_min(1e-10) / 448 - if use_ue8m0_scales: - reference_scales = reference_scales.log2().ceil().exp2() - reference_q = (x.float().view(*shape[:-1], groups, group_size) / reference_scales[..., None]).to(q.dtype) - assert scales.shape == (*shape[:-1], groups) - torch.testing.assert_close(scales, reference_scales, rtol=0, atol=0) - torch.testing.assert_close(q.float(), reference_q.view_as(x).float(), rtol=0, atol=0) - - -@pytest.mark.parametrize("use_sgl", [False, True]) -@pytest.mark.parametrize("rows", [1, 17]) -@pytest.mark.parametrize("groups", [1, 2, 3, 4, 5, 8]) -@pytest.mark.parametrize("group_size", [64, 128]) -def test_packed_ue8m0_scales(monkeypatch, rows, groups, group_size, use_sgl): - if use_sgl and not activation.HAS_SGL_KERNEL: - pytest.skip("requires SGL kernel") - monkeypatch.setattr(activation, "HAS_SGL_KERNEL", use_sgl) - torch.manual_seed(20261008) - x = torch.randn(rows, groups * group_size, device="cuda", dtype=torch.bfloat16) - x[0].zero_() - if rows > 1: - x[1].fill_(1e-12) - x[-1, -group_size:].fill_(2.0 ** 16) - q, packed_scales = activation.per_token_group_quant_fp8(x, group_size, use_ue8m0_scales=True, use_packed_ue8m0=True) - - assert packed_scales.dtype == torch.int32 - assert packed_scales.shape == (rows, (groups + 3) // 4) - assert packed_scales.stride() == (1, (rows + 3) // 4 * 4) - if groups % 4: - assert torch.count_nonzero(packed_scales[:, -1].to(torch.int64) >> (groups % 4 * 8)) == 0 - scales = _unpack_scales(packed_scales, groups) - amax = x.float().reshape(rows, groups, group_size).abs().amax(-1) - reference_scales = torch.exp2(torch.ceil(torch.log2(amax.clamp_min(1e-10) / 448.0))) - reference_q = (x.float().reshape(rows, groups, group_size) / reference_scales[..., None]).to(q.dtype) - torch.testing.assert_close(scales, reference_scales, rtol=0, atol=0) - torch.testing.assert_close(q.float(), reference_q.reshape_as(x).float(), rtol=0, atol=0) - - -@pytest.mark.parametrize("experts", [None, 2]) -@pytest.mark.parametrize("use_ue8m0_scales", [False, True]) -def test_weight_quantization_partial_blocks(monkeypatch, experts, use_ue8m0_scales): - monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") - torch.manual_seed(20261008) - shape = (129, 257) if experts is None else (experts, 129, 257) - x = torch.randn(shape, device="cuda", dtype=torch.bfloat16) - x[..., :128, :128].zero_() - x[..., 128:, 128:256].fill_(1e-12) - q, scales = weight.weight_quant(x, use_ue8m0_scales=use_ue8m0_scales) - - blocks = F.pad(x.float(), (0, 127, 0, 127)).reshape(-1, 2, 128, 3, 128) - amax = blocks.abs().amax(dim=(2, 4)).reshape_as(scales) - reference_scales = amax.clamp_min(1e-4) / 448.0 if use_ue8m0_scales else amax / 448.0 - if use_ue8m0_scales: - reference_scales = torch.exp2(torch.ceil(torch.log2(reference_scales.double()))).float() - denom = reference_scales if use_ue8m0_scales else reference_scales + 1e-6 - expanded = denom.repeat_interleave(128, dim=-2).repeat_interleave(128, dim=-1)[..., :129, :257] - torch.testing.assert_close(scales, reference_scales, rtol=0, atol=0) - torch.testing.assert_close(q.float(), (x.float() / expanded).to(q.dtype).float(), rtol=0, atol=0) - - -@pytest.mark.parametrize("env_value", ["1", "0"]) -def test_ue8m0_env_preserves_unquantized_method(monkeypatch, env_value): - monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") - monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", env_value) - method = Quantcfg({"n_layer": 1}).get_quant_method(0, "q_proj") - assert method.method_name == "none" - - -@pytest.mark.skipif(not deepgemm.HAS_DEEPGEMM, reason="requires DeepGEMM") -@pytest.mark.parametrize("sm100", [False, True]) -def test_deepgemm_selects_scale_format_for_machine(monkeypatch, sm100, packed_weight_layout): - monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", "1") - monkeypatch.setattr(deepgemm, "is_sm100_gpu", lambda: sm100) - method = Quantcfg({"n_layer": 1}, quant_type="fp8w8a8-b128-deepgemm").get_quant_method(0, "q_proj") - weight_pack, _ = method.create_weight([128], 512, torch.bfloat16, 0) - method.load_weight(torch.ones((128, 512), dtype=torch.bfloat16, device="cuda"), weight_pack) - seen = [] - - def check_gemm_inputs(a, b, out): - seen.append(a[1].dtype) - assert b[1].dtype == (torch.int32 if sm100 else torch.float32) - scales = _unpack_scales(a[1], 4) if sm100 else a[1] - torch.testing.assert_close(scales, torch.full_like(scales, 2.0 ** -8), rtol=0, atol=0) - weight_scales = _unpack_scales(b[1], 4) if sm100 else b[1] - torch.testing.assert_close(weight_scales, torch.full_like(weight_scales, 2.0 ** -8), rtol=0, atol=0) - out.zero_() - - # Exercise format selection and real quantization on Hopper; SM100 GEMM needs a separate machine. - monkeypatch.setattr(deepgemm, "_deepgemm_fp8_nt", check_gemm_inputs) - method.apply( - torch.ones((3, 512), dtype=torch.bfloat16, device="cuda"), weight_pack, use_custom_tensor_mananger=False - ) - method.apply( - torch.ones((3, 512), dtype=torch.bfloat16, device="cuda"), weight_pack, use_custom_tensor_mananger=False - ) - assert seen == [torch.int32 if sm100 else torch.float32] * 2 - assert len(packed_weight_layout) == int(sm100) - - -@pytest.mark.skipif(not deepgemm.HAS_DEEPGEMM, reason="requires DeepGEMM") -@pytest.mark.parametrize("sm100", [False, True]) -@pytest.mark.parametrize("ue8m0", [False, True]) -@pytest.mark.parametrize("prequantized", [False, True]) -@pytest.mark.parametrize("experts", [None, 2]) -def test_weight_scale_loading_preserves_fused_views( - monkeypatch, packed_weight_layout, sm100, ue8m0, prequantized, experts -): - monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", str(int(ue8m0))) - monkeypatch.setattr(deepgemm, "is_sm100_gpu", lambda: sm100) - method = Quantcfg({"n_layer": 1}, quant_type="fp8w8a8-b128-deepgemm").get_quant_method(0, "q_proj") - args = dict(out_dims=[128, 129], in_dim=640, dtype=torch.bfloat16, device_id=0) - full, parts = ( - method.create_weight(**args) if experts is None else method.create_moe_weight(**args, num_experts=experts) - ) - packed = sm100 and ue8m0 - expected_parts = [] - torch.manual_seed(20261008) - for part in parts: - raw = torch.randn_like(part.weight, dtype=torch.bfloat16) - raw[..., 128:256] *= 8 - q, scales = weight.weight_quant(raw, use_ue8m0_scales=ue8m0) - loads = ( - [(part, raw, q, scales)] - if experts is None - else [(part.get_expert(i), raw[i], q[i], scales[i]) for i in range(experts)] - ) - for dest, bf16, fp8, sf in loads: - if prequantized: - method.load_weight(fp8, dest) - method.load_weight_scale(sf.cpu(), dest) - else: - method.load_weight(bf16, dest) - assert all(dest.load_ok) - torch.testing.assert_close(part.weight.float(), q.float(), rtol=0, atol=0) - expected = ( - scales.index_select(-2, torch.arange(part.weight.shape[-2], device="cuda") // 128) if packed else scales - ) - flat = part.weight_scale.reshape(-1, part.weight_scale.shape[-1]) - actual = _unpack_scales(flat, 5).reshape_as(expected) if packed else part.weight_scale - torch.testing.assert_close(actual, expected, rtol=0, atol=0) - expected_parts.append(expected) - - expected = torch.cat(expected_parts, dim=-2) - assert full.weight_scale.dtype == (torch.int32 if packed else torch.float32) - if packed: - assert full.weight_scale.shape[-2:] == (257, 2) - assert full.weight_scale.stride()[-2:] == (1, 260) - actual = _unpack_scales(full.weight_scale.reshape(-1, 2), 5).reshape_as(expected) - else: - actual = full.weight_scale - torch.testing.assert_close(actual, expected, rtol=0, atol=0) - assert len(packed_weight_layout) == (2 * (experts or 1) if packed else 0) - - -@pytest.mark.skipif(not deepgemm.HAS_DEEPGEMM, reason="requires DeepGEMM") -@pytest.mark.parametrize("need_trans", [False, True]) -@pytest.mark.parametrize("k", [512, 640]) -def test_tp_moe_consumes_packed_weight_scales(monkeypatch, packed_weight_layout, need_trans, k): - from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe import grouped_matmul - - monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", "1") - monkeypatch.setattr(deepgemm, "is_sm100_gpu", lambda: True) - method = Quantcfg({"n_layer": 1}, quant_type="fp8w8a8-b128-deepgemm").get_quant_method(0, "q_proj") - full, _ = method.create_moe_weight([256], k, torch.bfloat16, 0, 2) - torch.manual_seed(20261008) - raw = torch.randn((2, 256, k), device="cuda", dtype=torch.bfloat16) - q, scales = weight.weight_quant(raw, use_ue8m0_scales=True) - method.load_weight(raw, full) - x = torch.randn((17, k), device="cuda", dtype=torch.bfloat16) - counts = torch.tensor([9, 8], device="cuda", dtype=torch.int32) - indices = torch.zeros((2, 17), device="cuda", dtype=torch.int32) - indices[0, :9] = torch.arange(9, device="cuda") - indices[1, :8] = torch.arange(9, 17, device="cuda") - outputs = [] - for sf in (scales, full.weight_scale): - out = torch.empty((17, 256), device="cuda", dtype=torch.bfloat16) - grouped_matmul( - 17, - x, - None, - counts, - indices, - torch.ones((2, 17), device="cuda"), - q, - sf, - 1, - out, - mul_routed_weight=True, - use_fp8_w8a8=True, - run_config=dict( - BLOCK_SIZE_M=16, - BLOCK_SIZE_N=32, - BLOCK_SIZE_K=64, - GROUP_SIZE_M=1, - num_warps=4, - num_stages=2, - NEED_TRANS=need_trans, - ), - ) - outputs.append(out) - torch.testing.assert_close(outputs[0], outputs[1], rtol=0, atol=0) - - -@pytest.mark.skipif(not deepgemm.HAS_DEEPGEMM, reason="requires DeepGEMM") -@pytest.mark.parametrize("prequantized", [False, True]) -@pytest.mark.parametrize("env_value", [None, "1", "0"]) -@pytest.mark.parametrize("scale_fmt", ["no_config", None, "ue8m0", "float32"]) -@pytest.mark.parametrize("rows", [1, 17, 32]) -def test_deepgemm_load_and_apply(monkeypatch, rows, scale_fmt, env_value, prequantized): - monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") - if env_value is None: - monkeypatch.delenv("LIGHTLLM_USE_UE8M0_SCALES", raising=False) - else: - monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", env_value) - torch.manual_seed(20261008) - x = torch.randn(rows, 1024, device="cuda", dtype=torch.bfloat16) - w = torch.randn(256, 1024, device="cuda", dtype=torch.bfloat16) - config = {"n_layer": 1} - if scale_fmt != "no_config": - config["quantization_config"] = {"quant_method": "fp8", "weight_block_size": [128, 128]} - if scale_fmt is not None: - config["quantization_config"]["scale_fmt"] = scale_fmt - method = Quantcfg(config, quant_type="fp8w8a8-b128-deepgemm").get_quant_method(0, "q_proj") - use_ue8m0_scales = scale_fmt == "ue8m0" or env_value == "1" - weight_pack, _ = method.create_weight([256], 1024, torch.bfloat16, 0) - assert method.use_ue8m0_scales == use_ue8m0_scales - assert method.use_packed_ue8m0 == (use_ue8m0_scales and deepgemm.is_sm100_gpu()) - monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", "0" if use_ue8m0_scales else "1") - if prequantized: - qweight, scales = weight.weight_quant(w, use_ue8m0_scales=use_ue8m0_scales) - method.load_weight(qweight, weight_pack) - method.load_weight_scale(scales, weight_pack) - else: - method.load_weight(w, weight_pack) - assert method.use_ue8m0_scales == use_ue8m0_scales - weight_scales = weight_pack.weight_scale - if method.use_packed_ue8m0: - weight_scales = _unpack_scales(weight_scales, 8)[::128] - assert torch.all(weight_scales > 0) - log_scales = torch.log2(weight_scales) - assert torch.equal(log_scales, log_scales.round()) == use_ue8m0_scales - - quantize_activation = deepgemm.per_token_group_quant_fp8 - - def check_activation_scales(*args, **kwargs): - result = quantize_activation(*args, **kwargs) - packed = use_ue8m0_scales and deepgemm.is_sm100_gpu() - assert kwargs["use_packed_ue8m0"] == packed - assert result[1].dtype == (torch.int32 if packed else torch.float32) - if not packed: - log_scales = torch.log2(result[1]) - assert torch.equal(log_scales, log_scales.round()) == use_ue8m0_scales - return result - - monkeypatch.setattr(deepgemm, "per_token_group_quant_fp8", check_activation_scales) - - out = method.apply(x, weight_pack, use_custom_tensor_mananger=False) - assert method.use_ue8m0_scales == use_ue8m0_scales - # Later calls reuse the choice made when creating the weight. - monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", "0" if use_ue8m0_scales else "1") - method.apply(x, weight_pack, out=out, use_custom_tensor_mananger=False) - reference = x.float() @ w.float().T - nrmse = (out.float() - reference).square().mean().sqrt() / reference.square().mean().sqrt() - cosine = F.cosine_similarity(out.float().flatten(), reference.flatten(), dim=0) - assert torch.isfinite(out).all() - assert nrmse.item() < 0.08 - assert cosine.item() > 0.995 - - -@pytest.mark.skipif(not deepgemm.HAS_DEEPGEMM, reason="requires DeepGEMM") -@pytest.mark.parametrize("prepacked", [False, True]) -@pytest.mark.parametrize("sm100", [False, True]) -def test_redundancy_expert_update_packs_scales(monkeypatch, packed_weight_layout, prepacked, sm100): - from types import SimpleNamespace - from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.ep_redundancy import ( - FusedMoeWeightEPAutoRedundancy, - ) - - monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", "1") - monkeypatch.setattr(deepgemm, "is_sm100_gpu", lambda: sm100) - method = Quantcfg({"n_layer": 1}, quant_type="fp8w8a8-b128-deepgemm").get_quant_method(0, "q_proj") - w13, _ = method.create_moe_weight([128, 128], 512, torch.bfloat16, 0, 4) - w2, _ = method.create_moe_weight([512], 128, torch.bfloat16, 0, 4) - update = object.__new__(FusedMoeWeightEPAutoRedundancy) - update.redundancy_expert_num = 2 - update.redundancy_expert_ids = [6, 7] - update._ep_w = SimpleNamespace( - w13=w13, - w2=w2, - quant_method=method, - redundancy_expert_ids_tensor=torch.zeros(2, device="cuda", dtype=torch.int64), - ) - expected = [] - for dest in (w13, w2): - dest.weight.zero_() - dest.weight_scale.zero_() - shape = (2,) + dest.weight.shape[-2:] - raw = torch.randn(shape, device="cuda", dtype=torch.bfloat16) - q, sf = weight.weight_quant(raw, use_ue8m0_scales=True) - if prepacked: - source, _ = method.create_moe_weight([shape[-2]], shape[-1], torch.bfloat16, 0, 2) - method.load_weight(raw, source) - sf = source.weight_scale - expected.append([q, sf]) - update.w13, update.w2 = expected - before = len(packed_weight_layout) - update.commit() - assert len(packed_weight_layout) == before + (2 if sm100 and not prepacked else 0) - for dest, (q, sf) in zip((w13, w2), expected): - assert torch.count_nonzero(dest.weight[:2].float()) == 0 - assert torch.count_nonzero(dest.weight_scale[:2]) == 0 - torch.testing.assert_close(dest.weight[2:].float(), q.float(), rtol=0, atol=0) - if sm100: - groups = (q.shape[-1] + 127) // 128 - actual = _unpack_scales(dest.weight_scale[2:].reshape(-1, dest.weight_scale.shape[-1]), groups) - expected_sf = ( - _unpack_scales(sf.reshape(-1, sf.shape[-1]), groups) - if prepacked - else sf.repeat_interleave(128, dim=-2).reshape(-1, groups) - ) - else: - actual, expected_sf = dest.weight_scale[2:], sf - torch.testing.assert_close(actual, expected_sf, rtol=0, atol=0) - assert update._ep_w.redundancy_expert_ids_tensor.tolist() == [6, 7] - - -@pytest.mark.skipif(not deepgemm.HAS_DEEPGEMM, reason="requires DeepGEMM") -@pytest.mark.parametrize("is_prefill", [False, True]) -@pytest.mark.parametrize("sm100", [False, True]) -@pytest.mark.parametrize("scale_source", ["none", "env", "config"]) -def test_ep_moe_rounds_scales_and_preserves_logical_block_size(monkeypatch, is_prefill, sm100, scale_source): - from lightllm.common.basemodel.triton_kernel.fused_moe import grouped_fused_moe_ep as ep - from lightllm.utils.envs_utils import get_env_use_ue8m0_scales - - monkeypatch.delenv("LIGHTLLM_USE_UE8M0_SCALES", raising=False) - assert not get_env_use_ue8m0_scales() - config = {"n_layer": 1} - if scale_source == "env": - monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", "1") - elif scale_source == "config": - config["quantization_config"] = {"quant_method": "fp8", "weight_block_size": [128, 128], "scale_fmt": "ue8m0"} - use_ue8m0_scales = scale_source != "none" - monkeypatch.setattr(deepgemm, "is_sm100_gpu", lambda: sm100) - method = Quantcfg(config, quant_type="fp8w8a8-b128-deepgemm").get_quant_method(0, "q_proj") - assert get_env_use_ue8m0_scales() == use_ue8m0_scales - w13, _ = method.create_moe_weight([128, 128], 512, torch.bfloat16, 0, 2) - w2, _ = method.create_moe_weight([512], 128, torch.bfloat16, 0, 2) - x = torch.randn((3, 512), device="cuda", dtype=torch.bfloat16) - q, sf = ep.quantize_fused_experts_input(x, w13, method) - assert q.shape == x.shape and sf.shape == (3, 4) - assert sf.dtype == torch.float32 - expected_scales = x.float().reshape(3, 4, 128).abs().amax(-1).clamp_min(1e-10) / 448 - if use_ue8m0_scales: - expected_scales = torch.exp2(torch.ceil(torch.log2(expected_scales))) - expected_q = (x.float().reshape(3, 4, 128) / expected_scales.unsqueeze(-1)).to(q.dtype).reshape_as(q) - torch.testing.assert_close(q.float(), expected_q.float(), atol=0, rtol=0) - torch.testing.assert_close(sf, expected_scales) - seen = [] - monkeypatch.setattr(ep.dist_group_manager, "ep_buffer", object(), raising=False) - - def check_dispatch(**kwargs): - seen.append(kwargs["block_size_k"]) - assert ( - kwargs["w1_scale"].dtype - == kwargs["w2_scale"].dtype - == (torch.int32 if sm100 and use_ue8m0_scales else torch.float32) - ) - return x - - monkeypatch.setattr(ep.prefill_decode_sync, "fused_experts_impl", check_dispatch) - ep.fused_experts( - x, - w13, - w2, - torch.ones((3, 1), device="cuda"), - torch.zeros((3, 1), device="cuda", dtype=torch.int64), - num_experts=2, - quant_method=method, - is_prefill=is_prefill, - ) - assert seen == [128] diff --git a/unit_tests/common/test_autotuner.py b/unit_tests/common/test_autotuner.py deleted file mode 100644 index 547246cfe5..0000000000 --- a/unit_tests/common/test_autotuner.py +++ /dev/null @@ -1,486 +0,0 @@ -import json -from contextlib import nullcontext - -import pytest -import torch -from frozendict import frozendict - -from lightllm.common.triton_utils import autotuner as autotuner_module -from lightllm.common.triton_utils.autotuner import AutotuneKernelType, AutotuneLevel, Autotuner, autotune - - -@pytest.fixture(autouse=True) -def autotune_environment(monkeypatch): - monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", None) - monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: AutotuneLevel.ADAPTIVE_AUTOTUNE) - monkeypatch.setattr(autotuner_module.dist, "is_initialized", lambda: False) - monkeypatch.setattr(autotuner_module.KernelConfigs, "get_config_file_name", lambda params: "configs.json") - - -def make_kernel(tmp_path, monkeypatch, name, kernel_type=None, warmup_all_exist_config=True): - calls = [] - benchmarks = [] - options = {} if kernel_type is None else {"kernel_type": kernel_type} - - @autotune( - kernel_name=name, - configs_gen_func=lambda: [{"block": 1}, {"block": 2}], - static_key_func=lambda: {"dtype": "test"}, - run_key_func=lambda size: size, - warmup_all_exist_config=warmup_all_exist_config, - **options, - ) - def kernel(size, run_config=None): - calls.append((size, run_config)) - return run_config - - def bench(size, run_config): - benchmarks.append((size, run_config)) - return 1.0 / run_config["block"] - - cache_dir = tmp_path / name - cache_dir.mkdir() - kernel._cache_dir = str(cache_dir) - monkeypatch.setattr(kernel, "_bench", bench) - return kernel, calls, benchmarks, cache_dir / "configs.json" - - -@pytest.mark.parametrize("repeat_override,retry_override", [(None, None), (8, 2)]) -def test_short_kernel_benchmark_sampling_options(monkeypatch, repeat_override, retry_override): - calls, replays = [], [] - - @autotune( - kernel_name="short_kernel_sampling", - configs_gen_func=lambda: [], - static_key_func=lambda: {}, - run_key_func=lambda size: size, - benchmark_repeats=32, - benchmark_retries=5, - ) - def kernel(size, run_config=None): - calls.append(size) - - class Stream: - def synchronize(self): - pass - - class Graph: - def replay(self): - replays.append(None) - - class Event: - def record(self): - pass - - def synchronize(self): - pass - - def elapsed_time(self, other): - return 0.032 - - monkeypatch.setattr(torch.cuda, "current_stream", Stream) - monkeypatch.setattr(torch.cuda, "Stream", Stream) - monkeypatch.setattr(torch.cuda, "CUDAGraph", Graph) - monkeypatch.setattr(torch.cuda, "graph", lambda *args, **kwargs: nullcontext()) - monkeypatch.setattr(torch.cuda, "Event", lambda **kwargs: Event()) - actual = kernel._bench(16, n_repeat=repeat_override, n_retries=retry_override) - repeats = 32 if repeat_override is None else repeat_override - retries = 5 if retry_override is None else retry_override - assert len(calls) == repeats + 1 # Initial warmup plus calls captured in the graph. - assert len(replays) == retries - assert actual == pytest.approx(0.032 / repeats) - - -@pytest.mark.parametrize("repeats,retries", [(0, 3), (3, 0)]) -def test_benchmark_sampling_requires_positive_counts(repeats, retries): - with pytest.raises(ValueError, match="must be positive"): - autotune( - kernel_name="invalid_sampling", - configs_gen_func=lambda: [], - static_key_func=lambda: {}, - run_key_func=lambda size: size, - benchmark_repeats=repeats, - benchmark_retries=retries, - )(lambda size, run_config=None: None) - - -def test_key_defaults_preserve_explicit_arguments_and_required_checks(): - @autotune( - kernel_name="key_defaults", - configs_gen_func=lambda: [], - static_key_func=lambda window: {"window": window}, - run_key_func=lambda size: size, - ) - def kernel(size, window=(-1, -1), run_config=None): - return run_config - - assert kernel._static_key(8) == {"window": (-1, -1)} - assert kernel._static_key(8, (511, 0)) == {"window": (511, 0)} - assert kernel._static_key(8, window=(1023, 0)) == {"window": (1023, 0)} - assert kernel._static_key(8, window=None) == {"window": None} - with pytest.raises(KeyError, match="Missing argument 'size'"): - kernel._run_key() - - -@pytest.mark.parametrize("level", [AutotuneLevel.ADAPTIVE_AUTOTUNE, AutotuneLevel.FORCE_AUTOTUNE]) -def test_two_warmup_phases_only_persist_matching_kernel_configs(tmp_path, monkeypatch, level): - monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: level) - general, _, general_benchmarks, general_cache = make_kernel(tmp_path, monkeypatch, "general") - decode, decode_calls, decode_benchmarks, decode_cache = make_kernel( - tmp_path, monkeypatch, "decode", AutotuneKernelType.DECODE_ATTENTION - ) - - with Autotuner.autotune_warmup(): - assert general(8) == {"block": 2} - assert decode(8) is None - assert len(general_benchmarks) == 2 - assert decode_benchmarks == [] - assert decode_calls == [(8, None)] - assert not decode_cache.exists() - general_cache_before = general_cache.read_bytes() - - with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - assert general(16) == {"block": 2} - assert decode(16) == {"block": 2} - assert len(general_benchmarks) == 2 - assert len(decode_benchmarks) == 2 - assert general_cache.read_bytes() == general_cache_before - assert json.loads(general_cache.read_text()) == {"8": {"block": 2}} - assert json.loads(decode_cache.read_text()) == {"16": {"block": 2}} - - -def test_other_phase_warms_history_without_tuning_and_later_uses_new_config(tmp_path, monkeypatch): - decode, calls, benchmarks, cache_file = make_kernel( - tmp_path, monkeypatch, "decode", AutotuneKernelType.DECODE_ATTENTION - ) - cache_file.write_text(json.dumps({"8": {"block": 1}, "32": {"block": 2}})) - cache_before = cache_file.read_bytes() - - with Autotuner.autotune_warmup(): - assert decode(16) == {"block": 1} - # Historical configs are warmed in either phase, but only the matching phase searches new configs. - assert calls == [(16, {"block": 1}), (16, {"block": 2}), (16, {"block": 1})] - assert benchmarks == [] - assert cache_file.read_bytes() == cache_before - - with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - assert decode(16) == {"block": 2} - assert len(benchmarks) == 2 - assert decode(16) == {"block": 2} - assert json.loads(cache_file.read_text())["16"] == {"block": 2} - - -@pytest.mark.parametrize("level", [AutotuneLevel.ADAPTIVE_AUTOTUNE, AutotuneLevel.FORCE_AUTOTUNE]) -def test_excluded_kernel_does_not_enter_distributed_tuning(tmp_path, monkeypatch, level): - general, calls, benchmarks, cache_file = make_kernel(tmp_path, monkeypatch, "general") - monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: level) - monkeypatch.setattr(autotuner_module.dist, "is_initialized", lambda: True) - monkeypatch.setattr(autotuner_module, "get_global_rank", lambda: 0) - monkeypatch.setattr(autotuner_module, "get_global_world_size", lambda: 2) - - def unexpected_collective(*args, **kwargs): - pytest.fail("An excluded kernel must not enter autotuning collectives") - - monkeypatch.setattr(general, "_get_autotune_group", unexpected_collective) - monkeypatch.setattr(autotuner_module.dist, "all_gather_object", unexpected_collective) - with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - assert general(16) is None - assert calls == [(16, None)] - assert benchmarks == [] - assert not cache_file.exists() - - -@pytest.mark.parametrize("level", [0, 1, 2, 3]) -@pytest.mark.parametrize("kernel_type", [AutotuneKernelType.GENERAL, AutotuneKernelType.DECODE_ATTENTION]) -def test_matching_phase_respects_kernel_autotune_policy(tmp_path, monkeypatch, level, kernel_type): - monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: level) - kernel, _, benchmarks, cache_file = make_kernel(tmp_path, monkeypatch, kernel_type.value, kernel_type) - cache_file.write_text(json.dumps({"16": {"block": 1}})) - cache_before = cache_file.read_bytes() - with Autotuner.autotune_warmup(kernel_type): - result = kernel(16) - if level == AutotuneLevel.FORCE_AUTOTUNE and kernel_type == AutotuneKernelType.GENERAL: - assert result == {"block": 2} - assert len(benchmarks) == 2 - assert json.loads(cache_file.read_text()) == {"16": {"block": 2}} - else: - assert result == (None if level == AutotuneLevel.CLOSE_AUTOTUNE else {"block": 1}) - assert benchmarks == [] - assert cache_file.read_bytes() == cache_before - - -@pytest.mark.parametrize("kernel_type", [AutotuneKernelType.GENERAL, AutotuneKernelType.DECODE_ATTENTION]) -def test_force_autotune_reuses_decode_configs_across_layers(tmp_path, monkeypatch, kernel_type): - monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: AutotuneLevel.FORCE_AUTOTUNE) - kernel, _, benchmarks, cache_file = make_kernel(tmp_path, monkeypatch, kernel_type.value, kernel_type) - with Autotuner.autotune_warmup(kernel_type): - for size in [8, 16]: - benchmarks.clear() - assert kernel(size) == {"block": 2} - assert len(benchmarks) == 2 # A new run key still needs tuning. - for _ in range(3): - assert kernel(size) == {"block": 2} - assert len(benchmarks) == (2 if kernel_type == AutotuneKernelType.DECODE_ATTENTION else 8) - assert json.loads(cache_file.read_text()) == {"8": {"block": 2}, "16": {"block": 2}} - - -def test_history_is_warmed_on_load_or_during_autotune_warmup(tmp_path, monkeypatch): - kernel, calls, benchmarks, cache_file = make_kernel(tmp_path, monkeypatch, "general") - cache_file.write_text(json.dumps({"8": {"block": 1}, "32": {"block": 2}})) - assert kernel(16) == {"block": 1} - assert calls == [(16, {"block": 1}), (16, {"block": 2}), (16, {"block": 1})] - calls.clear() - with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - assert kernel(16) == {"block": 1} - assert calls == [(16, {"block": 1})] - calls.clear() - assert kernel(16) == {"block": 1} - assert calls == [(16, {"block": 1})] - assert benchmarks == [] - - -def test_repeated_configs_are_skipped_after_success(tmp_path, monkeypatch): - kernel, calls, benchmarks, cache_file = make_kernel(tmp_path, monkeypatch, "general") - config = {"block": 1, "warps": 4} - cache_file.write_text(json.dumps({"8": config, "16": config, "32": {"block": 2}})) - with Autotuner.autotune_warmup(): - assert kernel(16) == config - assert calls == [(16, config), (16, {"block": 2}), (16, config)] - - calls.clear() - with Autotuner.autotune_warmup(): - assert kernel(16) == config - with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - assert kernel(16) == config - assert kernel(16) == config - assert calls == [(16, config)] * 3 - assert benchmarks == [] - - -@pytest.mark.parametrize("level", [0, 1, 2]) -@pytest.mark.parametrize("phase", [None, AutotuneKernelType.GENERAL, AutotuneKernelType.DECODE_ATTENTION]) -def test_disabled_history_warmup_loads_cache_without_extra_state_updates(tmp_path, monkeypatch, level, phase): - monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: level) - calls = [] - - @autotune( - kernel_name="stateful_kernel", - kernel_type=AutotuneKernelType.DECODE_ATTENTION, - configs_gen_func=lambda: [{"block": 1}, {"block": 2}], - static_key_func=lambda: {}, - run_key_func=lambda state: state.numel(), - warmup_all_exist_config=False, - ) - def kernel(state, run_config=None): - calls.append(run_config) - state.add_(run_config["block"]) - return state - - kernel._cache_dir = str(tmp_path) - cache_file = tmp_path / "configs.json" - cache_file.write_text(json.dumps({"4": {"block": 1}, "8": {"block": 2}})) - cache_before = cache_file.read_bytes() - - def unexpected_warmup(*args, **kwargs): - pytest.fail("Disabling history warmup must skip its execution and argument cloning") - - monkeypatch.setattr(kernel, "kernel_warmup", unexpected_warmup) - monkeypatch.setattr(kernel, "_mutate_args_clone", unexpected_warmup) - state = torch.zeros(4) - scope = nullcontext() if phase is None else Autotuner.autotune_warmup(phase) - with scope: - for _ in range(2): - assert kernel(state) is state - assert calls == [{"block": 1}, {"block": 1}] - torch.testing.assert_close(state, torch.full((4,), 2.0)) - assert cache_file.read_bytes() == cache_before - assert not kernel.warmuped_configs_set - - -@pytest.mark.parametrize("kernel_type", [AutotuneKernelType.GENERAL, AutotuneKernelType.DECODE_ATTENTION]) -def test_disabled_history_warmup_still_searches_new_configs(tmp_path, monkeypatch, kernel_type): - kernel, calls, benchmarks, cache_file = make_kernel( - tmp_path, monkeypatch, kernel_type.value, kernel_type, warmup_all_exist_config=False - ) - with Autotuner.autotune_warmup(kernel_type): - for size in [8, 16]: - calls.clear() - assert kernel(size) == {"block": 2} - assert calls == [(size, {"block": 2})] - calls.clear() - assert kernel(size) == {"block": 2} - assert calls == [(size, {"block": 2})] - assert len(benchmarks) == 4 - assert json.loads(cache_file.read_text()) == {"8": {"block": 2}, "16": {"block": 2}} - assert not kernel.warmuped_configs_set - - -def test_failed_warmup_retries_only_during_autotune_warmup(tmp_path, monkeypatch): - kernel, calls, _, cache_file = make_kernel(tmp_path, monkeypatch, "general") - cache_file.write_text(json.dumps({"8": {"block": 1}, "32": {"block": 2}})) - - def shape_sensitive_kernel(size, run_config=None): - calls.append((size, run_config)) - if size < 16 and run_config["block"] == 2: - raise RuntimeError("This config requires a larger input") - return run_config - - monkeypatch.setattr(kernel, "fn", shape_sensitive_kernel) - with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - assert kernel(8) == {"block": 1} - assert calls == [(8, {"block": 1}), (8, {"block": 2}), (8, {"block": 1})] - calls.clear() - assert kernel(16) == {"block": 1} - assert calls == [(16, {"block": 1})] - calls.clear() - with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - assert kernel(16) == {"block": 1} - assert calls == [(16, {"block": 2}), (16, {"block": 1})] - calls.clear() - with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - assert kernel(16) == {"block": 1} - assert calls == [(16, {"block": 1})] - - -def test_new_configs_are_warmed_on_later_calls_and_state_is_not_persisted(tmp_path, monkeypatch): - kernel, calls, benchmarks, cache_file = make_kernel(tmp_path, monkeypatch, "general") - cache_file.write_text(json.dumps({"8": {"block": 1}})) - with Autotuner.autotune_warmup(): - assert kernel(8) == {"block": 1} - calls.clear() - with Autotuner.autotune_warmup(): - assert kernel(16) == {"block": 2} - assert calls == [(16, {"block": 2})] - assert len(benchmarks) == 2 - assert json.loads(cache_file.read_text()) == {"8": {"block": 1}, "16": {"block": 2}} - - calls.clear() - assert kernel(16) == {"block": 2} - assert calls == [(16, {"block": 2})] - calls.clear() - with Autotuner.autotune_warmup(): - assert kernel(16) == {"block": 2} - assert calls == [(16, {"block": 2}), (16, {"block": 2})] - calls.clear() - with Autotuner.autotune_warmup(): - assert kernel(16) == {"block": 2} - assert calls == [(16, {"block": 2})] - - # A new autotuner loading the same file must warm its configs again. - reloaded, reload_calls, _, _ = make_kernel(tmp_path, monkeypatch, "reloaded") - reloaded._cache_dir = str(cache_file.parent) - assert reloaded(16) == {"block": 2} - assert len(reload_calls) == 3 - assert {call[1]["block"] for call in reload_calls[:2]} == {1, 2} - reload_calls.clear() - with Autotuner.autotune_warmup(): - assert reloaded(16) == {"block": 2} - assert reload_calls == [(16, {"block": 2})] - - -def test_warmup_preserves_mutated_input_and_skips_repeated_execution(): - executions = [] - - @autotune( - kernel_name="mutating_kernel", - configs_gen_func=lambda: [{"block": 1}], - static_key_func=lambda: {}, - run_key_func=lambda state: state.numel(), - mutates_args=["state"], - ) - def kernel(state, run_config=None): - executions.append(run_config) - state.add_(run_config["block"]) - - state = torch.zeros(4) - static_key = frozendict({}) - kernel.kernel_warmup(static_key, state, run_config={"block": 1}) - torch.testing.assert_close(state, torch.zeros(4)) - - kernel.kernel_warmup(static_key, state, run_config={"block": 1}) - assert executions == [{"block": 1}] - torch.testing.assert_close(state, torch.zeros(4)) - - -def test_explicit_config_bypasses_tuning(tmp_path, monkeypatch): - kernel, calls, benchmarks, cache_file = make_kernel(tmp_path, monkeypatch, "general") - with Autotuner.autotune_warmup(): - assert kernel(16, run_config={"block": 3}) == {"block": 3} - assert calls == [(16, {"block": 3})] - assert benchmarks == [] - assert not cache_file.exists() - - -@pytest.mark.parametrize("keyword_input", [False, True]) -@pytest.mark.parametrize("level", [AutotuneLevel.ADAPTIVE_AUTOTUNE, AutotuneLevel.FORCE_AUTOTUNE]) -def test_rebuild_inputs_only_for_benchmarking(tmp_path, monkeypatch, keyword_input, level): - monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: level) - rebuilds = [] - executions = [] - benchmarks = [] - original = torch.zeros(2) - - def rebuild(state, size, run_config=None): - rebuilds.append(state) - return (torch.ones(size),), {"size": size} - - @autotune( - kernel_name="rebuild_inputs", - kernel_type=AutotuneKernelType.DECODE_ATTENTION, - configs_gen_func=lambda: [{"block": 1}, {"block": 2}], - static_key_func=lambda state: {"input_size": state.numel()}, - run_key_func=lambda size: size, - rebuild_input_func=rebuild, - ) - def kernel(state, size, run_config=None): - executions.append(state) - return state - - def bench(state, size, run_config): - benchmarks.append(state) - torch.testing.assert_close(state, torch.ones(8)) - assert size == 8 - return 1.0 / run_config["block"] - - kernel._cache_dir = str(tmp_path) - monkeypatch.setattr(kernel, "_bench", bench) - args, kwargs = ((), {"state": original, "size": 8}) if keyword_input else ((original, 8), {}) - - assert kernel(*args, **kwargs) is original - with Autotuner.autotune_warmup(): - assert kernel(*args, **kwargs) is original - with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - assert kernel(*args, **kwargs, run_config={"block": 3}) is original - assert not rebuilds - assert kernel(*args, **kwargs) is original - assert kernel(*args, **kwargs) is original # Exact cache hit only warms historical configs. - - assert kernel(*args, **kwargs) is original - assert len(rebuilds) == 1 and rebuilds[0] is original - assert len(benchmarks) == 2 and benchmarks[0] is benchmarks[1] - assert all(state is original for state in executions) - assert kernel.cached_configs == {frozendict({"input_size": 2}): {"8": {"block": 2}}} - torch.testing.assert_close(original, torch.zeros(2)) - - -def test_default_api_and_nested_phase_restore_after_exception(): - assert not Autotuner.is_autotune_warmup() - assert not Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.GENERAL) - assert not Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.DECODE_ATTENTION) - Autotuner.start_autotune_warmup() - assert Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.GENERAL) - assert not Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.DECODE_ATTENTION) - with pytest.raises(RuntimeError, match="test failure"): - with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - assert Autotuner.is_autotune_warmup() - assert Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.DECODE_ATTENTION) - assert not Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.GENERAL) - with Autotuner.autotune_warmup(): - assert Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.GENERAL) - assert Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.DECODE_ATTENTION) - raise RuntimeError("test failure") - assert Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.GENERAL) - Autotuner.end_autotune_warmup() - assert not Autotuner.is_autotune_warmup() - assert not Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.GENERAL) - assert not Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.DECODE_ATTENTION) diff --git a/unit_tests/common/test_cache_copy_control_flow.py b/unit_tests/common/test_cache_copy_control_flow.py deleted file mode 100644 index 26fe317052..0000000000 --- a/unit_tests/common/test_cache_copy_control_flow.py +++ /dev/null @@ -1,98 +0,0 @@ -import ast -import __future__ -from pathlib import Path -import sys -from types import SimpleNamespace -from unittest.mock import Mock - -import pytest - - -ROOT = Path(__file__).resolve().parents[2] - - -def load_methods(path, class_name, method_names, optimize=0): - tree = ast.parse((ROOT / path).read_text()) - cls = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == class_name) - cls.body = [node for node in cls.body if isinstance(node, ast.FunctionDef) and node.name in method_names] - module = ast.Module(body=[cls], type_ignores=[]) - namespace = {"logger": Mock()} - exec(compile(module, path, "exec", flags=__future__.annotations.compiler_flag, optimize=optimize), namespace) - return namespace[class_name] - - -@pytest.mark.parametrize("optimize", [0, 1, 2]) -@pytest.mark.parametrize("prefix", [0, 64]) -def test_cpu_cache_restore_allocates_before_read(optimize, prefix): - cache_class = load_methods( - "lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py", - "MultiLevelKvCacheModule", - {"read_req_from_cpu"}, - optimize, - ) - req = SimpleNamespace(cur_kv_len=prefix, hold_kv_len=prefix) - indices = list(range(prefix)) - - def allocate(count): - indices.extend(range(req.hold_kv_len, req.hold_kv_len + count)) - req.hold_kv_len += count - return indices[-count:] - - def get_indices(start, end): - selected = indices[start:end] - return SimpleNamespace(values=selected, numel=lambda: len(selected)) - - req.kv_cache_alloc_need = Mock(side_effect=lambda length: length - req.hold_kv_len) - req.alloc_kv_mem = Mock(side_effect=allocate) - req.token_index = SimpleNamespace(get_gpu_token_indexes=get_indices) - operator = SimpleNamespace(read_req=Mock(), get_page_layouts=lambda: []) - cache = cache_class() - cache.backend = SimpleNamespace( - model=SimpleNamespace(att_state_manager=SimpleNamespace(get_cpu_cache_operator=lambda: operator)), - is_hybrid_att_model=False, - ) - cache.cpu_cache_client = SimpleNamespace(cpu_kv_cache_tensor=object()) - - cache.read_req_from_cpu(req, [0], [256], 0) - - req.alloc_kv_mem.assert_called_once_with(256 - prefix) - operator.read_req.assert_called_once() - assert operator.read_req.call_args.kwargs["token_indices"].values == list(range(256)) - assert req.cur_kv_len == req.hold_kv_len == 256 - - -@pytest.mark.parametrize("return_fail_index", [False, True]) -@pytest.mark.parametrize("success", [False, True]) -def test_native_batch_return_formats(monkeypatch, return_fail_index, success): - copier = load_methods( - "lightllm/common/cuda_wrapper.py", - "PinnedH2DBatchCopier", - {"_copy_on_stream", "_try_native_batch_h2d"}, - ) - copier._native_batch_available = True - copier._protect_pinned_sources = Mock() - error = 0 if success else SimpleNamespace(name="cudaErrorInvalidValue") - result = (error, 0) if return_fail_index else (error,) - runtime = SimpleNamespace( - cudaMemcpyAttributes=SimpleNamespace, - cudaMemcpySrcAccessOrder=SimpleNamespace(cudaMemcpySrcAccessOrderStream=1), - cudaError_t=SimpleNamespace(cudaSuccess=0), - cudaMemcpyBatchAsync=Mock(return_value=result), - ) - monkeypatch.setitem(sys.modules, "cuda", SimpleNamespace()) - monkeypatch.setitem(sys.modules, "cuda.bindings", SimpleNamespace(runtime=runtime)) - source = SimpleNamespace(data_ptr=lambda: 100, numel=lambda: 256, element_size=lambda: 4) - target = SimpleNamespace(data_ptr=lambda: 200, copy_=Mock()) - stream = SimpleNamespace(cuda_stream=7) - - copier._copy_on_stream([target], [source], stream) - - runtime.cudaMemcpyBatchAsync.assert_called_once() - assert runtime.cudaMemcpyBatchAsync.call_args.args[:4] == ([200], [100], [1024], 1) - assert copier._native_batch_available is success - if success: - target.copy_.assert_not_called() - copier._protect_pinned_sources.assert_called_once_with([source], stream) - else: - target.copy_.assert_called_once_with(source, non_blocking=True) - copier._protect_pinned_sources.assert_not_called() diff --git a/unit_tests/common/test_deepseek4_paged_cache.py b/unit_tests/common/test_deepseek4_paged_cache.py index f5bd92e08a..06de22959d 100644 --- a/unit_tests/common/test_deepseek4_paged_cache.py +++ b/unit_tests/common/test_deepseek4_paged_cache.py @@ -86,36 +86,6 @@ def cache(monkeypatch, tmp_path, request): get_env_start_args.cache_clear() -def test_attention_input_params_select_token_buffers(cache): - manager, requests = cache - for layer_index, slot in ((0, 0), (1, 1), (2, None)): - swa, compressed_kv = requests.get_att_input_params(layer_index) - assert swa.data_ptr() == requests.req_to_swa.flatten(1, 2)[layer_index].data_ptr() - if slot is None: - assert compressed_kv is None - else: - token_buffer = manager.get_att_input_params(layer_index) - assert token_buffer is manager.kv_buffer[slot] - assert token_buffer.shape[:2] == (manager.size + manager.page_size, 1) - assert compressed_kv.untyped_storage().data_ptr() == token_buffer.data_ptr() - - -def test_big_page_pool_reserves_cpu_cache_staging_slots(cache): - manager, _ = cache - pool = manager.big_page_buffers - args = get_env_start_args() - page_tokens = args.linear_att_hash_page_size * args.linear_att_page_block_num - capacity = (manager.size + page_tokens - 1) // page_tokens - assert pool.size == capacity + 2 and pool.keep_num == 2 - assert manager.CPU_CACHE_BIG_PAGE_LOAD_TEMP_BUFFER_ID == capacity - assert manager.CPU_CACHE_BIG_PAGE_OFFLOAD_TEMP_BUFFER_ID == capacity + 1 - slots = pool.alloc_state_cache(capacity) - assert slots == list(range(capacity)) - assert pool.alloc_one_state_cache() is None - pool.free_state_cache(slots) - assert pool.get_free_cache_num() == capacity - - def test_copy_mem_to_mem_preserves_complete_block_indices(cache): manager, _ = cache # 压缩块内的 KV/scale/indexer 不可拆分；允许按非连续、乱序的完整块复制。 @@ -160,47 +130,6 @@ def test_cache_io_copies_individual_token_slots(cache): assert buffer[destination[2:]].eq(0).all() -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, - ) - ) - # 扩展压缩 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]) -def test_pd_memory_reservation_matches_allocated_bytes(cache, monkeypatch, page_size): - manager, _ = cache - args = get_env_start_args() - monkeypatch.setattr(args, "run_mode", "decode") - monkeypatch.setattr(args, "pd_kv_page_num", 3, raising=False) - monkeypatch.setattr(args, "pd_kv_page_size", page_size, raising=False) - reserved = manager.get_pd_kv_move_buffer_size() - buffer = manager.alloc_paged_kv_move_buffer(3, page_size) - assert reserved == buffer.nbytes - assert buffer.shape == manager.get_paged_kv_move_buffer_shape(3, page_size) - - def _packed_swa(buffer): return buffer.view(buffer.shape[0], -1, 74880) if buffer.shape[-1] == 585 else buffer @@ -373,173 +302,6 @@ def _reserve(manager, requests, length): return req_idx, slots -@pytest.mark.parametrize("ratio", [4, 128]) -@pytest.mark.parametrize("cache", [256, 512, 768, 1024], indirect=True) -def test_noncontiguous_pages_obey_logical_compression_boundaries(cache, ratio): - manager, requests = cache - page_size = manager.page_size - req_idx, allocated = _reserve(manager, requests, 3 * page_size) - slots = allocated.view(3, page_size)[torch.tensor([2, 0, 1], device="cuda")].reshape(-1) - requests.req_to_token_indexs[req_idx, : slots.numel()] = slots - positions = torch.tensor( - [0, 2, 3, 126, 127, 254, 255, page_size - 1, page_size, 2 * page_size - 1, 2 * page_size], device="cuda" - ) - indexes = torch.empty((len(positions), slots.numel() // ratio), dtype=torch.int32, device="cuda") - lengths = torch.empty_like(positions, dtype=torch.int32) - build_compressed_kv_indices( - torch.full_like(positions, req_idx), positions, requests.req_to_token_indexs, ratio, indexes, lengths - ) - group_slots = slots[ratio - 1 :: ratio] // ratio - for row, position in enumerate(positions.tolist()): - count = (position + 1) // ratio - torch.testing.assert_close(indexes[row, :count], group_slots[:count]) - assert indexes[row, count:].eq(-1).all() - assert lengths[row].item() == max(1, count) - manager.free(allocated) - assert manager.allocator.can_use_mem_size == manager.size - - -def test_indexer_writer_only_writes_closed_groups(cache): - manager, requests = cache - _, slots = _reserve(manager, requests, 512) - slots = slots.view(2, 256).flip(0).reshape(-1) - keys = torch.randn((259, 128), dtype=torch.bfloat16, device="cuda") - positions = torch.arange(259, dtype=torch.int32, device="cuda") - _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, 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. - selected_keys = keys[closed].float() - scale = selected_keys.abs().amax(dim=-1, keepdim=True).clamp_min(1e-4) / 448.0 - quantized = (selected_keys / scale).clamp(-448.0, 448.0).to(torch.float8_e4m3fn) - expected = torch.cat((quantized.view(torch.uint8), scale.view(torch.uint8)), dim=-1) - torch.testing.assert_close(actual, expected, rtol=0, atol=0) - # Positions 256..258 cannot publish their reserved next compressed slot. - assert _read_packed_rows(_compressed_kv_pools(manager)[1], 0, slots[256:257] // 4).eq(0).all() - - -def test_speculative_retry_reuses_swa_without_freeing_token_pages(cache): - manager, requests = cache - req_idx, slots = _reserve(manager, requests, 512) - - def cpu(data): - return torch.tensor(data, dtype=torch.int32) - - _prepare_prefill(requests, cpu([req_idx]), cpu([0]), cpu([254])) - for sequences in ([255, 256, 257, 258], [255, 256, 257, 258], [256, 257, 258, 259]): - swa_slots = _swa_slots(requests, req_idx, torch.tensor(sequences, device="cuda") - 1) - assert swa_slots.unique().numel() == len(sequences) - assert manager.allocator.can_use_mem_size == manager.size - 512 - requests.free([req_idx], slots) - assert manager.allocator.can_use_mem_size == manager.size - - -@pytest.mark.parametrize("model_kind", ["target", "mtp", "dspark"]) -def test_model_selects_held_slots_without_allocating_swa(cache, model_kind): - from lightllm.common.basemodel.batch_objs import ModelInput - from lightllm.models.deepseek_v4.model import DeepseekV4TpPartModel - from lightllm.models.deepseek_v4_mtp.model import DeepseekV4MTPModel - from lightllm.models.deepseek_v4_dspark.model import DeepseekV4DSparkModel - - manager, requests = cache - req_idx, slots = _reserve(manager, requests, 512) - _prepare_swa(requests, req_idx, 0, 256) - model_cls = {"target": DeepseekV4TpPartModel, "mtp": DeepseekV4MTPModel, "dspark": DeepseekV4DSparkModel}[ - 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, - max_q_seq_len=1, - max_kv_seq_len=257, - b_req_idx=torch.tensor([req_idx], dtype=torch.int32), - b_seq_len=torch.tensor([257], dtype=torch.int32), - b_mtp_index=torch.zeros(1, dtype=torch.int32), - b_position_delta=torch.zeros(1, dtype=torch.int32), - b_shared_seq_len=torch.zeros(1, dtype=torch.int32), - b_shared_radix_node_id=torch.full((1,), -1, dtype=torch.int64), - multimodal_params=[{"images": [], "audios": []}], - ) - model_input.to_cuda() - for _ in range(2): - torch.testing.assert_close(model._select_mem_indexes(model_input), slots[256:257]) - assert manager.allocator.can_use_mem_size == manager.size - 512 - requests.free([req_idx], slots) - - -@pytest.mark.parametrize("cache", [(256, 0), (256, 5)], indirect=True) -def test_prefill_chunk_keeps_all_new_swa_rows(cache): - manager, requests = cache - req_idx, slots = _reserve(manager, requests, 2048) - _prepare_prefill(requests, torch.tensor([req_idx]), torch.tensor([0]), torch.tensor([2048])) - positions = torch.arange(2048, device="cuda") - swa_slots = _swa_slots(requests, req_idx, positions) - assert swa_slots.unique().numel() == 2048 - packed = torch.randint(0, 256, (2048, 584), dtype=torch.uint8, device="cuda") - _write_swa(requests, 0, req_idx, positions, packed) - # Checkpoints can still read any boundary of the complete chunk here. - torch.testing.assert_close(_read_packed_rows(_swa_buffer(requests, req_idx), 0, swa_slots), packed) - _finish_prefill(requests, [req_idx]) - tail = positions[-requests.sliding_window :] - torch.testing.assert_close(_read_swa(requests, 0, req_idx, tail), packed[tail]) - requests.free([req_idx], slots) - - -def test_prefill_checkpoint_reads_precede_window_retention(cache): - """The chunk-end checkpoint and late draft KV must survive a long prefill.""" - manager, requests = cache - req, slots = _reserve(manager, requests, 2048) - _prepare_swa(requests, req, 0, 2048) - positions = torch.arange(2048, device="cuda") - packed = torch.randint(0, 256, (2048, 584), dtype=torch.uint8, device="cuda") - _write_swa(requests, 0, req, positions, packed) - _state_buffer(requests, "c4").uniform_() - _state_buffer(requests, "c4_indexer").uniform_() - live_tail = requests.req_to_c4_state.clone() - states = requests.create_small_page_cache_manager(1) - index = states.alloc_one_state_cache() - requests.save_state(_checkpoint_req(req, 2048), index, states) - window, _, _ = states.get_state_cache(index) - torch.testing.assert_close(window[0].cuda(), packed[-requests.sliding_window :], rtol=0, atol=0) - for boundary in range(256, 2049, 256): - torch.testing.assert_close( - _read_swa(requests, 0, req, positions[boundary - requests.sliding_window : boundary]), - packed[boundary - requests.sliding_window : boundary], - rtol=0, - atol=0, - ) - - # A DSpark context commit writes KV after the target model has returned. - _write_swa(requests, 2, req, positions, packed.flip(0)) - _finish_prefill(requests, [req]) - # Reuse the scratch: only fixed request windows and CPU snapshots remain. - requests._test_backend.get_swa_prefill_workspace(requests._test_backend.model)[0].zero_() - requests.save_state(_checkpoint_req(req, 2048), index, states) - saved_swa = states.get_state_cache(index)[0] - torch.testing.assert_close(saved_swa[0].cuda(), packed[-requests.sliding_window :], rtol=0, atol=0) - tail = positions[-requests.sliding_window :] - for layer, expected in ((0, packed), (2, packed.flip(0))): - torch.testing.assert_close( - _read_swa(requests, layer, req, tail), - expected[tail], - rtol=0, - atol=0, - ) - torch.testing.assert_close(requests.req_to_c4_state, live_tail) - requests.free([req], slots) - states.free_state_cache([index]) - - def test_non_aligned_chunks_preserve_window_and_checkpoint_history(cache): manager, requests = cache req, slots = _reserve(manager, requests, 3072) @@ -564,51 +326,6 @@ def test_non_aligned_chunks_preserve_window_and_checkpoint_history(cache): requests.free([req], slots) -@pytest.mark.parametrize("cache", [256, 512, 768, 1024], indirect=True) -def test_compressed_kv_copy_preserves_packed_data_and_scale_regions(cache): - manager, requests = cache - source, src = _reserve(manager, requests, 2 * manager.page_size) - destination, dst = _reserve(manager, requests, 2 * manager.page_size) - requests.req_to_swa.random_(0, 256) - swa_state = requests.req_to_swa.clone() - src = src.view(2, manager.page_size).flip(0).reshape(-1) - for pool in _compressed_kv_pools(manager): - pool.random_(0, 256) - manager.operator.copy_mem_to_mem(src_mem_index=src, dst_mem_index=dst) - for pool in _compressed_kv_pools(manager): - torch.testing.assert_close(pool[:, src[::256].long() // 256], pool[:, dst[::256].long() // 256], rtol=0, atol=0) - torch.testing.assert_close(requests.req_to_swa, swa_state) - requests.free([source, destination], torch.cat([src, dst])) - - -@pytest.mark.parametrize("cache", [256, 512, 768, 1024], indirect=True) -def test_hold_page_is_separate_from_request_rings(cache): - manager, requests = cache - hold = requests.req_to_token_indexs[requests.HOLD_REQUEST_ID] - torch.testing.assert_close( - hold[: manager.page_size], - torch.arange(manager.size, manager.size + manager.page_size, dtype=torch.int32, device="cuda"), - ) - assert ( - _swa_blocks(requests, requests.HOLD_REQUEST_ID, torch.arange(32, device="cuda")) - .eq(requests.HOLD_SWA_SLOT // 128) - .all() - ) - assert (hold // 4).lt(_compressed_kv_pools(manager)[0].shape[1] * 64).all() - assert (hold // 128).lt(_compressed_kv_pools(manager)[2].shape[1] * 2).all() - free_tokens = manager.allocator.can_use_mem_size - reqs = [requests.alloc(), requests.alloc()] - assert set(_request_pages(requests, reqs[0])).isdisjoint(_request_pages(requests, reqs[1])) - assert ( - _swa_blocks(requests, torch.tensor(reqs, device="cuda")[:, None], torch.arange(32, device="cuda")[None, :]) - .lt(requests.HOLD_SWA_SLOT // 128) - .all() - ) - assert manager.allocator.can_use_mem_size == free_tokens - for req in reqs: - requests.free_req(req) - - @pytest.mark.parametrize("length", [256, 512, 2048]) def test_shared_compressed_kv_restores_private_continuation(cache, length): manager, requests = cache @@ -714,1348 +431,3 @@ def test_pd_kv_and_att_state_roundtrip(cache, length): prefill.free([source], src) requests.free([destination], dst) assert manager.allocator.can_use_mem_size == manager.size - - -@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires two CUDA devices") -@pytest.mark.parametrize("length", [257, 2051]) -def test_pd_transport_completion_covers_peer_gpu_state(cache, length): - import ctypes - - transport, _ = cache - with torch.cuda.device(1): - peer = DeepseekV4MemoryManager(8192, torch.bfloat16, 1, 512, 3, transport.cache_config) - requests = DeepseekV4ReqManager(2, 4096, None, peer.cache_config, mtp_step=get_env_start_args().mtp_step) - requests.bind_mem_manager(peer) - source, src = _reserve(peer, requests, length) - destination, dst = _reserve(peer, requests, length) - _randomize_compressed_kv(peer) - requests.req_to_swa.flatten(1, 2).random_(0, 256) - for buffer in ( - _state_buffer(requests, "c4"), - _state_buffer(requests, "c4_indexer"), - _state_buffer(requests, "c128"), - ): - buffer.uniform_() - for name in ( - "req_to_swa", - "req_to_c4_state", - "req_to_c4_indexer_state", - "req_to_c128_state", - ): - setattr(peer, name, getattr(requests, name)) - torch.cuda.synchronize() - peer_stream = torch.cuda.Stream() - consumer_stream = torch.cuda.Stream() - copy_stream = torch.cuda.Stream() - # CUDA IPC enables peer access in production. A same-process test needs - # to establish that access before Triton dereferences the peer pointer. - 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): - transport.write_mem_to_page_kv_move_buffer( - src[start : start + transport.page_size].tolist(), 0, 1, mem_managers, 1 - ) - torch.cuda.current_stream().synchronize() # The network consumes the completed packed page. - transport.read_page_kv_move_buffer_to_mem( - dst[start : start + transport.page_size].tolist(), 0, 1, mem_managers, 1 - ) - torch.cuda.current_stream().synchronize() # Return the page to the pool only after completion. - with torch.cuda.stream(peer_stream), torch.cuda.stream(copy_stream): - transport.write_mem_to_page_kv_move_buffer([], 0, 1, mem_managers, 1, "att_state", source) - with torch.cuda.device(1): - torch.cuda._sleep(20_000_000) - transport.read_page_kv_move_buffer_to_mem([], 0, 1, mem_managers, 1, "att_state", destination) - completed = torch.cuda.Event() - completed.record(copy_stream) - completed.synchronize() - # 在独立流验证逻辑窗口；无效 MTP 槽位不属于 PD 数据。 - 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), - ): - 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])) - - -@pytest.mark.parametrize("ratio,layer,is_indexer", [(4, 0, False), (128, 1, False), (4, 0, True)]) -@pytest.mark.parametrize("length", [512, 2048]) -@pytest.mark.parametrize("cache", [(256, 0), (256, 1), (256, 5)], indirect=True) -def test_compressor_checkpoint_continuation_matches_full_prefill(cache, ratio, layer, is_indexer, length): - from lightllm.models.deepseek_v4.layer_infer.compressor import fused_compress - from lightllm.models.deepseek_v4.workspace import DeepseekV4Workspace - - manager, requests = cache - workspace = DeepseekV4Workspace(SimpleNamespace(tp_world_size_=1, batch_max_tokens=length)) - workspace.c128_compress_scratch = torch.empty((1, length, 512), dtype=torch.float32, device="cuda") - source, src = _reserve(manager, requests, length) - destination, dst = _reserve(manager, requests, length) - _prepare_swa(requests, source, 0, length) - head_dim = requests.cache_config.indexer_head_dim if is_indexer else requests.cache_config.head_dim - scores = torch.randn((length, head_dim * (4 if ratio == 4 else 2)), dtype=torch.float32, device="cuda") - cos = torch.ones((length, 32), dtype=torch.float32, device="cuda") - sin = torch.zeros_like(cos) - weight = torch.ones(head_dim, dtype=torch.float32, device="cuda") - - def compress(req_idx, slots, start, end, decode=False): - positions = torch.arange(start, end, dtype=torch.int32, device="cuda") - width = end - start - state = SimpleNamespace( - dsv4_workspace=workspace, - microbatch_index=0, - 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, - _dsv4_token_to_batch_idx=None if decode else torch.zeros(width, dtype=torch.int32, device="cuda"), - b_req_idx=torch.full((width if decode else 1,), req_idx, dtype=torch.int32, device="cuda"), - b_mtp_index=torch.arange(width, dtype=torch.int32, device="cuda") - if decode - else torch.zeros(1, dtype=torch.int32, device="cuda"), - b_seq_len=positions + 1 if decode else torch.tensor([end], dtype=torch.int32, device="cuda"), - b_ready_cache_len=None if decode else torch.tensor([start], dtype=torch.int32, device="cuda"), - b_q_start_loc=None if decode else torch.zeros(1, dtype=torch.int32, device="cuda"), - dsv4_swa_write_slots=_swa_slots(requests, req_idx, positions), - dsv4_swa_runtime_write_slots=torch.where( - positions >= end - requests.swa_capacity_tokens, - _swa_slots(requests, req_idx, positions, prefill=False), - -1, - ) - if not decode - else None, - ) - output = fused_compress( - kv_score=scores[start:end], - infer_state=state, - 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, - ) - if is_indexer and req_idx == destination: - closed_groups = (positions + 1) % ratio == 0 - torch.testing.assert_close(output[closed_groups], reference[start:end][closed_groups], rtol=0, atol=0) - return output - - reference = compress(source, src, 0, length) - states = requests.create_small_page_cache_manager(1) - index = states.alloc_one_state_cache() - # A short final chunk must preserve the previous rows needed by the checkpoint tail. - checkpoint_len = length // 2 - for start, end in ((0, checkpoint_len - 2), (checkpoint_len - 2, checkpoint_len)): - _prepare_swa(requests, source, start, end) - compress(source, src, start, end) - requests.save_state(_checkpoint_req(source, checkpoint_len), index, states) - requests.restore_state(SimpleNamespace(req_idx=destination, cur_kv_len=checkpoint_len), states, index) - verify_width = requests.c4_state_tokens - 8 + 1 - for start in range(checkpoint_len, length - verify_width, 256): - end = min(start + 256, length - verify_width) - _prepare_swa(requests, destination, start, end) - compress(destination, dst, start, end) - _finish_prefill(requests, [destination]) - # Rejected speculative rows are written again in the same held token page. - compress(destination, dst, length - verify_width, length, decode=True) - compress(destination, dst, length - verify_width, length, decode=True) - if not is_indexer: - pool = _compressed_kv_pools(manager)[0] if ratio == 4 else _compressed_kv_pools(manager)[2] - torch.testing.assert_close( - _read_packed_rows(pool, 0, src[checkpoint_len + ratio - 1 : length : ratio].long() // ratio), - _read_packed_rows(pool, 0, dst[checkpoint_len + ratio - 1 : length : ratio].long() // ratio), - rtol=0, - atol=0, - ) - requests.free([source, destination], torch.cat([src, dst])) - states.free_state_cache([index]) - - -def test_swa_index_cuda_graph_reads_updated_private_pages(cache): - from lightllm.models.deepseek_v4.triton_kernel.build_swa_indices import build_swa_indices - - manager, requests = cache - source, src = _reserve(manager, requests, 512) - destination, dst = _reserve(manager, requests, 512) - _prepare_prefill(requests, torch.tensor([source, destination]), torch.zeros(2), torch.full((2,), 512)) - req_ids = torch.tensor([source, requests.HOLD_REQUEST_ID], dtype=torch.int32, device="cuda") - positions = torch.tensor([255, 1], dtype=torch.int32, device="cuda") - indexes = torch.empty((2, 128), dtype=torch.int32, device="cuda") - lengths = torch.empty(2, dtype=torch.int32, device="cuda") - write_slots = torch.empty_like(lengths) - - def build(): - build_swa_indices( - req_ids, - positions, - requests, - indexes, - lengths, - write_slots, - prefill_state=requests._test_prefill_states[source], - ) - - stream = torch.cuda.Stream() - stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(stream): - build() - torch.cuda.current_stream().wait_stream(stream) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - build() - req_ids[0] = destination - positions[0] = 256 - graph.replay() - expected = _swa_slots(requests, destination, torch.arange(256, 128, -1, dtype=torch.int32, device="cuda")) - torch.testing.assert_close(indexes[0], expected) - assert write_slots[0].item() == expected[0].item() - assert write_slots[1].item() == requests._test_prefill_states[source].HOLD_SWA_SLOT + 1 - requests.free([source, destination], torch.cat([src, dst])) - - -@pytest.mark.parametrize( - "dspark,cache", - [(False, (256, 0)), (False, (256, 1)), (False, (256, 5)), (True, (256, 1)), (True, (256, 5))], - indirect=["cache"], -) -def test_swa_ring_attention_matches_full_cache_after_speculative_retries(cache, dspark): - from lightllm.common.basemodel.attention.nsa import dsv4_fp8_flashmla_sparse as attention - from lightllm.models.deepseek_v4.triton_kernel.build_swa_indices import build_swa_indices - from lightllm.models.deepseek_v4.triton_kernel.build_dspark_swa_indices import build_dspark_swa_indices - - torch.manual_seed(1593) - manager, requests = cache - source, src = _reserve(manager, requests, 4096) - destination, dst = _reserve(manager, requests, 4096) - _prepare_swa(requests, destination, 0, 512) - # The reference retains every position in the other forward workspace. - _prepare_swa(requests, source, 0, 4096, microbatch_index=1) - values = torch.randn((4096, 512), dtype=torch.bfloat16, device="cuda") * 0.1 - packed = torch.zeros((4096, 584), dtype=torch.uint8, device="cuda") - packed[:, :448] = values[:, :448].to(torch.float8_e4m3fn).view(torch.uint8) - packed[:, 448:576] = values[:, 448:].contiguous().view(torch.uint8) - packed[:, 576:583] = 127 - positions = torch.arange(4096, device="cuda") - _write_swa_rows(requests, 2, source, positions, packed) - _write_swa(requests, 2, destination, positions[:512], packed[:512]) - _finish_prefill(requests, [destination]) - - mtp_step = requests.c4_state_tokens - 8 - token_num = mtp_step if dspark else mtp_step + 1 - width = 256 if dspark else 128 - req_idx = torch.full((token_num,), destination, dtype=torch.int32, device="cuda") - query_positions = torch.arange(512, 512 + token_num, dtype=torch.int32, device="cuda") - indices = torch.empty((token_num, width), dtype=torch.int32, device="cuda") - lengths = torch.empty(token_num, dtype=torch.int32, device="cuda") - write_slots = torch.empty_like(lengths) - - def build(): - if dspark: - build_dspark_swa_indices( - 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, - ) - else: - build_swa_indices(req_idx, query_positions, requests, indices, lengths, write_slots) - - backend = attention.DeepseekV4FlashMlaFp8SparseAttBackend( - type( - "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, - ), - "layers_infer": [SimpleNamespace(layer_num_=2)], - "tp_rank_": 0, - "tp_world_size_": 1, - "req_manager": requests, - }, - )() - ) - state_info = SimpleNamespace(mem_manager=manager) - state = backend.create_att_decode_state(state_info) - state.init_state() - params = dict( - compress_ratio=0, - head_dim_v=512, - softmax_scale=512 ** -0.5, - swa_indices=indices.unsqueeze(1), - swa_lengths=lengths, - attn_sink=torch.zeros(64, device="cuda"), - ) - control = attention.AttControl(nsa_decode=True, nsa_decode_dict=params) - q = torch.randn((token_num, 64, 512), dtype=torch.bfloat16, device="cuda") * 0.1 - packed_kv = requests.get_att_input_params(2) - - def forward(): - build() - return state.decode_att(q, packed_kv, None, att_control=control) - - stream = torch.cuda.Stream() - stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(stream): - forward() - torch.cuda.current_stream().wait_stream(stream) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph, stream=stream): - actual = forward() - - reference_state = backend.create_att_decode_state(state_info) - reference_state.init_state() - 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], - ), - 2, - ) - committed_end = 512 - # Both sides of page boundaries, repeated retries, and several complete ring wraps. - for start in (512, 513, 638, 639, 640, 640, 641, 767, 768, 1023, 1279, 1280, 1919, 1920, 2047): - # Advance intervening committed KV in small steps. A large parallel - # prefill would instead borrow pages before writing the whole chunk. - for gap_start in range(committed_end, start, 128): - gap = positions[gap_start : min(start, gap_start + 128)] - _write_swa_rows(requests, 2, source, gap, packed[gap]) - _write_swa_rows(requests, 2, destination, gap, packed[gap]) - query_positions.copy_(torch.arange(start, start + token_num, dtype=torch.int32)) - current = packed[start : start + token_num].clone() - current[:, :448] = ( - (values[start : start + token_num, :448] * (1 + start % 7)).to(torch.float8_e4m3fn).view(torch.uint8) - ) - logical_positions = query_positions.long() - _write_swa_rows(requests, 2, source, logical_positions, current) - _write_swa_rows(requests, 2, destination, logical_positions, current) - history = torch.arange(max(0, start - 128), start, device="cuda") - graph.replay() - if dspark: - visible = torch.cat((history.flip(0), logical_positions)) - reference_indices = _swa_slots(requests, source, visible).expand(token_num, -1).contiguous() - reference_lengths = torch.full_like(lengths, visible.numel()) - else: - visible = logical_positions[:, None] - torch.arange(128, device="cuda") - reference_indices = _swa_slots(requests, source, visible) - reference_lengths = torch.full_like(lengths, 128) - reference_indices = torch.nn.functional.pad( - reference_indices, (0, width - reference_indices.shape[1]), value=-1 - ) - reference_indices = reference_indices.to(torch.int32) - reference_params = dict(params, swa_indices=reference_indices.unsqueeze(1), swa_lengths=reference_lengths) - expected = reference_state.decode_att( - q, reference_kv, None, att_control=attention.AttControl(nsa_decode=True, nsa_decode_dict=reference_params) - ) - torch.testing.assert_close(actual, expected, rtol=0, atol=0, msg=f"SWA mismatch at position {start}") - torch.testing.assert_close(write_slots, _swa_slots(requests, destination, logical_positions).to(torch.int32)) - # Accept one position; speculative future rows stay dirty until the next write. - committed_end = start + 1 - requests.free([source, destination], torch.cat([src, dst])) - - -@pytest.mark.parametrize( - "cache,hash_page_size", [(256, 256), (256, 512), (512, 512), (768, 768), (1024, 1024)], indirect=["cache"] -) -@pytest.mark.parametrize("hit_extra_page", [False, True]) -def test_hybrid_radix_hit_fork_pause_abort_and_eviction(cache, monkeypatch, hash_page_size, hit_extra_page): - from sortedcontainers import SortedDict - from lightllm.server.router.model_infer.infer_batch import InferReq, g_infer_context, CacheTier - from lightllm.server.router.dynamic_prompt.hybrid_att_radix_cache import HybridAttPagedRadixCache - from lightllm.utils.kv_cache_utils import compute_token_list_hash - - manager, requests = cache - 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.manager) - args = get_env_start_args() - args.linear_att_hash_page_size = hash_page_size - args.linear_att_page_block_num = block_num - tail_len = big_page_tokens + hash_page_size - 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, - }.items(): - monkeypatch.setattr(g_infer_context, name, value, raising=False) - - def request(req_idx, total_len): - req = InferReq.__new__(InferReq) - req.req_idx, req.args = req_idx, args - req.cur_kv_len, req.hold_kv_len, req.cur_output_len = 0, 0, 0 - req.shared_kv_node = None - req.tail_small_page_buffer_id = None - req.hybrid_len_to_big_page_id = SortedDict() - req.hybrid_cache_len = (total_len - 1) // hash_page_size * hash_page_size - req.mm_input_spans = [] - req.cache_tiers = {CacheTier.GPU} - req.sampling_param = SimpleNamespace(disable_prompt_cache=False) - req.prompt_selected_logprobs = SimpleNamespace(copy_capture_slots_if_needed=lambda **kwargs: None) - tokens = list(range(total_len)) - hashes = compute_token_list_hash(tokens, hash_page_size) - req.shm_req = SimpleNamespace( - input_len=total_len, - shm_prompt_ids=SimpleNamespace(arr=tokens), - hybrid_token_hash_list=SimpleNamespace(size=len(hashes), get_all=lambda: hashes), - shm_cur_kv_len=0, - prompt_cache_len=0, - ) - req.get_chuncked_input_token_len = req.get_chuncked_input_token_len_for_hybrid_att - req.req_manager = requests - req.args = args - if not hasattr(args, "req_token_index_mode"): - args.req_token_index_mode = "pin_mem_page_cache" - from lightllm.server.router.model_infer.req_token_index import create_req_token_index - - req.token_index = create_req_token_index(req, torch.cuda.current_stream()) - return req - - source, held = _reserve(manager, requests, tail_len + 1) - req = request(source, tail_len + 1) - req.token_index.append_token_indexes(held.cpu()) - req.hold_kv_len = held.numel() - for end in (big_page_tokens, tail_len): - assert req.get_chuncked_input_token_len() == end - _prepare_swa(requests, source, req.cur_kv_len, end) - for pool in (requests.req_to_swa.flatten(1, 2), *_compressed_kv_pools(manager)): - pool.random_(0, 256) - _state_buffer(requests, "c4").uniform_() - _state_buffer(requests, "c4_indexer").uniform_() - g_infer_context.save_hybrid_state_to_cache(torch.tensor([source], device="cuda"), [req]) - req.cur_kv_len = end - _prepare_swa(requests, source, req.cur_kv_len, tail_len + 1) - req.cur_kv_len = tail_len + 1 - freed = [] - g_infer_context.free_a_req_mem(freed, req) - requests.memory.allocator.free(torch.cat(freed)) - requests.free_req(source) - assert radix.get_tree_total_tokens_num() == tail_len - - forks = [request(requests.alloc(), hit_len + 1) for _ in range(2)] - for fork in forks: - from lightllm.server.router.model_infer.req_token_index import ReqTokenIndex - - fork._hybrid_match_radix_cache() - ReqTokenIndex.submit_pending_uploads([fork.token_index]) - assert fork.cur_kv_len == hit_len - assert fork.hold_kv_len == hit_len - assert fork.shared_kv_node.node_prefix_total_len == big_page_tokens - if hit_extra_page: - assert requests.req_to_token_indexs[fork.req_idx, big_page_tokens].item() != held[big_page_tokens].item() - restored = requests.req_to_token_indexs[fork.req_idx, :hit_len] - for pool, ratio in zip(_compressed_kv_pools(manager), (4, 4, 128)): - torch.testing.assert_close( - _read_packed_rows(pool, 0, held[ratio - 1 : hit_len : ratio].long() // ratio), - _read_packed_rows(pool, 0, restored[ratio - 1 : hit_len : ratio].long() // ratio), - ) - assert set(_request_pages(requests, forks[0].req_idx)).isdisjoint(_request_pages(requests, forks[1].req_idx)) - # Pause can publish a reusable prefix; abort with GPU placement disabled only dereferences it. - forks[1].cache_tiers.clear() - for fork in forks: - freed = [] - g_infer_context.free_a_req_mem(freed, fork) - requests.memory.allocator.free(torch.cat(freed)) - requests.free_req(fork.req_idx) - radix.free_radix_cache_to_get_enough_token(manager.size) - assert manager.allocator.can_use_mem_size == manager.size - assert ( - manager.big_page_buffers.get_free_cache_num() - == manager.big_page_buffers.size - manager.big_page_buffers.keep_num - ) - assert small.get_free_cache_num() == small.size - - -@pytest.mark.parametrize("page_cap", [3, 130]) -def test_c4_indexer_page_table_uses_each_query_request_and_causal_length(cache, page_cap): - from lightllm.models.deepseek_v4.triton_kernel.build_c4_indexer_page_table import build_c4_indexer_page_table - - manager, requests = cache - first, first_slots = _reserve(manager, requests, 512) - second, second_slots = _reserve(manager, requests, 768) - req_ids = [second, first, second, requests.HOLD_REQUEST_ID, first] - lengths = [1, 0, 129, 64, 65] - token_req_idx = torch.tensor(req_ids, dtype=torch.int32, device="cuda") - c4_lengths = torch.tensor(lengths, dtype=torch.int32, device="cuda") - storage = torch.full((len(req_ids) * page_cap + 7,), -17, dtype=torch.int32, device="cuda") - out = storage[:-7].view(len(req_ids), page_cap) - actual = build_c4_indexer_page_table( - token_req_idx, - c4_lengths, - page_cap * 64, - requests.req_to_token_indexs, - requests.HOLD_REQUEST_ID, - out=out, - ) - expected = torch.zeros_like(out) - request_slots = {first: first_slots, second: second_slots} - for row, (req_idx, length) in enumerate(zip(req_ids, lengths)): - if req_idx == requests.HOLD_REQUEST_ID: - continue - compressed_slots = request_slots[req_idx][3 : length * 4 : 4].long() // 4 - pages = compressed_slots[::64] // 64 - expected[row, : pages.numel()] = pages.to(torch.int32) - torch.testing.assert_close(actual, expected, rtol=0, atol=0) - assert actual.data_ptr() == out.data_ptr() - assert storage[-7:].eq(-17).all() - empty = build_c4_indexer_page_table( - token_req_idx[:0], - c4_lengths[:0], - page_cap * 64, - requests.req_to_token_indexs, - requests.HOLD_REQUEST_ID, - ) - assert empty.shape == (0, page_cap) - requests.free([first, second], torch.cat((first_slots, second_slots))) - - -@pytest.mark.parametrize("mode", ["decode", "prefill"]) -def test_c4_indexer_metadata_scoring_and_graph_replay(cache, monkeypatch, mode): - from lightllm.models.deepseek_v4 import infer_struct as infer_struct_module - from lightllm.models.deepseek_v4.layer_infer import indexer as indexer_module - - torch.manual_seed(1593) - monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_DP", "0") - monkeypatch.setenv("LIGHTLLM_DP_WORLD_SIZE", "1") - manager, requests = cache - req_idx, slots = _reserve(manager, requests, 4096) - heads, tokens = 64, 5 - config = dict( - compress_ratios=[4, 128, 0], - num_attention_heads=heads, - index_topk=512, - index_head_dim=128, - qk_rope_head_dim=64, - index_n_heads=heads, - sliding_window=128, - head_dim=512, - rms_norm_eps=1e-6, - ) - 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") - state.b_seq_len = torch.full_like(state.b_req_idx, 4096) - state.position_ids = torch.arange(3068, 3068 + tokens, device="cuda") - state.dsv4_sparse_req_idx = torch.full_like(state.position_ids, req_idx) - state._dsv4_token_to_batch_idx = torch.zeros(tokens, dtype=torch.int64, device="cuda") - # Exercise chunk scheduling without requiring a long prompt or large logits allocation. - monkeypatch.setattr(infer_struct_module, "C4_PREFILL_LOGITS_BUDGET_BYTES", 2 * 1024 * 4) - keys = (torch.randn((1024, 128), device="cuda") * 0.1).to(torch.float8_e4m3fn) - scales = torch.ones((1024, 1), dtype=torch.float32, device="cuda") - packed = torch.cat((keys.view(torch.uint8), scales.view(torch.uint8)), dim=1) - c4_slots = slots[3::4].long() // 4 - _write_packed_rows(_compressed_kv_pools(manager)[1], 0, c4_slots, packed) - q = (torch.randn((tokens, heads, 128), device="cuda") * 0.125).to(torch.float8_e4m3fn) - weights = torch.rand((tokens, heads), device="cuda") - - def select(): - return indexer._c4_indices(state, q, weights) - - stream = torch.cuda.Stream() - stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(stream): - select() - state.c4_indexer_metadata = None - torch.cuda.current_stream().wait_stream(stream) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph, stream=stream): - captured, captured_lengths = select() - # Change both causal lengths and Q inputs; graph metadata must be recomputed. - state.position_ids.add_(256) - state.b_seq_len.fill_(3584) - q.copy_((q.float() * 1.25).to(q.dtype)) - graph.replay() - actual = captured.clone() - actual_lengths = captured_lengths.clone() - state.c4_indexer_metadata = None - eager, eager_lengths = select() - torch.testing.assert_close(actual.sort(dim=-1).values, eager.sort(dim=-1).values, rtol=0, atol=0) - torch.testing.assert_close(actual_lengths, eager_lengths, rtol=0, atol=0) - torch.testing.assert_close(actual_lengths, torch.full_like(actual_lengths, 512)) - inverse_slots = torch.full((_compressed_kv_pools(manager)[1].shape[1] * 64,), -1, device="cuda", dtype=torch.long) - inverse_slots[c4_slots] = torch.arange(1024, device="cuda") - logical_indices = inverse_slots[actual[:, 0].long()] - scores = ((q.float() @ keys.float().T).relu() * weights[:, :, None]).sum(1) - valid_lengths = (state.position_ids + 1) // 4 - for row in range(tokens): - selected = logical_indices[row] - assert ((selected >= 0) & (selected < valid_lengths[row])).all() - selected_scores = scores[row, selected].sort(descending=True).values - reference_scores = scores[row, : valid_lengths[row]].topk(512).values - torch.testing.assert_close(selected_scores, reference_scores, rtol=2e-3, atol=2e-3) - requests.free([req_idx], slots) - - -@pytest.mark.parametrize("is_prefill", [False, True]) -@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 - - manager, requests = cache - first, first_slots = _reserve(manager, requests, 512) - 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") - model._cos_cached_sliding = model._sin_cached_sliding = rope - model._cos_cached_compress = model._sin_cached_compress = rope - - 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") - state.b_req_idx = torch.tensor([first, second], dtype=torch.int32, device="cuda") - state.b_seq_len = torch.tensor(lengths, dtype=torch.int32, device="cuda") - if is_prefill: - state.b_ready_cache_len = state.b_seq_len - 1 - state.prefill_att_state = backend.create_att_prefill_state(state) - else: - state.decode_att_state = backend.create_att_decode_state(state) - state.init_some_extra_state(model) - return state - - captured_state = make_state([130, 258], 0) - other_state = make_state([130, 258], 1) - 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) - 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) - other_state.dsv4_c4_indices.fill_(22) - assert captured_state.dsv4_c4_indices.eq(11).all() - - stream = torch.cuda.Stream() - stream.wait_stream(torch.cuda.current_stream()) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph, stream=stream): - swa_output = captured_state.dsv4_swa_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: - captured_state.copy_for_prefill_cuda_graph(updated) - else: - captured_state.copy_for_cuda_graph(updated) - graph.replay() - torch.testing.assert_close(swa_output, updated.dsv4_swa_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))) - - -@pytest.mark.parametrize("microbatch_index", [0, 1]) -@pytest.mark.parametrize("mode", ["proposal", "commit", "prefill"]) -def test_dspark_swa_metadata_is_built_once(cache, monkeypatch, microbatch_index, mode): - from unittest.mock import Mock - - import lightllm.models.deepseek_v4.triton_kernel.build_swa_indices as swa_kernel - import lightllm.models.deepseek_v4_dspark.infer_struct as dspark_state - from lightllm.models.deepseek_v4.workspace import DeepseekV4Workspace - from lightllm.models.deepseek_v4_dspark.infer_struct import DeepseekV4DSparkInferStateInfo - from lightllm.common.basemodel.attention.nsa.dsv4_fp8_flashmla_sparse import DeepseekV4FlashMlaFp8SparseAttBackend - - manager, requests = cache - req_idx, slots = _reserve(manager, requests, 512) - args = get_env_start_args() - 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, - ) - model = type("TestModel", (), {})() - vars(model).update( - config=config, - batch_max_tokens=3, - max_req_num=requests.max_request_num, - max_seq_length=4096, - graph=None, - 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", - layers_infer=[SimpleNamespace(layer_num_=i) for i in range(3)], - ) - backend = DeepseekV4FlashMlaFp8SparseAttBackend(model) - model.block_size = 3 - model.dsv4_workspace = DeepseekV4Workspace(model) - rope = torch.zeros((512, 32), device="cuda") - model._cos_cached_sliding = model._sin_cached_sliding = rope - model._cos_cached_compress = model._sin_cached_compress = rope - state = DeepseekV4DSparkInferStateInfo() - 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") - state.b_req_idx = torch.tensor([req_idx], dtype=torch.int32, device="cuda") - state.b_seq_len = torch.tensor([260], dtype=torch.int32, device="cuda") - state.b_ready_cache_len = torch.tensor([257], dtype=torch.int32, device="cuda") - state.prefill_att_state = backend.create_att_prefill_state(state) - else: - state.decode_att_state = backend.create_att_decode_state(state) - state.b_req_idx = torch.full((3,), req_idx, dtype=torch.int32, device="cuda") - state.b_seq_len = torch.arange(258, 261, dtype=torch.int32, device="cuda") - if mode == "commit": - state.mtp_draft_input_hiddens = torch.zeros((3, 1), device="cuda") - state.max_kv_seq_len = 260 - causal_builder = Mock(wraps=swa_kernel.build_swa_indices) - draft_builder = Mock(wraps=dspark_state.build_dspark_swa_indices) - monkeypatch.setattr(swa_kernel, "build_swa_indices", causal_builder) - monkeypatch.setattr(dspark_state, "build_dspark_swa_indices", draft_builder) - state.init_some_extra_state(model) - proposal = mode == "proposal" - assert causal_builder.call_count == int(not proposal) - assert draft_builder.call_count == int(proposal) - assert state.dsv4_swa_indices.shape == (3, 256 if proposal else 128) - torch.testing.assert_close( - state.dsv4_swa_lengths, torch.full_like(state.dsv4_swa_lengths, 131 if proposal else 128) - ) - expected_slots = _get_swa_token_indices( - requests, req_idx, state.position_ids, prefill_state=state.prefill_att_state if state.is_prefill else None - ).to(torch.int32) - torch.testing.assert_close(state.dsv4_swa_write_slots, expected_slots) - requests.free([req_idx], slots) - - -def test_empty_dspark_commit_is_local_but_proposal_keeps_common_padding_path(monkeypatch): - from lightllm.models.deepseek_v4.model import DeepseekV4TpPartModel - from lightllm.models.deepseek_v4_dspark.model import DeepseekV4DSparkModel - - model = DeepseekV4DSparkModel.__new__(DeepseekV4DSparkModel) - model_input = SimpleNamespace(batch_size=0, mtp_draft_input_hiddens=torch.empty((0, 8), device="cuda")) - # A KV-only empty commit is local; a full proposal still takes the common - # model path and participates in DP collectives through HOLD padding. - output = model._decode(model_input) - assert output.logits.shape == (0, 1) - assert output.logits.device == model_input.mtp_draft_input_hiddens.device - - proposal_calls = [] - monkeypatch.setattr(DeepseekV4TpPartModel, "_decode", lambda self, inputs: proposal_calls.append(inputs)) - model_input.mtp_draft_input_hiddens = None - model._decode(model_input) - assert proposal_calls == [model_input] - - -@pytest.mark.parametrize( - "length,prefix", [(768, 0), (768, 256), (2816, 0), (2816, 2304), (4096, 0), (4096, 256), (4096, 2304)] -) -@pytest.mark.parametrize("exhaust_small_pool", [False, True]) -@pytest.mark.parametrize("token_index_mode", ["pin_mem_cache", "pin_mem_page_cache", "pure_gpu"]) -def test_standard_cpu_cache_offload_and_load_use_saved_checkpoints( - cache, monkeypatch, length, prefix, exhaust_small_pool, token_index_mode -): - """Offload after the runtime ring is overwritten, then resume through the common loader.""" - from collections import deque - 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 unit_tests.common.dsv4_test_utils import DeepseekV4StateCacheManager - - manager, requests = cache - source, src = _reserve(manager, requests, length) - requests.create_small_page_cache_manager(4) - page_tokens = get_env_start_args().cpu_cache_token_page_size - ends = list(range(page_tokens, length + 1, page_tokens)) - if length % page_tokens: - ends.append(length) - source_checkpoints = {} - tail = None - expected_checkpoints = [] - for end in ends: - # Populate a runtime state at this endpoint without generating new tokens. - _prepare_swa(requests, source, end, end) - value = 11 + end // 256 - requests.req_to_swa.flatten(1, 2).fill_(value) - _state_buffer(requests, "c4").fill_(value) - _state_buffer(requests, "c4_indexer").fill_(value) - state_pool = requests.small_page_buffers if end % page_tokens else manager.big_page_buffers - slot = state_pool.alloc_one_state_cache() - requests.save_state(_checkpoint_req(source, end), slot, state_pool) - torch.cuda.synchronize() - expected_checkpoints.append(state_pool.buffer[slot].clone()) - if end % page_tokens: - 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) - _state_buffer(requests, "c4_indexer").fill_(77) - _randomize_compressed_kv(manager) - pages = torch.full( - (len(ends) + 2, 1, 1, 1, manager.cache_config.get_cpu_cache_big_page_bytes()), - 221, - dtype=torch.uint8, - pin_memory=True, - ) - published, released = [], [] - client = SimpleNamespace( - cpu_kv_cache_tensor=pages, - lock=SimpleNamespace(acquire_sleep1ms=lambda: None, release=lambda: None), - allocate_pages=lambda hashes, **kwargs: (list(range(1, len(hashes) + 1)), [False] * len(hashes)), - update_pages_status_to_ready=lambda **kwargs: published.append(kwargs), - deref_pages=lambda page_list: released.extend(page_list), - ) - req = SimpleNamespace( - req_idx=source, - req_id=42, - cur_kv_len=length, - hold_kv_len=src.numel(), - shared_kv_node=None, - hybrid_len_to_big_page_id=source_checkpoints, - hybrid_cache_len=length, - tail_small_page_buffer_id=tail, - cpu_cache_offload_task_status=InferReq._CpuCacheTaskStatus.NOT_STARTED, - cache_tiers=(CacheTier.CPU,), - shm_req=SimpleNamespace( - group_req_id=42, - request_id=42, - input_len=length + 1, - token_hash_list=SimpleNamespace(get_all=lambda: list(range(len(ends)))), - token_hash_page_len_list=SimpleNamespace(get_all=lambda: ends), - ), - ) - attach_req_token_index(req, requests, src.cpu(), mode=token_index_mode) - # 已命中的大页来自 radix，其余完整大页来自当前请求。 - shared_page_ids = [source_checkpoints.pop(end) for end in range(page_tokens, prefix + 1, page_tokens)] - 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", - SimpleNamespace( - get_big_page_ids_by_node=lambda node: shared_page_ids.copy(), - free_one_small_page_buffer=lambda: None, - free_radix_cache_to_get_enough_token=lambda need: None, - ), - ) - monkeypatch.setattr(context, "is_hybrid_att_model", True) - monkeypatch.setattr(context, "get_overlap_stream", lambda: producer) - monkeypatch.setattr(context, "get_can_alloc_token_num", lambda: manager.allocator.can_use_mem_size) - monkeypatch.setattr(common.dist, "broadcast_object_list", lambda *args, **kwargs: None) - monkeypatch.setattr(common.dist, "barrier", lambda **kwargs: None) - module = common.MultiLevelKvCacheModule.__new__(common.MultiLevelKvCacheModule) - module.args = SimpleNamespace( - page_size=manager.page_size, - cpu_cache_token_page_size=page_tokens, - linear_att_hash_page_size=256, - 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), - ) - module.cpu_cache_client = client - module.offload_tasks = deque() - module.load_tasks = deque() - module.page_index_buffer = torch.empty(8, dtype=torch.int32, device="cuda") - module.page_ready_buffer = torch.empty(8, dtype=torch.bool, device="cuda") - module.gloo_group = module.filter_group = module.init_sync_group = None - module.need_sync_compute_stream = lambda: False - assert module.offload_finished_reqs_to_cpu_cache([req]) == [] - assert req.cpu_cache_offload_task_status.is_running() - torch.cuda.synchronize() - module.update_cpu_cache_offload_task_states() - assert module.offload_finished_reqs_to_cpu_cache([req]) == [req] - assert published == [ - dict( - page_list=list(range(1, len(ends) + 1)), - deref=True, - disk_offload_enable=False, - token_num_in_page_list=length, - ) - ] - flat = pages.view(pages.shape[0], -1) - for i, expected in enumerate(expected_checkpoints, 1): - torch.testing.assert_close( - flat[i, manager.cache_config.get_cpu_cache_full_att_bytes() :], expected, rtol=0, atol=0 - ) - if tail is not None: - torch.testing.assert_close( - manager.big_page_buffers.buffer[manager.CPU_CACHE_BIG_PAGE_OFFLOAD_TEMP_BUFFER_ID], - offload_before, - rtol=0, - atol=0, - ) - - destination = requests.alloc() - dst = torch.empty(0, dtype=torch.int32, device="cuda") - if prefix: - dst = manager.alloc(prefix).cuda() - requests.req_to_token_indexs[destination, :prefix] = dst - manager.operator.copy_mem_to_mem(src[:prefix], dst) - restored = SimpleNamespace( - req_idx=destination, - cur_kv_len=prefix, - hold_kv_len=prefix, - mm_input_spans=[], - hybrid_cache_len=length, - hybrid_len_to_big_page_id={}, - tail_small_page_buffer_id=None, - sampling_param=SimpleNamespace(shm_param=SimpleNamespace(prompt_logprobs=-1)), - shm_req=SimpleNamespace( - input_len=length + 1, - disk_prompt_cache_len=0, - cpu_prompt_cache_len=0, - cpu_cache_match_page_indexes=SimpleNamespace(get_all=lambda: list(range(1, len(ends) + 1))), - token_hash_page_len_list=SimpleNamespace(get_all=lambda: ends), - ), - ) - attach_req_token_index(restored, requests, dst.cpu(), mode=token_index_mode) - if prefix >= page_tokens: - slot = manager.big_page_buffers.alloc_one_state_cache() - manager.big_page_buffers.buffer[slot].copy_(expected_checkpoints[0]) - restored.hybrid_len_to_big_page_id[page_tokens] = slot - - busy_small_slots = [] - if exhaust_small_pool: - busy_small_slots = requests.small_page_buffers.alloc_state_cache( - requests.small_page_buffers.get_free_cache_num() - ) - module.load_cpu_cache_to_reqs([restored]) - torch.cuda.synchronize() - module.update_cpu_cache_load_task_states() - assert restored.cur_kv_len == restored.hold_kv_len == length - assert released == list(range(1, len(ends) + 1)) - if length % page_tokens: - torch.testing.assert_close( - manager.big_page_buffers.buffer[manager.CPU_CACHE_BIG_PAGE_LOAD_TEMP_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( - pool[:, (src[::256] // 256).long()], pool[:, (dst[::256] // 256).long()], rtol=0, atol=0 - ) - check = DeepseekV4StateCacheManager(1, manager.cache_config) - _prepare_swa(requests, destination, length, length) - requests.save_state(_checkpoint_req(destination, length), 0, check) - torch.cuda.synchronize() - torch.testing.assert_close(check.buffer[0], expected_checkpoints[-1], rtol=0, atol=0) - requests.free([source, destination], torch.cat((src, dst))) - manager.big_page_buffers.free_state_cache( - shared_page_ids + list(source_checkpoints.values()) + list(restored.hybrid_len_to_big_page_id.values()) - ) - requests.small_page_buffers.free_state_cache( - [x for x in (tail, restored.tail_small_page_buffer_id) if x is not None] + busy_small_slots - ) - assert manager.allocator.can_use_mem_size == manager.size - assert ( - manager.big_page_buffers.get_free_cache_num() - == manager.big_page_buffers.size - manager.big_page_buffers.keep_num - ) - - -def test_prefill_fused_writer_preserves_old_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, requests = cache - req, held = _reserve(manager, requests, 2048) - reference = torch.zeros((1, 17, 74880), dtype=torch.uint8, device="cuda") - kv = torch.randn((2048, 512), dtype=torch.bfloat16, device="cuda") - weight = torch.randn(512, dtype=torch.bfloat16, device="cuda") - angles = 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 - layer_weight = SimpleNamespace(kv_norm_=SimpleNamespace(weight=weight)) - for start, end in ((0, 255), (255, 641), (641, 1100), (1100, 2048)): - _prepare_swa(requests, req, start, end) - positions = torch.arange(start, end, device="cuda", dtype=torch.int32) - indices = torch.empty((positions.numel(), 128), dtype=torch.int32, device="cuda") - lengths = torch.empty(positions.numel(), dtype=torch.int32, device="cuda") - writes = torch.empty_like(lengths) - runtime = torch.empty_like(lengths) - build_swa_indices( - torch.full_like(positions, req), - positions, - requests, - indices, - lengths, - writes, - prefill_state=requests._test_prefill_states[req], - swa_runtime_write_slots=runtime, - token_to_batch_idx=torch.zeros_like(positions), - b_seq_len=torch.tensor([end], device="cuda"), - ) - 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, - dsv4_swa_runtime_write_slots=runtime, - position_ids=positions, - ) - layer._norm_rope_and_cache_kv(kv[start:end], state, layer_weight) - fused_kv_norm_rope_to_swa(kv[start:end], weight, 1e-6, freqs, positions, positions, reference[0], 128) - visible = torch.arange(max(0, start - 127), end, device="cuda") - # 第八个 scale 字节仅用于对齐；workspace 复用时无需保持其初始值。 - torch.testing.assert_close( - _read_swa(requests, 2, req, visible)[:, :-1], - _read_packed_rows(reference, 0, visible)[:, :-1], - rtol=0, - atol=0, - ) - logical = positions[:, None] - torch.arange(128, device="cuda")[None, :] - valid = logical >= 0 - expected_slots = _swa_slots(requests, req, logical.clamp_min(0)) - torch.testing.assert_close(indices[valid], expected_slots[valid].to(torch.int32)) - _finish_prefill(requests, [req]) - tail = torch.arange(max(0, end - requests.sliding_window), end, device="cuda") - torch.testing.assert_close( - _read_swa(requests, 2, req, tail)[:, :-1], - _read_packed_rows(reference, 0, tail)[:, :-1], - rtol=0, - atol=0, - ) - requests.free([req], held) - - -def test_cpu_cache_preserves_already_published_pages(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 = [1] - readies = [True] - 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.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, - 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, - big_page_slot_ids=checkpoint_ids, - ) - torch.cuda.synchronize() - offset = 0 - 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() - assert pages[0].eq(221).all() and pages[2].eq(221).all() - manager.big_page_buffers.free_state_cache([state_id]) - requests.free([req], slots) - - -@pytest.mark.parametrize("prefix_len", range(4)) -@pytest.mark.parametrize("cached_first_image", [False, True]) -def test_image_metadata_and_embeddings_preserve_canonical_cache(cache, prefix_len, cached_first_image): - from lightllm.common.basemodel.triton_kernel.multimodal_emb import multimodal_emb - from lightllm.models.deepseek_v4.infer_struct import DeepseekV4InferStateInfo - from lightllm.models.deepseek_v4.tokenizer import DeepSeekV4Tokenizer - from lightllm.server.multimodal_params import MultimodalParams - from lightllm.utils.multimodal_chunk_utils import get_mm_input_spans - - tokenizer = DeepSeekV4Tokenizer( - SimpleNamespace(convert_tokens_to_ids=lambda token: 99), "/unused", {"vision_n_layers": 1} - ) - params = MultimodalParams(images=[{"type": "url", "data": "unused"} for _ in range(3)]) - token_base, canonical_length = 2 ** 62, 13 - for image in params.images: - image.token_id, image.token_num, image.start_index_in_embed_cache = token_base, canonical_length, 0 - input_ids = tokenizer.encode([7] * prefix_len + [99, 99, 8, 99, 8], params) - spans = get_mm_input_spans(input_ids, params.to_dict()) - expected_left, expected_right = [0] * len(input_ids), [0] * len(input_ids) - for start, end in spans: - image_start = input_ids.index(token_base + 3, start, end) - expected_left[image_start:end] = range(end - image_start) - expected_right[image_start:end] = range(end - image_start - 1, -1, -1) - - ready = spans[0][1] if cached_first_image else 0 - manager, requests = cache - req_idx, slots = _reserve(manager, requests, len(input_ids)) - backend = _test_backend(requests, token_capacity=128) - model = backend.model - model.has_vision = True - model.config["vision_max_n_token"] = 384 - model.dsv4_workspace = None - rope = torch.zeros((len(input_ids), 32), device="cuda") - model._cos_cached_sliding = model._sin_cached_sliding = rope - model._cos_cached_compress = model._sin_cached_compress = rope - 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") - state.b_ready_cache_len = torch.tensor([ready], dtype=torch.int32, device="cuda") - state.max_kv_seq_len = len(input_ids) - state.multimodal_params = [params.to_dict()] - state.prefill_att_state = backend.create_att_prefill_state(state) - state.init_some_extra_state(model) - assert state.dsv4_image_left.cpu().tolist() == expected_left[ready:] - assert state.dsv4_image_right.cpu().tolist() == expected_right[ready:] - - text_embeddings = torch.randn((100, 16), dtype=torch.bfloat16, device="cuda") - image_embeddings = torch.randn((canonical_length, 1, 16), dtype=torch.bfloat16, device="cuda") - out = torch.zeros((len(input_ids) - ready, 16), dtype=torch.bfloat16, device="cuda") - multimodal_emb( - out, - state.input_ids, - text_embeddings, - image_embeddings, - torch.tensor([canonical_length], dtype=torch.int64, device="cuda"), - torch.tensor([token_base], dtype=torch.int64, device="cuda"), - torch.tensor([0], dtype=torch.int64, device="cuda"), - 0, - 100, - 1, - ) - expected = torch.stack( - [ - text_embeddings[token] if token < 100 else image_embeddings[token - token_base, 0] - for token in input_ids[ready:] - ] - ) - torch.testing.assert_close(out, expected, rtol=0, atol=0) - requests.free([req_idx], slots) - - -@pytest.mark.parametrize("padding_tokens", [0, 3, 5]) -def test_prefill_metadata_does_not_read_cuda_scalars(cache, padding_tokens): - from lightllm.models.deepseek_v4.infer_struct import DeepseekV4InferStateInfo - from lightllm.models.deepseek_v4.workspace import DeepseekV4Workspace - - manager, requests = cache - backend = _test_backend(requests) - model = backend.model - model.has_vision = False - rope = torch.zeros((8, 32), device="cuda") - model._cos_cached_sliding = model._sin_cached_sliding = rope - model._cos_cached_compress = model._sin_cached_compress = rope - - model.dsv4_workspace = DeepseekV4Workspace(model) - real_tokens = 5 - padding_tokens - req_ids, lengths = ([0], [real_tokens]) if real_tokens else ([], []) - if padding_tokens: - req_ids.append(requests.HOLD_REQUEST_ID) - lengths.append(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") - infer_state.b_seq_len = torch.tensor(lengths, dtype=torch.int32, device="cuda") - infer_state.b_ready_cache_len = torch.zeros_like(infer_state.b_seq_len) - infer_state.max_kv_seq_len = max(lengths) - infer_state.prefill_att_state = backend.create_att_prefill_state(infer_state) - infer_state.init_some_extra_state(model) - - previous_debug_mode = torch.cuda.get_sync_debug_mode() - torch.cuda.set_sync_debug_mode("error") - try: - infer_state.init_some_extra_state(model) - finally: - torch.cuda.set_sync_debug_mode(previous_debug_mode) - assert infer_state.dsv4_sparse_req_idx.shape == (5,) - - -@pytest.mark.parametrize("window", [64, 128, 256, 512]) -@pytest.mark.parametrize("mtp_step", [0, 5]) -def test_configured_window_survives_chunks_and_checkpoint_restore(cache, window, mtp_step): - from dataclasses import replace - - manager, _ = cache - config = replace(manager.cache_config, sliding_window=window) - requests = DeepseekV4ReqManager(2, 4096, manager, config, mtp_step=mtp_step) - source, destination = requests.alloc(), requests.alloc() - packed = torch.randint(0, 256, (2048, config.kv_bytes), dtype=torch.uint8, device="cuda") - states = requests.create_small_page_cache_manager(1) - index = states.alloc_one_state_cache() - for start, end in ((0, 257), (257, 512), (512, 513), (513, 2048)): - _prepare_swa(requests, source, start, end) - positions = torch.arange(start, end, device="cuda") - _write_swa(requests, 0, source, positions, packed[start:end]) - if end % 256 == 0: - requests.save_state(_checkpoint_req(source, end), index, states) - requests.restore_state(SimpleNamespace(req_idx=destination, cur_kv_len=end), states, index) - positions = torch.arange(max(0, end - window), end, device="cuda") - actual = _read_swa(requests, 0, destination, positions) - torch.testing.assert_close(actual, packed[positions], rtol=0, atol=0) - _finish_prefill(requests, [source]) - positions = torch.arange(max(0, end - window), end, device="cuda") - actual = _read_swa(requests, 0, source, positions) - torch.testing.assert_close(actual, packed[positions], rtol=0, atol=0) - requests.free_req(source) - requests.free_req(destination) - states.free_state_cache([index]) - - -@pytest.mark.parametrize("mtp_step", [0, 3]) -def test_full_batch_unaligned_prefill_workspace_is_bounded_and_private(cache, mtp_step): - manager, _ = cache - request_count, chunk_tokens, ready = 8, 64, 382 - requests = DeepseekV4ReqManager(request_count, 4096, manager, manager.cache_config, mtp_step=mtp_step) - reqs = [requests.alloc() for _ in range(request_count)] - page_token_num = manager.cache_config.swa_page_token_num - first = (ready - requests.sliding_window + 1) // page_token_num * page_token_num - end = ready + chunk_tokens - required = request_count * ((end + page_token_num - 1) // page_token_num * page_token_num - first) - requests._test_backend = _test_backend(requests, request_count * chunk_tokens) - requests._test_prefill_states = {} - workspace, capacity = requests._test_backend.get_swa_prefill_workspace(requests._test_backend.model) - assert capacity >= required - assert workspace.untyped_storage().data_ptr() != requests.req_to_swa.untyped_storage().data_ptr() - workspace[:, capacity : 2 * capacity].fill_(173) - old_positions = torch.arange(ready - requests.sliding_window + 1, ready, device="cuda") - new_positions = torch.arange(ready, end, device="cuda") - for req in reqs: - _write_swa_rows( - requests, - 0, - req, - old_positions, - torch.full((old_positions.numel(), 584), req + 1, dtype=torch.uint8, device="cuda"), - prefill=False, - ) - _prepare_prefill( - requests, torch.tensor(reqs), torch.full((request_count,), ready), torch.full((request_count,), end) - ) - for req in reqs: - _write_swa( - requests, - 0, - req, - new_positions, - torch.full((chunk_tokens, 584), req + 17, dtype=torch.uint8, device="cuda"), - ) - for req in reqs: - positions = torch.arange(ready - requests.sliding_window + 1, end, device="cuda") - slots = _swa_slots(requests, req, positions) - assert slots.unique().numel() == positions.numel() - actual = _read_packed_rows(workspace, 0, slots) - expected = torch.where(positions[:, None] < ready, req + 1, req + 17).expand_as(actual).to(torch.uint8) - torch.testing.assert_close(actual, expected, rtol=0, atol=0) - assert workspace[:, capacity : 2 * capacity].eq(173).all() - _finish_prefill(requests, reqs) - for req in reqs: - positions = torch.arange(end - requests.sliding_window, end, device="cuda") - actual = _read_swa(requests, 0, req, positions) - expected = torch.where(positions[:, None] < ready, req + 1, req + 17).expand_as(actual).to(torch.uint8) - torch.testing.assert_close(actual, expected, rtol=0, atol=0) - requests.free_req(req) diff --git a/unit_tests/common/test_hybrid_req_manager_pd_sharing.py b/unit_tests/common/test_hybrid_req_manager_pd_sharing.py deleted file mode 100644 index ba349372ca..0000000000 --- a/unit_tests/common/test_hybrid_req_manager_pd_sharing.py +++ /dev/null @@ -1,145 +0,0 @@ -import multiprocessing -import traceback -from multiprocessing.reduction import ForkingPickler - -import pytest -import torch - -from lightllm.common.attention_state.bounded_att import GDNState, KDAState -from lightllm.server.router.model_infer.mode_backend.pd.p2p_fix import reduce_tensor -from unit_tests.common.test_linear_state import make_state - - -class _SerializationFailure: - def __reduce_ex__(self, protocol): - raise RuntimeError("injected serialization failure") - - -def _assert_received_manager_roundtrip(payload, state_name, mode, expected): - resources = ForkingPickler.loads(payload) - state = getattr(resources, state_name) - assert state.big_pages_pool is state.small_pages_pool is None - assert resources.full_att_allocator is None - layout = resources.get_pd_operator().get_pd_bounded_att_page_layouts()[state_name] - block = torch.empty((1, layout.byte_end), dtype=torch.uint8, device="cuda") - state.write_to_pd_page(1, block, 0, 8, layout) - state.read_from_pd_page(block, 2, 0, 8, layout) - width = state.config.conv_width - ssm_index = 2 if mode == "replay" else 2 * (state.mtp_step + 1) - actual = (state.conv_state[:, 2, :, :width], state.ssm_state[:, ssm_index]) - for tensor, values in zip(actual, expected): - torch.testing.assert_close(tensor, torch.tensor(values, dtype=tensor.dtype, device="cuda"), rtol=0, atol=0) - assert state.mtp_state_index[2] == 0 - assert state.mtp_state_index[1] == 1 - if mode == "replay": - assert state.cursors[2] == 0 - assert state.cursors[1] == 1 - torch.cuda.synchronize() - - -def _receive_pd_managers(connection): - torch.cuda.set_device(0) - while True: - message = connection.recv() - if message is None: - return - try: - _assert_received_manager_roundtrip(*message) - connection.send(None) - except Exception: - connection.send(traceback.format_exc()) - - -@pytest.fixture(scope="module") -def pd_receiver(): - context = multiprocessing.get_context("spawn") - sender, receiver = context.Pipe() - process = context.Process(target=_receive_pd_managers, args=(receiver,)) - process.start() - receiver.close() - try: - yield sender, process - finally: - if process.is_alive(): - sender.send(None) - process.join(timeout=30) - if process.is_alive(): - process.terminate() - process.join(timeout=5) - sender.close() - assert process.exitcode == 0 - - -@pytest.mark.parametrize("cls", [GDNState, KDAState]) -@pytest.mark.parametrize("mode", ["native", "replay"]) -def test_pd_export_copy_keeps_worker_pools_and_runtime_identity(cls, mode): - state, resources, _ = make_state(cls=cls, mode=mode, device="meta") - state.big_pages_pool, state.small_pages_pool = object(), object() - before = vars(state).copy() - exported = resources.get_pd_operator().create_ipc_copy() - copy = getattr(exported, state.state_name) - assert copy.big_pages_pool is copy.small_pages_pool is None - assert copy.conv_state is state.conv_state - assert copy.ssm_state is state.ssm_state - if mode == "replay": - assert copy.cursors is state.cursors - assert copy.keys is state.keys - assert copy.mtp_state_index is state.mtp_state_index - assert all(vars(state)[name] is value for name, value in before.items()) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA IPC and pinned memory required") -@pytest.mark.parametrize("cls", [GDNState, KDAState]) -@pytest.mark.parametrize("mode", ["native", "replay"]) -@pytest.mark.parametrize("fail_serialization", [False, True]) -def test_shared_attention_state_roundtrip_preserves_worker_resources( - runtime_start_args, monkeypatch, pd_receiver, cls, mode, fail_serialization -): - state, resources, _ = make_state(cls=cls, mode=mode) - state.initialize_page_pools(big_page_capacity=2, small_page_capacity=2) - state.conv_state.copy_(torch.arange(5, dtype=torch.bfloat16, device="cuda")) - for slot in range(state.ssm_state.shape[1]): - state.ssm_state[:, slot].fill_(slot + 1) - state.mtp_state_index[1], state.mtp_state_index[2] = 1, 2 - if mode == "replay": - state.raw_keys.zero_() - state.raw_values.zero_() - state.gates.fill_(-0.5) - state.betas.zero_() - state.cursors[1], state.cursors[2] = 1, 3 - width = state.config.conv_width - if mode == "replay": - expected = [ - state.conv_state[:, 1, :, :width].cpu().tolist(), - state.ssm_state[:, 1].cpu().tolist(), - ] - else: - expected = [ - state.conv_state[:, 1, :, :width].cpu().tolist(), - state.ssm_state[:, state.mtp_step + 1].cpu().tolist(), - ] - before = vars(state).copy() - source_slot = 1 if mode == "replay" else state.mtp_step + 1 - source_conv, source_ssm = state.conv_state[:, 1].clone(), state.ssm_state[:, source_slot].clone() - snapshots = [] - for pool in (state.big_pages_pool, state.small_pages_pool): - for tensor in (pool.conv_state, pool.ssm_state): - tensor.fill_(3) - snapshots.append((tensor, tensor.data_ptr(), tensor.clone())) - monkeypatch.setitem(ForkingPickler._extra_reducers, torch.Tensor, reduce_tensor) - exported = resources.get_pd_operator().create_ipc_copy() - if fail_serialization: - with pytest.raises(RuntimeError, match="injected serialization failure"): - ForkingPickler.dumps((_SerializationFailure(), exported)) - else: - sender, receiver = pd_receiver - sender.send((ForkingPickler.dumps(exported).tobytes(), state.state_name, mode, expected)) - assert sender.poll(30), f"PD receiver did not respond; exit code: {receiver.exitcode}" - error = sender.recv() - assert error is None, error - assert all(vars(state)[name] is value for name, value in before.items()) - torch.testing.assert_close(state.conv_state[:, 1], source_conv, rtol=0, atol=0) - torch.testing.assert_close(state.ssm_state[:, source_slot], source_ssm, rtol=0, atol=0) - for tensor, pointer, contents in snapshots: - assert tensor.data_ptr() == pointer - torch.testing.assert_close(tensor, contents, rtol=0, atol=0) diff --git a/unit_tests/common/test_linear_state.py b/unit_tests/common/test_linear_state.py deleted file mode 100644 index e538002a17..0000000000 --- a/unit_tests/common/test_linear_state.py +++ /dev/null @@ -1,334 +0,0 @@ -from dataclasses import replace -from importlib import import_module - -import pytest -import torch - -from lightllm.common.attention_state import AttentionStateManager -from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState, KDAState - - -def make_state(cls=GDNState, mode="native", step=2, dtype=torch.float32, device="cuda", config=None): - config = config or ConvSSMStateConfig( - logical_layer_ids=(0, 1), - global_k_heads=1, - global_v_heads=1 if cls is KDAState else 2, - key_dim=32, - value_dim=32, - conv_width=3, - conv_dtype=torch.bfloat16, - ssm_dtype=dtype, - ) - if mode == "replay": - from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState - from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDAState - - cls = ReplayKDAState if cls is KDAState else ReplayGDNState - state = cls(config, 4, mtp_step=step, device=device, ssm_replay_config={"cache_len": 4}) - else: - state = cls(config, 4, mtp_step=step, device=device) - resources = AttentionStateManager(request_capacity=4) - resources.register_state("main", state) - resources.finish_state_registration() - layout = resources.get_pd_operator().get_pd_bounded_att_page_layouts()[state.state_name] - return state, resources, layout - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("cls", [GDNState, KDAState]) -@pytest.mark.parametrize("mode,step", [("native", 0), ("native", 2), ("replay", 0), ("replay", 2)]) -@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -def test_checkpoint_and_pd_export_same_accepted_prefix_and_restore_reused_slot( - runtime_start_args, cls, mode, step, dtype -): - state, _, layout = make_state(cls, mode, step, dtype) - state.initialize_page_pools(big_page_capacity=2, small_page_capacity=2) - src, dst, offset = 1, 2, min(1, step) - state.conv_state.copy_(torch.arange(3 + step, device="cuda").to(torch.bfloat16)) - for slot in range(state.ssm_state.shape[1]): - state.ssm_state[:, slot].fill_(slot + 1) - if state.mtp_enabled: - state.mtp_state_index[src] = offset - if mode == "native": - expected_conv = state.conv_state[:, src, ..., :3].clone() - expected_ssm = state.ssm_state[:, src * (state.mtp_step + 1)].clone() - else: - # prefill 页只保存基准窗口和 canonical SSM，不按 MTP 偏移取窗口，也不折回 history。 - expected_conv = state.conv_state[:, src, ..., :3].clone() - expected_ssm = state.ssm_state[:, src].clone() - if mode == "replay": - state.raw_keys.zero_() - state.raw_values.zero_() - state.gates.fill_(-0.5) - state.betas.zero_() - state.cursors[src] = 1 - before = state.ssm_state.clone() - block = torch.empty((1, layout.byte_end), dtype=torch.uint8, device="cuda") - state.write_to_pd_page(src, block, 0, 8, layout) - for kind in ("small", "big"): - state.save_bounded_att_runtime_to_pages([src], [8], kind, [0]) - torch.cuda.synchronize() - for pool in (state.small_pages_pool, state.big_pages_pool): - torch.testing.assert_close(pool.conv_state[0], expected_conv.cpu(), rtol=0, atol=0) - torch.testing.assert_close(pool.ssm_state[0], expected_ssm.cpu(), rtol=2e-5, atol=2e-5) - assert torch.equal(state.ssm_state, before) - if mode == "replay": - assert state.cursors[src].item() == 1 - - for restore in ( - lambda: state.restore_bounded_att_runtime_from_pages([dst], [8], "small", [0]), - lambda: state.restore_bounded_att_runtime_from_pages([dst], [8], "big", [0]), - lambda: state.read_from_pd_page(block, dst, 0, 8, layout), - lambda: state.copy_runtime_state_from_dp( - src_req_indices=[src], dst_req_indices=[dst], remote_dp_bounded_state=state - ), - ): - if state.mtp_enabled: - state.mtp_state_index[dst] = step - if mode == "replay": - state.cursors[dst] = 3 - restore() - actual_conv = state.conv_state[:, dst, ..., :3] - if mode == "replay": - actual_ssm = state.ssm_state[:, dst] - else: - actual_ssm = state.ssm_state[:, dst * (state.mtp_step + 1)] - torch.testing.assert_close(actual_conv, expected_conv, rtol=0, atol=0) - torch.testing.assert_close(actual_ssm, expected_ssm, rtol=2e-5, atol=2e-5) - if state.mtp_enabled: - assert state.mtp_state_index[dst].item() == 0 - if mode == "replay": - assert state.cursors[dst].item() == 0 - assert state.cursors[src].item() == 1 - state.reset_bounded_att_state([dst]) - assert not state.conv_state[:, dst].any() - if mode == "replay": - assert not state.ssm_state[:, dst].any() - else: - start = dst * (state.mtp_step + 1) - assert not state.ssm_state[:, start : start + state.mtp_step + 1].any() - - -def copy_on_cpu(*, copies): - for source, destination in copies: - destination.copy_(source) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("cls", [GDNState, KDAState]) -@pytest.mark.parametrize("mode", ["native", "replay"]) -@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -@pytest.mark.parametrize("destination", ["cpu", "cuda"]) -def test_checkpoint_batch_export(runtime_start_args, cls, mode, dtype, destination): - state, _, _ = make_state(cls, mode, 2, dtype) - state.initialize_page_pools(3, device=destination) - state.conv_state.copy_(torch.arange(5, device="cuda").to(torch.bfloat16)) - for slot in range(state.ssm_state.shape[1]): - state.ssm_state[:, slot].fill_(slot + 1) - state.mtp_state_index.copy_(torch.tensor([0, 2, 1, 0], device="cuda")) - if mode == "replay": - state.raw_keys.zero_() - state.raw_values.zero_() - state.gates.fill_(-0.5) - state.betas.zero_() - state.cursors.copy_(torch.tensor([1, 2, 3, 0], device="cuda")) - requests = [2, 0, 1] - width = state.config.conv_width - if mode == "replay": - expected = [(state.conv_state[:, req, :, :width].clone(), state.ssm_state[:, req].clone()) for req in requests] - else: - expected = [ - ( - state.conv_state[:, req, :, :width].clone(), - state.ssm_state[:, req * (state.mtp_step + 1)].clone(), - ) - for req in requests - ] - before = state.ssm_state.clone() - history = getattr(state, "runtime", state) - cursors = history.cursors.clone() if mode == "replay" else None - stream = torch.cuda.Stream() - stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(stream): - state.save_bounded_att_runtime_to_pages(requests, [8] * 3, "big", [1, 2, 0]) - stream.synchronize() - pool = state.big_pages_pool - for slot, (conv, ssm) in zip([1, 2, 0], expected): - torch.testing.assert_close(pool.conv_state[slot].cpu(), conv.cpu(), atol=0, rtol=0) - torch.testing.assert_close(pool.ssm_state[slot].cpu(), ssm.cpu(), atol=0, rtol=0) - assert torch.equal(state.ssm_state, before) - if cursors is not None: - assert torch.equal(state.cursors, cursors) - target, _, _ = make_state(cls, mode, 2, dtype) - if target.mtp_enabled: - target.mtp_state_index.fill_(2) - target.copy_runtime_state_from_dp(requests, [0, 1, 2], state) - for req, (conv, ssm) in enumerate(expected): - actual_conv = target.conv_state[:, req, :, : target.config.conv_width] - if mode == "replay": - actual_ssm = target.ssm_state[:, req] - else: - actual_ssm = target.ssm_state[:, req * (target.mtp_step + 1)] - torch.testing.assert_close(actual_conv, conv, atol=0, rtol=0) - torch.testing.assert_close(actual_ssm, ssm, atol=0, rtol=0) - assert torch.equal(state.ssm_state, before) - if cursors is not None: - assert torch.equal(state.cursors, cursors) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_native_checkpoint_scratch_is_bounded(runtime_start_args): - config = ConvSSMStateConfig( - logical_layer_ids=(0, 1, 2, 3), - global_k_heads=1, - global_v_heads=2, - key_dim=128, - value_dim=128, - conv_width=3, - conv_dtype=torch.bfloat16, - ssm_dtype=torch.float32, - ) - state = GDNState(config, 65, mtp_step=2, device="cuda") - state.initialize_page_pools(64) - requests = list(range(64)) - state.save_bounded_att_runtime_to_pages(requests, [128] * 64, "big", requests) - torch.cuda.synchronize() - before = torch.cuda.memory_allocated() - torch.cuda.reset_peak_memory_stats() - state.save_bounded_att_runtime_to_pages(requests, [128] * 64, "big", requests) - torch.cuda.synchronize() - assert torch.cuda.max_memory_allocated() - before < 1024 * 1024 - - -def test_pd_block_roundtrip_preserves_replicated_heads_and_accepted_mtp_state(monkeypatch): - # Exercise layout and accepted-state selection independently of the CUDA copy kernel. - monkeypatch.setattr( - import_module("lightllm.common.attention_state.bounded_att.gdn.gdn_state"), - "copy_bounded_att_tensors", - copy_on_cpu, - ) - config = ConvSSMStateConfig( - logical_layer_ids=(0,), - global_k_heads=1, - global_v_heads=2, - key_dim=2, - value_dim=3, - conv_width=3, - conv_dtype=torch.float32, - ssm_dtype=torch.float32, - tp_world_size=4, - ) - states, layouts, snapshots = [], [], [] - src, dst = 1, 2 - for rank in range(4): - state, _, layout = make_state(device="cpu", config=replace(config, tp_rank=rank)) - state.conv_state.fill_(-1) - state.ssm_state.fill_(-2) - state.mtp_state_index.fill_(1) - state.conv_state[:, src, :2, :3].fill_(10 + rank) - state.conv_state[:, src, 2:4, :3].fill_(20 + rank) - state.conv_state[:, src, 4:, :3].fill_(30 + rank) - state.ssm_state[:, src * (state.mtp_step + 1)].fill_(40 + rank) - states.append(state) - layouts.append(layout) - snapshots.append((state.conv_state.clone(), state.ssm_state.clone())) - size = layouts[0].byte_end - block = torch.full((1, size + 16), 0xA5, dtype=torch.uint8) - for state, layout in zip(states, layouts): - state.write_to_pd_page(src, block, 0, 8, layout) - expected_conv = torch.cat( - [torch.full((1, channels, 3), value) for channels, value in ((2, 10), (2, 20), (3, 30), (3, 32))], - dim=1, - ).float() - expected_ssm = torch.stack((torch.full((1, 2, 3), 40), torch.full((1, 2, 3), 42)), dim=1).float() - conv_bytes = expected_conv.numel() * 4 - torch.testing.assert_close(block[0, :conv_bytes].view(torch.float32).view_as(expected_conv), expected_conv) - torch.testing.assert_close(block[0, conv_bytes:size].view(torch.float32).view_as(expected_ssm), expected_ssm) - assert torch.all(block[0, size:] == 0xA5) - exported = block.clone() - other_requests = torch.tensor([0, 1, 3]) - for rank, (state, layout, (conv_before, ssm_before)) in enumerate(zip(states, layouts, snapshots)): - torch.testing.assert_close(state.conv_state, conv_before) - torch.testing.assert_close(state.ssm_state, ssm_before) - state.read_from_pd_page(block, dst, 0, 8, layout) - expected_local_conv = torch.cat( - (expected_conv[:, :4], expected_conv[:, 4 + rank // 2 * 3 : 7 + rank // 2 * 3]), dim=1 - ) - torch.testing.assert_close(state.conv_state[:, dst, ..., :3], expected_local_conv) - torch.testing.assert_close(state.ssm_state[:, dst * 3], expected_ssm[:, rank // 2 : rank // 2 + 1]) - assert state.mtp_state_index.tolist() == [1, 1, 0, 1] - torch.testing.assert_close(state.conv_state[:, other_requests], conv_before[:, other_requests]) - other_slots = (other_requests[:, None] * 3 + torch.arange(3)).flatten() - torch.testing.assert_close(state.ssm_state[:, other_slots], ssm_before[:, other_slots]) - assert torch.equal(block, exported) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("cls", [GDNState, KDAState]) -@pytest.mark.parametrize("step", [0, 2]) -@pytest.mark.parametrize("restore_kind", ["checkpoint", "pd", "dp"]) -def test_replay_restored_slot_decodes_like_clean_slot(runtime_start_args, cls, step, restore_kind): - from unit_tests.common.basemodel.triton_kernel.linear_att.replay_history import replay_layer - - torch.manual_seed(1690) - state, _, layout = make_state(cls=cls, mode="replay", step=step) - state.initialize_page_pools(big_page_capacity=1) - state.ssm_state.normal_(mean=0, std=0.02) - state.conv_state.normal_(mean=0, std=0.1) - block = torch.empty((1, layout.byte_end), dtype=torch.uint8, device="cuda") - state.save_bounded_att_runtime_to_pages(req_indices=[0], state_token_lens=[8], page_kind="big", page_slot_ids=[0]) - state.write_to_pd_page(req_idx=0, pd_page_buffers=block, pd_page_index=0, state_token_len=8, layout=layout) - hv, kd, vd = state.ssm_state.shape[-3:] - kh = state.num_key_heads - width = step + 1 - q = torch.randn(width, kh, kd, device="cuda", dtype=torch.bfloat16).repeat(2, 1, 1) - k = torch.randn_like(q[:width]).repeat(2, 1, 1) - v = torch.randn(width, hv, vd, device="cuda", dtype=torch.bfloat16).repeat(2, 1, 1) - a = torch.full((2 * width, hv * kd if state.kda else hv), -3.0, device="cuda") - b = torch.zeros(2 * width, hv, device="cuda") - alog = torch.zeros(hv, device="cuda") - bias = torch.zeros(a.shape[-1], device="cuda") - reqs = torch.tensor([1, 2], dtype=torch.int32, device="cuda") - cu = torch.arange(3, dtype=torch.int32, device="cuda") * width - - def decode(): - outputs = [] - for layer in state.config.logical_layer_ids: - outputs.append( - replay_layer( - state=state, - layer_idx=layer, - q=q, - k=k, - v=v, - a=a, - b=b, - a_log=alog, - bias=bias, - reqs=reqs, - cu_seqlens=cu, - run_config={"BV": 32, "num_warps": 1, "num_stages": 3}, - ) - ) - state.finish_decode_token(logical_layer_id=layer, reqs=reqs, cu_seqlens=cu) - return outputs - - decode() - assert state.cursors[1].item() == width - state.reset_bounded_att_state(req_indices=[2]) - state.restore_bounded_att_runtime_from_pages( - req_indices=[2], state_token_lens=[8], page_kind="big", page_slot_ids=[0] - ) - if restore_kind == "checkpoint": - state.restore_bounded_att_runtime_from_pages( - req_indices=[1], state_token_lens=[8], page_kind="big", page_slot_ids=[0] - ) - elif restore_kind == "pd": - state.read_from_pd_page(pd_page_buffers=block, req_idx=1, pd_page_index=0, state_token_len=8, layout=layout) - else: - state.copy_runtime_state_from_dp(src_req_indices=[0], dst_req_indices=[1], remote_dp_bounded_state=state) - assert state.cursors[1].item() == 0 - for _ in range(3): - for output in decode(): - torch.testing.assert_close(output[:width], output[width:], rtol=0, atol=0) - torch.testing.assert_close(state.ssm_state[:, 1], state.ssm_state[:, 2], rtol=0, atol=0) diff --git a/unit_tests/common/test_linear_state_contract.py b/unit_tests/common/test_linear_state_contract.py deleted file mode 100644 index fbfc4a43a9..0000000000 --- a/unit_tests/common/test_linear_state_contract.py +++ /dev/null @@ -1,241 +0,0 @@ -import pytest -import torch - -from lightllm.common.attention_state.bounded_att import KDAState -from unit_tests.common.test_linear_state import make_state -from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState -from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import advance_cursor -from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState -from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDAState -from lightllm.models.glm5_next.config import Glm5NextResourceConfig -from lightllm.server.core.objs.start_args_type import StartArgs - - -@pytest.fixture -def linear_config(): - return ConvSSMStateConfig( - logical_layer_ids=(0, 1, 3, 4, 6, 7), - global_k_heads=1, - global_v_heads=2, - key_dim=32, - value_dim=32, - conv_width=3, - conv_dtype=torch.bfloat16, - ssm_dtype=torch.float32, - ) - - -@pytest.mark.parametrize("layer,row", [(0, 0), (1, 1), (3, 2), (4, 3), (6, 4), (7, 5)]) -def test_layer_buffers_follow_packed_linear_layers(linear_config, layer, row): - owner = GDNState(linear_config, request_capacity=1, device="cpu") - for i in range(6): - owner.conv_state[i].fill_(i) - owner.ssm_state[i].fill_(i + 10) - actual_conv, actual_ssm = owner.get_layer_view(layer) - assert (actual_conv == row).all() - assert (actual_ssm == row + 10).all() - - -@pytest.mark.parametrize("layer", [-1, 2, 5, 8]) -def test_linear_mapping_rejects_full_attention_and_out_of_range_layers(linear_config, layer): - owner = GDNState(linear_config, request_capacity=1, device="cpu") - with pytest.raises(KeyError): - owner.get_layer_view(layer) - - -def test_model_configuration_retains_nondefault_kda_gate_bound(): - 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], - "gate_lower_bound": -2.5, - }, - }, - StartArgs(tp=4, data_type="bfloat16"), - ) - assert config.gate_lower_bound == -2.5 - - -def test_kda_backend_receives_gate_bound_at_construction(linear_config): - owner = ReplayKDAState(linear_config, 3, mtp_step=2, gate_lower_bound=-2.5, device="cpu") - assert owner.gate_lower_bound == -2.5 - - -@pytest.mark.parametrize("cls", [GDNState, KDAState]) -@pytest.mark.parametrize("mtp_step,expected", [(0, "inline"), (2, "precompute")]) -def test_replay_defaults_to_precompute_for_speculative_decode(linear_config, cls, mtp_step, expected): - replay_cls = ReplayKDAState if cls is KDAState else ReplayGDNState - owner = replay_cls(linear_config, request_capacity=2, mtp_step=mtp_step, device="cpu") - history = getattr(owner, "runtime", owner) - assert history.compute_mode == expected - assert history.verify_width == mtp_step + 1 - - -@pytest.mark.parametrize("cls", [GDNState, KDAState]) -def test_mtp0_decode_keeps_explicit_precompute(linear_config, cls): - replay_cls = ReplayKDAState if cls is KDAState else ReplayGDNState - owner = replay_cls( - linear_config, - request_capacity=2, - mtp_step=0, - device="cpu", - ssm_replay_config={"cache_len": 8, "compute_mode": "precompute"}, - ) - assert owner.verify_width == 1 - assert owner.compute_mode == "precompute" - assert owner.capacity == 8 - - -class _CaptureKernel: - def __init__(self): - self.args = None - - def __getitem__(self, grid): - def launch(*args, **kwargs): - self.args = args - self.kwargs = kwargs - - return launch - - -@pytest.mark.parametrize( - "module_name,cls,bv,precompute", - [ - ("replay_gdn_decode", ReplayGDNState, 32, True), - ("replay_gdn_decode", ReplayGDNState, 8, False), - ("replay_kda_decode", ReplayKDAState, 32, True), - ("replay_kda_decode", ReplayKDAState, 8, False), - ], -) -def test_mtp0_precompute_reaches_replay_kernel(monkeypatch, linear_config, module_name, cls, bv, precompute): - import importlib - - decode = importlib.import_module(f"lightllm.common.basemodel.attention.linear.{module_name}") - captured = _CaptureKernel() - monkeypatch.setattr(decode, "_decode_replay", captured) - owner = cls( - linear_config, - request_capacity=2, - mtp_step=0, - device="cpu", - ssm_replay_config={"cache_len": 8, "compute_mode": "precompute"}, - ) - heads = owner.num_key_heads - q = torch.zeros(1, heads, linear_config.key_dim) - v = torch.zeros(1, owner.ssm_state.shape[-3], linear_config.value_dim) - gate = v.shape[-3] * linear_config.key_dim if cls is ReplayKDAState else v.shape[-3] - run_config = {"BV": bv, "num_warps": 1, "num_stages": 1} - if cls is ReplayKDAState: - decode.replay( - state_ssm=owner.ssm_state[0], - state_keys=owner.keys[0], - state_deltas=owner.deltas[0], - state_gates=owner.gates[0], - state_raw_keys=owner.raw_keys[0], - state_raw_values=owner.raw_values[0], - state_betas=owner.betas[0], - state_cursors=owner.cursors, - state_tile_sync_mark=owner.tile_sync_mark, - q=q, - k=q, - v=v, - a=torch.zeros(1, gate), - b=torch.zeros(1, v.shape[-3]), - a_log=torch.zeros(v.shape[-3]), - bias=torch.zeros(gate), - reqs=torch.zeros(1, dtype=torch.int32), - cu_seqlens=torch.arange(2, dtype=torch.int32), - state_verify_width=owner.verify_width, - state_compute_mode=owner.compute_mode, - state_gate_lower_bound=owner.gate_lower_bound, - run_config=run_config, - ) - else: - decode.replay( - owner, - 0, - q, - q, - v, - torch.zeros(1, gate), - torch.zeros(1, v.shape[-3]), - torch.zeros(v.shape[-3]), - torch.zeros(gate), - torch.zeros(1, dtype=torch.int32), - torch.arange(2, dtype=torch.int32), - run_config=run_config, - ) - assert captured.kwargs["PRECOMPUTE_PROJECTIONS"] is precompute - assert captured.kwargs["KDA"] is (cls is ReplayKDAState) - - -@pytest.mark.parametrize("cls", [GDNState, KDAState]) -def test_native_commit_selects_mtp_slot(cls): - if not torch.cuda.is_available(): - pytest.skip("CUDA required") - owner, _, _ = make_state(cls, step=2) - accepted = torch.tensor([3], dtype=torch.int32, device="cuda") - owner.commit_bounded_att_state(torch.tensor([1], device="cuda"), accepted, accepted) - assert owner.mtp_state_index.tolist() == [0, 2, 0, 0] - - -@pytest.mark.parametrize("cls", [GDNState, KDAState]) -def test_replay_fused_acceptance_keeps_cursors_inside_mode(monkeypatch, cls): - import sys - - owner, _, _ = make_state(cls, mode="replay", step=3, device="cpu") - history = getattr(owner, "runtime", owner) - calls = [] - monkeypatch.setattr(sys.modules[type(history).__module__], "commit_ssm_prefix", lambda *args: calls.append(args)) - widths = torch.tensor([1, 2, 3]) - owner.commit_bounded_att_state(torch.tensor([2, 0, 1]), torch.tensor([1, 2, 3]), widths) - assert len(calls) == 1 - assert calls[0][3] is widths - assert calls[0][4] is history.cursors - assert calls[0][0] is owner.mtp_state_index - - -def _commit_prefix_cpu(indices, reqs, lengths, req_verify_width, replay_cursors) -> None: - """CPU 上按请求行核对索引和 cursor。正式提交算子只跑 CUDA。""" - - indices[reqs] = (lengths - 1).to(indices.dtype) - hold = replay_cursors.numel() - 1 - for row in range(reqs.numel()): - req = int(reqs[row]) - if req == hold: - continue - rejected = int(req_verify_width[row]) - int(lengths[row]) - replay_cursors[req] = int(replay_cursors[req]) - rejected - - -@pytest.mark.parametrize("device", ["cpu", "cuda"]) -def test_acceptance_keeps_request_order_with_strided_inputs(device): - if device == "cuda" and not torch.cuda.is_available(): - pytest.skip("CUDA required") - owner, _, _ = make_state(mode="replay", step=3, device=device) - owner.mtp_state_index.fill_(2) - history = getattr(owner, "runtime", owner) - history.cursors.zero_() - reqs = torch.tensor([2, -99, 0, -99, 3, -99], device=device)[::2] - lengths = torch.tensor([1, -99, 3, -99, 2, -99], dtype=torch.int32, device=device)[::2] - if device == "cuda": - flat_lengths = lengths.contiguous() - cu = torch.zeros(flat_lengths.numel() + 1, dtype=torch.int32, device=device) - cu[1:] = torch.cumsum(flat_lengths, 0) - advance_cursor(history, reqs.contiguous(), cu) - owner.commit_bounded_att_state(reqs, lengths, lengths) - else: - # 正式提交只接受 CUDA。CPU 用例只核对 stride 下的请求顺序。 - history.cursors[0] = 3 - history.cursors[2] = 1 - _commit_prefix_cpu(owner.mtp_state_index, reqs, lengths, lengths, history.cursors) - assert owner.mtp_state_index.tolist() == [2, 2, 0, 1] - assert history.cursors.tolist() == [3, 0, 1, 0] diff --git a/unit_tests/common/test_mem_manager_ipc.py b/unit_tests/common/test_mem_manager_ipc.py deleted file mode 100644 index 4008cd6686..0000000000 --- a/unit_tests/common/test_mem_manager_ipc.py +++ /dev/null @@ -1,110 +0,0 @@ -from multiprocessing.reduction import ForkingPickler - -import pytest -import torch -import torch.multiprocessing as mp - -from lightllm.common.attention_state import AttentionStateManager -from lightllm.common.attention_state.allocators import full_att_kv_allocator as allocator_module -from lightllm.common.attention_state.full_att import GQAAttentionConfig, GQAFullAttentionState - - -class _Counter: - def __init__(self, *args): - pass - - def set_value(self, value): - self.value = value - - -def _make_manager(monkeypatch, size, device="cpu"): - monkeypatch.setattr(allocator_module, "SharedInt", _Counter) - monkeypatch.setattr(allocator_module, "get_current_rank_in_node", lambda: 0) - manager = AttentionStateManager(request_capacity=2) - state = GQAFullAttentionState( - GQAAttentionConfig(logical_layer_groups=((0,),), global_head_num=1, head_dim=4, dtype=torch.float32), - device=device, - ) - state.allocate_storage(4) - manager.register_state("main", state) - manager.finish_state_registration() - manager.full_att_allocator = allocator_module.FullAttKVAllocator(size, pin_memory=False) - mapping = torch.arange(4, dtype=torch.int32, device=device).reshape(1, 4) - return manager, mapping - - -def _export(manager, mapping, kind): - operator = manager.get_dp_shared_operator() if kind == "dp" else manager.get_pd_operator() - return ForkingPickler.dumps((operator.create_ipc_copy(), mapping)) - - -@pytest.mark.parametrize("kind", ["dp", "pd"]) -def test_ipc_descriptor_size_does_not_scale_with_local_allocator(monkeypatch, kind): - small, small_mapping = _make_manager(monkeypatch, 256) - large, large_mapping = _make_manager(monkeypatch, 131072) - small_payload = _export(small, small_mapping, kind) - large_payload = _export(large, large_mapping, kind) - assert abs(len(large_payload) - len(small_payload)) < 128 - - -@pytest.mark.parametrize("kind", ["dp", "pd"]) -def test_ipc_roundtrip_retains_shared_buffers_and_local_owner(monkeypatch, kind): - source, mapping = _make_manager(monkeypatch, 256) - allocator = source.full_att_allocator - exported, exported_mapping = ForkingPickler.loads(_export(source, mapping, kind)) - assert source.full_att_allocator is allocator - assert allocator.alloc(1).tolist() == [0] - assert exported.full_att_allocator is None - assert exported.get_dp_shared_operator().owner is exported - assert exported.get_pd_operator().owner is exported - state = exported.get_full_attention_states()[0] - state.k_storage.fill_(17) - state.v_storage.fill_(19) - exported_mapping.fill_(23) - source_state = source.get_full_attention_states()[0] - assert torch.all(source_state.k_storage == 17) - assert torch.all(source_state.v_storage == 19) - assert torch.all(mapping == 23) - - -def _update_exported_buffers(payload, connection): - try: - exported, mapping = ForkingPickler.loads(payload) - assert exported.full_att_allocator is None - state = exported.get_full_attention_states()[0] - state.k_storage.fill_(17) - state.v_storage.fill_(19) - mapping.fill_(23) - if mapping.is_cuda: - torch.cuda.synchronize() - connection.send("complete") - finally: - connection.close() - - -@pytest.mark.parametrize("kind", ["dp", "pd"]) -@pytest.mark.parametrize("device", ["cpu", "cuda"]) -def test_ipc_buffers_remain_shared_between_processes(monkeypatch, device, kind): - source, mapping = _make_manager(monkeypatch, 256, device) - allocator = source.full_att_allocator - context = mp.get_context("spawn") - receiver, sender = context.Pipe(duplex=False) - process = context.Process(target=_update_exported_buffers, args=(_export(source, mapping, kind).tobytes(), sender)) - try: - process.start() - sender.close() - assert receiver.poll(60), "IPC consumer did not complete" - assert receiver.recv() == "complete" - process.join(10) - assert process.exitcode == 0 - assert source.full_att_allocator is allocator - state = source.get_full_attention_states()[0] - assert torch.all(state.k_storage == 17) - assert torch.all(state.v_storage == 19) - assert torch.all(mapping == 23) - finally: - if process.is_alive(): - process.terminate() - process.join(10) - receiver.close() - sender.close() diff --git a/unit_tests/common/test_pinned_memory_non_blocking_copy.py b/unit_tests/common/test_pinned_memory_non_blocking_copy.py deleted file mode 100644 index 594157560b..0000000000 --- a/unit_tests/common/test_pinned_memory_non_blocking_copy.py +++ /dev/null @@ -1,39 +0,0 @@ -import pytest -import torch -from torch.profiler import ProfilerActivity, profile - - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") - - -def _profile_request_table_copy(non_blocking: bool): - token_num = 1024 * 1024 - mem_indexes = torch.arange(token_num, dtype=torch.int32, pin_memory=True) - req_to_token_indexes = torch.empty((1, token_num), dtype=torch.int32, device="cuda") - - # 先完成分配等准备操作，避免其 CUDA runtime 调用混入待验证区间。 - torch.cuda.synchronize() - with profile( - activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], - acc_events=True, - ) as prof: - req_to_token_indexes[0].copy_(mem_indexes, non_blocking=non_blocking) - if non_blocking: - # 异步路径需要在读取结果前显式等待；这里会记录为 cudaDeviceSynchronize， - # 与 copy_ 内部可能产生的 cudaStreamSynchronize 可以明确区分。 - torch.cuda.synchronize() - - runtime_calls = {event.key for event in prof.key_averages()} - assert torch.equal(req_to_token_indexes[0].cpu(), mem_indexes) - return runtime_calls - - -def test_pinned_cpu_to_gpu_non_blocking_copy_avoids_internal_stream_sync(): - blocking_calls = _profile_request_table_copy(non_blocking=False) - non_blocking_calls = _profile_request_table_copy(non_blocking=True) - - # 两条路径底层都提交异步 H2D memcpy；同步路径随后在 copy_ 内部等待当前 stream。 - assert "cudaMemcpyAsync" in blocking_calls - assert "cudaMemcpyAsync" in non_blocking_calls - assert "cudaStreamSynchronize" in blocking_calls - assert "cudaStreamSynchronize" not in non_blocking_calls diff --git a/unit_tests/common/test_pinned_memory_stream.py b/unit_tests/common/test_pinned_memory_stream.py index f87ee97b51..064d2cb894 100644 --- a/unit_tests/common/test_pinned_memory_stream.py +++ b/unit_tests/common/test_pinned_memory_stream.py @@ -50,24 +50,6 @@ def create_pool(*args, **kwargs): submit(lambda: None).result(timeout=10) -def test_empty_batch_and_fallback_do_not_create_release_pool(monkeypatch): - copier = cuda_wrapper.PinnedH2DBatchCopier - monkeypatch.setattr(copier, "_release_pool", None) - monkeypatch.setattr(copier, "_native_batch_available", False) - monkeypatch.setattr(torch.cuda, "current_stream", lambda: SimpleNamespace(cuda_stream=0)) - - def unexpected_pool(*args, **kwargs): - pytest.fail("Empty batches and copy_ fallback must not create the source release pool") - - monkeypatch.setattr(cuda_wrapper, "ThreadPoolExecutor", unexpected_pool) - cuda_wrapper.cuda_memcpy_pinned_h2d_batch_async([], []) - source = torch.arange(8, dtype=torch.int32) - target = torch.empty_like(source) - cuda_wrapper.cuda_memcpy_pinned_h2d_batch_async([target], [source]) - assert copier._release_pool is None - assert torch.equal(target, source) - - def test_release_pool_is_initialized_once_for_concurrent_callers(monkeypatch): copier = cuda_wrapper.PinnedH2DBatchCopier monkeypatch.setattr(copier, "_release_pool", None) @@ -98,82 +80,6 @@ def get_pool(): assert creations == [{"max_workers": 1, "thread_name_prefix": "lightllm-pinned-release"}] -@requires_cuda -@pytest.mark.parametrize("native_batch", [False, True]) -def test_default_stream_copy_orders_reads_and_reuses_upload_stream(monkeypatch, native_batch): - runtime = pytest.importorskip("cuda.bindings.runtime") - copier = cuda_wrapper.PinnedH2DBatchCopier - monkeypatch.setattr(copier, "_native_batch_available", native_batch) - stream = torch.cuda.default_stream() - source = torch.arange(8192, dtype=torch.int32, pin_memory=True) - target = torch.empty(8192, dtype=torch.int32, device="cuda") - native_copy = runtime.cudaMemcpyBatchAsync - submitted_streams, observed = [], [] - - def record_batch(*args): - submitted_streams.append(args[-1]) - return native_copy(*args) - - monkeypatch.setattr(runtime, "cudaMemcpyBatchAsync", record_batch) - with torch.cuda.stream(stream): - for _ in range(2): - torch.cuda._sleep(20_000_000) - target.fill_(-1) - cuda_wrapper.cuda_memcpy_pinned_h2d_batch_async([target], [source]) - assert torch.cuda.current_stream() == stream - observed.append(target.clone()) - for result in observed: - assert torch.equal(result.cpu(), source) - if native_batch: - assert len(submitted_streams) == 2 - assert submitted_streams[0] == submitted_streams[1] != 0 - else: - assert submitted_streams == [] - - -@requires_cuda -def test_raw_batch_records_one_completion_event_without_small_copy(monkeypatch, release_futures): - runtime = pytest.importorskip("cuda.bindings.runtime") - source = torch.arange(2048, dtype=torch.int32, pin_memory=True) - another_source = torch.arange(2048, dtype=torch.int32, pin_memory=True).add_(10000) - sources = [source[128:384], source[1000:1256], another_source[17:273]] - target = torch.full((1024,), -1, dtype=torch.int32, device="cuda") - targets = [target[2:258], target[300:556], target[600:856]] - expected = torch.full((1024,), -1, dtype=torch.int32) - expected[2:258], expected[300:556], expected[600:856] = sources - stream = torch.cuda.Stream() - stream.wait_stream(torch.cuda.default_stream()) - calls = [] - native_copy, event_record = runtime.cudaMemcpyBatchAsync, torch.cuda.Event.record - - def record_native(*args): - calls.append(("batch", args[3], args[-1])) - return native_copy(*args) - - def record_event(event, on_stream=None): - calls.append(("event", 1, on_stream.cuda_stream)) - return event_record(event, on_stream) - - def unexpected_copy(*args, **kwargs): - pytest.fail("Raw batch source protection must not submit a small copy_ or enter per-tensor fallback") - - def unexpected_sync(*args, **kwargs): - pytest.fail("Normal raw batches must not synchronize a CUDA stream on the submitting thread") - - monkeypatch.setattr(runtime, "cudaMemcpyBatchAsync", record_native) - monkeypatch.setattr(torch.cuda.Event, "record", record_event) - monkeypatch.setattr(torch.Tensor, "copy_", unexpected_copy) - with torch.cuda.stream(stream), monkeypatch.context() as no_sync: - no_sync.setattr(torch.cuda.Stream, "synchronize", unexpected_sync) - no_sync.setattr(torch.cuda, "synchronize", unexpected_sync) - cuda_wrapper.cuda_memcpy_pinned_h2d_batch_async(targets, sources) - assert len(release_futures) == 1 - release_futures[0].result(timeout=10) - assert cuda_wrapper.PinnedH2DBatchCopier._native_batch_available - assert calls == [("batch", 3, stream.cuda_stream), ("event", 1, stream.cuda_stream)] - assert torch.equal(target.cpu(), expected) - - @requires_cuda def test_background_wait_holds_tensor_references_without_blocking_submitter(monkeypatch, release_futures): pytest.importorskip("cuda.bindings.runtime") @@ -254,74 +160,6 @@ def test_async_copy_keeps_freed_source_storage_out_of_reuse(monkeypatch, native_ assert torch.equal(target.cpu(), expected) -@requires_cuda -@pytest.mark.parametrize("failure", ["record", "pool", "submit"]) -def test_event_setup_failure_propagates_without_sync_or_fallback(monkeypatch, failure): - pytest.importorskip("cuda.bindings.runtime") - source = torch.arange(1024, dtype=torch.int32, pin_memory=True) - target = torch.empty_like(source, device="cuda") - stream = torch.cuda.Stream() - stream.wait_stream(torch.cuda.default_stream()) - synchronize = torch.cuda.Stream.synchronize - - def fail_setup(*args, **kwargs): - raise RuntimeError(f"test {failure} failure") - - def unexpected_sync(*args, **kwargs): - pytest.fail("Source protection setup errors must propagate without synchronous cleanup") - - def unexpected_copy(*args, **kwargs): - pytest.fail("copy_ fallback must not hide source protection failure after raw batch submission") - - if failure == "record": - monkeypatch.setattr(torch.cuda.Event, "record", fail_setup) - elif failure == "pool": - monkeypatch.setattr(cuda_wrapper.PinnedH2DBatchCopier, "_release_pool", None) - monkeypatch.setattr(cuda_wrapper, "ThreadPoolExecutor", fail_setup) - else: - monkeypatch.setattr(cuda_wrapper.PinnedH2DBatchCopier._get_release_pool(), "submit", fail_setup) - monkeypatch.setattr(torch.cuda.Stream, "synchronize", unexpected_sync) - monkeypatch.setattr(torch.Tensor, "copy_", unexpected_copy) - try: - with torch.cuda.stream(stream): - torch.cuda._sleep(20_000_000) - with pytest.raises(RuntimeError, match=f"test {failure} failure"): - cuda_wrapper.cuda_memcpy_pinned_h2d_batch_async([target], [source]) - finally: - # 测试捕获注入异常后仍要继续运行，自行持有源并等待在途传输结束。 - synchronize(stream) - assert cuda_wrapper.PinnedH2DBatchCopier._native_batch_available - assert torch.equal(target.cpu(), source) - - -@requires_cuda -def test_pytorch_fallback_needs_no_completion_event_or_pool_task(monkeypatch, release_futures): - monkeypatch.setattr(cuda_wrapper.PinnedH2DBatchCopier, "_native_batch_available", False) - source = torch.arange(1024, dtype=torch.int32, pin_memory=True) - target = torch.empty_like(source, device="cuda") - stream = torch.cuda.Stream() - stream.wait_stream(torch.cuda.default_stream()) - copy_sizes = [] - copy = torch.Tensor.copy_ - - def unexpected_event(*args, **kwargs): - pytest.fail("PyTorch copy_ already protects source storage and must not create an extra completion event") - - def record_copy(dst, src, *args, **kwargs): - copy_sizes.append(src.numel() * src.element_size()) - return copy(dst, src, *args, **kwargs) - - monkeypatch.setattr(torch.cuda, "Event", unexpected_event) - monkeypatch.setattr(torch.Tensor, "copy_", record_copy) - with torch.cuda.stream(stream): - cuda_wrapper.cuda_memcpy_pinned_h2d_batch_async([target], [source]) - stream.synchronize() - assert not cuda_wrapper.PinnedH2DBatchCopier._native_batch_available - assert release_futures == [] - assert copy_sizes == [source.numel() * source.element_size()] - assert torch.equal(target.cpu(), source) - - @requires_cuda def test_source_release_pool_is_reused_across_batches(release_futures): pytest.importorskip("cuda.bindings.runtime") diff --git a/unit_tests/common/test_req_manager_hold_page.py b/unit_tests/common/test_req_manager_hold_page.py deleted file mode 100644 index a5125f9326..0000000000 --- a/unit_tests/common/test_req_manager_hold_page.py +++ /dev/null @@ -1,27 +0,0 @@ -from types import SimpleNamespace - -import pytest -import torch - -from lightllm.common.req_manager import ReqManager - - -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) - req_manager.page_size = 4 - req_manager.bind_hold_page((100, 101, 102, 103)) - - 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() - - -def test_hold_request_indexes_require_all_requests_released(): - req_manager = ReqManager.__new__(ReqManager) - req_manager.req_list = SimpleNamespace(is_all_free=lambda: False) - - with pytest.raises(AssertionError, match="all requests are released"): - req_manager.bind_hold_page((100, 101, 102, 103)) diff --git a/unit_tests/common/test_req_manager_page.py b/unit_tests/common/test_req_manager_page.py index 4671576018..dd0da891c1 100644 --- a/unit_tests/common/test_req_manager_page.py +++ b/unit_tests/common/test_req_manager_page.py @@ -191,70 +191,3 @@ def _make_scheduler_req(manager, req_idx, current, total): req.kv_cache_alloc_need = MethodType(InferReq.kv_cache_alloc_need, req) req.alloc_kv_mem = MethodType(InferReq.alloc_kv_mem, req) return req - - -@pytest.mark.parametrize( - "hold_kv_len, available, batch_max_tokens, expected_alloc_sizes, second_req_wait_pause", - [(4, 8, 3, [4], False), (8, 0, 3, [], False), (4, 4, 6, [4], True)], -) -def test_prefill_scheduler_checks_compute_and_kv_budgets_separately( - monkeypatch, hold_kv_len, available, batch_max_tokens, expected_alloc_sizes, second_req_wait_pause -): - manager, alloc, backend, allocations = _make_scheduler(monkeypatch, 4, available, batch_max_tokens) - reqs = [_make_scheduler_req(manager, i, 3, lambda: 6) for i in range(2)] - for req in reqs: - req.alloc_kv_mem(hold_kv_len) - allocations.clear() - backend._filter_not_ready_reqs = lambda req_ids: reqs - - prefill_reqs, decode_reqs, prefill_tokens = backend._get_classed_reqs(req_ids=[0, 1]) - - assert prefill_reqs == [reqs[0]] - assert decode_reqs == [] - assert prefill_tokens == 3 - assert allocations == 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 - assert reqs[1].wait_pause is second_req_wait_pause - torch.cuda.synchronize() - assert torch.equal(manager.req_to_token_indexs[0, :8].cpu(), reqs[0].token_index.get_cpu_token_indexes(0, 8)) - - -def test_decode_scheduler_publishes_batched_request_table(monkeypatch): - manager, alloc, backend, allocations = _make_scheduler(monkeypatch, 4, 12) - req = _make_scheduler_req(manager, 0, 4, lambda: 5) - req.alloc_kv_mem(4) - allocations.clear() - backend._filter_not_ready_reqs = lambda req_ids: [req] - - prefill_reqs, decode_reqs, prefill_tokens = backend._get_classed_reqs(req_ids=[0]) - - assert prefill_reqs == [] - assert decode_reqs == [req] - assert prefill_tokens == 0 - assert req.hold_kv_len == 16 - assert allocations == [12] - torch.cuda.synchronize() - assert manager.req_to_token_indexs[0, :16].tolist() == list(range(16)) - - -@pytest.mark.parametrize("page_size, expected_alloc_token_num", [(1, 9), (3, 9), (4, 12), (7, 14), (8, 8)]) -def test_decode_small_page_preallocation_reuses_one_allocation_for_multiple_steps( - monkeypatch, page_size, expected_alloc_token_num -): - manager, alloc, backend, allocations = _make_scheduler(monkeypatch, page_size, 32) - req = _make_scheduler_req(manager, 0, page_size, lambda: req.cur_kv_len + 1) - req.alloc_kv_mem(page_size) - allocations.clear() - backend._filter_not_ready_reqs = lambda req_ids: [req] - - for _ in range(expected_alloc_token_num): - prefill_reqs, decode_reqs, prefill_tokens = backend._get_classed_reqs(req_ids=[0]) - assert prefill_reqs == [] - assert decode_reqs == [req] - assert prefill_tokens == 0 - req.cur_kv_len += 1 - - assert req.hold_kv_len == page_size + expected_alloc_token_num - assert allocations == [expected_alloc_token_num] diff --git a/unit_tests/models/deepseek2/test_rope_repeat.py b/unit_tests/models/deepseek2/test_rope_repeat.py deleted file mode 100644 index 2f578e03eb..0000000000 --- a/unit_tests/models/deepseek2/test_rope_repeat.py +++ /dev/null @@ -1,25 +0,0 @@ -import torch -import time -import pytest -from lightllm.models.deepseek2.triton_kernel.repeat_rope import repeat_rope - - -def test_torch_cat(): - source = torch.randn((100, 1, 1077), device="cuda") - dest = torch.randn((100, 7, 1077), device="cuda") - - repeat_rope(dest, source) - torch.equal(dest[:, 0, :], source) - torch.equal(dest[:, -1, :], source) - - source = torch.randn((100, 1, 128), device="cuda") - dest = torch.randn((100, 64, 128), device="cuda") - - repeat_rope(dest, source) - torch.equal(dest[:, 0, :], source) - torch.equal(dest[:, -1, :], source) - return - - -if __name__ == "__main__": - pytest.main() diff --git a/unit_tests/models/deepseek3_2/triton_kernel/test_hadamard_transform.py b/unit_tests/models/deepseek3_2/triton_kernel/test_hadamard_transform.py deleted file mode 100644 index 8a54d6d9fd..0000000000 --- a/unit_tests/models/deepseek3_2/triton_kernel/test_hadamard_transform.py +++ /dev/null @@ -1,83 +0,0 @@ -import pytest -import torch -import triton - -from lightllm.models.deepseek3_2.triton_kernel.hadamard_transform import hadamard_transform - - -TP = 8 -INDEX_N_HEADS = 64 -INDEX_HEAD_DIM = 128 -TP_INDEX_N_HEADS = INDEX_N_HEADS // TP -SCALE = INDEX_HEAD_DIM ** -0.5 - - -def _get_sgl_kernel_hadamard_transform(): - if not torch.cuda.is_available(): - pytest.skip("CUDA is required for hadamard_transform comparison") - try: - from sgl_kernel import hadamard_transform as sgl_hadamard_transform - except ImportError: - pytest.skip("sgl_kernel.hadamard_transform is not available") - return sgl_hadamard_transform - - -def _bench(fn, x): - ms = triton.testing.do_bench_cudagraph(lambda: fn(x, scale=SCALE), return_mode="median") - return ms, fn(x, scale=SCALE) - - -@pytest.mark.parametrize("tokens", [1, 16, 128, 512, 1024, 2048, 4096, 8192, 16384]) -def test_hadamard_transform_matches_sgl_kernel_deepseek_v32_shapes(tokens): - sgl_hadamard_transform = _get_sgl_kernel_hadamard_transform() - - q = torch.randn(tokens, TP_INDEX_N_HEADS, INDEX_HEAD_DIM, dtype=torch.bfloat16, device="cuda") - k = torch.randn(tokens, INDEX_HEAD_DIM, dtype=torch.bfloat16, device="cuda") - - q_expected = sgl_hadamard_transform(q, scale=SCALE) - q_actual = hadamard_transform(q, scale=SCALE) - k_expected = sgl_hadamard_transform(k, scale=SCALE) - k_actual = hadamard_transform(k, scale=SCALE) - torch.cuda.synchronize() - - assert torch.equal(q_actual, q_expected) - assert torch.equal(k_actual, k_expected) - - -def test_hadamard_transform_perf_report_deepseek_v32_shapes(): - sgl_hadamard_transform = _get_sgl_kernel_hadamard_transform() - - print( - "\nDeepSeek-V3.2 per-rank shapes with tp=8:" - "\n q: [tokens, 8, 128]" - "\n k: [tokens, 128]" - "\n\ntokens | q_diff | k_diff | sgl_q ms | tri_q ms | sgl_k ms | tri_k ms | tri(q+k) ms | slowdown q+k" - ) - - for tokens in [1, 16, 128, 512, 1024, 2048, 4096, 8192, 16384]: - q = torch.randn(tokens, TP_INDEX_N_HEADS, INDEX_HEAD_DIM, dtype=torch.bfloat16, device="cuda") - k = torch.randn(tokens, INDEX_HEAD_DIM, dtype=torch.bfloat16, device="cuda") - - q_expected = sgl_hadamard_transform(q, scale=SCALE) - q_actual = hadamard_transform(q, scale=SCALE) - k_expected = sgl_hadamard_transform(k, scale=SCALE) - k_actual = hadamard_transform(k, scale=SCALE) - torch.cuda.synchronize() - - q_diff = (q_expected.float() - q_actual.float()).abs().max().item() - k_diff = (k_expected.float() - k_actual.float()).abs().max().item() - sgl_q_ms, _ = _bench(sgl_hadamard_transform, q) - tri_q_ms, _ = _bench(hadamard_transform, q) - sgl_k_ms, _ = _bench(sgl_hadamard_transform, k) - tri_k_ms, _ = _bench(hadamard_transform, k) - sgl_sum_ms = sgl_q_ms + sgl_k_ms - tri_sum_ms = tri_q_ms + tri_k_ms - - print( - f"{tokens:6d} | {q_diff:6.1g} | {k_diff:6.1g} | " - f"{sgl_q_ms:8.4f} | {tri_q_ms:8.4f} | {sgl_k_ms:8.4f} | {tri_k_ms:8.4f} | " - f"{tri_sum_ms:11.4f} | {tri_sum_ms / sgl_sum_ms:10.2f}x" - ) - - assert q_diff == 0 - assert k_diff == 0 diff --git a/unit_tests/models/deepseek3_2/triton_kernel/test_topk_index_to_mem_index.py b/unit_tests/models/deepseek3_2/triton_kernel/test_topk_index_to_mem_index.py deleted file mode 100644 index 5e3aaf38d0..0000000000 --- a/unit_tests/models/deepseek3_2/triton_kernel/test_topk_index_to_mem_index.py +++ /dev/null @@ -1,24 +0,0 @@ -import torch -import pytest -from lightllm.models.deepseek3_2.triton_kernel.topk_index_to_mem_index import trans_topk_index_to_mem_index - - -def test_trans_topk_index_to_mem_index(): - """Test trans_topk_index_to_mem_index converts topk indices to memory indices correctly.""" - batch_size = 1 - topk = 2048 - - # Create topk_index tensor with some valid indices and some -1 (padding) - topk_index = torch.zeros((batch_size, topk), dtype=torch.int32, device="cuda") - topk_index[:, 0:2048] = torch.arange(0, 2048, dtype=torch.int32, device="cuda") - - # Create ragged_mem_index lookup table - ragged_mem_index = torch.arange(0, 2048, dtype=torch.int32, device="cuda") + 10 - - topk_mem_index = trans_topk_index_to_mem_index(topk_index, ragged_mem_index) - - assert torch.equal(topk_mem_index, (torch.arange(0, 2048, dtype=torch.int32, device="cuda") + 10).view(1, -1)) - - -if __name__ == "__main__": - pytest.main() diff --git a/unit_tests/models/deepseek_v4/test_c128_workspace.py b/unit_tests/models/deepseek_v4/test_c128_workspace.py deleted file mode 100644 index c7fcdca099..0000000000 --- a/unit_tests/models/deepseek_v4/test_c128_workspace.py +++ /dev/null @@ -1,548 +0,0 @@ -"""C128 capacity, two-stage compression alignment and graph replay regressions.""" - -import bisect -from types import SimpleNamespace as NS, MethodType -from unittest.mock import patch -import pytest -import torch -from lightllm.common.basemodel.basemodel import TpPartBaseModel -from lightllm.common.basemodel.mtp_manager import MtpManager -from lightllm.common.basemodel.cuda_graph import CudaGraph -from lightllm.common.basemodel import cuda_graph as graph_impl -from lightllm.models.deepseek_v4.workspace import DeepseekV4Workspace -from lightllm.server.core.objs.start_args_type import StartArgs -from lightllm.server.router.model_infer.mode_backend.generic_pre_process import ( - prepare_decode_inputs, - overlap_prepare_decode_inputs, -) - - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") - - -def model_for( - batch=256, - requests=256, - step=3, - tp=1, - overlap=False, - prefill_overlap=False, - graph_requests=64, - graph=True, - dynamic=False, -): - args = StartArgs( - batch_max_tokens=batch, - running_max_req_size=requests, - mtp_mode="dspark" if step else None, - mtp_step=step, - mtp_dynamic_verify=dynamic, - enable_decode_microbatch_overlap=overlap, - enable_prefill_microbatch_overlap=prefill_overlap, - enable_tpsp_mix_mode=overlap or prefill_overlap or tp > 1, - ) - mtp = MtpManager.__new__(MtpManager) - mtp.args = args - model = NS( - batch_max_tokens=batch, - max_req_num=requests, - tp_world_size_=tp, - args=args, - mtp_manager=mtp, - is_mtp_draft_model=False, - disable_cudagraph=not graph, - config={"head_dim": 512}, - req_manager=NS(HOLD_REQUEST_ID=requests), - ) - model._align_decode_batch_size = MethodType(TpPartBaseModel._align_decode_batch_size, model) - graph_requests = min(graph_requests, requests) - model.graph_max_batch_size = model._align_decode_batch_size( - (graph_requests // 2 if overlap else graph_requests) * mtp.get_decode_tokens_per_request(False) - ) - return model - - -def reqs_for(count, step): - return [ - NS( - req_idx=i, - mtp_step=step, - cur_kv_len=124, - get_cur_total_len=lambda: 125, - multimodal_params={"images": [], "audios": []}, - get_radix_cache_shared_len=lambda: 0, - shared_kv_node=None, - ) - for i in range(count) - ] - - -def inputs_for(model, count): - reqs = reqs_for(count, model.args.mtp_step) - if model.args.enable_decode_microbatch_overlap: - result = overlap_prepare_decode_inputs(reqs) - return [result[0], result[3]] - return [prepare_decode_inputs(reqs)[0]] - - -def graph_buckets(model): - with patch.object(graph_impl, "get_env_start_args", lambda: model.args): - width = model.mtp_manager.get_decode_tokens_per_request(False) - alignment = model.mtp_manager.get_decode_batch_alignment(False) - return CudaGraph.gen_cuda_graph_batch_sizes( - alignment, 32 * width, 8 * alignment, model.graph_max_batch_size, model.tp_world_size_ - ) - - -@pytest.mark.parametrize("overlap", [False, True]) -def test_reported_overflow_reproduction_is_now_covered(overlap): - m = model_for(overlap=overlap) - w = DeepseekV4Workspace(m) - w.init_c128_compress_scratch(m) - assert ( - w.c128_compress_scratch.shape == (2, 512, 512) if overlap else w.c128_compress_scratch.shape == (1, 1024, 512) - ) - for region, input in enumerate(inputs_for(m, 256)): - position = input.b_seq_len - 1 - boundaries = torch.where((position + 1) % 128 == 0)[0] - assert int(boundaries[boundaries >= 256][0]) == 259 - scratch = w.get_c128_compress_scratch(input.batch_size, region) - assert int(boundaries.max()) < scratch.shape[0] - assert w.token_capacity == 256 - - -@pytest.mark.parametrize("requests", [17, 256, 257]) -@pytest.mark.parametrize("step", [0, 3, 5]) -@pytest.mark.parametrize("tp", [1, 8]) -@pytest.mark.parametrize("overlap", [False, True]) -@pytest.mark.parametrize("graph", [False, True]) -def test_real_decode_mtp_tp_and_graph_padding_fit(requests, step, tp, overlap, graph): - m = model_for(batch=9, requests=requests, step=step, tp=tp, overlap=overlap, graph=graph) - w = DeepseekV4Workspace(m) - w.init_c128_compress_scratch(m) - buckets = graph_buckets(m) if graph else [] - for count in [0, 1, requests // 2, requests]: - for region, input in enumerate(inputs_for(m, count)): - rows = m._align_decode_batch_size(max(1, input.batch_size)) - if buckets and rows <= buckets[-1]: - rows = buckets[bisect.bisect_left(buckets, rows)] - input.input_ids = torch.zeros(input.batch_size, dtype=torch.int64) - padded = TpPartBaseModel._create_padded_decode_model_input(m, input, rows) - assert padded.batch_size == rows - assert w.get_c128_compress_scratch(rows, region).shape == (rows, 512) - for rows in buckets: - assert w.get_c128_compress_scratch(rows, 0).shape[0] == rows - assert w.token_capacity == ((9 + tp - 1) // tp) * tp - - -@pytest.mark.parametrize("dynamic", [False, True]) -def test_dynamic_verify_plans_for_maximum_width_and_odd_microbatch(dynamic): - m = model_for(batch=1, requests=17, step=5, tp=8, overlap=True, graph=True, dynamic=dynamic) - w = DeepseekV4Workspace(m) - w.init_c128_compress_scratch(m) - assert w.get_c128_compress_scratch(m._align_decode_batch_size(9 * 6), 0).shape[0] >= 54 - assert w.token_capacity == 8 - - -def test_graph_max_bucket_can_exceed_eager_bound(): - m = model_for(batch=1, requests=17, step=3, tp=8) - m.graph_max_batch_size = 256 # lower-level model API can explicitly request a larger graph bucket - w = DeepseekV4Workspace(m) - w.init_c128_compress_scratch(m) - assert w.get_c128_compress_scratch(256, 0).shape[0] == 256 - - -def test_prefill_arena_keeps_its_own_capacity(): - m = model_for(batch=4097, requests=256, step=3, tp=8, prefill_overlap=True) - w = DeepseekV4Workspace(m) - w.init_c128_compress_scratch(m) - assert w.c128_compress_scratch.shape == (2, 4104, 512) - w.init_flashmla_prefill_q(2, 4, 512, torch.bfloat16) - assert w.get_flashmla_prefill_q(4104).shape == (4104, 4, 512) - with pytest.raises(ValueError, match="capacity"): - w.get_flashmla_prefill_q(4105) - - -@pytest.mark.parametrize("overlap", [False, True]) -def test_checked_accessor_rejects_bad_rows_and_regions(overlap): - m = model_for(overlap=overlap) - w = DeepseekV4Workspace(m) - w.init_c128_compress_scratch(m) - for index in [-1, 2 if overlap else 1]: - with pytest.raises(ValueError, match="region"): - w.get_c128_compress_scratch(1, index) - for rows in [-1, w.c128_compress_scratch.shape[1] + 1]: - with pytest.raises(ValueError, match="capacity"): - w.get_c128_compress_scratch(rows, 0) - assert w.get_c128_compress_scratch(0, 0).shape == (0, 512) - - -def compress_case(workspace, batch, region=0, boundary=True): - """Four verify rows per request, with known C128 history and packed output.""" - from lightllm.models.deepseek_v4.layer_infer.compressor import fused_compress - - steps = torch.arange(4, device="cuda", dtype=torch.int32).repeat(batch) - positions = (124 if boundary else 2) + steps - request_ids = torch.arange(batch, device="cuda", dtype=torch.int32).repeat_interleave(4) - runtime = torch.randn(batch + 1, 131, 1024, device="cuda") * 0.4 - values = torch.randn(batch * 4, 1024, device="cuda") * 0.4 - packed = torch.full((batch + 1, 1168), 85, dtype=torch.uint8, device="cuda") - state = NS(config=NS(window_size=128), get_layer_view=lambda **_: runtime) - compressed = NS(config=NS(block_token_size=256), get_layer_view=lambda **_: NS(kv=packed)) - infer = NS( - position_ids=positions, - mem_index=request_ids * 256 + steps, - b_req_idx=request_ids, - b_mtp_index=steps, - b_seq_len=positions + 1, - b_ready_cache_len=None, - b_q_start_loc=None, - _dsv4_token_to_batch_idx=None, - is_prefill=False, - microbatch_index=region, - dsv4_workspace=workspace, - get_compressor_state=lambda **_: state, - get_compressed_kv_state=lambda **_: compressed, - ) - norm = torch.ones(512, device="cuda", dtype=torch.bfloat16) - cos = torch.ones(256, 32, device="cuda") - sin = torch.zeros_like(cos) - - def run(): - return fused_compress( - kv_score=values, - infer_state=infer, - layer_idx=0, - norm_weight=norm, - eps=1e-6, - head_dim=512, - qk_rope_head_dim=64, - compress_ratio=128, - cos_table=cos, - sin_table=sin, - ) - - return NS(run=run, infer=infer, runtime=runtime, values=values, packed=packed, batch=batch, region=region) - - -def check_compressed_reference(case, history, workspace): - current = case.values.view(case.batch, 4, 1024) - window = torch.cat((history[: case.batch, :124], current), dim=1) - expected = (window[..., :512] * window[..., 512:].softmax(dim=1)).sum(dim=1) - actual = workspace.get_c128_compress_scratch(case.batch * 4, case.region)[3::4] - torch.testing.assert_close(actual, expected, rtol=1e-5, atol=5e-7) - torch.testing.assert_close(case.runtime[: case.batch, 124:128], current, rtol=0, atol=0) - normalized = expected * torch.rsqrt(expected.square().mean(dim=1, keepdim=True) + 1e-6) - rope = case.packed[: case.batch, 448:576].contiguous().view(torch.bfloat16) - torch.testing.assert_close(rope.float(), normalized[:, -64:].bfloat16().float(), rtol=0.008, atol=2e-6) - # Only the first packed slot belongs to a boundary; the second one stays untouched. - assert (case.packed[: case.batch, 576:1152] == 85).all() - assert (case.packed[: case.batch, 1160:1168] == 85).all() - assert (case.packed[case.batch] == 85).all() - - -@pytest.mark.parametrize("overlap", [False, True]) -def test_c128_kernel_preserves_guards_regions_and_reference(overlap): - torch.manual_seed(1013) - model = model_for(overlap=overlap, graph=False) - workspace = DeepseekV4Workspace(model) - workspace.init_c128_compress_scratch(model) - regions, rows, head = workspace.c128_compress_scratch.shape - arena = torch.full((regions, rows + 2, head), 17.0, device="cuda") - workspace.c128_compress_scratch = arena[:, 1:-1] - batch = 128 if overlap else 256 - cases = [compress_case(workspace, batch, region) for region in range(regions)] - histories = [case.runtime.clone() for case in cases] - streams = [torch.cuda.Stream() for _ in cases] - for stream in streams: - stream.wait_stream(torch.cuda.current_stream()) - for case, stream in zip(cases, streams): - with torch.cuda.stream(stream): - case.run() - for stream in streams: - torch.cuda.current_stream().wait_stream(stream) - for case, history in zip(cases, histories): - check_compressed_reference(case, history, workspace) - assert (arena[:, 0] == 17).all() and (arena[:, -1] == 17).all() - - -def test_c128_capture_without_boundary_then_replay_with_boundary(): - torch.manual_seed(1013) - workspace = DeepseekV4Workspace(model_for(graph=True)) - workspace.init_c128_compress_scratch(model_for(graph=True)) - case = compress_case(workspace, 256, boundary=False) - for _ in range(2): - case.run() - torch.cuda.synchronize() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - case.run() - history = case.runtime.clone() - case.infer.position_ids.copy_(124 + case.infer.b_mtp_index) - case.infer.b_seq_len.copy_(case.infer.position_ids + 1) - graph.replay() - check_compressed_reference(case, history, workspace) - - -@pytest.mark.parametrize("boundary", [False, True]) -def test_short_scratch_is_rejected_before_any_compression_launch(monkeypatch, boundary): - from lightllm.models.deepseek_v4.layer_infer import compressor - - model = model_for(graph=False) - workspace = DeepseekV4Workspace(model) - workspace.init_c128_compress_scratch(model) - workspace.c128_compress_scratch = torch.empty((1, 256, 512), device="cuda") - case = compress_case(workspace, 256, boundary=boundary) - - class RejectLaunch: - def __getitem__(self, grid): - pytest.fail("unsafe compression kernel was launched") - - for name in ("_compress_window_kernel", "_compress_norm_rope_pack_kernel"): - monkeypatch.setattr(compressor, name, RejectLaunch()) - with pytest.raises(ValueError, match="capacity"): - case.run() - - -def reference_window(values, state, position, seq_len, ready_len, q_start, req_idx): - """Per-channel window softmax, matching the compress stage one boundary row at a time.""" - start = position - 127 - kv_rows = [] - score_rows = [] - for pos in range(start, start + 128): - if pos < 0 or pos >= seq_len: - kv_rows.append(torch.zeros(512, device=values.device)) - score_rows.append(torch.full((512,), float("-inf"), device=values.device)) - elif pos >= ready_len: - row = values[q_start + (pos - ready_len)] - kv_rows.append(row[:512]) - score_rows.append(row[512:]) - else: - row = state[req_idx, pos % state.shape[1]] - kv_rows.append(row[:512]) - score_rows.append(row[512:]) - kv = torch.stack(kv_rows) - weight = torch.softmax(torch.stack(score_rows), dim=0) - return (kv * weight).sum(dim=0) - - -def reference_norm(compressed, weight, eps): - variance = compressed.square().sum() / compressed.numel() - return compressed * torch.rsqrt(variance + eps) * weight.float() - - -def reference_rope(normed, position, cos, sin): - rope = normed[-64:].float().view(32, 2) - even, odd = rope[:, 0], rope[:, 1] - row = (position // 128) * 128 - rotated = torch.stack((even * cos[row] - odd * sin[row], odd * cos[row] + even * sin[row]), dim=-1) - return rotated.reshape(-1).to(torch.bfloat16) - - -def packed_rope(packed, slot): - page, token_in_page = divmod(int(slot), 2) - start = token_in_page * 576 + 448 - return packed[page, start : start + 128].contiguous().view(torch.bfloat16) - - -def stage_spec(scenario): - """Decode and prefill layouts whose boundary rows exercise the window scratch and both page slots.""" - cos = torch.randn(256, 32, device="cuda") - sin = torch.randn(256, 32, device="cuda") - norm = torch.randn(512, device="cuda").to(torch.bfloat16) - prefill = scenario.startswith("prefill") - if scenario == "decode": - batch = 2 - steps = torch.arange(4, device="cuda", dtype=torch.int32).repeat(batch) - positions = 124 + steps - request_ids = torch.arange(batch, device="cuda", dtype=torch.int32).repeat_interleave(4) - mem_index = request_ids * 256 + steps - runtime = torch.randn(batch + 1, 131, 1024, device="cuda") - values = torch.randn(batch * 4, 1024, device="cuda") - b_req_idx, b_seq_len, b_mtp = request_ids, positions + 1, steps - token_to_batch = b_ready = b_q_start = None - elif scenario == "decode_second_slot": - batch = 2 - steps = torch.arange(4, device="cuda", dtype=torch.int32).repeat(batch) - positions = 124 + steps - request_ids = torch.arange(batch, device="cuda", dtype=torch.int32).repeat_interleave(4) - mem_index = request_ids * 256 + 128 + steps - runtime = torch.randn(batch + 1, 131, 1024, device="cuda") - values = torch.randn(batch * 4, 1024, device="cuda") - b_req_idx, b_seq_len, b_mtp = request_ids, positions + 1, steps - token_to_batch = b_ready = b_q_start = None - elif scenario == "prefill_current": - positions = torch.arange(128, device="cuda", dtype=torch.int32) - mem_index = positions - runtime = torch.randn(1, 131, 1024, device="cuda") - values = torch.randn(128, 1024, device="cuda") - token_to_batch = torch.zeros(128, dtype=torch.int32, device="cuda") - b_req_idx = torch.zeros(1, dtype=torch.int32, device="cuda") - b_seq_len = torch.tensor([128], dtype=torch.int32, device="cuda") - b_ready = torch.zeros(1, dtype=torch.int32, device="cuda") - b_q_start = torch.zeros(1, dtype=torch.int32, device="cuda") - b_mtp = torch.zeros(1, dtype=torch.int32, device="cuda") - else: - positions = torch.arange(64, 192, device="cuda", dtype=torch.int32) - mem_index = 256 + positions - runtime = torch.randn(1, 131, 1024, device="cuda") - values = torch.randn(128, 1024, device="cuda") - token_to_batch = torch.zeros(128, dtype=torch.int32, device="cuda") - b_req_idx = torch.zeros(1, dtype=torch.int32, device="cuda") - b_seq_len = torch.tensor([192], dtype=torch.int32, device="cuda") - b_ready = torch.tensor([64], dtype=torch.int32, device="cuda") - b_q_start = torch.zeros(1, dtype=torch.int32, device="cuda") - b_mtp = torch.zeros(1, dtype=torch.int32, device="cuda") - packed = torch.full((int(mem_index.max()) // 128 + 2, 1168), 85, dtype=torch.uint8, device="cuda") - return NS( - positions=positions, - mem_index=mem_index, - runtime=runtime, - values=values, - packed=packed, - cos=cos, - sin=sin, - norm=norm, - eps=1e-6, - prefill=prefill, - token_to_batch=token_to_batch, - b_req_idx=b_req_idx, - b_seq_len=b_seq_len, - b_mtp=b_mtp, - b_ready=b_ready, - b_q_start=b_q_start, - ) - - -def attach_stage_case(workspace, spec, runtime, packed): - from lightllm.models.deepseek_v4.layer_infer.compressor import fused_compress - - state = NS(config=NS(window_size=128), get_layer_view=lambda **_: runtime) - compressed = NS(config=NS(block_token_size=256), get_layer_view=lambda **_: NS(kv=packed)) - infer = NS( - position_ids=spec.positions, - mem_index=spec.mem_index, - b_req_idx=spec.b_req_idx, - b_mtp_index=spec.b_mtp, - b_seq_len=spec.b_seq_len, - b_ready_cache_len=spec.b_ready, - b_q_start_loc=spec.b_q_start, - _dsv4_token_to_batch_idx=spec.token_to_batch, - is_prefill=spec.prefill, - microbatch_index=0, - dsv4_workspace=workspace, - get_compressor_state=lambda **_: state, - get_compressed_kv_state=lambda **_: compressed, - ) - - def run(): - return fused_compress( - kv_score=spec.values, - infer_state=infer, - layer_idx=0, - norm_weight=spec.norm, - eps=spec.eps, - head_dim=512, - qk_rope_head_dim=64, - compress_ratio=128, - cos_table=spec.cos, - sin_table=spec.sin, - ) - - return run - - -def launch_fused_oracle(spec, runtime, packed): - import triton - from lightllm.models.deepseek_v4.triton_kernel.compressor import _fused_compress_norm_rope_kernel - - slots = torch.where((spec.positions + 1) % 128 == 0, spec.mem_index.reshape(-1) // 128, -1) - _fused_compress_norm_rope_kernel[(spec.positions.shape[0],)]( - spec.values, - spec.values.stride(0), - spec.values.stride(1), - runtime.flatten(0, 1), - spec.positions, - spec.token_to_batch, - spec.b_req_idx, - spec.b_mtp, - spec.b_seq_len, - spec.b_ready, - spec.b_q_start, - slots, - spec.norm, - spec.eps, - spec.cos, - spec.cos.stride(0), - spec.cos.stride(1), - spec.sin, - spec.sin.stride(0), - spec.sin.stride(1), - packed, - HEAD_DIM=512, - CANDIDATE_DIM=512, - STATE_DIM=1024, - COMPRESS_RATIO=128, - WINDOW_SIZE=128, - IS_C4=False, - IS_PREFILL=spec.prefill, - STATE_CAPACITY_TOKENS=runtime.shape[1], - ROPE_HEAD_DIM=64, - FP8_MAX=torch.finfo(torch.float8_e4m3fn).max, - AMAX_MIN=1e-4, - NOPE_DIM=448, - QUANT_BLOCK=64, - SCALE_BYTES=8, - PAGE_SIZE=2, - BYTES_PER_PAGE=packed.stride(0), - BLOCK=triton.next_power_of_2(512), - IS_IN_INDEXER=False, - num_warps=4, - ) - - -@pytest.mark.parametrize("scenario", ["decode", "decode_second_slot", "prefill_current", "prefill_split"]) -def test_two_stages_match_fused_kernel_and_reference(scenario): - torch.manual_seed(1013) - spec = stage_spec(scenario) - model = model_for(graph=False) - workspace = DeepseekV4Workspace(model) - workspace.init_c128_compress_scratch(model) - staged_runtime, oracle_runtime = spec.runtime.clone(), spec.runtime.clone() - staged_packed, oracle_packed = spec.packed.clone(), spec.packed.clone() - history = spec.runtime.clone() - workspace.c128_compress_scratch.fill_(17) - - attach_stage_case(workspace, spec, staged_runtime, staged_packed)() - launch_fused_oracle(spec, oracle_runtime, oracle_packed) - torch.cuda.synchronize() - assert torch.equal(staged_packed, oracle_packed) - - positions = spec.positions - boundary = (positions + 1) % 128 == 0 - rows = positions.shape[0] - compressed = workspace.get_c128_compress_scratch(rows, 0) - assert torch.equal(compressed[~boundary], torch.full_like(compressed[~boundary], 17)) - assert torch.equal(workspace.c128_compress_scratch[0, rows], torch.full((512,), 17, device="cuda")) - - for token in boundary.nonzero(as_tuple=False).flatten().tolist(): - if spec.prefill: - ready = int(spec.b_ready[int(spec.token_to_batch[token])]) - q_start = int(spec.b_q_start[int(spec.token_to_batch[token])]) - seq_len = int(spec.b_seq_len[int(spec.token_to_batch[token])]) - req_idx = int(spec.b_req_idx[int(spec.token_to_batch[token])]) - else: - ready = int(positions[token] - spec.b_mtp[token]) - q_start = token - int(spec.b_mtp[token]) - seq_len = int(spec.b_seq_len[token]) - req_idx = int(spec.b_req_idx[token]) - expected = reference_window(spec.values, history, int(positions[token]), seq_len, ready, q_start, req_idx) - torch.testing.assert_close(compressed[token], expected, rtol=1e-5, atol=1e-6) - normed = reference_norm(compressed[token], spec.norm, spec.eps) - slot = int(spec.mem_index[token] // 128) - torch.testing.assert_close( - packed_rope(staged_packed, slot).float(), - reference_rope(normed, int(positions[token]), spec.cos, spec.sin).float(), - rtol=0, - atol=0, - ) diff --git a/unit_tests/models/deepseek_v4/test_kernels.py b/unit_tests/models/deepseek_v4/test_kernels.py deleted file mode 100644 index bf29e0e70b..0000000000 --- a/unit_tests/models/deepseek_v4/test_kernels.py +++ /dev/null @@ -1,321 +0,0 @@ -import pytest -import torch - -from lightllm.models.deepseek_v4.triton_kernel import norm_rope, topk_softplus_sqrt, indexer_topk -from lightllm.models.deepseek_v4.triton_kernel.rotary_emb import inverse_rotary_emb - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") - - -def freqs_and_positions(batch, dtype=torch.int64): - torch.manual_seed(13) - angles = torch.randn(128, 32, device="cuda") - freqs = torch.polar(torch.ones_like(angles), angles) - positions = torch.randint(0, 128, (batch,), device="cuda", dtype=dtype) - return freqs, positions - - -def rope_reference(x, freqs, positions): - pairs = torch.view_as_complex(x[..., -64:].float().reshape(*x.shape[:-1], 32, 2).contiguous()) - rotation = freqs[positions.long()] - while rotation.ndim < pairs.ndim: - rotation = rotation.unsqueeze(1) - return torch.view_as_real(pairs * rotation).flatten(-2) - - -@pytest.mark.parametrize("tokens,heads,dim_stride", [(0, 8, 1), (1, 3, 1), (7, 16, 2), (129, 64, 1)]) -@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) -def test_inverse_rotary_emb_strides_and_unwritten_channels(tokens, heads, dim_stride, dtype): - freqs, positions = freqs_and_positions(tokens) - padded = torch.randn(tokens, heads + 2, 512 * dim_stride, device="cuda", dtype=dtype) - before = padded.clone() - attention = padded[:, :heads, ::dim_stride] - expected = rope_reference(attention, freqs.conj(), positions).to(dtype) - rotation = freqs[positions] - # Independent table strides: complex real view for cos, contiguous sin. - inverse_rotary_emb(attention[..., -64:], rotation.real, rotation.imag.contiguous()) - tolerance = 0.01 if dtype == torch.bfloat16 else 1e-5 - torch.testing.assert_close(attention[..., -64:], expected, rtol=tolerance, atol=tolerance) - torch.testing.assert_close(attention[..., :-64], before[:, :heads, ::dim_stride][..., :-64], rtol=0, atol=0) - torch.testing.assert_close(padded[:, heads:], before[:, heads:], rtol=0, atol=0) - if dim_stride == 2: - torch.testing.assert_close(padded[..., 1::2], before[..., 1::2], rtol=0, atol=0) - - -def test_inverse_rotary_emb_cuda_graph_replay(): - freqs, positions = freqs_and_positions(7) - rotation = freqs[positions] - source = torch.randn(7, 5, 512, device="cuda", dtype=torch.bfloat16) - output = torch.empty_like(source) - - def run(): - inverse_rotary_emb(output[..., -64:], rotation.real, rotation.imag) - - output.copy_(source) - run() - torch.cuda.synchronize() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - run() - for _ in range(2): - source.normal_() - positions.random_(0, 128) - rotation.copy_(freqs[positions]) - output.copy_(source) - graph.replay() - expected = source.clone() - expected[..., -64:] = rope_reference(source, freqs.conj(), positions).bfloat16() - torch.testing.assert_close(output, expected, rtol=0.01, atol=0.01) - - -@pytest.mark.parametrize("batch,heads", [(0, 8), (1, 8), (7, 16), (128, 16)]) -@pytest.mark.parametrize("position_dtype", [torch.int32, torch.int64]) -def test_q_norm_rope_strided_and_bf16_rounding(batch, heads, position_dtype): - freqs, pos = freqs_and_positions(batch, position_dtype) - q = torch.randn(batch, heads + 2, 512, device="cuda", dtype=torch.bfloat16)[:, :heads] - padded = torch.full((batch, heads + 4, 512), 9, device="cuda", dtype=torch.bfloat16) - out = padded[:, :heads] - norm_rope.fused_q_norm_rope(q, out, 1e-6, freqs, pos) - normalized = (q.float() * torch.rsqrt(q.float().square().mean(-1, keepdim=True) + 1e-6)).bfloat16() - expected = normalized.clone() - expected[..., -64:] = rope_reference(normalized, freqs, pos).bfloat16() - torch.testing.assert_close(out, expected, rtol=0.01, atol=0.02) - assert (padded[:, heads:] == 9).all() - # In-place outputs must also work without changing the normalization input early. - inplace = q.contiguous().clone() - norm_rope.fused_q_norm_rope(inplace, inplace, 1e-6, freqs, pos) - torch.testing.assert_close(inplace, expected, rtol=0.01, atol=0.02) - - -@pytest.mark.parametrize("batch", [1, 7, 129]) -@pytest.mark.parametrize("position_dtype", [torch.int32, torch.int64]) -def test_k_norm_rope_cache_layout_and_unwritten_bytes(batch, position_dtype): - freqs, pos = freqs_and_positions(batch, position_dtype) - kv = torch.randn(batch, 640, device="cuda", dtype=torch.bfloat16)[:, :512] - kv[0].zero_() - weight = torch.randn(512, device="cuda", dtype=torch.bfloat16) - slots = torch.randperm(384, device="cuda")[:batch].int() - cache = torch.full((3, 74880), 165, device="cuda", dtype=torch.uint8) - norm_rope.fused_kv_norm_rope_to_swa(kv, weight, 1e-6, freqs, pos, slots, cache, 128) - values = kv.float() * torch.rsqrt(kv.float().square().mean(-1, keepdim=True) + 1e-6) * weight.float() - grouped = values[:, :448].reshape(batch, 7, 64) - scale_raw = grouped.abs().amax(-1).clamp_min(1e-4) / 448 - bits = scale_raw.contiguous().view(torch.int32) - exponent = ((bits >> 23) & 255) + ((bits & 0x7FFFFF) != 0).int() - scales = torch.ldexp(torch.ones_like(scale_raw), exponent - 127) - expected_fp8 = (grouped / scales[..., None]).clamp(-448, 448).to(torch.float8_e4m3fn) - expected_rope = rope_reference(values, freqs, pos).bfloat16() - pages = slots.long() // 128 - offsets = slots.long() % 128 - data_offsets = offsets[:, None] * 576 + torch.arange(576, device="cuda") - data = cache[pages[:, None], data_offsets] - torch.testing.assert_close( - data[:, :448].contiguous().view(torch.float8_e4m3fn).float(), - expected_fp8.reshape(batch, 448).float(), - rtol=0.13, - atol=0.03125, - ) - torch.testing.assert_close(data[:, 448:].contiguous().view(torch.bfloat16), expected_rope, rtol=0.01, atol=0.02) - scale_offsets = 128 * 576 + offsets[:, None] * 8 + torch.arange(7, device="cuda") - torch.testing.assert_close(cache[pages[:, None], scale_offsets], exponent.byte(), rtol=0, atol=0) - written = torch.zeros_like(cache, dtype=torch.bool) - written[pages[:, None], data_offsets] = True - written[pages[:, None], scale_offsets] = True - assert (cache[~written] == 165).all() - - -@pytest.mark.parametrize("batch,heads", [(0, 64), (1, 3), (7, 64), (128, 64)]) -def test_indexer_rope_hadamard_quant_and_preallocated_outputs(batch, heads): - freqs, pos = freqs_and_positions(batch) - q = torch.randn(batch, heads, 128, device="cuda", dtype=torch.bfloat16) - weights = torch.randn(batch, heads, device="cuda", dtype=torch.bfloat16) - if batch: - q[0, 0].zero_() - output = torch.empty_like(q, dtype=torch.float8_e4m3fn) - out_weights = torch.empty(batch, heads, 1, device="cuda") - got, got_weights = norm_rope.fused_q_indexer_rope_hadamard_quant(q, weights, 0.01, freqs, pos, output, out_weights) - assert got is output and got_weights is out_weights - rotated = q.float() - rotated[..., -64:] = rope_reference(q, freqs, pos) - hadamard = torch.ones(1, 1, device="cuda") - for _ in range(7): - hadamard = torch.cat((torch.cat((hadamard, hadamard), 1), torch.cat((hadamard, -hadamard), 1)), 0) - transformed = (rotated @ hadamard) * (128 ** -0.5) - scale = transformed.abs().amax(-1, keepdim=True).clamp_min(1e-4) / 448 - expected = (transformed / scale).clamp(-448, 448).to(torch.float8_e4m3fn) - torch.testing.assert_close(got.float(), expected.float(), rtol=0.13, atol=0.03125) - torch.testing.assert_close(got_weights, weights.float().unsqueeze(-1) * 0.01 * scale, rtol=4e-6, atol=1e-8) - - -def router_reference(logits, factor, bias, tokens, table, image_bias, image_start): - scores = torch.sqrt(logits.clamp_min(0) + torch.log(1 + torch.exp(-logits.abs()))) - scores = torch.nan_to_num(scores, nan=0) - selection = scores + (bias if bias is not None else 0) - if image_bias is not None: - image = tokens >= image_start - selection = torch.where(image[:, None], scores + image_bias, selection) - indices = selection.argsort(dim=-1, descending=True, stable=True)[:, :6] - if table is not None: - hash_ids = table[tokens.clamp_max(table.shape[0] - 1)] - if image_bias is None: - indices = hash_ids - else: - indices = torch.where(image[:, None], indices, hash_ids) - weights = scores.gather(1, indices) - total = weights.sum(-1, keepdim=True) - return weights * factor / torch.where(total > 0, total, 1), indices - - -@pytest.mark.parametrize("experts", [256, 384]) -@pytest.mark.parametrize("mode", ["bias", "hash", "vision", "hash_vision"]) -@pytest.mark.parametrize("rows,distribution", [(1, "normal"), (513, "normal"), (17, "ties"), (513, "underflow")]) -def test_router_modes_scaling_ties_and_underflow(experts, mode, rows, distribution): - torch.manual_seed(23) - logits = torch.randn(rows, experts, device="cuda") - bias = torch.randn(experts, device="cuda") * 0.1 - image_bias = torch.randn_like(bias) * 0.1 if "vision" in mode else None - if distribution == "ties": - logits.zero_() - bias.zero_() - if image_bias is not None: - image_bias.zero_() - elif distribution == "underflow": - logits.fill_(-100) - tokens = torch.randint(0, 32, (rows,), device="cuda") - table = torch.stack([torch.randperm(experts, device="cuda")[:6] for _ in range(32)]) if "hash" in mode else None - if image_bias is not None: - tokens[::2] = 100_000 + torch.arange(0, rows, 2, device="cuda") - correction = None if table is not None else bias - weights = torch.empty(rows, 6, device="cuda") - indices = torch.empty(rows, 6, device="cuda", dtype=torch.int64) - topk_softplus_sqrt.topk_softplus_sqrt(weights, indices, logits, 2.5, correction, tokens, table, image_bias, 32) - ref_weights, ref_indices = router_reference(logits, 2.5, correction, tokens, table, image_bias, 32) - torch.testing.assert_close(indices, ref_indices, rtol=0, atol=0) - torch.testing.assert_close(weights, ref_weights, rtol=5e-6, atol=1e-6) - - -def assert_topk(scores, lengths, pages, out, raw, page_size): - for row, length in enumerate(lengths.tolist()): - count = min(out.shape[1], length) - chosen = raw[row, :count].long() - assert (raw[row, count:] == -1).all() and (out[row, count:] == -1).all() - assert chosen.unique().numel() == count - assert ((chosen >= 0) & (chosen < length)).all() - torch.testing.assert_close( - out[row, :count], (pages[row, chosen // page_size] * page_size + chosen % page_size).int(), rtol=0, atol=0 - ) - torch.testing.assert_close( - scores[row, chosen].sort(descending=True).values, scores[row, :length].topk(count).values, rtol=0, atol=0 - ) - - -@pytest.mark.parametrize("width", [128, 512, 640, 8192, 262144]) -@pytest.mark.parametrize("page_size", [1, 64, 128]) -def test_topk_valid_prefix_raw_indices_and_padding(width, page_size): - torch.manual_seed(31) - scores = torch.randn(8, width + 128, device="cuda")[:, :width] - lengths = torch.tensor( - [0, 1, min(511, width), min(512, width), min(513, width), width // 3, width - 3, width], - device="cuda", - dtype=torch.int32, - ) - scores.masked_fill_(torch.arange(width, device="cuda")[None, :] >= lengths[:, None], float("nan")) - pages = torch.randint(0, 4096, (8, (width + page_size - 1) // page_size), device="cuda", dtype=torch.int32) - out = torch.empty(8, 512, device="cuda", dtype=torch.int32) - raw = torch.empty_like(out) - indexer_topk.indexer_topk(scores, lengths, pages, out, page_size, raw) - assert_topk(scores, lengths, pages, out, raw, page_size) - # The allocation-free in-place translation must select the same values too. - indexer_topk.indexer_topk(scores, lengths, pages, out, page_size) - valid = out >= 0 - assert valid.sum(-1).tolist() == lengths.clamp_max(512).tolist() - - -@pytest.mark.parametrize( - "width,rows,page_size", [(513, 1, 64), (640, 2, 128), (1024, 1, 1), (1024, 4, 64), (1536, 2, 128), (2048, 2, 64)] -) -@pytest.mark.parametrize("distribution", ["normal", "ties", "concentrated"]) -def test_short_topk_rank_selection_and_optional_raw(width, rows, distribution, page_size): - torch.manual_seed(43) - scores = torch.randn(rows, width + 1, device="cuda")[:, :width] - if distribution == "ties": - scores.fill_(1) - elif distribution == "concentrated": - scores.mul_(1e-6).add_(1) - lengths = torch.full((rows,), width, device="cuda", dtype=torch.int32) - if rows > 1: - lengths[0] = 513 - scores.masked_fill_(torch.arange(width, device="cuda")[None, :] >= lengths[:, None], float("nan")) - pages = torch.randint(0, 4096, (rows, (width + page_size - 1) // page_size), device="cuda", dtype=torch.int32) - out = torch.empty(rows, 512, device="cuda", dtype=torch.int32) - raw = torch.empty_like(out) - indexer_topk.indexer_topk(scores, lengths, pages, out, page_size, raw) - assert_topk(scores, lengths, pages, out, raw, page_size) - expected = out.clone() - indexer_topk.indexer_topk(scores, lengths, pages, out, page_size) - torch.testing.assert_close(out, expected, rtol=0, atol=0) - - -@pytest.mark.parametrize("rows,width", [(8, 8192), (1, 1024), (2, 2048), (4, 1024)]) -def test_topk_cuda_graph_replay_changing_lengths(rows, width): - scores = torch.randn(rows, width, device="cuda") - lengths = torch.zeros(rows, device="cuda", dtype=torch.int32) - pages = torch.randint(0, 4096, (rows, width // 64), device="cuda", dtype=torch.int32) - out = torch.empty(rows, 512, device="cuda", dtype=torch.int32) - raw = torch.empty_like(out) - fn = lambda: indexer_topk.indexer_topk(scores, lengths, pages, out, 64, raw) - fn() - torch.cuda.synchronize() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - fn() - for shift in range(8): - scores.normal_() - lengths.copy_( - torch.tensor([0, 1, 511, 512, 513, 1024, 4096, width], device="cuda", dtype=torch.int32) - .clamp_max(width) - .roll(shift)[:rows] - ) - scores.masked_fill_(torch.arange(width, device="cuda")[None, :] >= lengths[:, None], float("nan")) - graph.replay() - assert_topk(scores, lengths, pages, out, raw, 64) - - -def test_norm_cuda_graph_dependent_launch_and_replay(): - freqs, positions = freqs_and_positions(7) - source = torch.randn(7, 8, 512, device="cuda", dtype=torch.bfloat16) - first = torch.empty_like(source) - second = torch.empty_like(source) - - def run(): - norm_rope.fused_q_norm_rope(source, first, 1e-6, freqs, positions) - norm_rope.fused_q_norm_rope(first, second, 1e-6, freqs, positions) - - run() - torch.cuda.synchronize() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - run() - for _ in range(3): - source.normal_() - positions.random_(0, 128) - graph.replay() - expected = source.clone() - for _ in range(2): - expected = ( - expected.float() * torch.rsqrt(expected.float().square().mean(-1, keepdim=True) + 1e-6) - ).bfloat16() - expected[..., -64:] = rope_reference(expected, freqs, positions).bfloat16() - torch.testing.assert_close(second, expected, rtol=0.02, atol=0.03) - - -@pytest.mark.parametrize("topk", [128, 256, 1024]) -@pytest.mark.parametrize("width", [64, 640, 4096]) -def test_topk_uses_configured_output_capacity(topk, width): - scores = torch.randn(2, width, device="cuda") - lengths = torch.tensor([width // 2, width], dtype=torch.int32, device="cuda") - pages = torch.randint(0, 1000, (2, (width + 63) // 64), dtype=torch.int32, device="cuda") - out = torch.empty((2, topk), dtype=torch.int32, device="cuda") - raw = torch.empty_like(out) - indexer_topk.indexer_topk(scores, lengths, pages, out, 64, raw) - assert_topk(scores, lengths, pages, out, raw, 64) diff --git a/unit_tests/models/deepseek_v4/test_model.py b/unit_tests/models/deepseek_v4/test_model.py deleted file mode 100644 index 8394d18b48..0000000000 --- a/unit_tests/models/deepseek_v4/test_model.py +++ /dev/null @@ -1,300 +0,0 @@ -from types import SimpleNamespace - -import pytest -import torch - -from lightllm.models.deepseek_v4.model import DeepseekV4TpPartModel -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) - model.config = dict( - qk_rope_head_dim=64, - rope_theta=10000, - compress_rope_theta=160000, - max_position_embeddings=8193, - rope_scaling=dict(type="yarn", factor=16, original_max_position_embeddings=65536, beta_fast=32, beta_slow=1), - ) - model.max_seq_length = 8193 - model.layers_infer = [SimpleNamespace(compress_ratio=ratio, index_infer=SimpleNamespace()) for ratio in (0, 4, 128)] - model._init_to_get_rotary() - - positions = torch.tensor([0, 1, 127, 1023, 8192], device="cuda") - exponents = torch.arange(0, 64, 2, device="cuda", dtype=torch.float32) / 64 - angles = torch.outer(positions.float(), 1.0 / (10000 ** exponents)) - expected = torch.complex(angles.cos(), angles.sin()) - torch.testing.assert_close(model._freqs_cis_sliding[positions], expected, rtol=1e-6, atol=1e-6) - unscaled_compress = torch.outer(positions.float(), 1.0 / (160000 ** exponents)) - assert not torch.allclose( - model._freqs_cis_compress[positions], torch.complex(unscaled_compress.cos(), unscaled_compress.sin()) - ) - assert model.layers_infer[0].freqs_cis is model._freqs_cis_sliding - for layer in model.layers_infer[1:]: - assert layer.freqs_cis is model._freqs_cis_compress - - -@pytest.mark.parametrize( - "balance,ep,tpsp,tp,redundancy,error", - [ - (True, False, False, 2, 0, "enable_dp_prefill_balance"), - (False, True, True, 2, 0, "FFN and mHC token layouts"), - (False, True, False, 2, 2, "redundant experts"), - (False, True, True, 1, 0, None), - (False, True, False, 2, 0, None), - (False, False, True, 2, 0, None), - ], -) -def test_unsupported_parallel_layouts_fail_before_inference(monkeypatch, balance, ep, tpsp, tp, redundancy, error): - monkeypatch.setattr("lightllm.models.deepseek_v4.model.get_redundancy_expert_num", lambda: redundancy) - model = DeepseekV4TpPartModel.__new__(DeepseekV4TpPartModel) - model.load_way, model.tp_world_size_ = "HF", tp - model.config = {"num_attention_heads": 64, "o_groups": 8, "index_n_heads": 64} - model.args = SimpleNamespace(enable_dp_prefill_balance=balance, enable_ep_moe=ep, enable_tpsp_mix_mode=tpsp) - if error: - with pytest.raises(ValueError, match=error): - model._verify_params() - else: - model._verify_params() - - -def test_memory_check_rejects_sequence_larger_than_full_token_pool(): - model = DeepseekV4TpPartModel.__new__(DeepseekV4TpPartModel) - 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) - - with pytest.raises(AssertionError, match="max_total_token_num must be >= max_seq_length"): - model._check_mem_size() - - -def test_memory_check_still_requires_one_prefill_batch(): - model = DeepseekV4TpPartModel.__new__(DeepseekV4TpPartModel) - 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) - - with pytest.raises(AssertionError, match="greater than batch_max_tokens"): - model._check_mem_size() - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("deferred", [False, True]) -def test_draft_hidden_collector_materializes_mhc_streams(monkeypatch, deferred): - from lightllm.common.basemodel.hidden_collector import LayerHiddenCollector - from lightllm.common.basemodel.mtp_manager import MtpManager - from lightllm.models.deepseek_v4.model import DeepseekV4LayerHiddenCollector - - monkeypatch.setattr(LayerHiddenCollector, "_load_layer_ids", lambda self: frozenset([0])) - model = DeepseekV4TpPartModel.__new__(DeepseekV4TpPartModel) - model.config = {"hc_mult": 4, "hidden_size": 8} - model.layers_num = 2 - model.pre_infer = SimpleNamespace(_tpsp_allgather=lambda input, infer_state: input) - model.mtp_manager = MtpManager.__new__(MtpManager) - model.mtp_manager.args = SimpleNamespace(mtp_mode="dspark") - model._init_hidden_collector() - assert isinstance(model.hidden_collector_prototype, DeepseekV4LayerHiddenCollector) - collector = model.hidden_collector_prototype.new_instance() - - torch.manual_seed(6103) - residual = torch.randn((2, 4, 8), dtype=torch.bfloat16, device="cuda") - output = torch.randn((2, 8), dtype=torch.bfloat16, device="cuda") - post_mix = torch.randn((2, 4), device="cuda") - residual_mix = torch.randn((2, 4, 4), device="cuda") - streams = ( - torch.einsum("bsh,bst->bth", residual.float(), residual_mix) + output.float()[:, None] * post_mix[..., None] - ) - streams = streams.bfloat16() - hidden = (output, residual.flatten(1), residual_mix, post_mix) if deferred else streams.flatten(1) - - collector.add(layer_index=1, hidden=None) - collector.add(layer_index=0, hidden=hidden) - actual = collector.finish_output(SimpleNamespace(need_dp_prefill_balance=False)).spec_hidden - torch.testing.assert_close(actual, streams.mean(dim=1), rtol=0.02, atol=0.02) - assert not model.hidden_collector_prototype.layer_hiddens - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("tokens", [1, 17]) -@pytest.mark.parametrize("prefill", [False, True], ids=["decode", "prefill"]) -@pytest.mark.parametrize("truncated", [False, True], ids=["full", "truncated"]) -def test_mhc_layers_and_head_cuda_graph_replay(monkeypatch, tokens, prefill, truncated): - import triton - - from lightllm.models.deepseek_v4.layer_infer.post_layer_infer import DeepseekV4PostLayerInfer - from lightllm.models.llama.layer_infer.post_layer_infer import LlamaPostLayerInfer - from unit_tests.common.basemodel.triton_kernel.mhc.test_mhc import _hc_post_reference, _hc_pre_reference - - torch.manual_seed(6104) - triton.set_allocator(lambda size, alignment, stream: torch.empty(size, device="cuda", dtype=torch.int8)) - streams, hidden = 4, 4096 - cfg = {"hc_mult": streams, "hidden_size": hidden, "rms_norm_eps": 1e-5, "hc_eps": 1e-6} - residual = torch.randn(tokens, streams * hidden, device="cuda", dtype=torch.bfloat16) - layers = [] - weights = [] - for index in range(2): - layer = DeepseekV4TransformerLayerInfer.__new__(DeepseekV4TransformerLayerInfer) - layer.hc_mult, layer.embed_dim_ = streams, hidden - layer.eps_, layer.hc_eps, layer.sinkhorn_iters = cfg["rms_norm_eps"], cfg["hc_eps"], 20 - layer.is_last_layer = index == 1 and not truncated - layer.context_attention_forward = layer.token_attention_forward = lambda x, *_: x * 0.125 - layer._ffn = lambda x, *_: x * 0.25 - layers.append(layer) - weight = SimpleNamespace() - for prefix in ("attn", "ffn"): - setattr( - weight, - f"hc_{prefix}_fn_", - SimpleNamespace(weight=torch.randn(24, streams * hidden, device="cuda") * 0.005), - ) - setattr(weight, f"hc_{prefix}_scale_", SimpleNamespace(weight=torch.randn(3, device="cuda"))) - setattr(weight, f"hc_{prefix}_base_", SimpleNamespace(weight=torch.randn(24, device="cuda"))) - setattr( - weight, - f"{prefix}_norm_", - SimpleNamespace(weight=torch.randn(hidden, device="cuda", dtype=torch.bfloat16)), - ) - weights.append(weight) - - post = DeepseekV4PostLayerInfer.__new__(DeepseekV4PostLayerInfer) - post.alloc_tensor = torch.empty - post_weight = SimpleNamespace( - network_config_=cfg, - hc_head_fn_=SimpleNamespace(weight=torch.randn(streams, streams * hidden, device="cuda") * 0.005), - hc_head_scale_=SimpleNamespace(weight=torch.randn(1, device="cuda")), - hc_head_base_=SimpleNamespace(weight=torch.randn(streams, device="cuda")), - ) - # Keep real mHC calls and replace the logits path with its input. - monkeypatch.setattr(LlamaPostLayerInfer, "token_forward", lambda self, x, *_: x) - - def run(): - state = residual - for layer, weight in zip(layers, weights): - forward = layer.context_forward if prefill else layer.token_forward - state = forward(state, None, weight) - if isinstance(state, tuple): - x, current_residual, residual_mix, post_mix = state - assert x.shape == (tokens, hidden) - assert current_residual.shape == (tokens, streams * hidden) - assert residual_mix.shape == (tokens, streams, streams) - assert post_mix.shape == (tokens, streams) - assert isinstance(state, tuple) == truncated - return post.token_forward(state, None, post_weight) - - def reference(): - state = residual - for weight in weights: - for prefix, factor in (("attn", 0.125), ("ffn", 0.25)): - x, mix, post_mix = _hc_pre_reference( - state, - getattr(weight, f"hc_{prefix}_fn_").weight, - getattr(weight, f"hc_{prefix}_scale_").weight, - getattr(weight, f"hc_{prefix}_base_").weight, - streams, - cfg["rms_norm_eps"], - cfg["hc_eps"], - 20, - ) - normalized = x.float() * torch.rsqrt(x.float().square().mean(-1, keepdim=True) + cfg["rms_norm_eps"]) - normalized = (normalized * getattr(weight, f"{prefix}_norm_").weight.float()).bfloat16() - state = _hc_post_reference(normalized * factor, state, mix, post_mix, streams) - inv_rms = torch.rsqrt(state.float().square().mean(-1, keepdim=True) + cfg["rms_norm_eps"]) - gates = ( - torch.sigmoid( - (state.float() @ post_weight.hc_head_fn_.weight.T) * inv_rms * post_weight.hc_head_scale_.weight - + post_weight.hc_head_base_.weight - ) - + cfg["hc_eps"] - ) - return (state.view(tokens, streams, hidden).float() * gates[..., None]).sum(1).bfloat16() - - torch.testing.assert_close(run(), reference(), atol=0.06, rtol=0.04) - torch.cuda.synchronize() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - actual = run() - for _ in range(2): - residual.normal_() - graph.replay() - torch.testing.assert_close(actual, reference(), atol=0.06, rtol=0.04) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("tp_size,rank", [(1, 0), (2, 0), (2, 1)]) -def test_output_projection_matches_bf16_reference(monkeypatch, tp_size, rank): - from lightllm.common.basemodel.layer_weights.base_layer_weight import BaseLayerWeight - from lightllm.common.basemodel.layer_weights.meta_weights import platform_op - from lightllm.models.deepseek_v4.layer_weights.transformer_layer_weight import DeepseekV4TransformerLayerWeight - - monkeypatch.setattr( - platform_op, - "get_env_start_args", - lambda: SimpleNamespace(hardware_platform="cuda", enable_torch_fallback=False, enable_triton_fallback=False), - ) - monkeypatch.setenv("LIGHTLLM_DP_WORLD_SIZE", str(tp_size)) - monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_DP", str(rank)) - monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") - torch.manual_seed(6102) - layer = DeepseekV4TransformerLayerWeight.__new__(DeepseekV4TransformerLayerWeight) - BaseLayerWeight.__init__(layer) - layer.prefix, layer.data_type_ = "layers.0", torch.bfloat16 - layer.hidden, layer.n_heads, layer.head_dim = 128, 8, 64 - layer.q_lora_rank, layer.o_lora_rank, layer.o_groups = 128, 128, 2 - layer.get_quant_method = lambda name: None - layer._init_qkvo() - layer.experts = SimpleNamespace(quant_method=SimpleNamespace(weight_scale_suffix=None)) - - weight_name = "layers.0.attn.wo_a.weight" - scale_name = "layers.0.attn.wo_a.scale" - packed = torch.randn(256, 256, device="cuda").to(torch.float8_e4m3fn) - scales = torch.tensor([[0.125, 0.25], [0.5, 1]], device="cuda").to(torch.float8_e8m0fnu) - expected_weight = packed.float().unflatten(0, (-1, 128)).unflatten(-1, (-1, 128)) * scales[:, None, :, None].float() - expected_weight = expected_weight.flatten(2, 3).flatten(0, 1).bfloat16() - weight_b = torch.randn(128, 256, device="cuda", dtype=torch.bfloat16) * 0.1 - layer.load_hf_weights({weight_name: packed, scale_name: scales, "layers.0.attn.wo_b.weight": weight_b}) - assert layer.wo_a_.verify_load() - local_groups = layer.o_groups // tp_size - grouped_weight = expected_weight.view(2, 128, 256)[rank * local_groups : (rank + 1) * local_groups] - heads = layer.n_heads // tp_size - padded = torch.full((5, 64, 64), 9, device="cuda", dtype=torch.bfloat16) - attention = padded[:, :heads] - attention.normal_(std=0.1) - angles = torch.randn(5, 32, device="cuda") - cos, sin = angles.cos(), angles.sin() - pairs = torch.view_as_complex(attention.float().reshape(5, heads, 32, 2).contiguous()) - rotated = torch.view_as_real(pairs * torch.complex(cos, sin).conj()[:, None]).flatten(-2).bfloat16() - expected = torch.einsum("tgd,grd->tgr", rotated.reshape(5, local_groups, 256), grouped_weight).flatten(1) - expected = expected @ weight_b[:, rank * local_groups * 128 : (rank + 1) * local_groups * 128].T - - infer_layer = DeepseekV4TransformerLayerInfer.__new__(DeepseekV4TransformerLayerInfer) - infer_layer.qk_rope_head_dim, infer_layer.tp_groups, infer_layer.compress_ratio = 64, local_groups, 0 - infer_layer._tpsp_reduce = lambda input, infer_state: input - state = SimpleNamespace(position_cos_sliding=cos, position_sin_sliding=sin) - torch.testing.assert_close(infer_layer._get_o(attention, state, layer), expected, rtol=0.01, atol=0.01) - assert padded[:, heads:].eq(9).all() diff --git a/unit_tests/models/deepseek_v4/test_shared_expert_overlap.py b/unit_tests/models/deepseek_v4/test_shared_expert_overlap.py deleted file mode 100644 index c088d811eb..0000000000 --- a/unit_tests/models/deepseek_v4/test_shared_expert_overlap.py +++ /dev/null @@ -1,247 +0,0 @@ -"""Shared-expert overlap on the EP prefill and decode paths.""" - -from types import SimpleNamespace -import pytest -import torch -from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe_ep import ( - fused_experts_impl, - prefill_decode_sync as sync, -) -from lightllm.models.deepseek_v4.layer_infer import transformer_layer_infer as layer_module -from lightllm.models.deepseek_v4.layer_infer.transformer_layer_infer import DeepseekV4TransformerLayerInfer -from lightllm.models.deepseek_v4.layer_weights.transformer_layer_weight import DeepseekV4TransformerLayerWeight - - -def _layer(ep_mode): - layer = DeepseekV4TransformerLayerInfer.__new__(DeepseekV4TransformerLayerInfer) - layer.embed_dim_ = 4 - layer.layer_num_ = 0 - layer.enable_ep_moe = ep_mode - layer.swiglu_limit = 7 - layer._tpsp_allgather = lambda input, infer_state: input - ids = torch.zeros(3, 2, dtype=torch.int64) - layer._select_experts = lambda *args: (torch.ones(3, 2), ids) - return layer, ids - - -def _shared_recorder(events, scale=2): - def shared(*args, **kwargs): - x = kwargs["input"] if "input" in kwargs else args[0] - events.append("shared") - return x * scale - - return shared - - -def _moe_inputs(): - hidden = torch.ones(3, 4) - topk_weights = torch.ones(3, 2) - topk_idx = torch.zeros(3, 2, dtype=torch.int64) - w1 = torch.ones(2, 8, 4) - w2 = torch.ones(2, 4, 4) - return hidden, topk_weights, topk_idx, w1, w2 - - -def test_ep_ffn_passes_the_shared_expert_callback(monkeypatch): - monkeypatch.setattr(layer_module, "use_sm100_mega_moe", lambda method: False) - layer, _ids = _layer(True) - events = [] - layer._ffn_shared_expert = _shared_recorder(events) - - def routed(**kwargs): - events.append("routed") - return torch.full((3, 4), 3.0) + kwargs["shared_expert_func"]() - - experts = SimpleNamespace(quant_method=None, experts_with_topk=routed) - result = layer._ffn( - torch.ones(3, 4), - SimpleNamespace(is_prefill=False), - SimpleNamespace(experts=experts), - ) - assert events == ["routed", "shared"] - assert torch.all(result == 5) - - -def test_decode_runs_shared_expert_between_combine_send_and_recv(monkeypatch): - events = [] - hidden, topk_weights, topk_idx, w1, w2 = _moe_inputs() - import deep_ep - - monkeypatch.setattr(deep_ep.ElasticBuffer, "capture", staticmethod(lambda: events.append("event") or "event")) - monkeypatch.setattr(sync, "per_token_group_quant_fp8", lambda *args, **kwargs: (hidden, hidden)) - monkeypatch.setattr(sync, "get_mk_alignment_for_contiguous_layout", lambda: 128) - monkeypatch.setattr(sync, "get_deepep_num_max_dispatch_tokens_per_rank_decode", lambda: 16) - monkeypatch.setattr(sync, "get_ep_num_sms", lambda overlap_with_compute: 8) - - def dispatch(*args, **kwargs): - events.append("dispatch") - handle = SimpleNamespace( - psum_num_recv_tokens_per_expert=torch.tensor([3, 3]), - expert_alignment=128, - recv_src_metadata=torch.zeros(3, 4, dtype=torch.int32), - psum_num_recv_tokens_per_scaleup_rank=torch.tensor([3]), - ) - return (hidden, hidden), None, topk_weights, handle, None - - def gemm(**kwargs): - assert kwargs["expected_m"] == 3 # ceil(3 tokens * 8 ranks * topk 2 / 16 experts) - events.append("gemm") - return hidden - - def reduce(**kwargs): - events.append("reduce") - return hidden, kwargs["recv_src_metadata"] - - class _Event: - def current_stream_wait(self): - events.append("combine_recv") - - def combine(*args, **kwargs): - assert kwargs["previous_event"] == "event" - assert kwargs["async_with_compute_stream"] is True - events.append("combine_send") - return torch.full((3, 4), 3.0), None, _Event() - - monkeypatch.setattr(sync, "decode_masked_group_gemm", gemm) - monkeypatch.setattr(sync, "ep_reduce_decode_output", reduce) - result = fused_experts_impl( - hidden, - w1, - w2, - topk_weights, - topk_idx, - num_experts=16, - buffer=SimpleNamespace(num_ranks=8, dispatch=dispatch, combine=combine), - is_prefill=False, - use_fp8_w8a8=True, - use_fp8_all2all=True, - shared_expert_func=lambda: events.append("shared") or torch.full((3, 4), 2.0), - ) - assert events == ["dispatch", "gemm", "reduce", "event", "combine_send", "shared", "combine_recv"] - assert torch.all(result == 5) - - -def test_prefill_overlaps_shared_expert_with_combine(monkeypatch): - events = [] - hidden, topk_weights, topk_idx, w1, w2 = _moe_inputs() - import deep_ep - - monkeypatch.setattr(deep_ep.ElasticBuffer, "capture", staticmethod(lambda: events.append("event") or "event")) - monkeypatch.setattr(sync, "get_deepep_num_max_dispatch_tokens_per_rank_prefill", lambda: 16) - monkeypatch.setattr(sync, "per_token_group_quant_fp8", lambda *args, **kwargs: (hidden, hidden)) - monkeypatch.setattr(sync, "get_mk_alignment_for_contiguous_layout", lambda: 128) - monkeypatch.setattr(sync, "get_ep_num_sms", lambda overlap_with_compute: 8) - - def dispatch(*args, **kwargs): - events.append("dispatch") - handle = SimpleNamespace( - num_recv_tokens_per_expert_list=[0], - recv_src_metadata=torch.zeros(3, 4), - ) - return (hidden, hidden), None, topk_weights, handle, None - - class _Event: - def current_stream_wait(self): - events.append("combine_recv") - - def combine(*args, **kwargs): - assert kwargs["previous_event"] == "event" - assert kwargs["async_with_compute_stream"] is True - events.append("combine_send") - return torch.full((3, 4), 3.0), None, _Event() - - result = fused_experts_impl( - hidden, - w1, - w2, - topk_weights, - topk_idx, - num_experts=16, - buffer=SimpleNamespace(dispatch=dispatch, combine=combine), - is_prefill=True, - use_fp8_w8a8=True, - use_fp8_all2all=True, - shared_expert_func=lambda: events.append("shared") or torch.full((3, 4), 2.0), - ) - assert events == ["dispatch", "event", "combine_send", "shared", "combine_recv"] - assert torch.all(result == 5) - - -def test_rename_shared_expert_as_the_last_routed_expert(): - weight = DeepseekV4TransformerLayerWeight.__new__(DeepseekV4TransformerLayerWeight) - weight.prefix = "layers.0" - weight.n_routed_experts = 4 - weight.experts = SimpleNamespace( - quant_method=SimpleNamespace(weight_scale_suffix="weight_scale_inv", weight_zero_point_suffix=None) - ) - tensors = { - "layers.0.ffn.shared_experts.w1.weight": torch.ones(2), - "layers.0.ffn.shared_experts.w1.scale": torch.ones(1), - "layers.0.ffn.shared_experts.w2.weight": torch.ones(2), - "layers.0.ffn.shared_experts.w3.weight": torch.ones(2), - } - weight._rename_shared_experts(tensors) - assert torch.equal(tensors["layers.0.ffn.experts.4.w1.weight"], torch.ones(2)) - assert torch.equal(tensors["layers.0.ffn.experts.4.w1.scale"], torch.ones(1)) - assert "layers.0.ffn.experts.4.w2.weight" in tensors - assert "layers.0.ffn.experts.4.w3.weight" in tensors - - -def test_auto_enable_fused_shared_experts_for_v4_tp(monkeypatch): - from lightllm.utils import config_utils - - monkeypatch.setattr(config_utils, "get_model_type", lambda path: "deepseek_v4") - args = SimpleNamespace(model_dir="/model", enable_ep_moe=False, enable_fused_shared_experts=False) - config_utils.auto_set_fused_shared_experts(args) - assert args.enable_fused_shared_experts is True - - args.enable_ep_moe = True - args.enable_fused_shared_experts = False - config_utils.auto_set_fused_shared_experts(args) - assert args.enable_fused_shared_experts is False - - -def test_tp_fused_shared_expert_skips_the_standalone_ffn(monkeypatch): - monkeypatch.setattr(layer_module, "use_sm100_mega_moe", lambda method: False) - layer, _ids = _layer(False) - events = [] - layer._ffn_shared_expert = _shared_recorder(events) - reductions = [] - layer._tpsp_reduce = lambda input, infer_state: reductions.append(input.clone()) or input - - def routed(**kwargs): - events.append("routed") - return torch.full((3, 4), 3.0) - - experts = SimpleNamespace(quant_method=None, experts_with_topk=routed) - result = layer._ffn( - torch.ones(3, 4), - SimpleNamespace(is_prefill=False), - SimpleNamespace(experts=experts, num_fused_shared_experts=1), - ) - assert events == ["routed"] - assert torch.all(result == 3) - assert len(reductions) == 1 - assert torch.all(reductions[0] == 3) - - -@pytest.mark.parametrize("ep_mode", [False, True]) -def test_inplace_backends_compute_shared_expert_first(monkeypatch, ep_mode): - monkeypatch.setattr(layer_module, "use_sm100_mega_moe", lambda method: ep_mode) - layer, _ids = _layer(ep_mode) - events = [] - x = torch.ones(3, 4) - layer._ffn_shared_expert = _shared_recorder(events) - - def routed(**kwargs): - events.append("routed") - kwargs["input_tensor"].fill_(9) - return torch.full((3, 4), 3.0) - - reductions = [] - layer._tpsp_reduce = lambda input, infer_state: reductions.append(input) or input - experts = SimpleNamespace(quant_method=None, experts_with_topk=routed) - result = layer._ffn(x, SimpleNamespace(is_prefill=False), SimpleNamespace(experts=experts)) - assert events == ["shared", "routed"] - assert torch.all(result == 5) - assert len(reductions) == (0 if ep_mode else 1) diff --git a/unit_tests/models/deepseek_v4/test_tokenizer.py b/unit_tests/models/deepseek_v4/test_tokenizer.py deleted file mode 100644 index b96db74cab..0000000000 --- a/unit_tests/models/deepseek_v4/test_tokenizer.py +++ /dev/null @@ -1,79 +0,0 @@ -import copy -from types import SimpleNamespace - -import pytest - -from lightllm.models.deepseek_v4.tokenizer import DeepSeekV4Tokenizer - - -IMAGE_TOKEN = "<｜deepseek_image｜>" - - -class _Tokenizer: - def convert_tokens_to_ids(self, token): - assert token == IMAGE_TOKEN - return 99 - - def encode(self, prompt, **kwargs): - ids = [] - for index, text in enumerate(prompt.split(IMAGE_TOKEN)): - if index: - ids.append(99) - ids.extend([7] * len(text)) - return ids - - -@pytest.fixture -def tokenizer(tmp_path): - # An empty model directory selects the bundled encoder used by default. - return DeepSeekV4Tokenizer(_Tokenizer(), str(tmp_path), {"vision_n_layers": 1}) - - -@pytest.mark.parametrize("tokenize", [False, True]) -def test_mixed_content_preserves_image_order_and_expands_placeholders(tokenizer, tokenize): - messages = [ - { - "role": "user", - "content": [ - {"type": "text", "text": "before"}, - {"type": "image_url", "image_url": {"url": "first"}}, - {"type": "text", "text": "between"}, - {"type": "image_url", "image_url": {"url": "second"}}, - {"type": "text", "text": "after"}, - ], - } - ] - original = copy.deepcopy(messages) - prompt = tokenizer.apply_chat_template(messages=messages, tokenize=tokenize) - if not tokenize: - assert f"before{IMAGE_TOKEN}between{IMAGE_TOKEN}after" in prompt - images = [SimpleNamespace(token_id=1000, token_num=13), SimpleNamespace(token_id=2000, token_num=17)] - ids = tokenizer.encode(prompt, SimpleNamespace(images=images)) - for image in images: - expected = list(range(image.token_id + image.start_idx % 4, image.token_id + image.token_num)) - assert ids[image.start_idx : image.start_idx + len(expected)] == expected - assert images[0].start_idx < images[1].start_idx - assert 99 not in ids - assert messages == original - - -@pytest.mark.parametrize( - "kwargs,thinking", - [ - ({"reasoning_effort": "none"}, False), - ({"reasoning_effort": "none", "thinking": False}, False), - ({"reasoning_effort": "none", "enable_thinking": False}, False), - ({"reasoning_effort": "none", "thinking": True}, True), - ({"reasoning_effort": "high"}, True), - ({"reasoning_effort": "max"}, True), - ({}, True), - ], -) -def test_reasoning_effort_uses_encoder_modes(tokenizer, kwargs, thinking): - prompt = tokenizer.apply_chat_template(messages=[{"role": "user", "content": "hello"}], **kwargs) - assert prompt.endswith("" if thinking else "") - - -def test_unknown_reasoning_effort_is_not_silently_replaced(tokenizer): - with pytest.raises(AssertionError, match="Invalid reasoning effort: low"): - tokenizer.apply_chat_template(messages=[{"role": "user", "content": "hello"}], reasoning_effort="low") diff --git a/unit_tests/models/deepseek_v4/test_vision.py b/unit_tests/models/deepseek_v4/test_vision.py deleted file mode 100644 index ad931ac2b5..0000000000 --- a/unit_tests/models/deepseek_v4/test_vision.py +++ /dev/null @@ -1,350 +0,0 @@ -import importlib.util -import io -import json -import os -import sys -from types import SimpleNamespace - -import numpy as np -import pytest -import torch -from PIL import Image -from torch import nn -from torch.nn.attention import SDPBackend, sdpa_kernel - -from lightllm.common.basemodel.attention_vit.sdpa.fp import SdpaVitAttBackend -from lightllm.models.deepseek_v4 import deepseek_v4_visual as visual -from lightllm.models.deepseek_v4 import image_processor as ours -from lightllm.server import visualserver - - -REFERENCE_DIR = "/mtc/models/DeepSeek-V4-Flash-Vision-Exp/inference" -REFERENCE_PROCESSOR_PATH = os.path.join(REFERENCE_DIR, "image_processor.py") -REFERENCE_VISION_PATH = os.path.join(REFERENCE_DIR, "vision.py") - -PATCH_SIZE = 14 -DOWNSAMPLE_RATIO = 3 -MAX_N_TOKEN = 384 -MIN_PIXELS = 147456 -MAX_WH_RATIO = 8 - -REFERENCE_ARGS = SimpleNamespace( - vision_patch_size=PATCH_SIZE, - vision_downsample_ratio=DOWNSAMPLE_RATIO, - vision_max_n_token=MAX_N_TOKEN, - vision_min_pixels=MIN_PIXELS, - vision_max_wh_ratio=MAX_WH_RATIO, -) - - -def _load_reference_processor(): - spec = importlib.util.spec_from_file_location("deepseek_v4_reference_image_processor", REFERENCE_PROCESSOR_PATH) - module = importlib.util.module_from_spec(spec) - sys.modules[spec.name] = module - spec.loader.exec_module(module) - return module - - -def _load_reference_vision(): - spec = importlib.util.spec_from_file_location("deepseek_v4_reference_vision", REFERENCE_VISION_PATH) - module = importlib.util.module_from_spec(spec) - sys.modules[spec.name] = module - spec.loader.exec_module(module) - return module - - -@pytest.fixture(scope="module") -def reference_processor(): - if not os.path.exists(REFERENCE_PROCESSOR_PATH): - pytest.skip("DeepSeek-V4 vision reference processor is not available") - return _load_reference_processor() - - -@pytest.fixture(scope="module") -def reference_vision(): - if not os.path.exists(REFERENCE_VISION_PATH): - pytest.skip("DeepSeek-V4 reference vision tower is not available") - return _load_reference_vision() - - -@pytest.mark.parametrize( - "height,width", - [(14, 14), (15, 30), (41, 512), (100, 100), (410, 756), (756, 42)], -) -def test_grid_tokens_matches_reference(reference_processor, height, width): - assert ours.grid_tokens(height, width, PATCH_SIZE, DOWNSAMPLE_RATIO) == reference_processor.grid_tokens( - height, width, PATCH_SIZE, DOWNSAMPLE_RATIO - ) - - -@pytest.mark.parametrize("start_pos", range(8)) -@pytest.mark.parametrize("n_llm_h,n_llm_w", [(1, 1), (2, 3), (3, 2), (4, 5), (7, 8)]) -def test_image_block_and_permutation_match_reference(reference_processor, n_llm_h, n_llm_w, start_pos): - expected_types, expected_perm = reference_processor.build_image_block(n_llm_h, n_llm_w, start_pos) - types, perm = ours.build_image_block(n_llm_h, n_llm_w, start_pos) - assert torch.equal(types, expected_types) - assert torch.equal(perm, expected_perm) - - -@pytest.mark.parametrize("prefix_len", [0, 1, 2, 3, 511, 512, 513, 514]) -@pytest.mark.parametrize("text_gap", range(4)) -@pytest.mark.parametrize("reuse_image", [False, True]) -def test_tokenizer_spans_match_reference_blocks(reference_processor, prefix_len, text_gap, reuse_image): - from lightllm.models.deepseek_v4.tokenizer import DeepSeekV4Tokenizer - from lightllm.server.multimodal_params import MultimodalParams - from lightllm.utils.multimodal_chunk_utils import adjust_mm_chunk_end, clamp_mm_cache_len, get_mm_input_spans - - tokenizer = DeepSeekV4Tokenizer( - SimpleNamespace(convert_tokens_to_ids=lambda token: 99), "/unused", {"vision_n_layers": 1} - ) - grids = [(1, 1), (3, 2), (1, 1)] - params = MultimodalParams(images=[{"type": "url", "data": "unused"} for _ in grids]) - prompt = [7] * prefix_len - expected_ids = prompt.copy() - expected_spans = [] - expected_cache = [] - for index, (image, (height, width)) in enumerate(zip(params.images, grids)): - canonical_types, _ = reference_processor.build_image_block(height, width, start_pos=0) - image.token_id = 2 ** 62 + (0 if reuse_image and index == 2 else index * 1024) - image.token_num = len(canonical_types) - image.start_index_in_embed_cache = 100 + (0 if reuse_image and index == 2 else index * 1024) - expected_cache.append((image.token_id, image.token_num, image.start_index_in_embed_cache)) - start = len(expected_ids) - actual_types, _ = reference_processor.build_image_block(height, width, start_pos=start) - suffix_start = len(canonical_types) - len(actual_types) - assert torch.equal(canonical_types[suffix_start:], actual_types) - expected_ids.extend(range(image.token_id + suffix_start, image.token_id + image.token_num)) - expected_spans.append((start, len(expected_ids))) - prompt.extend([99] + [8] * text_gap) - expected_ids.extend([8] * text_gap) - - input_ids = tokenizer.encode(prompt, params) - assert input_ids == expected_ids - assert [ - (image.token_id, image.token_num, image.start_index_in_embed_cache) for image in params.images - ] == expected_cache - spans = get_mm_input_spans(np.asarray(input_ids, dtype=np.int64), params.to_dict()) - assert spans == expected_spans - for candidate in range(len(input_ids) + 1): - expected_end = next((start for start, end in expected_spans if start < candidate < end), candidate) - assert clamp_mm_cache_len(candidate, spans) == expected_end - - for chunk_size in (1, 4, 16, 256, 512): - start = 0 - while start < len(input_ids): - end = adjust_mm_chunk_end(start, min(start + chunk_size, len(input_ids)), spans) - assert start < end <= len(input_ids) - assert not any(left < end < right for left, right in expected_spans) - start = end - - -@pytest.mark.parametrize( - "width,height", - [(800, 600), (100, 2000), (2000, 100), (50, 50), (13, 7), (384, 384)], -) -def test_image_geometry_and_pixels_match_reference(reference_processor, width, height): - rng = np.random.default_rng(width * 10000 + height) - image = Image.fromarray(rng.integers(0, 256, (height, width, 3), dtype=np.uint8)) - image_bytes = io.BytesIO() - image.save(image_bytes, format="PNG") - - expected = reference_processor.load_image({"data": image_bytes.getvalue()}, REFERENCE_ARGS) - actual = ours.load_image(image, REFERENCE_ARGS) - assert actual[1:] == expected[1:] - assert torch.equal(actual[0], expected[0]) - - n_vit_h, n_vit_w, n_llm_h, n_llm_w, token_num = ours.get_image_grid( - height, - width, - PATCH_SIZE, - DOWNSAMPLE_RATIO, - MAX_N_TOKEN, - MIN_PIXELS, - MAX_WH_RATIO, - ) - assert (n_vit_h, n_vit_w, n_llm_h, n_llm_w) == actual[1:] - reference_types, _ = reference_processor.build_image_block(n_llm_h, n_llm_w, start_pos=0) - assert token_num == len(reference_types) - assert token_num <= MAX_N_TOKEN - - -def _tiny_model(): - return visual.DeepseekV4VisionModel( - {"weight_dir": "/unused"}, - hidden_size=8, - vision_n_layers=1, - vision_dim=8, - vision_n_heads=2, - vision_inter_dim=12, - vision_patch_size=2, - vision_rope_theta=10000.0, - vision_downsample_ratio=1, - vision_max_n_token=64, - vision_min_pixels=0, - vision_max_wh_ratio=8, - ) - - -def _tiny_vision_args(): - return SimpleNamespace( - hidden_size=12, - dim=12, - vision_n_layers=2, - vision_dim=16, - vision_n_heads=4, - vision_inter_dim=24, - vision_patch_size=2, - vision_rope_theta=10000.0, - vision_downsample_ratio=2, - ) - - -@pytest.mark.parametrize("n_vit_h,n_vit_w", [(4, 4), (5, 3)]) -def test_vit_matches_reference(reference_vision, n_vit_h, n_vit_w, monkeypatch): - monkeypatch.setattr(visualserver, "VIT_ATTN_BACKEND", SdpaVitAttBackend) - torch.manual_seed(0) - args = _tiny_vision_args() - model = visual.ViT(args).eval() - previous_dtype = torch.get_default_dtype() - torch.set_default_dtype(torch.bfloat16) - try: - reference = reference_vision.ViT(args).eval() - finally: - torch.set_default_dtype(previous_dtype) - reference.load_state_dict(model.state_dict()) - patches = torch.randn( - n_vit_h * n_vit_w, - 3, - args.vision_patch_size, - args.vision_patch_size, - dtype=torch.bfloat16, - ) - with torch.inference_mode(), sdpa_kernel(SDPBackend.MATH): - actual = model(patches, n_vit_h, n_vit_w) - expected = reference(patches, n_vit_h, n_vit_w) - torch.testing.assert_close(actual, expected, rtol=0, atol=0) - - -@pytest.mark.parametrize("n_vit_h,n_vit_w", [(4, 4), (5, 3)]) -def test_aligner_matches_reference(reference_vision, n_vit_h, n_vit_w): - torch.manual_seed(0) - args = _tiny_vision_args() - model = visual.Aligner(args).eval() - previous_dtype = torch.get_default_dtype() - torch.set_default_dtype(torch.bfloat16) - try: - reference = reference_vision.Aligner(args).eval() - finally: - torch.set_default_dtype(previous_dtype) - reference.load_state_dict(model.state_dict()) - hidden = torch.randn(n_vit_h * n_vit_w, args.vision_dim, dtype=torch.bfloat16) - with torch.inference_mode(): - actual = model(hidden, n_vit_h, n_vit_w) - expected = reference(hidden, n_vit_h, n_vit_w) - torch.testing.assert_close(actual, expected, rtol=0, atol=0) - - -def test_visual_parameter_dtypes(): - model = _tiny_model() - for module in model.modules(): - if isinstance(module, nn.Linear): - assert module.weight.dtype == torch.bfloat16 - if module.bias is not None: - assert module.bias.dtype == torch.bfloat16 - elif isinstance(module, visual.RMSNorm): - assert module.weight.dtype == torch.float32 - assert model.image_start.dtype == torch.bfloat16 - assert model.image_pad.dtype == torch.bfloat16 - - -def test_load_model_reads_only_indexed_visual_weights(tmp_path, monkeypatch): - model = _tiny_model() - expected = { - name: torch.full_like(tensor, index + 1) for index, (name, tensor) in enumerate(model.state_dict().items()) - } - weight_map = {name: "visual.safetensors" for name in expected} - weight_map["layers.0.self_attn.weight"] = "language.safetensors" - (tmp_path / "model.safetensors.index.json").write_text(json.dumps({"weight_map": weight_map})) - - requested = [] - opened = [] - - class FakeSafetensors: - def __enter__(self): - return self - - def __exit__(self, exc_type, exc_value, traceback): - return False - - def get_tensor(self, key): - requested.append(key) - return expected[key] - - def fake_safe_open(path, framework, device): - opened.append(os.path.basename(path)) - assert framework == "pt" - assert device == "cpu" - return FakeSafetensors() - - monkeypatch.setattr(visual, "safe_open", fake_safe_open) - model.load_model(str(tmp_path)) - - assert opened == ["visual.safetensors"] - assert set(requested) == set(expected) - for name, tensor in model.state_dict().items(): - assert torch.equal(tensor, expected[name]) - - -def test_encode_returns_canonical_cache_block(monkeypatch): - model = _tiny_model() - with torch.no_grad(): - model.image_start.fill_(10) - model.image_pad.fill_(20) - model.image_newline.fill_(30) - model.image_end.fill_(40) - - class FakeVision(nn.Module): - def forward(self, patches, n_h, n_w): - return torch.zeros((n_h * n_w, 1), dtype=torch.bfloat16) - - class FakeAligner(nn.Module): - def forward(self, x, n_h, n_w): - values = torch.arange(n_h * n_w, dtype=torch.bfloat16) - return values[:, None].expand(-1, model.config.hidden_size) - - model.vision = FakeVision() - model.aligner = FakeAligner() - - image_bytes = io.BytesIO() - Image.new("RGB", (6, 4), color=(1, 2, 3)).save(image_bytes, format="PNG") - shm_names = [] - monkeypatch.setattr(visual, "get_shm_name_data", lambda uuid: f"{uuid}-data") - - def fake_read_shm(name): - shm_names.append(name) - return image_bytes.getvalue() - - monkeypatch.setattr(visual, "read_shm", fake_read_shm) - embeds, uuids, valid_ids = model.encode([SimpleNamespace(uuid="image-1")]) - - types, perm = ours.build_canonical_image_block(2, 3) - sentinel_table = torch.stack( - [ - model.image_start, - model.image_pad, - model.image_pad, - model.image_newline, - model.image_end, - ] - ) - expected = sentinel_table[types] - row_major = torch.arange(6, dtype=torch.bfloat16)[:, None].expand(-1, model.config.hidden_size) - expected[types == ours.IMAGE] = row_major[perm] - - assert shm_names == ["image-1-data"] - assert uuids == ["image-1"] - assert valid_ids == [[0, len(types)]] - assert torch.equal(embeds, expected) - assert torch.equal(types[:3], torch.full((3,), ours.IMAGE_PAD, dtype=torch.int64)) - assert types[3] == ours.IMAGE_START - assert types[-1] == ours.IMAGE_END diff --git a/unit_tests/models/deepseek_v4/test_vision_integration.py b/unit_tests/models/deepseek_v4/test_vision_integration.py deleted file mode 100644 index 571f993266..0000000000 --- a/unit_tests/models/deepseek_v4/test_vision_integration.py +++ /dev/null @@ -1,321 +0,0 @@ -from types import SimpleNamespace - -import pytest -import torch - - -VISION_CONFIG = { - "vision_n_layers": 32, - "vision_patch_size": 14, - "vision_downsample_ratio": 3, - "vision_max_n_token": 384, - "vision_min_pixels": 147456, - "vision_max_wh_ratio": 8, -} - - -class _Tokenizer: - def convert_tokens_to_ids(self, token): - assert token == "<｜deepseek_image｜>" - return 99 - - def encode(self, prompt, **kwargs): - return prompt - - -@pytest.mark.parametrize("prefix_len", range(8)) -def test_tokenizer_uses_position_dependent_canonical_suffix(prefix_len): - from lightllm.models.deepseek_v4.tokenizer import DeepSeekV4Tokenizer - from lightllm.server.multimodal_params import MultimodalParams - - tokenizer = DeepSeekV4Tokenizer(_Tokenizer(), "/unused", VISION_CONFIG) - params = MultimodalParams(images=[{"type": "url", "data": "unused"}]) - image = params.images[0] - image.token_id, image.token_num = 100_000, 13 - prompt = [7] * prefix_len + [99, 8] - - result = tokenizer.encode(prompt, params) - - cache_skip = prefix_len % 4 - assert result == [7] * prefix_len + list(range(image.token_id + cache_skip, image.token_id + 13)) + [8] - assert image.start_idx == prefix_len - image_end = prefix_len + image.token_num - cache_skip - assert result[image.start_idx : image_end] == list( - range(image.token_id + cache_skip, image.token_id + image.token_num) - ) - - from lightllm.utils.multimodal_chunk_utils import get_mm_input_spans - - assert get_mm_input_spans(result, params.to_dict()) == [(image.start_idx, image_end)] - - -def test_tokenizer_keeps_dsv4_on_one_dimensional_positions(): - from lightllm.models.deepseek_v4.tokenizer import DeepSeekV4Tokenizer - - tokenizer = DeepSeekV4Tokenizer(_Tokenizer(), "/unused", VISION_CONFIG) - image = SimpleNamespace(image_h=64, image_w=64, grid_thwd=None) - - assert tokenizer.get_image_token_length(image) > 0 - assert image.grid_thwd is None - - -def test_chunk_boundary_moves_before_complete_image_block(): - from lightllm.server.router.model_infer.infer_batch import InferReq - - req = InferReq.__new__(InferReq) - req.args = SimpleNamespace(chunked_prefill_size=8192, batch_max_tokens=8192) - req.cur_kv_len = 0 - req.cur_output_len = 0 - req.shm_req = SimpleNamespace(input_len=10_000) - req.mm_input_spans = [(8100, 8400)] - - assert req.get_chuncked_input_token_len() == 8100 - - req.cur_kv_len = 8100 - assert req.get_chuncked_input_token_len() == 10_000 - - -def test_chunk_boundary_between_adjacent_images_never_splits_the_previous_image(): - from lightllm.server.router.model_infer.infer_batch import InferReq - - req = InferReq.__new__(InferReq) - req.args = SimpleNamespace(chunked_prefill_size=640, batch_max_tokens=640) - req.cur_kv_len = 0 - req.cur_output_len = 0 - req.shm_req = SimpleNamespace(input_len=768) - req.mm_input_spans = [(0, 384), (384, 768)] - - assert req.get_chuncked_input_token_len() == 384 - - req.cur_kv_len = 384 - assert req.get_chuncked_input_token_len() == 768 - - -def test_atomic_image_block_can_exceed_chunk_size(): - from lightllm.server.router.model_infer.infer_batch import InferReq - - req = InferReq.__new__(InferReq) - req.args = SimpleNamespace(chunked_prefill_size=256, batch_max_tokens=1024) - req.cur_kv_len = 0 - req.cur_output_len = 0 - req.shm_req = SimpleNamespace(input_len=1000) - req.mm_input_spans = [(100, 800)] - - assert req.get_chuncked_input_token_len() == 100 - - req.cur_kv_len = 100 - assert req.get_chuncked_input_token_len() == 800 - - -def test_vision_model_allows_cpu_cache(): - from lightllm.models.deepseek_v4.model import DeepseekV4TpPartModel - from lightllm.server.core.objs import StartArgs - - model = SimpleNamespace( - load_way="HF", - tp_world_size_=1, - config={ - "num_attention_heads": 8, - "o_groups": 8, - "index_n_heads": 8, - "vision_n_layers": 1, - }, - args=StartArgs(enable_cpu_cache=True), - ) - - DeepseekV4TpPartModel._verify_params(model) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") -@pytest.mark.parametrize("is_hash", [False, True]) -def test_router_uses_bias_vl_only_for_image_tokens(is_hash): - from lightllm.models.deepseek_v4.layer_infer.transformer_layer_infer import DeepseekV4TransformerLayerInfer - - router = DeepseekV4TransformerLayerInfer.__new__(DeepseekV4TransformerLayerInfer) - router.has_vision = True - router.is_hash = is_hash - router.vocab_size = 32 - router.num_experts_per_tok = 6 - router.routed_scaling_factor = 1.5 - router.alloc_tensor = torch.empty - - torch.manual_seed(0) - logits = torch.randn((2, 256), dtype=torch.float32, device="cuda") - hash_table = torch.zeros((router.vocab_size, 6), dtype=torch.int64, device="cuda") - hash_table[2] = torch.tensor([1, 4, 7, 10, 13, 16], device="cuda") - text_bias = torch.zeros(256, dtype=torch.float32, device="cuda") - text_bias[20:26] = torch.tensor([60.0, 55.0, 50.0, 45.0, 40.0, 35.0], device="cuda") - vision_bias = torch.zeros(256, dtype=torch.float32, device="cuda") - vision_bias[200:206] = torch.tensor([60.0, 55.0, 50.0, 45.0, 40.0, 35.0], device="cuda") - layer_weight = SimpleNamespace( - n_routed_experts=256, - gate_weight_=SimpleNamespace(mm=lambda x, out: out.copy_(x)), - gate_tid2eid_=SimpleNamespace(weight=hash_table), - gate_bias_=SimpleNamespace(weight=text_bias), - gate_bias_vl_=SimpleNamespace(weight=vision_bias), - ) - infer_state = SimpleNamespace(is_prefill=True, input_ids=torch.tensor([2, 100_000], device="cuda")) - - weights, indices = router._select_experts(logits, infer_state, layer_weight) - - scores = torch.sqrt(torch.nn.functional.softplus(logits)) - text_indices = hash_table[infer_state.input_ids[0]] if is_hash else (scores[0] + text_bias).topk(6).indices - image_indices = (scores[1] + vision_bias).topk(6).indices - expected_indices = torch.stack((text_indices, image_indices)) - torch.testing.assert_close(indices, expected_indices) - expected = scores.gather(1, expected_indices) - expected = expected / expected.sum(dim=-1, keepdim=True) * 1.5 - torch.testing.assert_close(weights, expected, rtol=2e-5, atol=1e-6) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") -def test_build_image_visibility_scatter(): - from lightllm.models.deepseek_v4.triton_kernel.build_swa_indices import build_image_visibility - - image_spans = torch.tensor([[0, 1, 4], [0, 7, 3], [1, 6, 3]], dtype=torch.int32, device="cuda") - b_q_start_loc = torch.tensor([0, 6], dtype=torch.int32, device="cuda") - b_ready_cache_len = torch.tensor([2, 5], dtype=torch.int32, device="cuda") - b_q_seq_len = torch.tensor([6, 5], dtype=torch.int32, device="cuda") - image_left = torch.zeros(11, dtype=torch.int32, device="cuda") - image_right = torch.zeros(11, dtype=torch.int32, device="cuda") - - build_image_visibility( - image_spans, - b_q_start_loc, - b_ready_cache_len, - b_q_seq_len, - image_left, - image_right, - ) - - assert image_left.cpu().tolist() == [1, 2, 3, 0, 0, 0, 0, 0, 1, 2, 0] - assert image_right.cpu().tolist() == [2, 1, 0, 0, 0, 2, 0, 2, 1, 0, 0] - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") -def test_swa_index_adds_bidirectional_image_visibility(): - from lightllm.models.deepseek_v4.triton_kernel.build_swa_indices import build_swa_indices - - positions = torch.tensor([5, 8, 10], dtype=torch.int32, device="cuda") - req_idx = torch.zeros(3, dtype=torch.int32, device="cuda") - prefill_state = SimpleNamespace( - prefill_token_ranges=torch.tensor([[0, 128]], dtype=torch.int32, device="cuda"), - swa_prefill_workspace_offsets=torch.tensor([128], dtype=torch.int32, device="cuda"), - HOLD_SWA_SLOT=256, - ) - requests = SimpleNamespace( - config=SimpleNamespace(page_token_size=128), - swa_capacity_tokens=128, - HOLD_SWA_SLOT=128, - HOLD_REQUEST_ID=1, - ) - write_slots = torch.empty_like(positions) - output = torch.empty((3, 8), dtype=torch.int32, device="cuda") - lengths = torch.empty(3, dtype=torch.int32, device="cuda") - image_left = torch.tensor([0, 2, 4], dtype=torch.int32, device="cuda") - image_right = torch.tensor([0, 2, 0], dtype=torch.int32, device="cuda") - - build_swa_indices( - req_idx, - positions, - requests, - output, - lengths, - write_slots, - sliding_window=4, - image_left=image_left, - image_right=image_right, - prefill_state=prefill_state, - ) - - assert output.cpu().tolist() == [ - [133, 132, 131, 130, -1, -1, -1, -1], - [133, 134, 135, 136, 137, 138, -1, -1], - [134, 135, 136, 137, 138, -1, -1, -1], - ] - assert lengths.cpu().tolist() == [4, 6, 5] - - -@pytest.mark.parametrize( - "loaded_start,load_end,spans,expected", - [ - (0, 4096, [], 4096), - (0, 2048, [(2048, 2500)], 2048), - (0, 4096, [(3500, 4096)], 4096), - (0, 4096, [(3500, 4500)], 2048), - (0, 4096, [(1000, 2500), (3500, 4500)], 0), - (2304, 4096, [(3000, 4500)], 2304), - ], -) -def test_common_cpu_cache_load_never_splits_an_image(monkeypatch, loaded_start, load_end, spans, expected): - from lightllm.server.router.model_infer.mode_backend import multi_level_kv_cache as common - - ends = list(range(2048, load_end + 1, 2048)) - page_ids = list(range(10, 10 + len(ends))) - released, loaded = [], [] - req = SimpleNamespace( - req_idx=0, - cur_kv_len=loaded_start, - hold_kv_len=loaded_start, - mm_input_spans=spans, - sampling_param=SimpleNamespace(shm_param=SimpleNamespace(prompt_logprobs=-1)), - shm_req=SimpleNamespace( - input_len=8192, - disk_prompt_cache_len=0, - cpu_cache_match_page_indexes=SimpleNamespace(get_all=lambda: page_ids), - token_hash_page_len_list=SimpleNamespace(get_all=lambda: ends), - ), - ) - table = torch.arange(8192, dtype=torch.int32).view(1, -1) - - def allocate(alloc_token_num): - if alloc_token_num == 0: - return None - start = req.cur_kv_len - req.hold_kv_len += alloc_token_num - return torch.arange(start, start + alloc_token_num, dtype=torch.int32) - - req.args = SimpleNamespace(page_size=256) - req.token_index = SimpleNamespace(get_gpu_token_indexes=lambda start, end: table[0, start:end]) - req.kv_cache_alloc_need = lambda target: max(0, target - req.hold_kv_len) - req.alloc_kv_mem = allocate - - from collections import deque - - module = common.MultiLevelKvCacheModule.__new__(common.MultiLevelKvCacheModule) - module.load_tasks = deque() - module.args = SimpleNamespace(page_size=256) - module.backend = SimpleNamespace( - is_master_in_dp=True, - dp_world_size=1, - is_hybrid_att_model=False, - model=SimpleNamespace( - req_manager=SimpleNamespace( - req_to_token_indexs=table, alloc_token_indexes=allocate, wait_token_indexes=lambda: None - ), - att_state_manager=SimpleNamespace( - get_cpu_cache_operator=lambda: SimpleNamespace( - get_page_layouts=lambda: {}, - read_req=lambda **kwargs: loaded.append(req.cur_kv_len + kwargs["token_indices"].numel()), - ), - full_att_allocator=None, - ), - ), - ) - module.cpu_cache_client = SimpleNamespace( - 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) - monkeypatch.setattr(common.dist, "barrier", lambda **kwargs: None) - monkeypatch.setattr(torch.Tensor, "cuda", lambda self, **kwargs: self) - module.load_cpu_cache_to_reqs([req]) - for task in getattr(module, "load_tasks", ()): - task.event.synchronize() - module.update_cpu_cache_load_task_states() - assert req.cur_kv_len == req.hold_kv_len == expected - assert loaded == ([expected] if expected > loaded_start else []) - assert released == page_ids diff --git a/unit_tests/models/deepseek_v4/test_workspace.py b/unit_tests/models/deepseek_v4/test_workspace.py deleted file mode 100644 index 2bc4a56719..0000000000 --- a/unit_tests/models/deepseek_v4/test_workspace.py +++ /dev/null @@ -1,97 +0,0 @@ -from types import SimpleNamespace - -import pytest -import torch - -from lightllm.models.deepseek_v4.model import DeepseekV4TpPartModel -from lightllm.models.deepseek_v4.workspace import DeepseekV4C4PrefillWorkspace, DeepseekV4Workspace -from lightllm.models.deepseek_v4_dspark.model import DeepseekV4DSparkModel -from lightllm.models.deepseek_v4_mtp.model import DeepseekV4MTPModel - - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") - - -@pytest.mark.parametrize("draft_class", [DeepseekV4MTPModel, DeepseekV4DSparkModel]) -def test_draft_initialization_preserves_captured_prefill_buffer(monkeypatch, draft_class): - args = SimpleNamespace(llm_kv_type="fp8kv_dsa") - monkeypatch.setattr("lightllm.models.deepseek_v4.model.get_env_start_args", lambda: args) - - def backend_class(**kwargs): - return lambda model: SimpleNamespace(tp_q_head_num=2, padded_q_head_num=4) - - monkeypatch.setattr("lightllm.models.deepseek_v4.model.get_nsa_prefill_att_backend_class", backend_class) - monkeypatch.setattr("lightllm.models.deepseek_v4.model.get_nsa_decode_att_backend_class", backend_class) - main = DeepseekV4TpPartModel.__new__(DeepseekV4TpPartModel) - main.batch_max_tokens, main.tp_world_size_, main.data_type = 8, 1, torch.bfloat16 - main.config = {"head_dim": 512, "block_size": 3, "dspark_noise_token_id": 1} - main.layers_infer = [SimpleNamespace(compress_ratio=0)] - main.dsv4_workspace = DeepseekV4Workspace(main) - for name in ( - "_freqs_cis_sliding", - "_freqs_cis_compress", - "_cos_cached_sliding", - "_sin_cached_sliding", - "_cos_cached_compress", - "_sin_cached_compress", - ): - setattr(main, name, torch.ones((1, 32), device="cuda")) - main._init_att_backend() - buffer = main.dsv4_workspace.flashmla_prefill_q - buffer[:, :2].fill_(2) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - result = main.dsv4_workspace.flashmla_prefill_q + 1 - - draft = draft_class.__new__(draft_class) - draft.main_model, draft.config, draft.data_type = main, main.config, main.data_type - draft.layers_infer = [SimpleNamespace(compress_ratio=0)] - draft.context_wkv_weight = object() - draft._init_custom() - draft._init_att_backend() - - assert draft.dsv4_workspace is main.dsv4_workspace - assert main.dsv4_workspace.flashmla_prefill_q is buffer - assert draft.layers_infer[0].flashmla_q_head_num_ == 4 - main.dsv4_workspace.flashmla_prefill_q[:, :2].fill_(4) - graph.replay() - torch.testing.assert_close(result[:, :2], torch.full_like(result[:, :2], 5), rtol=0, atol=0) - torch.testing.assert_close(result[:, 2:], torch.ones_like(result[:, 2:]), rtol=0, atol=0) - - -@pytest.mark.parametrize("batch_max_tokens,tp,capacity", [(4097, 8, 4104), (4096, 8, 4096), (1, 8, 8), (3, 1, 3)]) -def test_workspace_covers_tp_padding_and_checks_q_capacity(batch_max_tokens, tp, capacity): - from lightllm.models.deepseek_v4.triton_kernel.norm_rope import fused_q_norm_rope - - model = SimpleNamespace( - batch_max_tokens=batch_max_tokens, - tp_world_size_=tp, - config={"index_n_heads": 2, "index_head_dim": 128}, - ) - workspace = DeepseekV4Workspace(model) - assert workspace.token_capacity == capacity - workspace.init_flashmla_prefill_q(2, 4, 512, torch.bfloat16) - workspace.init_c4_prefill_aux(model) - q = workspace.get_flashmla_prefill_q(capacity) - q_input = torch.ones((capacity, 2, 512), dtype=torch.bfloat16, device="cuda") - freqs = torch.ones((1, 32), dtype=torch.complex64, device="cuda") - positions = torch.zeros(capacity, dtype=torch.int32, device="cuda") - fused_q_norm_rope(q_input, q[:, :2], 1e-6, freqs, positions) - torch.testing.assert_close(q[:, :2], q_input, rtol=0, atol=0) - assert q[:, 2:].count_nonzero() == 0 - assert workspace.c4_prefill_aux.indexer_q_outputs(capacity)[0].shape[0] == capacity - with pytest.raises(ValueError, match="capacity"): - workspace.get_flashmla_prefill_q(capacity + 1) - - -@pytest.mark.parametrize("overlap", [False, True]) -def test_c4_aux_is_one_workspace(overlap): - model = SimpleNamespace( - batch_max_tokens=4, - tp_world_size_=1, - config={"index_n_heads": 2, "index_head_dim": 8}, - args=SimpleNamespace(enable_prefill_microbatch_overlap=overlap), - ) - workspace = DeepseekV4Workspace(model) - workspace.init_c4_prefill_aux(model) - assert isinstance(workspace.c4_prefill_aux, DeepseekV4C4PrefillWorkspace) diff --git a/unit_tests/models/glm5_next/test_cache.py b/unit_tests/models/glm5_next/test_cache.py deleted file mode 100644 index 6da740353d..0000000000 --- a/unit_tests/models/glm5_next/test_cache.py +++ /dev/null @@ -1,270 +0,0 @@ -from unit_tests.common.att_resources.model_helpers import build_hybrid_resources -import dataclasses -from types import SimpleNamespace - -import pytest -import torch - -from lightllm.common.attention_state.allocators import FullAttKVAllocator - -from lightllm.common.attention_state.full_att import ( - PackedMLAFullAttentionState, - PackedMLAAttentionConfig, - PackedMLALogicalConfig, -) -from lightllm.common.req_manager import ReqManager -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 - - -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], - }, - }, - StartArgs(tp=4, data_type="bfloat16"), - ) - # 模型配置保留原始 latent 宽度；packed 行宽由 State 创建入口计算。 - assert config.full_att_head_dim == 512 - - -@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("mode", ["native", "replay"]) -def test_hybrid_checkpoint_restore_and_packed_kv_copy( - monkeypatch, - small_page, - tp_world_size, - mtp_step, - mode, -): - 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, - data_type="bfloat16", - linear_att_hash_page_size=4, - linear_att_page_block_num=2, - cpu_cache_token_page_size=8, - mtp_step=mtp_step, - mtp_mode="eagle_with_att" if mtp_step else None, - ) - set_unique_server_name(args) - get_env_start_args.cache_clear() - set_env_start_args(dataclasses.asdict(args)) - config = Glm5NextResourceConfig( - tp_world_size=tp_world_size, - full_att_all_num_kv_heads=1, - full_att_head_dim=512, - global_linear_k_heads=2 * tp_world_size, - global_linear_v_heads=2 * tp_world_size, - head_linear_k_dim=128, - head_linear_v_dim=128, - conv_kernel_size=4, - conv_state_dtype=torch.bfloat16, - ssm_state_dtype=torch.float32, - full_attention_interval=4, - all_layer_num=4, - ) - full_layer_ids = config.get_full_att_layer_indices() - if mtp_step: - full_layer_ids += (config.all_layer_num,) - full = PackedMLAFullAttentionState( - PackedMLAAttentionConfig( - 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, - ) - full.allocate_storage(16) - req = ReqManager(3, 16) - 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 = build_hybrid_resources( - config=config, - request_capacity=4, - mtp_step=mtp_step, - ssm_state_mode=mode, - 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, - ) - full = main_full - linear = resources.main_kda_conv_bf16_ssm_fp32_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) - 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( - 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) - 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_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]) - linear.conv_state[:, 0].normal_() - linear.ssm_state[:, 0].normal_() - conv = linear.conv_state[:, 0, ..., :3].clone() - ssm = linear.ssm_state[:, 0].clone() - if mode == "replay": - linear.raw_keys.zero_() - linear.raw_values.zero_() - linear.gates.fill_(-0.5) - linear.betas.zero_() - linear.cursors[0] = 1 - linear.cursors[2] = 3 - if mtp_step: - linear.mtp_state_index[2] = mtp_step - resources.save_bounded_att_runtime_to_pages( - req_indices=[0], - state_token_lens=[8], - page_kind=kind, - page_slot_ids=[slot], - ) - torch.cuda.synchronize() - linear.conv_state[:, 0].zero_() - linear.ssm_state[:, 0].zero_() - 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], - ) - torch.cuda.synchronize() - assert torch.equal(linear.conv_state[:, 2, ..., :3], conv) - restored_slot = 2 if mode == "replay" else 2 * (linear.mtp_step + 1) - assert torch.equal(linear.ssm_state[:, restored_slot], ssm) - if mode == "replay": - assert linear.cursors[2] == 0 - if mtp_step: - assert linear.mtp_state_index[2] == 0 - # 整组边界不依赖旧窗口，恢复大小页无需改动无重叠窗口。 - assert torch.equal(tail.buffer, unchanged) - resources.reset_bounded_att_state([2]) - assert not linear.conv_state[:, 2].any() - if mode == "replay": - assert not linear.ssm_state[:, 2].any() - else: - assert not linear.ssm_state[:, 2 * (linear.mtp_step + 1) : 3 * (linear.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 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) - packed_bytes[:, 0].random_(0, 256) - full.copy_full_att_tokens(torch.tensor([0], device="cuda"), torch.tensor([7], device="cuda")) - assert torch.equal(packed_bytes[:, 0], packed_bytes[:, 7]) - - -@pytest.mark.parametrize("gate_type", ["silu", "sigmoid"]) -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( - 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=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=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( - 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( - 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( - actual=out, - expected=expected.bfloat16(), - atol=0.008, - rtol=0.008, - ) diff --git a/unit_tests/models/glm5_next/test_compact_topk_modes.py b/unit_tests/models/glm5_next/test_compact_topk_modes.py deleted file mode 100644 index 973ef62c07..0000000000 --- a/unit_tests/models/glm5_next/test_compact_topk_modes.py +++ /dev/null @@ -1,65 +0,0 @@ -import pytest -import torch -from triton.runtime.jit import JITFunction - -from lightllm.models.glm5_next.triton_kernel.fp8_prefill import _compact_topk_indices_kernel - - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") - - -@pytest.mark.parametrize("index_dtype", [torch.int32, torch.int64]) -@pytest.mark.parametrize("aligned_chunks", [False, True]) -def test_compact_topk_preserves_layout_and_reuses_kernel(monkeypatch, index_dtype, aligned_chunks): - torch.manual_seed(1525) - misses, hashes = [], set() - original = JITFunction._do_compile - - def observe(self, *args, **kwargs): - if self is _compact_topk_indices_kernel: - misses.append(1) - return original(self, *args, **kwargs) - - monkeypatch.setattr(JITFunction, "_do_compile", observe) - # Exercise aligned production chunks and unaligned chunk boundaries separately. - chunk_starts = [0, 2097152, 4194304, 0] if aligned_chunks else [127, 8193, 8193, 127] - for index in range(4): - batch = [1, 3, 33, 65][index] - low = chunk_starts[index] - high = low + [1, 2305, 16385, 30000][index] - topk = torch.randint(-7, 20000, (batch, 2176), device="cuda", dtype=index_dtype) - topk[:, ::17] = -1 - base = max(0, low - 32) - starts = torch.arange(batch, device="cuda", dtype=torch.int32) * 31 + base - # Include positive indices outside either end of the current chunk. - topk[0, 1:6] = torch.tensor( - [low - 2 - base, low - 1 - base, low - base, high - 1 - base, high - base], - device="cuda", - dtype=torch.int32, - ) - output = torch.full_like(topk, -99999) - lengths = torch.full((batch,), -987, device="cuda", dtype=torch.int32) - before_misses = len(misses) - compiled = _compact_topk_indices_kernel[(batch,)]( - topk, - starts, - output, - lengths, - low, - high, - topk.stride(0), - topk.stride(1), - TOPK=2176, - BLOCK=4096, - ) - if hashes: - assert len(misses) == before_misses - hashes.add(compiled.hash) - ragged = topk + starts[:, None] - valid = (topk >= 0) & (ragged >= low) & (ragged < high) - assert torch.equal(lengths, valid.sum(1).int()) - for row in range(batch): - expected = ragged[row][valid[row]] - low - assert torch.equal(output[row, : expected.numel()], expected) - assert (output[row, expected.numel() :] == -99999).all() - assert len(hashes) == 1 diff --git a/unit_tests/models/glm5_next/test_direct_nsa_indices.py b/unit_tests/models/glm5_next/test_direct_nsa_indices.py deleted file mode 100644 index d0e6dcb9bd..0000000000 --- a/unit_tests/models/glm5_next/test_direct_nsa_indices.py +++ /dev/null @@ -1,257 +0,0 @@ -import weakref -from types import SimpleNamespace - -import pytest -import torch - -from lightllm.common.basemodel.attention.base_att import AttControl -from lightllm.common.basemodel.attention.nsa.glm5_next import ( - Glm5NextSparseAttBackend, - Glm5NextSparseDecodeState, -) -from lightllm.common.basemodel.attention.nsa.fp8_glm5_next import ( - Fp8Glm5NextSparseAttBackend, - Fp8Glm5NextSparseDecodeState, -) -from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.common.attention_state.full_att.packed_mla.fp8_e8m0_packed_mla_write import write_fp8_e8m0_packed_kv -from lightllm.common.basemodel.attention.nsa import fp8_glm5_next -from lightllm.models.glm5_next.triton_kernel.kpool import expand_topk, prefill_nsa_params - - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") - - -def tensor(values): - return torch.tensor(values, device="cuda", dtype=torch.int32) - - -@pytest.mark.parametrize("counts", [(3, 0, 1, 2), (257, 0, 1, 513), (1,) * 129]) -def test_prefill_metadata_reuses_existing_prefixes_during_graph_replay(counts): - batch, queries = len(counts), sum(counts) - positions = torch.empty(queries, device="cuda", dtype=torch.int32) - req = torch.empty(batch, device="cuda", dtype=torch.int32) - cu_q, cu_kv = torch.empty(batch + 1, device="cuda", dtype=torch.int32), torch.empty( - batch + 1, device="cuda", dtype=torch.int32 - ) - - def update(counts, phase): - q_prefix, kv_prefix, visible, starts, requests = [0], [0], [], [], [] - req_values = [(batch - i + phase) % (batch + 1) for i in range(batch)] - for i, count in enumerate(counts): - prefix = 1000000 + i * 7 + phase - visible.extend(range(prefix + 1, prefix + count + 1)) - starts.extend([kv_prefix[-1]] * count) - requests.extend([req_values[i]] * count) - q_prefix.append(q_prefix[-1] + count) - kv_prefix.append(kv_prefix[-1] + prefix + count) - positions.copy_(tensor(visible) - 1) - req.copy_(tensor(req_values)) - cu_q.copy_(tensor(q_prefix)) - cu_kv.copy_(tensor(kv_prefix)) - return tensor(visible), tensor(requests), tensor(starts) - - expected = update(counts, 0) - actual = prefill_nsa_params(positions, req, cu_q, cu_kv) - for observed, reference in zip(actual, expected): - torch.testing.assert_close(observed, reference, rtol=0, atol=0) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - actual = prefill_nsa_params(positions, req, cu_q, cu_kv) - expected = update(counts[::-1], 11) - graph.replay() - for observed, reference in zip(actual, expected): - torch.testing.assert_close(observed, reference, rtol=0, atol=0) - - -@pytest.mark.parametrize("dense", [False, True]) -def test_topk_direct_request_lookup_preserves_mtp_hold_and_strided_tables(dense): - table = torch.arange(8 * 70, device="cuda", dtype=torch.int32).reshape(8, 70)[::2, ::2] - req = tensor([2, 2, 0, 3, 3]) - lengths = tensor([5, 7, 9, 2, 2] if dense else [17, 18, 9, 2, 2]) - groups = tensor([[3, 1], [1, 0], [0, 1], [0, 0], [0, 0]]) - - def check(actual): - mem, relative = actual - expected_mem, expected_relative = torch.full_like(mem, -1), torch.full_like(relative, -1) - for i, (request, length) in enumerate(zip(req.tolist(), lengths.tolist())): - tokens = ( - list(range(length)) - if dense - else [group * 4 + j for group in groups[i].tolist()[: min(length // 4, 2)] for j in range(4)] - + list(range(length // 4 * 4, length)) - ) - expected_relative[i, : len(tokens)] = tensor(tokens) - expected_mem[i, : len(tokens)] = table[request, tensor(tokens).long()] - torch.testing.assert_close(mem, expected_mem, rtol=0, atol=0) - torch.testing.assert_close(relative, expected_relative, rtol=0, atol=0) - - check(expand_topk(groups, lengths, req, table, 8, dense=dense)) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - actual = expand_topk(groups, lengths, req, table, 8, dense=dense) - req.copy_(tensor([0, 0, 2, 3, 3])) - table.add_(10000) - lengths[:2].add_(1) - graph.replay() - check(actual) - - -@pytest.mark.parametrize("fp8", [False, True]) -def test_decode_aliases_infer_state_without_ragged_allocation(monkeypatch, fp8): - import lightllm.common.basemodel.triton_kernel.gen_nsa_ks_ke as legacy - - def forbidden(*args, **kwargs): - raise AssertionError("GLM must not generate or copy a full ragged index") - - monkeypatch.setattr(legacy, "gen_nsa_ks_ke", forbidden) - monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", str(torch.cuda.current_device())) - model = type( - "Model", - (), - { - "config": {"index_kpool": 4, "index_topk": 2048}, - "batch_max_tokens": 4, - "max_req_num": 1, - "max_seq_length": 32, - "att_state_role": "decode_test", - "args": SimpleNamespace(enable_prefill_microbatch_overlap=False, run_mode="decode"), - }, - )() - backend = (Fp8Glm5NextSparseAttBackend if fp8 else Glm5NextSparseAttBackend)(model) - assert not hasattr(backend, "ragged_mem_buffers") - infer = InferStateInfo() - infer.b_req_idx = tensor([2, 2, 0, 3, 3]) - infer.b_seq_len = tensor([999999, 1000000, 65536, 2, 2]) - cls = Fp8Glm5NextSparseDecodeState if fp8 else Glm5NextSparseDecodeState - state = cls(backend=backend, _infer_state_weak_ref=weakref.ref(infer)) - state.init_state() - assert state.q_kv_seq_len is infer.b_seq_len - assert state.q_req_idx is infer.b_req_idx - assert not hasattr(state, "ragged_mem_index") and not hasattr(state, "ks") - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - state.init_state() - infer.b_seq_len[:3].sub_(100) - infer.b_req_idx[:3].copy_(tensor([0, 0, 1])) - graph.replay() - assert state.q_kv_seq_len is infer.b_seq_len and state.q_req_idx is infer.b_req_idx - - -@pytest.mark.parametrize( - "queries,heads,total_kv,chunk_size", - [ - (1, 64, 65536, 2097152), - (2, 64, 65536, 2097152), - (3, 8, 65536, 2097152), - (9, 64, 512, 2097152), - (9, 64, 131073, 2097152), - (4, 16, 40, 3), - (4, 64, 40, 7), - ], -) -def test_fp8_prefill_reads_request_table_without_ragged_buffer(monkeypatch, queries, heads, total_kv, chunk_size): - monkeypatch.setattr(fp8_glm5_next, "KV_GATHER_CHUNK_SIZE", chunk_size) - - torch.manual_seed(1685) - batch = min(queries, 2) - counts = [queries] if batch == 1 else [queries // 2, queries - queries // 2] - seqs = [total_kv] if batch == 1 else [total_kv // 2, total_kv - total_kv // 2] - req = tensor([2, 0][:batch]) - table = torch.full((4, 2 * (max(seqs) + 7)), -1, device="cuda", dtype=torch.int32)[:, ::2] - physical = torch.randperm(total_kv + 32, device="cuda", dtype=torch.int32)[:total_kv] - offset = 0 - for request, length in zip(req.tolist(), seqs): - table[request, :length] = physical[offset : offset + length] - offset += length - q = torch.randn(queries, heads, 512, device="cuda", dtype=torch.bfloat16) - source = torch.randn(total_kv + 32, 1, 512, device="cuda", dtype=torch.bfloat16) - packed = torch.zeros(total_kv + 32, 1, 720, device="cuda", dtype=torch.uint8) - write_fp8_e8m0_packed_kv(source, torch.arange(total_kv + 32, device="cuda", dtype=torch.int32), packed) - prefill_kv = torch.randn(queries, 1, 512, device="cuda", dtype=torch.bfloat16) - visible, query_req, current_positions, q_prefix, kv_prefix = [], [], [], [0], [0] - for request, length, count in zip(req.tolist(), seqs, counts): - visible.extend(range(length - count + 1, length + 1)) - query_req.extend([request] * count) - current_positions.extend(range(length - count, length)) - q_prefix.append(q_prefix[-1] + count) - kv_prefix.append(kv_prefix[-1] + length) - query_req, current_positions = tensor(query_req), tensor(current_positions) - mem_index = table[query_req.long(), current_positions.long()] - topk = torch.full((queries, 2176), -1, device="cuda", dtype=torch.int32) - for i, length in enumerate(visible): - topk[i, :7] = tensor([0, length - 1, 0, 1, length // 2, length - 1, -1]) - topk_mem = torch.empty_like(topk) - - def refresh_physical(): - topk_mem.copy_(table[query_req.long()[:, None], topk.clamp_min(0).long()]) - topk_mem.masked_fill_(topk < 0, -1) - mem_index.copy_(table[query_req.long(), current_positions.long()]) - - refresh_physical() - infer = InferStateInfo() - infer.mem_index, infer.b_req_idx = mem_index, req - infer.position_ids = tensor(visible) - 1 - infer.b1_cu_q_seq_len, infer.b1_cu_kv_seq_len = tensor(q_prefix), tensor(kv_prefix) - infer.total_token_num = total_kv - infer.req_manager = SimpleNamespace(req_to_token_indexs=table) - monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", str(torch.cuda.current_device())) - model = type( - "Model", - (), - { - "batch_max_tokens": queries, - "prefill_graph": None, - "max_req_num": batch, - "max_seq_length": max(seqs), - "att_state_role": "indices_test", - "config": {"index_topk": 2048}, - "args": SimpleNamespace(enable_prefill_microbatch_overlap=False, run_mode="normal"), - }, - )() - backend = Fp8Glm5NextSparseAttBackend(model) - state = backend.create_att_prefill_state(infer) - state.init_state() - assert not hasattr(state, "ragged_mem_index") and not hasattr(backend, "ragged_mem_buffers") - control = AttControl( - nsa_prefill=True, - nsa_prefill_dict={ - "topk_indices": topk, - "topk_mem_indices": topk_mem, - "prefill_cache_kv": prefill_kv, - "softmax_scale": 512 ** -0.5, - "kv_lora_rank": 512, - }, - ) - - def run(): - return state.prefill_att(q, packed[..., :584], None, control) - - def check(output): - scales = packed[..., 576:583].view(torch.float8_e8m0fnu).float().reshape(-1, 7, 1) - values = packed[..., :448].view(torch.float8_e4m3fn).float().reshape(-1, 7, 64) * scales - kv = torch.cat((values.reshape(-1, 1, 448).bfloat16(), packed[..., 448:576].view(torch.bfloat16)), -1) - kv[mem_index.long()] = prefill_kv - expected = [] - for i in range(queries): - keys = kv[topk_mem[i][topk_mem[i] >= 0].long(), 0].float() - expected.append((q[i].float() @ keys.T * 512 ** -0.5).softmax(-1) @ keys) - torch.testing.assert_close(output.float(), torch.stack(expected), atol=0.012, rtol=0.015) - - check(run()) - workspace = backend._prefill_workspaces[0] - if workspace is not None: - assert set(workspace) == {"gathered_kv", "chunk_topk_indices", "chunk_cache_seqlens", "cu_seqlens_q"} - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - output = run() - graph.replay() - check(output) - q.normal_() - prefill_kv.normal_() - for request, length in zip(req.tolist(), seqs): - table[request, :length].copy_(table[request, :length].roll(1)) - topk[0].fill_(-1) - refresh_physical() - graph.replay() - check(output) diff --git a/unit_tests/models/glm5_next/test_fp8_compile_cache.py b/unit_tests/models/glm5_next/test_fp8_compile_cache.py deleted file mode 100644 index 6a07b19811..0000000000 --- a/unit_tests/models/glm5_next/test_fp8_compile_cache.py +++ /dev/null @@ -1,94 +0,0 @@ -import pytest -import torch - -from lightllm.models.glm5_next.triton_kernel import ( - fp8_prefill as fp8, -) - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") - - -def record_kernel_hashes(monkeypatch, kernel): - hashes = [] - original = kernel.run - - def run(*args, **kwargs): - compiled = original(*args, **kwargs) - hashes.append(compiled.hash) - return compiled - - monkeypatch.setattr(kernel, "run", run) - return hashes - - -@pytest.mark.parametrize("transposed_lse", [False, True]) -def test_online_attention_merge_reuses_kernel_across_query_counts(monkeypatch, transposed_lse): - hashes = record_kernel_hashes(monkeypatch, fp8._merge_att_output_kernel) - for num_q_tokens in [1, 3, 8, 16, 17, 33, 64]: - output = torch.randn(num_q_tokens, 4, 512, dtype=torch.float32, device="cuda") - lse = torch.randn(4, num_q_tokens, dtype=torch.float32, device="cuda") - chunk_output = torch.randn_like(output) - chunk_lse = ( - torch.randn(num_q_tokens, 4, dtype=torch.float32, device="cuda").T - if transposed_lse - else torch.randn_like(lse) - ) - expected_lse = torch.logaddexp(lse, chunk_lse) - expected = output * torch.exp(lse - expected_lse).T.unsqueeze(-1) + chunk_output * torch.exp( - chunk_lse - expected_lse - ).T.unsqueeze(-1) - fp8._merge_att_output_kernel[(num_q_tokens, 1)]( - output, - lse, - chunk_output, - chunk_lse, - chunk_lse.stride(1), - chunk_lse.stride(0), - q_head_num=4, - q_token_num=num_q_tokens, - BLOCK_HEADS=4, - num_warps=4, - ) - torch.testing.assert_close(output, expected, rtol=2e-5, atol=2e-6) - torch.testing.assert_close(lse, expected_lse, rtol=2e-5, atol=2e-6) - assert len(hashes) == 7 and len(set(hashes)) == 1 - - -def test_prefill_kv_gather_reuses_kernel_across_token_counts(monkeypatch): - hashes = record_kernel_hashes(monkeypatch, fp8._gather_prefill_kv_kernel) - cached_kv = torch.randn(64, 512, dtype=torch.bfloat16, device="cuda") - packed_kv = torch.zeros(64, 1, 584, dtype=torch.uint8, device="cuda") - packed_kv[:, 0, :448] = cached_kv[:, :448].to(torch.float8_e4m3fn).view(torch.uint8) - packed_kv[:, 0, 448:576] = cached_kv[:, 448:].contiguous().view(torch.uint8) - packed_kv[:, 0, 576:583] = 127 - cached_kv[:, :448] = cached_kv[:, :448].to(torch.float8_e4m3fn).to(torch.bfloat16) - table = torch.arange(64, dtype=torch.int32, device="cuda").view(1, -1) - requests = torch.zeros(1, dtype=torch.int32, device="cuda") - cu_q = torch.tensor([0, 1], dtype=torch.int32, device="cuda") - for num_tokens in [1, 3, 8, 16, 17, 32, 33]: - cu_kv = torch.tensor([0, num_tokens], dtype=torch.int32, device="cuda") - prefill_kv = torch.randn(1, 1, 512, dtype=torch.bfloat16, device="cuda") - output = torch.empty(num_tokens, 1, 512, dtype=torch.bfloat16, device="cuda") - fp8._gather_prefill_kv_kernel[(num_tokens // 4 + 1,)]( - packed_kv, - prefill_kv, - output, - table, - requests, - cu_kv, - cu_q, - 0, - num_tokens, - packed_kv.stride(0), - prefill_kv.stride(0), - prefill_kv.stride(2), - req_to_token_indexs_stride_b=table.stride(0), - req_to_token_indexs_stride_s=table.stride(1), - batch_size=1, - BLOCK_ROWS=4, - num_warps=4, - ) - expected = cached_kv[:num_tokens].clone() - expected[-1] = prefill_kv[0, 0] - torch.testing.assert_close(output[:, 0], expected, rtol=0, atol=0) - assert len(hashes) == 7 and len(set(hashes)) == 1 diff --git a/unit_tests/models/glm5_next/test_fp8_kv.py b/unit_tests/models/glm5_next/test_fp8_kv.py deleted file mode 100644 index a6a5c7c9f2..0000000000 --- a/unit_tests/models/glm5_next/test_fp8_kv.py +++ /dev/null @@ -1,1073 +0,0 @@ -import weakref -from types import SimpleNamespace - -import pytest -import torch -import triton - -from lightllm.common.basemodel.attention.base_att import AttControl -from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.common.basemodel.attention.nsa.glm5_next import ( - Glm5NextSparsePrefillState, - Glm5NextSparseDecodeState, -) -from lightllm.common.basemodel.attention.nsa.fp8_glm5_next import ( - Fp8Glm5NextSparseAttBackend, - Fp8Glm5NextSparsePrefillState, - Fp8Glm5NextSparseDecodeState, -) -from lightllm.models.glm5_next.triton_kernel.fp8_sparse_decode import ( - can_use_fp8_sparse_direct_decode, - 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.packed_mla.fp8_e8m0_packed_mla_write import write_fp8_e8m0_packed_kv -from lightllm.common.basemodel.attention.nsa import fp8_glm5_next -from lightllm.models.glm5_next.triton_kernel.fp8_prefill import ( - _gather_prefill_kv_kernel, -) - - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") - - -@pytest.fixture(autouse=True) -def current_cuda_device(monkeypatch): - monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", str(torch.cuda.current_device())) - - -def _make_infer_state(**attributes): - state = InferStateInfo() - vars(state).update(attributes) - return state - - -def _make_workspace_backend(batch_max_tokens, overlap=False): - backend = object.__new__(Fp8Glm5NextSparseAttBackend) - backend.model = SimpleNamespace( - batch_max_tokens=batch_max_tokens, - prefill_graph=None, - config={"index_topk": 2048, "index_kpool": 4}, - max_req_num=1, - max_seq_length=fp8_glm5_next.KV_GATHER_CHUNK_SIZE, - att_state_role="unit_test", - args=SimpleNamespace(enable_prefill_microbatch_overlap=overlap, run_mode="normal"), - ) - backend._prefill_workspaces = backend.allocate_prefill_workspaces(backend.model) - return backend - - -class _PrefillCase: - def __init__(self, seq_lens, q_lens, heads, width=2176, capacity=None, backend=None, softmax_scale=0.0625): - device = "cuda" - tensor = lambda values: torch.tensor(values, device=device, dtype=torch.int32) - self.seq_lens, self.q_lens = seq_lens, q_lens - self.requests = list(range(2 * len(seq_lens) - 1, 0, -2)) - total, queries = sum(seq_lens), sum(q_lens) - capacity = capacity or total + 17 - self.req_to_token_indexs = torch.full( - (2 * len(seq_lens) + 1, max(seq_lens) + 7), -1, device=device, dtype=torch.int32 - ) - slots = torch.randperm(capacity, device=device, dtype=torch.int32)[:total] - cu_q, cu_kv, q_req_idx, positions = [0], [0], [], [] - for req, length, count in zip(self.requests, seq_lens, q_lens): - self.req_to_token_indexs[req, :length] = slots[cu_kv[-1] : cu_kv[-1] + length] - q_req_idx.extend([req] * count) - positions.extend(range(length - count, length)) - cu_q.append(cu_q[-1] + count) - cu_kv.append(cu_kv[-1] + length) - self.q_req_idx, self.position_ids = tensor(q_req_idx).long(), tensor(positions).long() - self.q = torch.randn(queries, heads, 512, device=device, dtype=torch.bfloat16) - self.source_kv = torch.randn(capacity, 1, 512, device=device, dtype=torch.bfloat16) - self.packed_kv = torch.zeros(capacity, 1, 720, device=device, dtype=torch.uint8) - self.all_slots = torch.arange(capacity, device=device, dtype=torch.int32) - write_fp8_e8m0_packed_kv(self.source_kv, self.all_slots, self.packed_kv) - self.prefill_cache_kv = torch.randn(queries, 1, 1024, device=device, dtype=torch.bfloat16)[..., ::2] - visible = tensor(positions) + 1 - self.topk_indices = torch.randint(max(seq_lens), (queries, width), device=device, dtype=torch.int32) - self.topk_indices.remainder_(visible[:, None]) - self.topk_indices.masked_fill_( - torch.arange(width, device=device)[None, :] >= visible.clamp_max(2048)[:, None], -1 - ) - self.topk_indices[:, ::19] = -1 - self.topk_indices[:, 0] = visible - 1 - self.topk_indices[:, 1] = visible - 1 - self.topk_mem_indices = torch.empty(queries, width * 2, device=device, dtype=torch.int32)[:, ::2] - self.infer_state = _make_infer_state( - is_prefill=True, - req_manager=SimpleNamespace(req_to_token_indexs=self.req_to_token_indexs), - b_req_idx=tensor(self.requests), - b_seq_len=tensor(seq_lens), - b1_cu_q_seq_len=tensor(cu_q), - b1_cu_kv_seq_len=tensor(cu_kv), - position_ids=tensor(positions), - total_token_num=total, - mem_index=torch.empty(queries, device=device, dtype=torch.int32), - ) - self.backend = backend or _make_workspace_backend(queries) - _, self.att_state = self.new_state() - self.control = AttControl( - nsa_prefill=True, - nsa_prefill_dict={ - "topk_indices": self.topk_indices, - "topk_mem_indices": self.topk_mem_indices, - "prefill_cache_kv": self.prefill_cache_kv, - "softmax_scale": softmax_scale, - "kv_lora_rank": 512, - }, - ) - - def new_state(self, **overrides): - if overrides: - infer = _make_infer_state(**vars(self.infer_state)) - vars(infer).update(overrides) - else: - infer = self.infer_state - state = self.backend.create_att_prefill_state(infer) - state.init_state() - return infer, state - - def refresh(self): - self.infer_state.mem_index.copy_(self.req_to_token_indexs[self.q_req_idx, self.position_ids]) - self.topk_mem_indices.copy_( - self.req_to_token_indexs[self.q_req_idx[:, None], self.topk_indices.clamp_min(0).long()] - ) - self.topk_mem_indices.masked_fill_(self.topk_indices < 0, -1) - self.control.nsa_prefill_dict["topk_indices"] = self.topk_indices - - def run(self, state=None): - self.refresh() - active = self.att_state if state is None else state - return active.prefill_att(self.q, self.packed_kv[..., :584], None, self.control) - - def check(self, actual): - values = _dequantize_kv_cache(self.packed_kv) - values[self.infer_state.mem_index.long()] = self.prefill_cache_kv - physical = self.req_to_token_indexs[self.q_req_idx[:, None], self.topk_indices.clamp_min(0).long()].masked_fill( - self.topk_indices < 0, -1 - ) - expected = _torch_sparse_attention(self.q, values, physical, self.control.nsa_prefill_dict["softmax_scale"]) - torch.testing.assert_close(actual.float(), expected, atol=0.012, rtol=0.015) - - def mutate(self): - self.q.normal_() - self.prefill_cache_kv.normal_() - self.source_kv.normal_() - write_fp8_e8m0_packed_kv(self.source_kv, self.all_slots, self.packed_kv) - for req, length in zip(self.requests, self.seq_lens): - self.req_to_token_indexs[req, :length].copy_(self.req_to_token_indexs[req, :length].roll(1)) - self.topk_indices.copy_(self.topk_indices.roll(1, dims=1)) - - -@pytest.mark.parametrize("draft", [False, True]) -@pytest.mark.parametrize("kv_type", ["None", "fp8kv_dsa"]) -def test_main_and_native_mtp_backend_selection(monkeypatch, draft, kv_type): - from lightllm.models.glm5_next.model import Glm5NextTpPartModel - from lightllm.models.glm5_next_mtp.model import Glm5NextMTPModel - - monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") - model = object.__new__(Glm5NextMTPModel if draft else Glm5NextTpPartModel) - model.args = SimpleNamespace(llm_kv_type=kv_type, enable_prefill_microbatch_overlap=False, run_mode="normal") - model.config = {"index_topk": 2048, "index_kpool": 4} - model.graph_max_batch_size = 4 - model.max_seq_length = 128 - model.max_req_num = 2 - model.batch_max_tokens = 4 - model.mtp_previous_draft_models = [] - model._init_att_backend() - infer_state = InferStateInfo() - prefill = model.prefill_att_backend.create_att_prefill_state(infer_state) - decode = model.decode_att_backend.create_att_decode_state(infer_state) - assert type(prefill) is (Fp8Glm5NextSparsePrefillState if kv_type == "fp8kv_dsa" else Glm5NextSparsePrefillState) - assert type(decode) is (Fp8Glm5NextSparseDecodeState if kv_type == "fp8kv_dsa" else Glm5NextSparseDecodeState) - infer_state.prefill_att_state = prefill - infer_state.decode_att_state = decode - infer_state.check_att_states() - assert prefill.infer_state is infer_state - assert decode.infer_state is infer_state - owner_ref = weakref.ref(infer_state) - del infer_state - assert owner_ref() is None - with pytest.raises(ReferenceError): - prefill.infer_state - with pytest.raises(ReferenceError): - decode.infer_state - - -def _dequantize_kv_cache(packed_kv): - values = packed_kv[..., :448].view(torch.float8_e4m3fn).float().reshape(-1, 7, 64) - scales = packed_kv[..., 576:583].view(torch.float8_e8m0fnu).float().reshape(-1, 7, 1) - tail = packed_kv[..., 448:576].view(torch.bfloat16) - return torch.cat(((values * scales).reshape(-1, 1, 448).bfloat16(), tail), dim=-1) - - -def _torch_sparse_attention(q, kv, indices, softmax_scale=0.0625): - outputs = [] - for row in range(q.shape[0]): - keys = kv[indices[row][indices[row] >= 0].long(), 0].float() - outputs.append((q[row].float() @ keys.T * softmax_scale).softmax(-1) @ keys) - return torch.stack(outputs) - - -def _capture_cuda_graph(run): - stream = torch.cuda.Stream() - stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(stream): - run() - torch.cuda.current_stream().wait_stream(stream) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - output = run() - return graph, output - - -@pytest.mark.parametrize("index_dtype", [torch.int32, torch.int64]) -@pytest.mark.parametrize("source_d_stride", [1, 2]) -def test_fp8_scatter_quantization_and_gather_prefill(index_dtype, source_d_stride): - torch.manual_seed(53) - # Noncontiguous KV from a fused projection, and a strided packed cache. - raw_source = torch.randn(11, 1, 512 * source_d_stride, device="cuda", dtype=torch.bfloat16) - source = raw_source if source_d_stride == 1 else raw_source[..., ::2] - source[0].zero_() - 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) - 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) - actual_scales = packed_kv[..., 576:583].view(torch.float8_e8m0fnu).float() - torch.testing.assert_close(actual_scales[mem_index], scales, rtol=0, atol=0) - assert torch.equal(packed_kv[mem_index, :, :448], values.view(torch.uint8).reshape(11, 1, 448)) - assert torch.equal(packed_kv[..., 448:576].view(torch.bfloat16)[mem_index], source[..., 448:]) - assert not packed_kv[mem_index, :, 583].any() - assert (packed_kv[..., 584:] == 71).all() - unused = torch.ones(31, device="cuda", dtype=torch.bool) - unused[mem_index] = False - assert (packed_kv[unused] == 71).all() - - # The packer is graph-safe for both source layouts and mutable destination mapping. - 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) - mem_index.copy_(mem_index.roll(1)) - graph.replay() - changed_groups = source[..., :448].float().reshape(11, 1, 7, 64) - changed_scales = torch.exp2(torch.ceil(torch.log2((changed_groups.abs().amax(-1) / 448).clamp_min(1e-4)))) - changed_values = (changed_groups / changed_scales.unsqueeze(-1)).to(torch.float8_e4m3fn) - assert torch.equal(packed_kv[mem_index, :, :448], changed_values.view(torch.uint8).reshape(11, 1, 448)) - - prefill_kv = source[8:].clone().add_(0.25) - gathered_kv = torch.empty((8, 1, 512), device="cuda", dtype=torch.bfloat16) - _gather_prefill_kv_kernel[(3,)]( - packed_kv[..., :584], - prefill_kv, - gathered_kv, - mem_index[None], - torch.zeros(1, device="cuda", dtype=torch.int32), - torch.tensor([0, 11], device="cuda", dtype=torch.int32), - torch.tensor([0, 3], device="cuda", dtype=torch.int32), - 2, - 8, - packed_kv.stride(0), - prefill_kv.stride(0), - prefill_kv.stride(2), - req_to_token_indexs_stride_b=11, - req_to_token_indexs_stride_s=1, - batch_size=1, - BLOCK_ROWS=4, - ) - expected = _dequantize_kv_cache(packed_kv) - expected[mem_index[8:].long()] = prefill_kv - assert torch.equal(gathered_kv, expected[mem_index[2:10].long()]) - - -@pytest.mark.parametrize("num_heads", [8, 16, 64, 128]) -def test_fp8_nope_decode_and_cuda_graph(num_heads): - if torch.cuda.get_device_capability()[0] != 9: - pytest.skip("Validated FlashMLA sparse decode target is Hopper") - import flash_mla - - torch.manual_seed(53) - 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) - 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) - indices = torch.randint(0, 4096, (9, 2176), device="cuda", dtype=torch.int32) - lanes = torch.arange(2176, device="cuda") - indices.masked_fill_(lanes[None, :] >= lengths[:, None], -1) - control = AttControl( - nsa_decode=True, nsa_decode_dict={"topk_mem_indices": indices, "softmax_scale": 0.0625, "kv_lora_rank": 512} - ) - infer_state = _make_infer_state() - state = Fp8Glm5NextSparseDecodeState( - _infer_state_weak_ref=weakref.ref(infer_state), flashmla_sched_meta=flash_mla.get_mla_metadata()[0] - ) - - def run(): - return state.decode_att((q, q[..., :0]), packed_kv[..., :584], None, control) - - def check(actual): - expected = _torch_sparse_attention(q, _dequantize_kv_cache(packed_kv), indices) - torch.testing.assert_close(actual.float(), expected, atol=0.012, rtol=0.015) - - check(run()) - graph, out = _capture_cuda_graph(run) - graph.replay() - check(out) - q.normal_() - source.normal_() - 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) - graph.replay() - check(out) - - -@pytest.mark.parametrize("cached", [False, True]) -@pytest.mark.parametrize("num_heads", [8, 16, 64, 128]) -@pytest.mark.parametrize("chunk_size", [3, 7, 32768]) -def test_fp8_nope_batched_prefill_and_graph(monkeypatch, cached, num_heads, chunk_size): - monkeypatch.setattr(fp8_glm5_next, "KV_GATHER_CHUNK_SIZE", chunk_size) - torch.manual_seed(53) - prefix = 4 if cached else 0 - case = _PrefillCase([prefix + 3] * 2, [3, 3], num_heads, width=128, capacity=64, softmax_scale=512 ** -0.5) - case.q = torch.randn(num_heads, 6, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1) - case.topk_indices[0].fill_(-1) - case.topk_indices[1, :2] = prefix - case.topk_indices[3].fill_(-1) - case.topk_indices[3, 0] = prefix - case.check(case.run()) - graph, output = _capture_cuda_graph(case.run) - graph.replay() - case.check(output) - case.mutate() - graph.replay() - case.check(output) - - -@pytest.mark.parametrize("num_heads", [8, 64]) -@pytest.mark.parametrize("chunk_size", [3, 32768]) -def test_fp8_uncached_prefill_reuses_current_kv(monkeypatch, num_heads, chunk_size): - from sgl_kernel.flash_attn import flash_attn_with_kvcache - - monkeypatch.setattr(fp8_glm5_next, "KV_GATHER_CHUNK_SIZE", chunk_size) - case = _PrefillCase([3, 3], [3, 3], num_heads, width=128) - case.prefill_cache_kv = torch.randn(6, 1, 1536, device="cuda", dtype=torch.bfloat16)[..., :512] - case.control.nsa_prefill_dict["prefill_cache_kv"] = case.prefill_cache_kv - - def check_kv_storage(*args, **kwargs): - assert kwargs["v_cache"].untyped_storage().data_ptr() == case.prefill_cache_kv.untyped_storage().data_ptr() - return flash_attn_with_kvcache(*args, **kwargs) - - monkeypatch.setattr("sgl_kernel.flash_attn.flash_attn_with_kvcache", check_kv_storage) - case.check(case.run()) - graph, output = _capture_cuda_graph(case.run) - case.mutate() - graph.replay() - case.check(output) - - -@pytest.mark.parametrize("num_heads", [8, 16, 64, 128]) -@pytest.mark.parametrize("softmax_scale", [0.0625, 512 ** -0.5]) -def test_fp8_chunked_prefill_sparse_rows_and_graph(monkeypatch, num_heads, softmax_scale): - monkeypatch.setattr(fp8_glm5_next, "KV_GATHER_CHUNK_SIZE", 129) - torch.manual_seed(53) - case = _PrefillCase([1025], [9], num_heads, capacity=2048, softmax_scale=softmax_scale) - lengths = torch.tensor([0, 1, 3, 127, 128, 511, 2048, 2049, 2051], device="cuda") - case.topk_indices.masked_fill_(torch.arange(2176, device="cuda")[None, :] >= lengths[:, None], -1) - case.topk_indices[1, 0] = 1016 - case.topk_indices[2, :3] = torch.tensor([2, 3, 4], device="cuda") - case.topk_indices[3, :127] = torch.randint(1016, 1025, (127,), device="cuda", dtype=torch.int32) - case.check(case.run()) - graph, output = _capture_cuda_graph(case.run) - graph.replay() - case.check(output) - case.mutate() - graph.replay() - case.check(output) - - -@pytest.mark.parametrize("num_q", [64, 129]) -@pytest.mark.parametrize("num_heads", [8, 64]) -def test_fp8_prefill_uses_flashattention_and_request_table_graph(monkeypatch, num_q, num_heads): - from sgl_kernel.flash_attn import flash_attn_with_kvcache - - torch.manual_seed(53100 + num_q) - case = _PrefillCase([32768, 32768], [num_q // 2, num_q - num_q // 2], num_heads, capacity=131072) - observed = [] - - def observed_flashattention(*args, **kwargs): - observed.append(kwargs["qv"].shape) - return flash_attn_with_kvcache(*args, **kwargs) - - def unexpected_flashmla(*args, **kwargs): - raise AssertionError("FP8 prefill must use its single FlashAttention backend") - - monkeypatch.setattr("sgl_kernel.flash_attn.flash_attn_with_kvcache", observed_flashattention) - monkeypatch.setattr("sgl_kernel.flash_mla.flash_mla_sparse_fwd", unexpected_flashmla) - case.check(case.run()) - assert observed == [(num_q, num_heads, 512)] - graph, output = _capture_cuda_graph(case.run) - case.mutate() - graph.replay() - case.check(output) - case.topk_indices.fill_(-1) - graph.replay() - assert torch.count_nonzero(output) == 0 - - -@pytest.mark.parametrize( - "num_q,num_heads", - [(1, 8), (2, 8), (8, 8), (1, 64), (2, 64), (3, 64), (4, 64), (5, 64), (6, 64), (7, 64), (8, 64)], -) -def test_fp8_cached_prefill_strided_holes_and_current(num_q, num_heads): - torch.manual_seed(53016 + num_q) - seqs = [65536 + value for value in [127, 521, 1031, 2053, 4099, 8191, 12289, 16381][:num_q]] - case = _PrefillCase(seqs, [1] * num_q, num_heads) - case.q = torch.randn(num_heads, num_q, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1) - backing = torch.full((num_q, 4352), -1, device="cuda", dtype=torch.int32) - backing[:, ::2] = case.topk_indices - case.topk_indices = backing[:, ::2] - output = case.run() - case.check(output) - assert output.is_contiguous() and torch.count_nonzero(output) > 0 - graph, output = _capture_cuda_graph(case.run) - graph.replay() - case.check(output) - case.mutate() - graph.replay() - case.check(output) - - -@pytest.mark.parametrize("num_q", [2, 3, 5, 8, 16, 32, 128]) -def test_fp8_cached_prefill_multi_query_and_graph(num_q): - torch.manual_seed(53200 + num_q) - case = _PrefillCase([131072, 131072], [num_q // 2, num_q - num_q // 2], 64) - case.check(case.run()) - for _ in range(2): - infer, state = case.new_state() - case.check(case.run(state)) - fresh_infer, fresh = case.new_state(mem_index=case.infer_state.mem_index.clone()) - case.check(case.run(fresh)) - graph, output = _capture_cuda_graph(case.run) - case.mutate() - graph.replay() - case.check(output) - case.topk_indices.fill_(-1) - graph.replay() - assert torch.count_nonzero(output) == 0 - - -def test_fp8_prefill_graph_with_empty_requests_and_changing_boundaries(): - case = _PrefillCase([65536] * 4, [0, 4, 0, 3], 64) - case.check(case.run()) - graph, output = _capture_cuda_graph(case.run) - replacement = _PrefillCase([65536] * 4, [2, 0, 5, 0], 64) - for name in ( - "q", - "packed_kv", - "prefill_cache_kv", - "req_to_token_indexs", - "q_req_idx", - "position_ids", - "topk_indices", - "topk_mem_indices", - ): - getattr(case, name).copy_(getattr(replacement, name)) - case.infer_state.b1_cu_q_seq_len.copy_(replacement.infer_state.b1_cu_q_seq_len) - case.infer_state.position_ids.copy_(replacement.infer_state.position_ids) - for name in ("q_kv_seq_len", "q_req_idx", "q_kv_start_loc"): - getattr(case.att_state, name).copy_(getattr(replacement.att_state, name)) - graph.replay() - case.check(output) - - -def test_fp8_one_query_prefill_and_decode_read_their_own_kv_sources(): - case = _PrefillCase([65536], [1], 64) - case.q.zero_() - case.prefill_cache_kv.fill_(1) - case.source_kv.fill_(2) - write_fp8_e8m0_packed_kv(case.source_kv, case.all_slots, case.packed_kv) - case.topk_indices.fill_(-1) - case.topk_indices[:, 0] = 65535 - case.refresh() - decode_infer = _make_infer_state(**vars(case.infer_state)) - decode_infer.is_prefill = False - decode_infer.is_cuda_graph_capture = True - decode_state = case.backend.create_att_decode_state(decode_infer) - decode_state.init_state() - control = AttControl( - nsa_decode=True, - nsa_decode_dict={ - "topk_mem_indices": case.topk_mem_indices, - "softmax_scale": 0.0625, - "kv_lora_rank": 512, - }, - ) - - def run(): - prefill = case.run() - decode = decode_state.decode_att((case.q, case.q[..., :0]), case.packed_kv[..., :584], None, control) - return prefill, decode - - prefill, decode = run() - torch.testing.assert_close(prefill, torch.ones_like(prefill), rtol=0, atol=0) - torch.testing.assert_close(decode, torch.full_like(decode, 2), rtol=0, atol=0) - graph, (prefill, decode) = _capture_cuda_graph(run) - case.prefill_cache_kv.fill_(3) - case.source_kv.fill_(4) - write_fp8_e8m0_packed_kv(case.source_kv, case.all_slots, case.packed_kv) - graph.replay() - torch.testing.assert_close(prefill, torch.full_like(prefill, 3), rtol=0, atol=0) - torch.testing.assert_close(decode, torch.full_like(decode, 4), rtol=0, atol=0) - - -def test_fp8_h64_b1_decode_uses_native_eager_and_direct_graph(): - if torch.cuda.get_device_capability()[0] != 9: - pytest.skip("Validated FlashMLA sparse decode target is Hopper") - import flash_mla - - torch.manual_seed(53300) - capacity = 4096 - 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) - 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) - control = AttControl( - nsa_decode=True, nsa_decode_dict={"topk_mem_indices": indices, "softmax_scale": 0.0625, "kv_lora_rank": 512} - ) - infer = _make_infer_state(is_cuda_graph_capture=False) - state = Fp8Glm5NextSparseDecodeState( - _infer_state_weak_ref=weakref.ref(infer), - decode_topk_bound=2051, - flashmla_sched_meta=flash_mla.get_mla_metadata()[0], - ) - - def run(): - return state.decode_att((q, q[..., :0]), packed[..., :584], None, control) - - expected = _torch_sparse_attention(q, _dequantize_kv_cache(packed), indices) - eager = run() - torch.testing.assert_close(eager.float(), expected, atol=0.012, rtol=0.015) - infer.is_cuda_graph_capture = True - graph, output = _capture_cuda_graph(run) - infer.is_cuda_graph_capture = False - graph.replay() - torch.testing.assert_close(output.float(), expected, atol=0.012, rtol=0.015) - - -@pytest.mark.parametrize("num_heads", [8, 16, 64, 128]) -def test_fp8_decode_qpad_kernel_strides_and_zero_padding(num_heads): - """The capture padding kernel reads a true D-stride and writes every padded head as zero.""" - torch.manual_seed(53016 + num_heads) - batch = 2 - # This has BMM-style B/H strides and a non-unit final dimension stride. - q = torch.randn(num_heads, batch, 1024, device="cuda", dtype=torch.bfloat16)[..., ::2].transpose(0, 1) - padded_heads = ((num_heads + 63) // 64) * 64 - head_dim = q.shape[-1] - block_size = triton.next_power_of_2(2 * head_dim) - padded = torch.empty((batch, 1, padded_heads, head_dim), device="cuda", dtype=torch.bfloat16) - _fused_strided_qpad_kernel[(batch, triton.cdiv(padded_heads * head_dim, block_size))]( - q, - padded, - q.stride(0), - q.stride(1), - q.stride(2), - num_heads, - padded_heads, - head_dim=head_dim, - BLOCK_SIZE=block_size, - num_warps=4, - ) - expected = torch.zeros_like(padded) - expected[:, 0, :num_heads] = q - assert torch.equal(padded, expected) - assert torch.count_nonzero(padded[:, 0, num_heads:]) == 0 - - -@pytest.mark.parametrize("num_heads", [8, 16, 64, 128]) -def test_fp8_decode_infer_state_graph_qpad_and_mutation(num_heads): - """Exercise the actual backend's eager and graph flags with strided BMM Q.""" - if torch.cuda.get_device_capability()[0] != 9: - pytest.skip("Validated FlashMLA sparse decode target is Hopper") - import flash_mla - - torch.manual_seed(53100 + num_heads) - batch, capacity = 2, 4096 - 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) - 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) - control = AttControl( - nsa_decode=True, - nsa_decode_dict={"topk_mem_indices": indices, "softmax_scale": 0.0625, "kv_lora_rank": 512}, - ) - infer_state = _make_infer_state(is_cuda_graph_capture=False) - state = Fp8Glm5NextSparseDecodeState( - _infer_state_weak_ref=weakref.ref(infer_state), flashmla_sched_meta=flash_mla.get_mla_metadata()[0] - ) - - def run(): - return state.decode_att((q, q[..., :0]), packed[..., :584], None, control) - - def reference(): - return _torch_sparse_attention(q, _dequantize_kv_cache(packed), indices) - - eager = run() - torch.testing.assert_close(eager.float(), reference(), atol=0.012, rtol=0.015) - assert eager.shape == q.shape - infer_state.is_cuda_graph_capture = True - graph, output = _capture_cuda_graph(run) - infer_state.is_cuda_graph_capture = False - graph.replay() - torch.testing.assert_close(output.float(), reference(), atol=0.012, rtol=0.015) - q.copy_(q.roll(1, dims=-1)) - saved = indices.clone() - indices[:, :2049].copy_(saved[:, :2049].roll(1, dims=1)) - assert torch.equal(indices[:, 2049:], saved[:, 2049:]) - graph.replay() - torch.testing.assert_close(output.float(), reference(), atol=0.012, rtol=0.015) - - -def test_fp8_sparse_direct_decode_factory_bound_guard(): - backend = object.__new__(Fp8Glm5NextSparseAttBackend) - backend.model = SimpleNamespace(config={"index_topk": 2048, "index_kpool": 4}) - infer_state = InferStateInfo() - assert backend.create_att_decode_state(infer_state).decode_topk_bound == 2051 - backend.model = SimpleNamespace(config={"index_topk": 2048, "index_kpool": 8}) - assert backend.create_att_decode_state(infer_state).decode_topk_bound is None - backend.model = SimpleNamespace(config={"index_topk": 2047, "index_kpool": 4}) - assert backend.create_att_decode_state(infer_state).decode_topk_bound is None - - -@pytest.mark.parametrize("batch", [1, 8]) -def test_fp8_sparse_direct_decode_actual_entry_edges_and_graph(monkeypatch, batch): - if torch.cuda.get_device_capability()[0] != 9: - pytest.skip("Validated FlashMLA sparse decode target is Hopper") - import flash_mla - - torch.manual_seed(53200 + batch) - capacity = 4096 - 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) - 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] - for row in range(batch): - count = counts[row] - if count: - indices[row, :count] = torch.randint(0, capacity, (count,), device="cuda", dtype=torch.int32) - if batch == 8: - indices[5, :64] = -1 # holes - indices[6, 64:128] = indices[6, :64] # duplicates - assert torch.equal(indices[:, 2051:], torch.full_like(indices[:, 2051:], -1)) - control = AttControl( - nsa_decode=True, nsa_decode_dict={"topk_mem_indices": indices, "softmax_scale": 0.0625, "kv_lora_rank": 512} - ) - state_owner = _make_infer_state(is_cuda_graph_capture=False) - state = Fp8Glm5NextSparseDecodeState( - _infer_state_weak_ref=weakref.ref(state_owner), - flashmla_sched_meta=flash_mla.get_mla_metadata()[0], - decode_topk_bound=2051, - ) - - def native_must_not_run(*args, **kwargs): - raise AssertionError("supported GLM-5.3 decode took native FlashMLA fallback") - - def run(): - return state.decode_att((q, q[..., :0]), packed[..., :584], None, control) - - assert not can_use_fp8_sparse_direct_decode(q, packed, indices.unsqueeze(1), 512, 2051) - expected = _torch_sparse_attention(q, _dequantize_kv_cache(packed), indices) - eager = run() - torch.testing.assert_close(eager.float(), expected, atol=0.02, rtol=0.02) - if batch == 1: - assert torch.count_nonzero(eager) > 0 - monkeypatch.setattr(flash_mla, "flash_mla_with_kvcache", native_must_not_run) - state.infer_state.is_cuda_graph_capture = True - graph, output = _capture_cuda_graph(run) - state.infer_state.is_cuda_graph_capture = False - graph.replay() - torch.testing.assert_close(output.float(), expected, atol=0.02, rtol=0.02) - q.copy_(q.roll(1, dims=-1)) - saved = indices.clone() - nonempty = (indices >= 0).sum(dim=1) - for row, count in enumerate(nonempty.tolist()): - if count: - indices[row, :count].copy_(saved[row, :count].roll(1)) - assert torch.equal(indices[:, 2051:], saved[:, 2051:]) - graph.replay() - expected = _torch_sparse_attention(q, _dequantize_kv_cache(packed), indices) - torch.testing.assert_close(output.float(), expected, atol=0.02, rtol=0.02) - indices.fill_(-1) - graph.replay() - assert torch.count_nonzero(output) == 0 - - -@pytest.mark.parametrize( - "num_heads,batch", - [ - (1, 1), - (4, 3), - (7, 5), - (8, 2), - (8, 3), - (8, 4), - (8, 5), - (16, 1), - (16, 3), - (16, 4), - (17, 3), - (32, 1), - (32, 2), - (32, 3), - (64, 1), - (64, 2), - (128, 1), - ], -) -def test_fp8_sparse_direct_decode_general_graph_factory_entry(monkeypatch, num_heads, batch): - """General head/token shapes use the direct path selected by the central predicate.""" - if torch.cuda.get_device_capability()[0] != 9: - pytest.skip("Validated FlashMLA sparse decode target is Hopper") - import flash_mla - - torch.manual_seed(53300 + batch) - capacity = 4096 - 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) - 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) - indices = torch.full((batch, 2176), -1, device="cuda", dtype=torch.int32) - counts = [([2051, 0, 2048, 2050, 1, 3, 127, 128][row % 8]) for row in range(batch)] - for row, count in enumerate(counts): - if count: - indices[row, :count] = torch.randint(0, capacity, (count,), device="cuda", dtype=torch.int32) - if batch >= 3: - indices[2, :32] = -1 # holes - if batch >= 4: - indices[3, 64:128] = indices[3, :64] # duplicates - assert torch.equal(indices[:, 2051:], torch.full_like(indices[:, 2051:], -1)) - control = AttControl( - nsa_decode=True, - nsa_decode_dict={"topk_mem_indices": indices, "softmax_scale": 0.0625, "kv_lora_rank": 512}, - ) - backend = object.__new__(Fp8Glm5NextSparseAttBackend) - backend.model = SimpleNamespace(config={"index_topk": 2048, "index_kpool": 4}) - infer_state = _make_infer_state(is_cuda_graph_capture=False) - state = backend.create_att_decode_state(infer_state) - state.flashmla_sched_meta = flash_mla.get_mla_metadata()[0] - assert state.decode_topk_bound == 2051 - native_state_owner = _make_infer_state(is_cuda_graph_capture=False) - native_state = Fp8Glm5NextSparseDecodeState( - _infer_state_weak_ref=weakref.ref(native_state_owner), - flashmla_sched_meta=flash_mla.get_mla_metadata()[0], - decode_topk_bound=None, - ) - - def run(current_state): - return current_state.decode_att((q, q[..., :0]), packed[..., :584], None, control) - - native_calls = [] - original_native = flash_mla.flash_mla_with_kvcache - - def observe_native(*args, **kwargs): - native_calls.append(True) - return original_native(*args, **kwargs) - - monkeypatch.setattr(flash_mla, "flash_mla_with_kvcache", observe_native) - eager = run(state) - native_output = run(native_state) - assert native_calls - assert torch.equal(eager, native_output) - expected = _torch_sparse_attention(q, _dequantize_kv_cache(packed), indices) - torch.testing.assert_close(eager.float(), expected, atol=0.02, rtol=0.02) - - def native_must_not_run(*args, **kwargs): - raise AssertionError("supported graph shape did not use the direct packed-KV path") - - monkeypatch.setattr(flash_mla, "flash_mla_with_kvcache", native_must_not_run) - state.infer_state.is_cuda_graph_capture = True - graph, output = _capture_cuda_graph(lambda: run(state)) - state.infer_state.is_cuda_graph_capture = False - graph.replay() - torch.testing.assert_close(output.float(), expected, atol=0.02, rtol=0.02) - q.copy_(q.roll(1, dims=-1)) - saved = indices.clone() - for row, count in enumerate((indices >= 0).sum(dim=1).tolist()): - if count: - indices[row, :count].copy_(saved[row, :count].roll(1)) - assert torch.equal(indices[:, 2051:], saved[:, 2051:]) - graph.replay() - torch.testing.assert_close( - output.float(), _torch_sparse_attention(q, _dequantize_kv_cache(packed), indices), atol=0.02, rtol=0.02 - ) - indices.fill_(-1) - graph.replay() - assert torch.count_nonzero(output) == 0 - - -def test_fp8_sparse_direct_decode_unknown_bound_fallback_and_strided_q(monkeypatch): - if torch.cuda.get_device_capability()[0] != 9: - pytest.skip("Validated FlashMLA sparse decode target is Hopper") - import flash_mla - - torch.manual_seed(53250) - capacity = 4096 - 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) - 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 - control = AttControl( - nsa_decode=True, nsa_decode_dict={"topk_mem_indices": late, "softmax_scale": 0.0625, "kv_lora_rank": 512} - ) - calls = [] - original = flash_mla.flash_mla_with_kvcache - - def native(*args, **kwargs): - calls.append(True) - return original(*args, **kwargs) - - monkeypatch.setattr(flash_mla, "flash_mla_with_kvcache", native) - unknown_owner = _make_infer_state(is_cuda_graph_capture=False) - unknown = Fp8Glm5NextSparseDecodeState( - _infer_state_weak_ref=weakref.ref(unknown_owner), flashmla_sched_meta=flash_mla.get_mla_metadata()[0] - ) - actual = unknown.decode_att((q, q[..., :0]), packed[..., :584], None, control) - torch.testing.assert_close( - actual.float(), _torch_sparse_attention(q, _dequantize_kv_cache(packed), late), atol=0.02, rtol=0.02 - ) - assert calls and not can_use_fp8_sparse_direct_decode( - q, packed[..., :584], late.unsqueeze(1), 512, None, is_cuda_graph_capture=True - ) - q_d2 = torch.randn(1, 8, 1024, device="cuda", dtype=torch.bfloat16)[..., ::2] - good = torch.full((1, 2176), -1, device="cuda", dtype=torch.int32) - good[:, :2051] = torch.randint(0, capacity, (1, 2051), device="cuda", dtype=torch.int32) - assert can_use_fp8_sparse_direct_decode( - q_d2, packed[..., :584], good.unsqueeze(1), 512, 2051, is_cuda_graph_capture=True - ) - assert not can_use_fp8_sparse_direct_decode( - q, packed[..., :584], good.unsqueeze(1), 256, 2051, is_cuda_graph_capture=True - ) - - -@pytest.mark.parametrize("num_heads,batch", [(1, 3), (7, 5), (16, 16), (32, 7), (65, 3), (128, 8)]) -@pytest.mark.parametrize("index_dtype", [torch.int32, torch.int64]) -def test_fp8_sparse_direct_decode_strided_inputs_and_packed_row_padding(monkeypatch, num_heads, batch, index_dtype): - """The kernel handles partial head tiles, larger batches and actual Q/index strides.""" - torch.manual_seed(53600 + num_heads + batch) - capacity = 4096 - source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) - packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) - 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) - indices[0].fill_(-1) - - def run(): - return fp8_sparse_direct_decode(q, packed, indices, 0.0625) - - def expected(): - return _torch_sparse_attention(q, _dequantize_kv_cache(packed), indices[:, 0]) - - if (num_heads, batch) in ((1, 3), (7, 5)): - infer = _make_infer_state(is_cuda_graph_capture=True) - state = Fp8Glm5NextSparseDecodeState(_infer_state_weak_ref=weakref.ref(infer), decode_topk_bound=2051) - - def native_must_not_run(*args, **kwargs): - raise AssertionError("supported strided decode took native fallback") - - monkeypatch.setattr("flash_mla.flash_mla_with_kvcache", native_must_not_run) - torch.testing.assert_close( - state.decode_att( - (q, q[..., :0]), - packed, - None, - AttControl( - nsa_decode=True, - nsa_decode_dict={"topk_mem_indices": indices, "softmax_scale": 0.0625, "kv_lora_rank": 512}, - ), - ).float(), - expected(), - atol=0.02, - rtol=0.02, - ) - - torch.testing.assert_close(run().float(), expected(), atol=0.02, rtol=0.02) - graph, output = _capture_cuda_graph(run) - graph.replay() - torch.testing.assert_close(output.float(), expected(), atol=0.02, rtol=0.02) - q.normal_() - indices[..., :2051].copy_(indices[..., :2051].roll(1, dims=0)) - graph.replay() - torch.testing.assert_close(output.float(), expected(), atol=0.02, rtol=0.02) - - -@pytest.mark.parametrize("index_dtype", [torch.int32, torch.int64]) -def test_fp8_sparse_direct_decode_central_predicate_accepts_strides_and_rejects_unaligned_kv(index_dtype): - q = torch.empty(1, 8, 1024, device="cuda", dtype=torch.bfloat16)[..., ::2] - packed = torch.empty(8, 1, 720, device="cuda", dtype=torch.uint8) - topk = torch.empty(1, 1, 4352, device="cuda", dtype=index_dtype)[..., ::2] - assert can_use_fp8_sparse_direct_decode(q, packed, topk, 512, 2051, is_cuda_graph_capture=True) - odd_rows = torch.empty(8, 1, 721, device="cuda", dtype=torch.uint8)[..., :584] - odd_base = torch.empty(8, 1, 722, device="cuda", dtype=torch.uint8)[..., 1:585] - strided_columns = torch.empty(8, 1, 1440, device="cuda", dtype=torch.uint8)[..., ::2] - for invalid in (odd_rows, odd_base, strided_columns): - assert not can_use_fp8_sparse_direct_decode(q, invalid, topk, 512, 2051, is_cuda_graph_capture=True) - - -@pytest.mark.parametrize("strided_input", ["q", "topk"]) -def test_fp8_sparse_direct_decode_offsets_beyond_int32(strided_input): - batch, heads, capacity = 4, 1, 64 - stride = 800_000_000 - shape = (batch, heads, 512) if strided_input == "q" else (batch, 1, 2176) - dtype = torch.bfloat16 if strided_input == "q" else torch.int32 - storage_size = (batch - 1) * stride + shape[-1] - required_bytes = storage_size * torch.empty((), dtype=dtype).element_size() - torch.cuda.empty_cache() - if torch.cuda.mem_get_info()[0] < required_bytes + 512 * 1024 ** 2: - pytest.skip("Large-stride regression requires sufficient free GPU memory") - # Allocate address space without initializing the gaps between logical rows. - storage = torch.empty(storage_size, device="cuda", dtype=dtype) - strided = storage.as_strided(shape, (stride, shape[-1], 1)) - q = strided if strided_input == "q" else torch.empty(batch, heads, 512, device="cuda", dtype=torch.bfloat16) - topk = strided if strided_input == "topk" else torch.empty(batch, 1, 2176, device="cuda", dtype=torch.int32) - q.copy_(torch.randn(q.shape, device="cuda", dtype=q.dtype)) - topk.fill_(-1) - 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) - 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]) - torch.testing.assert_close(actual.float(), expected, atol=0.02, rtol=0.02) - - -@pytest.mark.parametrize("num_heads,batch", [(16, 16), (32, 8), (64, 4), (128, 2)]) -def test_fp8_sparse_direct_decode_large_workload_retains_native_graph(monkeypatch, num_heads, batch): - if torch.cuda.get_device_capability()[0] != 9: - pytest.skip("Validated FlashMLA sparse decode target is Hopper") - import flash_mla - - 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) - 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) - assert not can_use_fp8_sparse_direct_decode(q, packed, topk, 512, 2051, is_cuda_graph_capture=True) - - def direct_must_not_run(*args, **kwargs): - raise AssertionError("larger workload took the slower direct path") - - monkeypatch.setattr( - "lightllm.models.glm5_next.triton_kernel.fp8_sparse_decode.fp8_sparse_direct_decode", direct_must_not_run - ) - infer = _make_infer_state(is_cuda_graph_capture=True) - state = Fp8Glm5NextSparseDecodeState( - _infer_state_weak_ref=weakref.ref(infer), - decode_topk_bound=2051, - flashmla_sched_meta=flash_mla.get_mla_metadata()[0], - ) - control = AttControl( - nsa_decode=True, nsa_decode_dict={"topk_mem_indices": topk, "softmax_scale": 0.0625, "kv_lora_rank": 512} - ) - run = lambda: state.decode_att((q, q[..., :0]), packed[..., :584], None, control) - expected = _torch_sparse_attention(q, _dequantize_kv_cache(packed), topk[:, 0]) - torch.testing.assert_close(run().float(), expected, atol=0.02, rtol=0.02) - graph, output = _capture_cuda_graph(run) - graph.replay() - torch.testing.assert_close(output.float(), expected, atol=0.02, rtol=0.02) - - -@pytest.mark.parametrize("num_heads", [8, 64]) -def test_fp8_prefill_workspace_reuse_mutable_inputs_and_graph_replay(num_heads): - torch.manual_seed(53450) - backend = _make_workspace_backend(66) - case = _PrefillCase([32768, 16384, 16384], [11, 11, 11], num_heads, backend=backend) - case.check(case.run()) - workspace = backend._prefill_workspaces[0] - assert set(workspace) == {"gathered_kv", "chunk_topk_indices", "chunk_cache_seqlens", "cu_seqlens_q"} - pointers = {name: value.data_ptr() for name, value in workspace.items()} - for _ in range(10): - infer, state = case.new_state() - case.check(case.run(state)) - case.mutate() - case.check(case.run()) - fresh = _PrefillCase([16384, 32768, 16384], [9, 11, 13], num_heads, capacity=131072, backend=backend) - fresh.check(fresh.run()) - assert {name: value.data_ptr() for name, value in backend._prefill_workspaces[0].items()} == pointers - for overrides in ({"prefill_cuda_graph_exe_list": []}, {"is_cuda_graph_capture": True}): - infer, state = case.new_state(**overrides) - case.check(case.run(state)) - graph, output = _capture_cuda_graph(case.run) - case.mutate() - graph.replay() - case.check(output) - - -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 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.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", - ) - state = create_packed_mla_state(config=config, 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_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_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() - 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/glm5_next/test_fused_mtp_norm.py b/unit_tests/models/glm5_next/test_fused_mtp_norm.py deleted file mode 100644 index c92ae5b86a..0000000000 --- a/unit_tests/models/glm5_next/test_fused_mtp_norm.py +++ /dev/null @@ -1,126 +0,0 @@ -from types import SimpleNamespace - -import pytest -import torch - -from lightllm.common.basemodel.triton_kernel.norm.rmsnorm import rmsnorm_forward -from lightllm.models.glm5_next_mtp.layer_infer import pre_layer_infer as mtp -from lightllm.models.glm5_next_mtp.triton_kernel.fused_eh_norm import fused_mtp_norm - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") - - -def reference( - embeddings: torch.Tensor, - previous_hidden: torch.Tensor, - weights: tuple[torch.Tensor, torch.Tensor, torch.Tensor], - eps: float, -) -> torch.Tensor: - rmsnorm_forward(x=previous_hidden, weight=weights[0], eps=eps, out=previous_hidden) - rmsnorm_forward(x=embeddings, weight=weights[1], eps=eps, out=embeddings) - rmsnorm_forward(x=previous_hidden, weight=weights[2], eps=eps, out=previous_hidden) - return torch.cat(tensors=(embeddings, previous_hidden), dim=-1) - - -@pytest.mark.parametrize("rows", [1, 4, 16, 64, 512, 8192]) -@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) -@pytest.mark.parametrize("scale", [1e-4, 1.0, 100.0]) -def test_norm_rounding_and_inplace_updates(rows: int, dtype: torch.dtype, scale: float) -> None: - torch.manual_seed(seed=5514) - hidden = 4096 - e = (torch.randn(size=(rows, hidden + 16), dtype=dtype, device="cuda") * scale)[:, :hidden] - p = (torch.randn(size=(rows, hidden + 16), dtype=dtype, device="cuda") * scale)[:, :hidden] - er, pr = e.clone(), p.clone() - weights = tuple(torch.randn(size=(hidden,), dtype=dtype, device="cuda") * 0.2 + 1 for _ in range(3)) - expected = reference(embeddings=er, previous_hidden=pr, weights=weights, eps=1e-5) - actual = fused_mtp_norm( - embeddings=e, - previous_hidden=p, - main_weight=weights[0], - embedding_weight=weights[1], - hidden_weight=weights[2], - eps=1e-5, - ) - torch.testing.assert_close(actual=actual, expected=expected, atol=0, rtol=0) - torch.testing.assert_close(actual=e, expected=er, atol=0, rtol=0) - torch.testing.assert_close(actual=p, expected=pr, atol=0, rtol=0) - - -def test_graph_replay_changes_both_inputs_and_preserves_strided_storage() -> None: - torch.manual_seed(seed=5515) - rows, hidden = 16, 4096 - e_storage = torch.randn(size=(rows, hidden * 2), device="cuda", dtype=torch.bfloat16) - p_storage = torch.randn(size=(rows, hidden * 2), device="cuda", dtype=torch.bfloat16) - e, p = e_storage[:, ::2], p_storage[:, ::2] - weights = tuple(torch.randn(size=(hidden,), device="cuda", dtype=e.dtype) * 0.2 + 1 for _ in range(3)) - fused_mtp_norm( - embeddings=e, - previous_hidden=p, - main_weight=weights[0], - embedding_weight=weights[1], - hidden_weight=weights[2], - eps=1e-5, - ) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(cuda_graph=graph): - actual = fused_mtp_norm( - embeddings=e, - previous_hidden=p, - main_weight=weights[0], - embedding_weight=weights[1], - hidden_weight=weights[2], - eps=1e-5, - ) - for scale in (0.01, 10.0): - e_storage.normal_().mul_(scale) - p_storage.normal_().mul_(scale) - untouched_e, untouched_p = e_storage[:, 1::2].clone(), p_storage[:, 1::2].clone() - er, pr = e.clone(), p.clone() - expected = reference(embeddings=er, previous_hidden=pr, weights=weights, eps=1e-5) - graph.replay() - torch.testing.assert_close(actual=actual, expected=expected, atol=0, rtol=0) - torch.testing.assert_close(actual=e, expected=er, atol=0, rtol=0) - torch.testing.assert_close(actual=p, expected=pr, atol=0, rtol=0) - assert torch.equal(input=e_storage[:, 1::2], other=untouched_e) - assert torch.equal(input=p_storage[:, 1::2], other=untouched_p) - - -@pytest.mark.parametrize("hidden", [4096, 1024]) -def test_model_integration(hidden: int) -> None: - torch.manual_seed(seed=5516) - e = torch.randn(size=(7, hidden), dtype=torch.bfloat16, device="cuda") - p = torch.randn_like(input=e) - er, pr = e.clone(), p.clone() - weights = tuple(torch.randn(size=(hidden,), dtype=e.dtype, device="cuda") * 0.2 + 1 for _ in range(3)) - - class Norm: - def __init__(self, weight: torch.Tensor) -> None: - self.weight = weight - - def __call__(self, input: torch.Tensor, eps: float, out: torch.Tensor) -> torch.Tensor: - return rmsnorm_forward(x=input, weight=self.weight, eps=eps, out=out) - - class Projection: - def __init__(self) -> None: - self.weight = torch.randn(size=(2 * hidden, 128), dtype=e.dtype, device=e.device) * 0.1 - - def mm(self, input_tensor: torch.Tensor) -> torch.Tensor: - return torch.mm(input=input_tensor, mat2=self.weight) - - layer = mtp.Glm5NextMTPPreLayerInfer.__new__(mtp.Glm5NextMTPPreLayerInfer) - layer.eps_ = 1e-5 - layer_weights = SimpleNamespace( - main_norm_weight_=Norm(weight=weights[0]), - enorm_weight_=Norm(weight=weights[1]), - hnorm_weight_=Norm(weight=weights[2]), - eh_proj_weight_=Projection(), - ) - expected = layer_weights.eh_proj_weight_.mm( - input_tensor=reference(embeddings=er, previous_hidden=pr, weights=weights, eps=1e-5) - ) - actual = layer._fuse_hidden( - input_embeddings=e, infer_state=SimpleNamespace(mtp_draft_input_hiddens=p), layer_weight=layer_weights - ) - torch.testing.assert_close(actual=actual, expected=expected, atol=0, rtol=0) - torch.testing.assert_close(actual=e, expected=er, atol=0, rtol=0) - torch.testing.assert_close(actual=p, expected=pr, atol=0, rtol=0) diff --git a/unit_tests/models/glm5_next/test_indexer_topk.py b/unit_tests/models/glm5_next/test_indexer_topk.py deleted file mode 100644 index 1a4ca0f184..0000000000 --- a/unit_tests/models/glm5_next/test_indexer_topk.py +++ /dev/null @@ -1,161 +0,0 @@ -from types import SimpleNamespace - -import pytest -import torch - -from lightllm.models.glm5_next.indexer import HAS_VLLM, Glm5NextNsaInfer - - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") - - -@pytest.fixture -def indexer(): - return Glm5NextNsaInfer( - 0, {"index_topk": 2048, "index_n_heads": 32, "index_head_dim": 128, "rms_norm_eps": 1e-5, "index_kpool": 4}, 1 - ) - - -def _require_vllm_topk(): - if not HAS_VLLM: - pytest.skip("vLLM top_k_per_row_decode required") - - -def _assert_topk(logits, lengths, indices): - topk = indices.shape[1] - valid = torch.arange(topk, device=logits.device)[None, :] < lengths[:, None] - assert (indices[~valid] == -1).all() - assert (indices[valid] >= 0).all() - assert (indices < lengths[:, None]).all() - for row, length in zip(indices, lengths.tolist()): - count = min(length, topk) - assert row[:count].unique().numel() == count - - got = logits.gather(1, indices.long().clamp_min(0)).masked_fill(~valid, -float("inf")) - positions = torch.arange(logits.shape[1], device=logits.device) - masked = logits.masked_fill(positions[None, :] >= lengths[:, None], -float("inf")) - # Compare values, since different equal-valued candidates are valid top-k. - torch.testing.assert_close(got.sort(descending=True).values, masked.topk(topk).values, rtol=0, atol=0) - - -@pytest.mark.parametrize("pools", [640, 8192, 65536, 262144]) -@pytest.mark.parametrize("distribution", ["normal", "concentrated", "ties"]) -def test_topk_variable_lengths(indexer, pools, distribution): - _require_vllm_topk() - torch.manual_seed(42) - # DeepGEMM logits can have a padded row stride, and output is a query chunk. - logits = torch.randn(8, pools + 128, device="cuda")[:, :pools] - if distribution == "concentrated": - logits.mul_(0.001).add_(1) - elif distribution == "ties": - logits.copy_(torch.randint(0, 4, logits.shape, device="cuda")) - lengths = torch.tensor([0, 1, 511, 512, 513, pools // 3, pools - 3, pools], device="cuda", dtype=torch.int32) - positions = torch.arange(pools, device="cuda") - logits.masked_fill_(positions[None, :] >= lengths[:, None], float("nan")) - output = torch.full((10, 512), -2, device="cuda", dtype=torch.int32) - indices = output[1:-1] - indexer.select_topk_indices(logits, lengths, indices) - _assert_topk(logits, lengths, indices) - assert (output[[0, -1]] == -2).all() - - -@pytest.mark.parametrize("pools", [8192, 262144]) -def test_topk_cuda_graph_with_changing_lengths(indexer, pools): - _require_vllm_topk() - logits = torch.zeros(8, pools, device="cuda") - lengths = torch.zeros(8, device="cuda", dtype=torch.int32) - indices = torch.empty(8, 512, device="cuda", dtype=torch.int32) - stream = torch.cuda.Stream() - stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(stream): - indexer.select_topk_indices(logits, lengths, indices) - torch.cuda.current_stream().wait_stream(stream) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - indexer.select_topk_indices(logits, lengths, indices) - - for shift in range(3): - logits.normal_().mul_(0.001).add_(1) - lengths.copy_( - torch.tensor([0, 1, 511, 512, 513, pools // 3, pools - 3, pools], device="cuda", dtype=torch.int32).roll( - shift - ) - ) - logits.masked_fill_(torch.arange(pools, device="cuda")[None, :] >= lengths[:, None], float("nan")) - graph.replay() - _assert_topk(logits, lengths, indices) - - -def test_topk_without_vllm(indexer, monkeypatch): - from lightllm.models.glm5_next import indexer as indexer_module - - monkeypatch.setattr(indexer_module, "HAS_VLLM", False) - logits = torch.randn(5, 640, device="cuda") - lengths = torch.tensor([0, 1, 511, 512, 639], device="cuda", dtype=torch.int32) - indices = torch.empty(5, 512, device="cuda", dtype=torch.int32) - indexer.select_topk_indices(logits, lengths, indices) - _assert_topk(logits, lengths, indices) - - -@pytest.mark.parametrize("mtp_size", [1, 3, 6]) -@pytest.mark.parametrize("max_pools", [1024, 262144]) -def test_paged_decode_matches_nonpaged_with_graph_and_mtp(indexer, monkeypatch, mtp_size, max_pools): - pytest.importorskip("deep_gemm") - _require_vllm_topk() - torch.manual_seed(27) - rows = 5 * mtp_size - storage = torch.empty(1024, 1, 584, device="cuda", dtype=torch.bfloat16) - packed = storage.view(torch.uint8)[:, :, -132:] - packed[:, 0, :128] = torch.randn(1024, 128, device="cuda").to(torch.float8_e4m3fn).view(torch.uint8) - scales = torch.pow(2.0, torch.randint(-3, 2, (1024,), device="cuda").float()) - packed[:, 0, 128:] = scales.view(torch.uint8).view(-1, 4) - table = torch.full((5, max_pools * 4), -1, device="cuda", dtype=torch.int32) - table[:, 3::4] = torch.randint(0, 1024, (5, max_pools), device="cuda", dtype=torch.int32) - req_idx = torch.tensor([1, 3, 0, 2, 4], device="cuda", dtype=torch.int32).repeat_interleave(mtp_size) - lengths = torch.zeros(rows, device="cuda", dtype=torch.int32) - q = torch.randn(rows, 32, 128, device="cuda").to(torch.float8_e4m3fn) - weights = torch.randn(rows, 32, device="cuda") * 0.05 - infer_state = SimpleNamespace( - b_req_idx=req_idx, - b_seq_len=lengths, - b1_cu_q_seq_len=torch.arange(rows + 1, device="cuda", dtype=torch.int32), - max_q_seq_len=1, - req_manager=SimpleNamespace(req_to_token_indexs=table), - ) - att_state = SimpleNamespace(q_kv_seq_len=lengths) - logits_outputs = [] - select_topk = indexer.select_topk_indices - - def capture_logits(logits, pool_lengths, indices): - logits_outputs.append(logits) - select_topk(logits, pool_lengths, indices) - - monkeypatch.setattr(indexer, "select_topk_indices", capture_logits) - - def run(): - return indexer._get_decode_indices(q, weights, packed, infer_state, att_state, max_pools) - - # Match startup capture with only empty HOLD rows, then replay real work. - run() - logits_outputs.clear() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - actual = run() - graph_logits = logits_outputs.pop() - for bases in ([0, 2047, 2051, max_pools * 4 - mtp_size + 1, 0], [7, 65, 511, 2053, 0]): - seqs = torch.tensor(bases, device="cuda", dtype=torch.int32)[:, None] - seqs = seqs + torch.arange(mtp_size, device="cuda", dtype=torch.int32)[None, :] - seqs[-1].zero_() - lengths.copy_(seqs.flatten()) - graph.replay() - indexer._get_prefill_indices(q, weights, packed, infer_state, att_state, max_pools) - reference_logits = logits_outputs.pop() - pool_lengths = lengths // 4 - valid = torch.arange(max_pools, device="cuda")[None, :] < pool_lengths[:, None] - torch.testing.assert_close(graph_logits[valid], reference_logits[valid], rtol=0, atol=0) - # Fragmented tables deliberately repeat keys; equal-score top-k ties - # may choose different indices. Check exact selected values and uniqueness. - _assert_topk(reference_logits, pool_lengths, actual) - assert (actual[-mtp_size:] == -1).all() - # Exercise request-slot reuse and speculative rollback on the same graph. - req_idx.copy_(req_idx.roll(mtp_size)) diff --git a/unit_tests/models/glm5_next/test_indexer_workspace.py b/unit_tests/models/glm5_next/test_indexer_workspace.py deleted file mode 100644 index 8360f184c4..0000000000 --- a/unit_tests/models/glm5_next/test_indexer_workspace.py +++ /dev/null @@ -1,78 +0,0 @@ -import weakref -from types import SimpleNamespace - -import pytest -import torch - -from lightllm.models.glm5_next.indexer import Glm5NextNsaInfer - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires the GLM GPU indexer kernels") -@pytest.mark.parametrize( - "max_pools,q_lengths,kv_lengths,expected_chunks", - [(640, [2047, 2048, 2049] * 4, [2304, 2432, 2560] * 4, 1), (262144, [31, 1026], [65536, 1048576], 3)], -) -def test_workspace_preserves_mixed_request_pool_selection( - monkeypatch, max_pools, q_lengths, kv_lengths, expected_chunks -): - """Chunked prefill bounds padded allocations and preserves top-k results.""" - import deep_gemm - - indexer = Glm5NextNsaInfer( - 0, {"index_kpool": 4, "index_topk": 2048, "index_n_heads": 32, "index_head_dim": 128, "rms_norm_eps": 1e-5}, 1 - ) - fp8_mqa_logits = deep_gemm.fp8_mqa_logits - select_topk = indexer.select_topk_indices - topk_scores = [] - previous_logits = None - chunk_count = 0 - - def track_logits_lifetime(*args, **kwargs): - nonlocal previous_logits, chunk_count - if previous_logits is not None: - assert previous_logits() is None, "the previous logits chunk is still alive" - logits = fp8_mqa_logits(*args, **kwargs) - allocated_bytes = logits.untyped_storage().nbytes() - assert allocated_bytes <= 512 * 1024 ** 2 - previous_logits = weakref.ref(logits) - chunk_count += 1 - return logits - - def capture_topk_scores(logits, lengths, indices): - select_topk(logits, lengths, indices) - scores = logits.gather(1, indices.long().clamp_min(0)).masked_fill(indices < 0, -float("inf")) - topk_scores.append(scores.sort(dim=1).values) - - monkeypatch.setattr(deep_gemm, "fp8_mqa_logits", track_logits_lifetime) - monkeypatch.setattr(indexer, "select_topk_indices", capture_topk_scores) - torch.manual_seed(20261029) - batch, tokens = len(q_lengths), sum(q_lengths) - packed_buffer = torch.empty((batch * max_pools, 1, 132), dtype=torch.uint8, device="cuda") - packed_buffer[:, 0, :128] = ( - torch.randn(batch * max_pools, 128, device="cuda").to(torch.float8_e4m3fn).view(torch.uint8) - ) - packed_buffer[:, 0, 128:] = torch.ones(batch * max_pools, device="cuda").view(torch.uint8).view(-1, 4) - req_to_token_indexs = torch.full((batch, max_pools * 4), -1, dtype=torch.int32, device="cuda") - req_to_token_indexs[:, 3::4] = torch.arange(batch * max_pools, dtype=torch.int32, device="cuda").view(batch, -1) - q_fp8 = torch.randn(tokens, 32, 128, device="cuda").to(torch.float8_e4m3fn) - weights = torch.randn(tokens, 32, device="cuda") * 0.05 - lengths = torch.cat([torch.arange(k - q + 1, k + 1, device="cuda") for q, k in zip(q_lengths, kv_lengths)]).int() - infer_state = SimpleNamespace( - b_req_idx=torch.arange(batch, dtype=torch.int32, device="cuda"), - b_seq_len=torch.tensor(kv_lengths, dtype=torch.int32, device="cuda"), - b1_cu_q_seq_len=torch.tensor([0] + q_lengths, device="cuda").cumsum(0, dtype=torch.int32), - max_q_seq_len=max(q_lengths), - req_manager=SimpleNamespace(req_to_token_indexs=req_to_token_indexs), - ) - att_state = SimpleNamespace(q_kv_seq_len=lengths) - indexer._get_prefill_indices(q_fp8, weights, packed_buffer, infer_state, att_state, max_pools) - assert chunk_count == expected_chunks - expected_scores = torch.cat(topk_scores) - topk_scores.clear() - chunk_size = min(indexer._get_prefill_chunk_size(max_pools) // 2, tokens // 2) - monkeypatch.setattr(indexer, "_get_prefill_chunk_size", lambda _: chunk_size) - actual = indexer._get_prefill_indices(q_fp8, weights, packed_buffer, infer_state, att_state, max_pools) - torch.testing.assert_close(torch.cat(topk_scores), expected_scores, rtol=0, atol=0) - valid = actual >= 0 - assert (valid.sum(dim=1) == (lengths // 4).clamp_max(indexer.index_topk // 4)).all().item() - assert ((actual == -1) | (valid & (actual < (lengths // 4)[:, None]))).all().item() diff --git a/unit_tests/models/glm5_next/test_kda_conv_qkv.py b/unit_tests/models/glm5_next/test_kda_conv_qkv.py deleted file mode 100644 index b677477519..0000000000 --- a/unit_tests/models/glm5_next/test_kda_conv_qkv.py +++ /dev/null @@ -1,186 +0,0 @@ -import weakref -from types import SimpleNamespace - -import pytest -import torch - -from lightllm.common.basemodel.attention.base_att import AttControl -from lightllm.common.basemodel.attention.linear import kda as backend_module -from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import causal_conv1d_fn -from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_qkv import ( - causal_conv1d_qkv, -) -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda import chunk_kda_with_fused_gate -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.l2norm import l2norm_fwd, USE_DEFAULT_FLA_NORM - - -pytestmark = pytest.mark.skipif( - not torch.cuda.is_available() or torch.cuda.get_device_capability() != (9, 0) or USE_DEFAULT_FLA_NORM, - reason="requires Hopper and the LightLLM L2 norm", -) - - -def reference(x, weight, states, starts, slots, warm, heads): - mixed = causal_conv1d_fn( - x.clone().T, - weight, - query_start_loc=starts, - cache_indices=slots, - has_initial_state=warm, - conv_states=states, - ).T - q, k, v = [part.view(1, x.shape[0], heads, 128) for part in mixed.chunk(3, dim=-1)] - return l2norm_fwd(q.contiguous()), l2norm_fwd(k.contiguous()), v.contiguous() - - -@pytest.mark.parametrize("heads", [1, 2, 8, 64]) -@pytest.mark.parametrize("state_width", [3, 4, 7]) -@pytest.mark.parametrize("lengths", [(2049, 2, 1024), (257, 1, 1), (1, 2, 3), (0, 1, 0, 2), (1,)]) -def test_conv_norm_preserves_outputs_and_history(heads, state_width, lengths): - torch.manual_seed(42) - channels = heads * 128 * 3 - opts = dict(device="cuda", dtype=torch.bfloat16) - projection = torch.randn(sum(lengths), channels + heads + 256, **opts) - original_projection = projection.clone() - x = projection[:, :channels] - weight = torch.randn(channels, 4, **opts) * 0.125 - states = torch.randn(5, channels, state_width, **opts) - expected_states = states.clone() - starts = torch.tensor([0, *lengths], device="cuda", dtype=torch.int32).cumsum(0, dtype=torch.int32) - slots = torch.tensor([1, 2, -1, 3][: len(lengths)], device="cuda", dtype=torch.int32) - cache_seqlens = torch.tensor([0, 7, 0, 1][: len(lengths)], device="cuda", dtype=torch.int32) - expected = reference(x, weight, expected_states[..., :3], starts, slots, cache_seqlens > 0, heads) - actual = causal_conv1d_qkv(x, weight, states[..., :3], starts, slots, cache_seqlens, max(lengths)) - for a, b in zip(actual, expected): - assert a.is_contiguous() - assert torch.equal(a.view(torch.uint8), b.view(torch.uint8)) - assert len({a.untyped_storage().data_ptr() for a in actual}) == 3 - assert torch.equal(states, expected_states) - assert torch.equal(projection, original_projection) - - -@pytest.mark.parametrize("lengths", [(3, 2, 4), (7, 1, 1)]) -def test_conv_norm_cuda_graph_changing_boundaries(lengths): - torch.manual_seed(42) - heads, channels = 2, 3 * 2 * 128 - opts = dict(device="cuda", dtype=torch.bfloat16) - x = torch.randn(9, channels + 258, **opts)[:, :channels] - weight = torch.randn(channels, 4, **opts) * 0.125 - initial = torch.randn(5, channels, 4, **opts) - states = initial.clone() - initial_boundaries = [0, lengths[0], lengths[0] + lengths[1], sum(lengths)] - starts = torch.tensor(initial_boundaries, device="cuda", dtype=torch.int32) - slots = torch.tensor([1, 2, -1], device="cuda", dtype=torch.int32) - cache_seqlens = torch.tensor([4, 0, 0], device="cuda", dtype=torch.int32) - causal_conv1d_qkv(x, weight, states[..., :3], starts, slots, cache_seqlens, max(lengths)) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - actual = causal_conv1d_qkv(x, weight, states[..., :3], starts, slots, cache_seqlens, max(lengths)) - for boundaries, cached in ((initial_boundaries, [4, 0, 0]), ([0, 4, 8, 9], [0, 7, 0])): - starts.copy_(torch.tensor(boundaries, device="cuda", dtype=torch.int32)) - cache_seqlens.copy_(torch.tensor(cached, device="cuda", dtype=torch.int32)) - states.copy_(initial) - expected_states = initial.clone() - expected = reference(x, weight, expected_states[..., :3], starts, slots, cache_seqlens > 0, heads) - graph.replay() - for a, b in zip(actual, expected): - assert torch.equal(a.view(torch.uint8), b.view(torch.uint8)) - assert torch.equal(states, expected_states) - - -@pytest.mark.parametrize("heads", [2, 8]) -@pytest.mark.parametrize("state_width", [3, 7]) -def test_conv_norm_chained_prefills(heads, state_width): - torch.manual_seed(42) - channels = 3 * heads * 128 - opts = dict(device="cuda", dtype=torch.bfloat16) - weight = torch.randn(channels, 4, **opts) * 0.125 - states = torch.randn(5, channels, state_width, **opts) - expected_states = states.clone() - slots = torch.tensor([1, 2, -1, 3], device="cuda", dtype=torch.int32) - cached = torch.tensor([0, 7, 0, 1], device="cuda", dtype=torch.int32) - for lengths in [(1, 2, 0, 3), (0, 1, 2, 0), (7, 0, 1, 2), (2, 1, 0, 1)]: - x = torch.randn(sum(lengths), channels + heads + 256, **opts)[:, :channels] - starts = torch.tensor([0, *lengths], device="cuda", dtype=torch.int32).cumsum(0, dtype=torch.int32) - expected = reference(x, weight, expected_states[..., :3], starts, slots, cached > 0, heads) - actual = causal_conv1d_qkv(x, weight, states[..., :3], starts, slots, cached, max(lengths)) - for a, b in zip(actual, expected): - assert torch.equal(a.view(torch.uint8), b.view(torch.uint8)) - assert torch.equal(states, expected_states) - cached += torch.tensor(lengths, device="cuda", dtype=torch.int32) - - -@pytest.mark.parametrize("lengths", [(129, 2), (129, 1, 1), (1,)]) -def test_backend_fused_conv_output_and_state(lengths): - torch.manual_seed(42) - heads, channels, tokens = 2, 3 * 2 * 128, sum(lengths) - opts = dict(device="cuda", dtype=torch.bfloat16) - projection = torch.randn(tokens, channels + heads + 256, **opts) - weight = torch.randn(channels, 4, **opts) * 0.125 - conv_states = torch.randn(4, channels, 4, **opts) - ssm_states = torch.randn(4, heads, 128, 128, device="cuda") * 0.01 - linear_state = SimpleNamespace(get_layer_view=lambda **kwargs: (conv_states, ssm_states)) - infer_state = type("State", (), {})() - infer_state.att_state_manager = SimpleNamespace(get_state=lambda **kwargs: linear_state) - infer_state.att_state_role = "main" - infer_state.b1_cu_q_seq_len = torch.tensor([0, *lengths], device="cuda", dtype=torch.int32).cumsum( - 0, dtype=torch.int32 - ) - infer_state.b_ready_cache_len = torch.arange(len(lengths), device="cuda", dtype=torch.int32) % 2 * 7 - infer_state.max_q_seq_len = max(lengths) - backend = SimpleNamespace( - head_dim=128, - tp_num_heads=heads, - tp_hidden_size=heads * 128, - conv_kernel_size=4, - lower_bound=-5.0, - ) - att = backend_module.KDAPrefillAttState(backend=backend, _infer_state_weak_ref=weakref.ref(infer_state)) - att.b_conv_buffer_idx = torch.arange(1, len(lengths) + 1, device="cuda", dtype=torch.int32) - att.b_ssm_buffer_idx = att.b_conv_buffer_idx - layer_weight = SimpleNamespace( - get_merged_kda_conv_weight=lambda: weight, - linear_A_log=SimpleNamespace(weight=torch.zeros(heads, device="cuda")), - linear_dt_bias=SimpleNamespace(weight=torch.zeros(heads * 128, device="cuda")), - ) - control = AttControl( - linear_att_prefill=True, - linear_att_prefill_dict=dict( - layer_weight=layer_weight, - layer_num=0, - mixed_qkv=projection[:, :channels], - raw_gate=torch.randn(tokens, heads * 128, **opts) * 0.1, - raw_beta=torch.randn(tokens, heads, **opts), - ), - ) - expected_conv, expected_ssm = conv_states.clone(), ssm_states.clone() - q, k, v = reference( - projection[:, :channels], - weight, - expected_conv[..., :3], - infer_state.b1_cu_q_seq_len, - att.b_conv_buffer_idx, - infer_state.b_ready_cache_len > 0, - heads, - ) - params = control.linear_att_prefill_dict - expected, final_state = chunk_kda_with_fused_gate( - q=q, - k=k, - v=v, - raw_g=params["raw_gate"].view(1, tokens, heads, 128), - beta=params["raw_beta"].float().sigmoid().view(1, tokens, heads), - A_log=layer_weight.linear_A_log.weight, - g_bias=layer_weight.linear_dt_bias.weight, - initial_state=ssm_states[att.b_ssm_buffer_idx].contiguous(), - output_final_state=True, - use_qk_l2norm_in_kernel=False, - cu_seqlens=infer_state.b1_cu_q_seq_len, - safe_gate=True, - lower_bound=-5.0, - ) - expected_ssm[att.b_ssm_buffer_idx] = final_state.to(ssm_states.dtype) - actual = att.prefill_att(None, None, None, att_control=control) - assert torch.equal(actual.view(torch.uint8), expected.view(torch.uint8)) - assert torch.equal(conv_states, expected_conv) - assert torch.equal(ssm_states, expected_ssm) diff --git a/unit_tests/models/glm5_next/test_kernels.py b/unit_tests/models/glm5_next/test_kernels.py deleted file mode 100644 index 49b2a08eb6..0000000000 --- a/unit_tests/models/glm5_next/test_kernels.py +++ /dev/null @@ -1,665 +0,0 @@ -import weakref -import dataclasses -import math -from types import SimpleNamespace - -import pytest -import torch -import triton - -from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.server.core.objs.start_args_type import StartArgs -from lightllm.utils.envs_utils import set_env_start_args -from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda import ( - chunk_kda_with_fused_gate, - fused_kda_gate_chunk_cumsum, -) -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, - gather_pools, - gather_paged_pools, - get_pool_ranges, - expand_topk, -) -from lightllm.models.glm5_next.triton_kernel.index_quant import hadamard_transform_quant_fp8 - - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") - - -@pytest.fixture(autouse=True) -def setup(): - torch.manual_seed(1525) - set_env_start_args(dataclasses.asdict(StartArgs())) - triton.set_allocator(lambda size, alignment, stream: torch.empty(size, device="cuda", dtype=torch.int8)) - - -@pytest.mark.parametrize("tokens", [1, 3, 65, 129]) -def test_kda_chunk_and_decode_match_recurrence(tokens): - heads, dim = 2, 128 - rand = lambda *shape: torch.randn(shape, device="cuda", dtype=torch.bfloat16) - q, k, v, gate = [rand(1, tokens, heads, dim) for _ in range(4)] - beta = rand(1, tokens, heads) - a = torch.randn(heads, device="cuda") - bias = torch.randn(heads, dim, device="cuda") - initial = torch.randn(1, heads, dim, dim, device="cuda") * 0.1 - expected = [] - state = initial[0].clone() - for i in range(tokens): - out, state = _reference_kda(q[0, i], k[0, i], v[0, i], gate[0, i], beta[0, i], a, bias, state) - expected.append(out) - expected = torch.stack(expected).unsqueeze(0) - actual, final = chunk_kda_with_fused_gate( - q=q, - k=k, - v=v.clone(), - raw_g=gate, - beta=beta.float().sigmoid(), - A_log=a, - g_bias=bias.flatten(), - initial_state=initial, - output_final_state=True, - use_qk_l2norm_in_kernel=True, - cu_seqlens=torch.tensor([0, tokens], dtype=torch.int32, device="cuda"), - safe_gate=True, - ) - torch.testing.assert_close(actual.float(), expected, atol=4e-3, rtol=3e-2) - torch.testing.assert_close(final[0], state, atol=8e-3, rtol=3e-2) - # Reuse a nonzero request slot; the neighboring requests must stay intact. - states = torch.randn(4, heads, dim, dim, device="cuda") - states[2] = initial[0] - unchanged = states[[0, 1, 3]].clone() - for i in range(tokens): - out, _ = fused_recurrent_kda( - q[:, i : i + 1], - k[:, i : i + 1], - v[:, i : i + 1], - gate[:, i : i + 1].reshape(1, 1, -1), - beta[:, i : i + 1], - a, - bias.flatten(), - states, - torch.tensor([2], device="cuda", dtype=torch.int32), - ) - torch.testing.assert_close(out[0, 0].float(), expected[0, i], atol=2e-3, rtol=1e-2) - torch.testing.assert_close(states[2], state, atol=2e-5, rtol=2e-4) - assert torch.equal(states[[0, 1, 3]], unchanged) - - -@pytest.mark.parametrize("seq_lens", [(65,), (3, 65, 129)]) -@pytest.mark.parametrize("safe_gate", [False, True]) -@pytest.mark.parametrize("strided", [False, True]) -def test_kda_gate_cumsum_packed_shape(seq_lens, safe_gate, strided): - heads, dim, chunk_size = 2, 128, 64 - if strided: - # All three input strides differ from contiguous [T, H, D]; empty_like - # allocates a contiguous output for this sliced, non-dense input view. - storage = torch.randn(sum(seq_lens) * 2, heads * 2, dim * 2, device="cuda", dtype=torch.bfloat16) - raw_g = storage[::2, ::2, ::2] - else: - raw_g = torch.randn(sum(seq_lens), heads, dim, device="cuda", dtype=torch.bfloat16) - a_log = torch.randn(heads, device="cuda") - bias = torch.randn(heads, dim, device="cuda") if len(seq_lens) > 1 else None - cu_seqlens = torch.tensor([0, *seq_lens], device="cuda", dtype=torch.int32).cumsum(0, dtype=torch.int32) - lower_bound = -3.0 - - actual = fused_kda_gate_chunk_cumsum( - raw_g, - A_log=a_log, - g_bias=bias.flatten() if bias is not None else None, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - safe_gate=safe_gate, - lower_bound=lower_bound, - ) - - gate_input = raw_g.float() + (bias if bias is not None else 0) - amplitude = a_log.exp()[None, :, None] - if safe_gate: - log_gate = lower_bound * torch.sigmoid(amplitude * gate_input) - else: - log_gate = -amplitude * torch.nn.functional.softplus(gate_input) - expected = torch.empty_like(log_gate) - seq_start = 0 - for seq_len in seq_lens: - for offset in range(0, seq_len, chunk_size): - start = seq_start + offset - end = seq_start + min(offset + chunk_size, seq_len) - expected[start:end] = log_gate[start:end].cumsum(0) / math.log(2) - seq_start += seq_len - - assert actual.shape == raw_g.shape - torch.testing.assert_close(actual, expected, atol=2e-5, rtol=1e-5) - - -@pytest.mark.parametrize( - "q_lens", - [ - (3, 2, 4), - (0, 1, 0, 1, 0), - (1,) * 64, - (257, 1, 513, 0), - (8129,) + (1,) * 63, - (3, 0, 1, 2) * 31 + (3, 0, 1), - (3, 0, 1, 2) * 32, - (3, 0, 1, 2) * 32 + (3,), - ], -) -def test_kpool_ranges_cuda_graph_with_changing_query_boundaries(q_lens): - # Move the longest query across the batch while replaying the same graph. - # Empty queries must not use the one-query-per-request decode shortcut. - tensor = lambda values: torch.tensor(values, device="cuda", dtype=torch.int32) - lengths = torch.empty(sum(q_lens), device="cuda", dtype=torch.int32) - cu_q_lens = torch.empty(len(q_lens) + 1, device="cuda", dtype=torch.int32) - max_pools = 4096 - - def set_inputs(counts, phase): - cu, visible_lengths, pool_starts = [0], [], [] - for batch, count in enumerate(counts): - prefix = batch * 7 + phase - visible_lengths.extend(range(prefix + 1, prefix + count + 1)) - pool_starts.extend([batch * max_pools] * count) - cu.append(cu[-1] + count) - lengths.copy_(tensor(visible_lengths)) - cu_q_lens.copy_(tensor(cu)) - expected_lengths = tensor([length // 4 for length in visible_lengths]) - expected_starts = tensor(pool_starts) - return expected_starts, expected_starts + expected_lengths, expected_lengths - - expected = set_inputs(q_lens, 0) - actual = get_pool_ranges(lengths, cu_q_lens, max(q_lens), max_pools) - for output, reference in zip(actual, expected): - assert torch.equal(output, reference) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - actual = get_pool_ranges(lengths, cu_q_lens, max(q_lens), max_pools) - expected = set_inputs(q_lens[::-1], 3) - graph.replay() - for output, reference in zip(actual, expected): - assert torch.equal(output, reference) - - -@pytest.mark.parametrize("cuda_graph", [False, True]) -def test_kpool_chunk_boundaries_and_fragmented_token_kv(cuda_graph): - # Unaligned chunks, reordered requests, and pools spanning old/new tails. - seqs = [9, 7] - ragged = torch.randperm(40, device="cuda", dtype=torch.int32)[: sum(seqs)] - req_idx = torch.tensor([1, 3], device="cuda", dtype=torch.int32) - table = torch.zeros(5, 12, device="cuda", dtype=torch.int32) - table[1, :9], table[3, :7] = ragged[:9], ragged[9:] - source = torch.randn(sum(seqs), 256, device="cuda", dtype=torch.bfloat16) - ape = torch.randn(4, 128, device="cuda") - packed_storage = torch.zeros(40, 1, 584, device="cuda", dtype=torch.bfloat16) - packed = packed_storage.view(torch.uint8)[:, :, -132:] - tail = torch.randn(5, 4, 256, device="cuda", dtype=torch.bfloat16) - unchanged = tail[[0, 2, 4]].clone() - tensor = lambda x: torch.tensor(x, device="cuda", dtype=torch.int32) - for chunks in [[(1, 0, 3), (3, 0, 2)], [(3, 2, 3), (1, 3, 9)], [(3, 3, 7)]]: - raw, lengths, starts, locations = [], [], [], [] - cu_q_lens = [0] - for req, first, end in chunks: - offset = 0 if req == 1 else 9 - raw.append(source[offset + first : offset + end]) - lengths.extend(range(first + 1, end + 1)) - starts.extend([len(locations)] * (end - first)) - locations.extend(table[req, :end].tolist()) - cu_q_lens.append(cu_q_lens[-1] + end - first) - args = ( - torch.cat(raw), - tail, - packed, - ape, - tensor(lengths), - tensor(locations)[tensor(starts) + tensor(lengths) - 1], - tensor([r for r, _, _ in chunks]), - tensor(cu_q_lens), - tensor([end for _, _, end in chunks]), - ) - max_q_len = max(end - first for _, first, end in chunks) - if cuda_graph: - # Compile before capture, then restore this chunk's input cache state. - saved_tail, saved_packed = tail.clone(), packed.clone() - compress_pools(*args, max_q_seq_len=max_q_len) - tail.copy_(saved_tail) - packed.copy_(saved_packed) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - compress_pools(*args, max_q_seq_len=max_q_len) - graph.replay() - else: - compress_pools(*args, max_q_seq_len=max_q_len) - for req, _, end in chunks: - offset = 0 if req == 1 else 9 - assert torch.equal(tail[req, : end % 4], source[offset + end // 4 * 4 : offset + end]) - assert torch.equal(tail[[0, 2, 4]], unchanged) - # Model the full KV copy performed by cache offload/load or request move. - packed_storage = packed_storage.clone() - packed = packed_storage.view(torch.uint8)[:, :, -132:] - keys, scales = gather_pools(packed, table, req_idx, torch.tensor(seqs, device="cuda", dtype=torch.int32), 3) - for batch, start in enumerate([0, 9]): - 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) - 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) - starts = torch.tensor([0, 9], device="cuda", dtype=torch.int32) - groups = torch.tensor([[1, 0], [0, -1]], device="cuda", dtype=torch.int32) - out, relative = expand_topk(groups, lengths, req_idx, table, index_topk=8) - for batch, expected in enumerate([[4, 5, 6, 7, 0, 1, 2, 3, 8], list(range(7))]): - assert relative[batch, : len(expected)].tolist() == expected - assert ( - out[batch, : len(expected)].tolist() - == ragged[starts[batch].item() + torch.tensor(expected, device="cuda")].tolist() - ) - assert (out[batch, len(expected) :] == -1).all() - - -def test_kpool_decode_cuda_graph_and_padding(): - # Replay one graph across all four tail phases, with reordered requests - # and repeated padding slots. Padding must not race on the hold state. - source = torch.randn(2, 8, 256, device="cuda", dtype=torch.bfloat16) - ape = torch.randn(4, 128, device="cuda") - tail = torch.zeros(5, 4, 256, device="cuda", dtype=torch.bfloat16) - storage = torch.zeros(32, 1, 584, device="cuda", dtype=torch.bfloat16) - packed = storage.view(torch.uint8)[:, :, -132:] - table = torch.randperm(32, device="cuda", dtype=torch.int32).view(4, 8) - raw = torch.zeros(4, 256, device="cuda", dtype=torch.bfloat16) - req_idx = torch.tensor([1, 3, 4, 4], device="cuda", dtype=torch.int32) - lengths = torch.ones(4, device="cuda", dtype=torch.int32) - mem_index = torch.zeros(4, device="cuda", dtype=torch.int32) - cu_q_lens = torch.arange(5, device="cuda", dtype=torch.int32) - - def run(): - compress_pools(raw, tail, packed, ape, lengths, mem_index, req_idx, cu_q_lens, lengths, max_q_seq_len=1) - - stream = torch.cuda.Stream() - stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(stream): - run() - torch.cuda.current_stream().wait_stream(stream) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - run() - tail.zero_() - for pos in range(8): - order = [pos % 2, 1 - pos % 2] - req_idx[:2] = torch.tensor([1 if i == 0 else 3 for i in order], device="cuda") - raw[:2] = source[order, pos] - lengths.fill_(pos + 1) - mem_index[:2] = table[order, pos] - graph.replay() - for batch, req in enumerate([1, 3]): - pending = (pos + 1) % 4 - assert torch.equal(tail[req, :pending], source[batch, (pos + 1) // 4 * 4 : pos + 1]) - 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) - 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()) - assert not tail[[0, 2, 4]].any() - - -@pytest.mark.parametrize("max_pools", [640, 262144]) -def test_gather_paged_pools_valid_pages_and_graph_replay(max_pools): - storage = torch.zeros(64, 1, 584, device="cuda", dtype=torch.bfloat16) - packed = storage.view(torch.uint8)[:, :, -132:] - source_keys = torch.randn(64, 128, device="cuda").to(torch.float8_e4m3fn).view(torch.uint8) - source_scales = torch.rand(64, device="cuda") + 0.1 - packed[:, 0, :128] = source_keys - packed[:, 0, 128:] = source_scales.view(torch.uint8).view(64, 4) - table = torch.full((4, max_pools * 4), -1, device="cuda", dtype=torch.int32) - locations = torch.randint(0, 64, (4, max_pools), device="cuda", dtype=torch.int32) - table[:, 3::4] = locations - req_idx = torch.tensor([2, 0, 3], device="cuda", dtype=torch.int32) - lengths = torch.tensor([max_pools, 3, 0], device="cuda", dtype=torch.int32) - - def check(pages, block_table): - page_bytes = pages.view(pages.shape[0], -1) - for row, (req, length) in enumerate(zip(req_idx.tolist(), lengths.tolist())): - page_ids = block_table[row, : triton.cdiv(length, 64)].long() - keys = page_bytes[page_ids, : 64 * 128].reshape(-1, 128) - scales = page_bytes[page_ids, 64 * 128 :].contiguous().view(torch.float32).flatten() - locs = locations[req, :length].long() - assert torch.equal(keys[:length], source_keys[locs]) - assert torch.equal(scales[:length], source_scales[locs]) - assert not keys[length:].any() - assert (scales[length:] == 1).all() - - check(*gather_paged_pools(packed, table, req_idx, lengths, max_pools)) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - pages, block_table = gather_paged_pools(packed, table, req_idx, lengths, max_pools) - # Empty HOLD rows initialize one dummy page, even with a 1M graph capacity. - pages.fill_(127) - lengths.zero_() - graph.replay() - page_bytes = pages.view(3, -1, 64 * 132) - assert not page_bytes[:, 0, : 64 * 128].any() - assert (page_bytes[:, 0, 64 * 128 :].contiguous().view(torch.float32) == 1).all() - assert (page_bytes[:, 1:] == 127).all() - for counts in ([513, 65, max_pools], [1, 0, 3]): - lengths.copy_(torch.tensor(counts, device="cuda", dtype=torch.int32)) - req_idx.copy_(torch.tensor([1, 3, 0], device="cuda", dtype=torch.int32)) - graph.replay() - check(pages, block_table) - - -@pytest.mark.parametrize("max_pools", [65535, 65536, 262144]) -def test_gather_pools_long_context_cuda_graph(max_pools): - # 1M-token graph warmup pads to 262144 pools even for empty hold requests. - # Also gather real pools at the end of that range from fragmented token KV. - storage = torch.zeros(64, 1, 584, device="cuda", dtype=torch.bfloat16) - packed = storage.view(torch.uint8)[:, :, -132:] - source_keys = torch.randn(64, 128, device="cuda").to(torch.float8_e4m3fn).view(torch.uint8) - source_scales = torch.rand(64, device="cuda") + 0.1 - packed[:, 0, :128] = source_keys - packed[:, 0, 128:] = source_scales.view(torch.uint8).view(64, 4) - table = torch.full((4, max_pools * 4), -1, device="cuda", dtype=torch.int32) - locations = torch.randint(0, 64, (4, max_pools), device="cuda", dtype=torch.int32) - table[:, 3::4] = locations - req_idx = torch.tensor([2, 0, 3], device="cuda", dtype=torch.int32) - seq_len = torch.tensor([max_pools * 4, 14, 2], device="cuda", dtype=torch.int32) - - def check(keys, scales): - keys = keys.view(torch.uint8).view(3, max_pools, 128) - scales = scales.view(3, max_pools) - for batch, (req, count) in enumerate([(2, max_pools), (0, 3), (3, 0)]): - locs = locations[req, :count].long() - assert torch.equal(keys[batch, :count], source_keys[locs]) - assert torch.equal(scales[batch, :count], source_scales[locs]) - assert not keys[batch, count:].any() - assert (scales[batch, count:] == 1).all() - - check(*gather_pools(packed, table, req_idx, seq_len, max_pools)) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - keys, scales = gather_pools(packed, table, req_idx, seq_len, max_pools) - # Replay with the same short hold lengths used by model startup. - original_lengths = seq_len.clone() - seq_len.fill_(2) - graph.replay() - assert not keys.view(torch.uint8).any() - assert (scales == 1).all() - seq_len.copy_(original_lengths) - graph.replay() - 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 - 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) - tail = torch.zeros(2, 4, 256, device="cuda", dtype=torch.bfloat16) - ape = torch.randn(4, 128, device="cuda") - lengths = torch.arange(1, 5, device="cuda", dtype=torch.int32) - zero = torch.zeros(4, device="cuda", dtype=torch.int32) - ragged = torch.tensor([0, 1, 2, loc], device="cuda", dtype=torch.int32) - cu_q_lens = torch.tensor([0, 4], device="cuda", dtype=torch.int32) - compress_pools(raw, tail, packed, ape, lengths, ragged, zero[:1], cu_q_lens, lengths[-1:], max_q_seq_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) - assert torch.equal(keys.view(torch.uint8), expected_key.view(torch.uint8)) - assert torch.equal(scales, expected_scale.flatten()) - - -@pytest.mark.parametrize("tp_world_size", [1, 4]) -def test_mhc_keeps_streams_through_decode_autotuning(monkeypatch, tp_world_size): - from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType - from lightllm.models.glm5_next.layer_infer.transformer_layer_infer import Glm5NextTransformerLayerInfer - from lightllm.models.glm5_next.model import Glm5NextTpPartModel - - hidden = 4096 - monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_DP", "0") - monkeypatch.setenv("LIGHTLLM_DP_WORLD_SIZE", str(tp_world_size)) - pre_infer = Glm5NextTpPartModel.pre_layer_infer_class({"hc_mult": 4}) - embeddings = torch.randn(1, hidden, device="cuda", dtype=torch.bfloat16) - input_ids = torch.zeros(1, device="cuda", dtype=torch.long) - pre_weight = SimpleNamespace(wte_weight_=lambda input_ids, alloc_func: embeddings[input_ids]) - pre_weight.wte_weight_.weight = embeddings - pre_weight.wte_weight_.tp_vocab_start_id = 0 - pre_weight.wte_weight_.tp_vocab_end_id = 1 - infer_state = SimpleNamespace(dist_group=None, multimodal_params=[]) - - def all_reduce(input_embeddings, **kwargs): - # TP communication must operate on the original embedding width. - assert input_embeddings.shape == (1, hidden) - input_embeddings.mul_(tp_world_size) - - monkeypatch.setattr("lightllm.models.llama.layer_infer.pre_layer_infer.all_reduce", all_reduce) - monkeypatch.setattr("lightllm.models.qwen_vl.layer_infer.pre_layer_infer.all_reduce", all_reduce) - expected_streams = (embeddings * tp_world_size).unsqueeze(1).expand(-1, 4, -1) - weight = SimpleNamespace( - att_norm_weight_=SimpleNamespace(weight=torch.ones(hidden, device="cuda", dtype=torch.bfloat16)), - ffn_norm_weight_=SimpleNamespace(weight=torch.ones(hidden, device="cuda", dtype=torch.bfloat16)), - ) - for prefix in ("attn", "ffn"): - setattr(weight, f"hc_{prefix}_fn", SimpleNamespace(weight=torch.randn(24, 4 * hidden, device="cuda") * 0.005)) - setattr(weight, f"hc_{prefix}_base", SimpleNamespace(weight=torch.zeros(24, device="cuda"))) - setattr(weight, f"hc_{prefix}_scale", SimpleNamespace(weight=torch.ones(3, device="cuda"))) - layer = object.__new__(Glm5NextTransformerLayerInfer) - layer.use_mhc = True - layer.embed_dim_, layer.mhc_streams = hidden, 4 - layer.num_hidden_layers, layer.autotune_layer_num = 5, 4 - layer.eps_, layer.hc_eps, layer.hc_sinkhorn_iters = 1e-5, 1e-6, 20 - layer.token_attention_forward = layer.context_attention_forward = layer._ffn = lambda x, *_: x * 0.1 - with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - x = pre_infer.token_forward(input_ids, infer_state, pre_weight) - torch.testing.assert_close(x.view(1, 4, hidden), expected_streams, atol=0, rtol=0) - for i in range(5): - layer.layer_num_ = i - x = layer.token_forward(x, None, weight) - assert x.shape == (1, hidden if i == 4 else 4 * hidden) - with Autotuner.autotune_warmup(): - x = pre_infer.context_forward(input_ids, infer_state, pre_weight) - torch.testing.assert_close(x.view(1, 4, hidden), expected_streams, atol=0, rtol=0) - for i in range(4): - layer.layer_num_ = i - x = layer.context_forward(x, None, weight) - assert x.shape == (1, hidden if i == 3 else 4 * hidden) - - -def test_shared_chunk_kernel_preserves_gdn_natural_log_decay(): - from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops import chunk_gated_delta_rule - - tokens, heads, dim = 67, 2, 128 - q, k, v = [torch.randn(1, tokens, heads, dim, device="cuda", dtype=torch.bfloat16) for _ in range(3)] - q = torch.nn.functional.normalize(q.float(), dim=-1).bfloat16() - k = torch.nn.functional.normalize(k.float(), dim=-1).bfloat16() - gate = -torch.rand(1, tokens, heads, device="cuda") - beta = torch.rand_like(gate) - state = torch.randn(1, heads, dim, dim, device="cuda") * 0.1 - initial = state.clone() - expected = [] - for i in range(tokens): - 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) - out, final = chunk_gated_delta_rule( - q, - k, - v, - gate, - beta, - initial_state=initial, - output_final_state=True, - cu_seqlens=torch.tensor([0, tokens], device="cuda", dtype=torch.int32), - ) - torch.testing.assert_close(out.float(), torch.stack(expected, 1), atol=4e-3, rtol=3e-2) - torch.testing.assert_close(final, state, atol=8e-3, rtol=3e-2) - - -@pytest.mark.parametrize("heads", [16, 64, 128]) -def test_nope_attention_native_512_and_cuda_graph(heads): - from lightllm.common.basemodel.attention.base_att import AttControl - from lightllm.common.basemodel.attention.nsa.glm5_next import Glm5NextSparsePrefillState, Glm5NextSparseDecodeState - - if torch.cuda.get_device_capability()[0] != 9: - pytest.skip("FA3 NoPE requires Hopper") - # The 72-element tail stores indexer data, not a zero RoPE embedding. - packed = torch.randn(4096, 1, 584, dtype=torch.bfloat16, device="cuda") - kv = packed[:, :, :512] - # Match the head-major projection's non-contiguous query layout. - q = torch.randn(heads, 9, 512, dtype=torch.bfloat16, device="cuda").transpose(0, 1) - lengths = torch.tensor([0, 1, 3, 127, 128, 511, 2048, 2049, 2051], dtype=torch.int32, device="cuda") - indexes = torch.randint(0, 4096, (9, 2176), dtype=torch.int32, device="cuda") - offsets = torch.arange(indexes.shape[1], device="cuda") - indexes.masked_fill_(offsets[None, :] >= lengths[:, None], -1) - - def reference(): - expected = [] - for i, length in enumerate(lengths.tolist()): - keys = kv[indexes[i, :length].long(), 0].float() - expected.append((q[i].float() @ keys.T * 0.0625).softmax(-1) @ keys) - return torch.stack(expected) - - expected = reference() - control = AttControl( - nsa_prefill_dict={"topk_mem_indices": indexes, "softmax_scale": 0.0625, "kv_lora_rank": kv.shape[-1]} - ) - prefill_state = Glm5NextSparsePrefillState() - prefill = prefill_state._nsa_prefill_att(q, kv, control) - torch.testing.assert_close(prefill.float(), expected, atol=0.012, rtol=0.015) - infer_state = InferStateInfo() - infer_state.b1_cu_q_seq_len = torch.arange(10, dtype=torch.int32, device="cuda") - infer_state.max_q_seq_len = 1 - decode = Glm5NextSparseDecodeState( - _infer_state_weak_ref=weakref.ref(infer_state), - nsa_cache_seqlens=lengths, - nsa_cu_seqlens_k_new=torch.nn.functional.pad(lengths.cumsum(0, dtype=torch.int32), (1, 0)), - ) - control.nsa_decode_dict = control.nsa_prefill_dict - out = decode._nsa_decode_att((q, q[..., :0]), kv, control) - torch.testing.assert_close(out.float(), expected, atol=0.012, rtol=0.015) - - stream = torch.cuda.Stream() - stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(stream): - prefill_state._nsa_prefill_att(q, kv, control) - decode._nsa_decode_att((q, q[..., :0]), kv, control) - torch.cuda.current_stream().wait_stream(stream) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - graph_prefill = prefill_state._nsa_prefill_att(q, kv, control) - graph_decode = decode._nsa_decode_att((q, q[..., :0]), kv, control) - graph.replay() - torch.testing.assert_close(graph_prefill.float(), expected, atol=0.012, rtol=0.015) - torch.testing.assert_close(graph_decode.float(), expected, atol=0.012, rtol=0.015) - - q.normal_() - packed.normal_() - lengths.copy_(lengths.roll(1)) - decode.nsa_cu_seqlens_k_new[1:].copy_(lengths.cumsum(0, dtype=torch.int32)) - indexes.random_(0, 4096) - indexes.masked_fill_(offsets[None, :] >= lengths[:, None], -1) - graph.replay() - expected = reference() - torch.testing.assert_close(graph_prefill.float(), expected, atol=0.012, rtol=0.015) - torch.testing.assert_close(graph_decode.float(), expected, atol=0.012, rtol=0.015) - - -@pytest.mark.parametrize("max_kv_seq_len", [2065, 1 << 20]) -def test_kpool_indexer_long_prefill_and_cached_decode(max_kv_seq_len): - from lightllm.models.glm5_next.indexer import Glm5NextNsaInfer - - tokens, heads, dim = 2065, 32, 128 - hidden = torch.randn(tokens, dim, device="cuda", dtype=torch.bfloat16) - q_weight = torch.randn(dim, heads * dim, device="cuda", dtype=torch.bfloat16) * 0.1 - ape = torch.randn(4, dim, device="cuda") - weights = SimpleNamespace( - wk_proj_=SimpleNamespace(mm=lambda x: x), - k_norm_=lambda x, eps: x, - index_kpool_compress_gate=SimpleNamespace(mm=lambda x: x * 0.25), - index_kpool_compress_ape=SimpleNamespace(weight=ape), - wq_b_proj_=SimpleNamespace(mm=lambda x: x @ q_weight), - weights_proj_=SimpleNamespace(mm=lambda x: torch.ones(x.shape[0], heads, device=x.device)), - ) - 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] - full_att_state = SimpleNamespace( - get_layer_view=lambda logical_layer_id: SimpleNamespace(indexer_k=storage.view(torch.uint8)[:, :, -132:]), - ) - infer = SimpleNamespace( - 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 "compressor_indexer" in state_type_filter - else 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]), - 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), - max_q_seq_len=tokens, - ) - state = SimpleNamespace( - q_kv_seq_len=torch.arange(1, tokens + 1, device="cuda", dtype=torch.int32), - q_req_idx=torch.zeros(tokens, device="cuda", dtype=torch.int32), - ) - indexer = Glm5NextNsaInfer( - 0, - {"index_topk": 2048, "index_n_heads": heads, "index_head_dim": dim, "rms_norm_eps": 1e-5, "index_kpool": 4}, - 1, - ) - _, 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() - 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. Stale window data is irrelevant at a complete group boundary. - prefix = 2060 - storage = storage.clone() - storage[ragged[prefix:].long()] = 0 - tail.fill_(19) - infer.is_prefill = False - infer.b1_cu_q_seq_len[1] = 1 - infer.max_q_seq_len = 1 - for pos in range(prefix, tokens): - infer.mem_index = ragged[pos : pos + 1] - infer.max_kv_seq_len = pos + 1 - infer.b_seq_len.fill_(pos + 1) - state.q_kv_seq_len = torch.tensor([pos + 1], device="cuda", dtype=torch.int32) - state.q_req_idx = infer.b_req_idx - _, decoded = indexer._get_indices(hidden[pos : pos + 1], hidden[pos : pos + 1], infer, state, weights) - valid = 2048 + (pos + 1) % 4 - assert set(decoded[0, :valid].tolist()) == set(full[pos, :valid].tolist()) - - 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) - query = (hidden[-1:] @ q_weight).view(heads, dim) - 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) - 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_kpool_gather.py b/unit_tests/models/glm5_next/test_kpool_gather.py deleted file mode 100644 index 9c25401a3f..0000000000 --- a/unit_tests/models/glm5_next/test_kpool_gather.py +++ /dev/null @@ -1,141 +0,0 @@ -"""Byte-exact gather and graph-replay coverage for GLM K-pool staging.""" - -import pytest -import torch -import triton - -from lightllm.models.glm5_next.triton_kernel import kpool -from lightllm.models.glm5_next.triton_kernel.kpool import gather_paged_pools, gather_pools - - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") - - -def _inputs(max_pools, requests=8): - # The packed view has both a nonzero byte offset and a larger row stride. - storage = torch.empty((128, 1, 584), dtype=torch.bfloat16, device="cuda") - packed = storage.view(torch.uint8)[:, :, -132:] - keys = torch.randn(128, 128, device="cuda").to(torch.float8_e4m3fn).view(torch.uint8) - scales = torch.pow(2.0, torch.randint(-4, 3, (128,), device="cuda")).float() - packed[:, 0, :128] = keys - packed[:, 0, 128:] = scales.view(torch.uint8).view(128, 4) - table_storage = torch.full((requests, max_pools * 4 + 16), -1, device="cuda", dtype=torch.int32) - table = table_storage[:, 4:-12] - locations = torch.randint(0, 128, (requests, max_pools), device="cuda", dtype=torch.int32) - table[:, 3::4] = locations - return packed, table, locations, keys, scales - - -@pytest.mark.parametrize("pools", [17, 33, 640, 65536, 262144]) -def test_gather_pools_bytes_scales_and_changing_graph_inputs(pools): - torch.manual_seed(29) - packed, table, locations, source_keys, source_scales = _inputs(pools) - req = torch.tensor([3, 1, 0, 2], device="cuda", dtype=torch.int32) - lengths = torch.zeros(4, device="cuda", dtype=torch.int32) - gather_pools(packed, table, req, lengths, pools) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - keys, scales = gather_pools(packed, table, req, lengths, pools) - for values in ([0, 3, pools * 4 - 3, pools * 4], [pools * 4, 0, 7, 65], [2, 2, 2, 2]): - lengths.copy_(torch.tensor(values, device="cuda", dtype=torch.int32).clamp_max(pools * 4)) - req.copy_(req.roll(1)) - graph.replay() - valid = torch.arange(pools, device="cuda")[None, :] < (lengths // 4)[:, None] - loc = locations[req.long()].long() - expected_keys = source_keys[loc].masked_fill(~valid[:, :, None], 0) - expected_scales = source_scales[loc].masked_fill(~valid, 1.0) - torch.testing.assert_close(keys.view(torch.uint8).view(4, pools, 128), expected_keys, rtol=0, atol=0) - torch.testing.assert_close(scales.view(4, pools), expected_scales, rtol=0, atol=0) - - -def _check_paged(pages, block_table, req, lengths, locations, source_keys, source_scales): - page_bytes = pages.view(pages.shape[0], -1) - for row, (slot, count) in enumerate(zip(req.tolist(), lengths.tolist())): - # Shared pages may contain the owner's later keys. Only the original - # per-query causal length is readable by the indexer. - page_ids = block_table[row, : triton.cdiv(count, 64)].long() - actual_keys = page_bytes[page_ids, : 64 * 128].reshape(-1, 128)[:count] - actual_scales = page_bytes[page_ids, 64 * 128 :].contiguous().view(torch.float32).flatten()[:count] - loc = locations[slot, :count].long() - torch.testing.assert_close(actual_keys, source_keys[loc], rtol=0, atol=0) - torch.testing.assert_close(actual_scales, source_scales[loc], rtol=0, atol=0) - - -@pytest.mark.parametrize("max_pools", [640, 4096, 262144]) -@pytest.mark.parametrize("reuse_group", [1, 2, 4, 8]) -def test_paged_reuse_mixed_requests_causal_lengths_and_graph_slot_recycling(max_pools, reuse_group): - torch.manual_seed(30) - packed, table, locations, source_keys, source_scales = _inputs(max_pools) - batch = reuse_group * 2 + 1 - shared = [1] * reuse_group + [2 + i % 2 for i in range(reuse_group)] + [4] - mixed = [1 + i % 3 for i in range(batch)] - req = torch.tensor(shared, device="cuda", dtype=torch.int32) - lengths = torch.tensor([max(0, max_pools - i * 63) for i in range(batch)], device="cuda", dtype=torch.int32) - gather_paged_pools(packed, table, req, lengths, max_pools, reuse_group=reuse_group) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - pages, block_table = gather_paged_pools(packed, table, req, lengths, max_pools, reuse_group=reuse_group) - identity = torch.arange(block_table.numel(), device="cuda", dtype=torch.int32).view_as(block_table) - scenarios = [ - (shared, [max(0, max_pools - i * 63) for i in range(batch)]), - (mixed, [max_pools if i % 2 else min(65, max_pools) for i in range(batch)]), - (shared, [min(i * 17, max_pools) for i in range(batch)]), - ([0] * batch, [0] * batch), - ([5] * reuse_group + mixed[reuse_group:], [max_pools] * batch), - ] - for slots, counts in scenarios: - req.copy_(torch.tensor(slots, device="cuda", dtype=torch.int32)) - lengths.copy_(torch.tensor(counts, device="cuda", dtype=torch.int32)) - # Recycle physical KV locations without changing any captured pointer. - locations.copy_(locations.roll(1, dims=1)) - table[:, 3::4] = locations - original_lengths = lengths.clone() - if not any(counts): - pages.fill_(127) - graph.replay() - torch.testing.assert_close(lengths, original_lengths, rtol=0, atol=0) - _check_paged(pages, block_table, req, lengths, locations, source_keys, source_scales) - if max(counts) < 4096 or reuse_group == 1: - torch.testing.assert_close(block_table, identity, rtol=0, atol=0) - elif slots == shared and reuse_group > 1: - valid_pages = triton.cdiv(max(counts[:reuse_group]), 64) - for row in range(1, reuse_group): - torch.testing.assert_close(block_table[row, :valid_pages], block_table[0, :valid_pages]) - if not any(counts): - dummy = pages.view(pages.shape[0], -1)[block_table[:, 0].long()] - assert not dummy[:, : 64 * 128].any() - assert (dummy[:, 64 * 128 :].contiguous().view(torch.float32) == 1).all() - - -@pytest.mark.parametrize("batch", [1, 3, 5]) -def test_paged_partial_groups_keep_valid_per_query_pages(batch): - packed, table, locations, source_keys, source_scales = _inputs(4097) - req = torch.ones(batch, device="cuda", dtype=torch.int32) - lengths = torch.tensor([4097 - i for i in range(batch)], device="cuda", dtype=torch.int32) - pages, block_table = gather_paged_pools(packed, table, req, lengths, 4097, reuse_group=4) - _check_paged(pages, block_table, req, lengths, locations, source_keys, source_scales) - if batch < 4: - expected = torch.arange(block_table.numel(), device="cuda", dtype=torch.int32).view_as(block_table) - torch.testing.assert_close(block_table, expected, rtol=0, atol=0) - - -def test_paged_batch_sizes_reuse_compiled_kernel(): - packed, table, locations, source_keys, source_scales = _inputs(4097) - req = torch.ones(16, device="cuda", dtype=torch.int32) - lengths = torch.arange(4097, 4081, -1, device="cuda", dtype=torch.int32) - kernel = kpool._gather_paged_pools - cache = kernel.device_caches[torch.cuda.current_device()][0] - # Warm the shared-group specialization, then vary both complete and partial - # groups. Grid dimensions must not create additional compiled variants. - gather_paged_pools(packed, table, req[:4], lengths[:4], 4097) - variants = len(cache) - for batch in (5, 8, 12, 16): - pages, block_table = gather_paged_pools(packed, table, req[:batch], lengths[:batch], 4097) - _check_paged(pages, block_table, req[:batch], lengths[:batch], locations, source_keys, source_scales) - assert len(cache) == variants - - -@pytest.mark.parametrize("reuse_group", [0, 3, 16]) -def test_reject_unsupported_reuse_group_before_allocating(reuse_group): - with pytest.raises(ValueError, match="reuse_group"): - gather_paged_pools(None, None, None, None, 4096, reuse_group=reuse_group) diff --git a/unit_tests/models/glm5_next/test_kpool_runtime_specialization.py b/unit_tests/models/glm5_next/test_kpool_runtime_specialization.py deleted file mode 100644 index 0338c98575..0000000000 --- a/unit_tests/models/glm5_next/test_kpool_runtime_specialization.py +++ /dev/null @@ -1,200 +0,0 @@ -import pytest -import torch -import triton -from triton.runtime.jit import JITFunction - -from lightllm.models.glm5_next.triton_kernel.index_quant import ( - hadamard_transform_quant_fp8, -) -from lightllm.models.glm5_next.triton_kernel.kpool import ( - _compress_pools, - _gather_pools, - _get_pool_ranges, - compress_pools, - gather_pools, -) - - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") - - -def _check_pool_ranges(batch, max_pools, single_query): - counts = [1] * batch if single_query else [0 if i % 7 == 3 else i % 3 + 1 for i in range(batch)] - cu = [0] - visible, owners = [], [] - for owner, count in enumerate(counts): - visible.extend(range(4, count + 4)) - owners.extend([owner] * count) - cu.append(len(visible)) - lengths = torch.tensor(visible, device="cuda", dtype=torch.int32) - cu_q_lens = torch.tensor(cu, device="cuda", dtype=torch.int32) - outputs = [torch.empty_like(lengths) for _ in range(3)] - tokens = lengths.numel() - blocks = triton.cdiv(tokens, 256) if single_query else tokens // 256 + batch - compiled = _get_pool_ranges[(blocks,)]( - lengths, cu_q_lens, *outputs, batch, tokens, max_pools, single_query, 256, num_warps=4 - ) - starts = torch.tensor(owners, device="cuda", dtype=torch.int32) * max_pools - expected = (starts, starts + lengths // 4, lengths // 4) - for actual, reference in zip(outputs, expected): - assert torch.equal(actual, reference) - return compiled.hash - - -@pytest.mark.parametrize("single_query", [False, True]) -def test_pool_ranges_reuse_compilation_across_runtime_sizes(single_query): - # The indexer rounds capacity up to 128 pools. Changing that capacity keeps - # its default alignment specialization while batch/token counts are dynamic. - hashes = { - _check_pool_ranges(batch, max_pools, single_query) - for batch, max_pools in [(1, 128), (3, 256), (16, 384), (17, 640), (127, 768), (128, 1024), (129, 1152)] - } - assert len(hashes) == 1 - - -@pytest.mark.parametrize("single_query", [False, True]) -@pytest.mark.parametrize("batch,max_pools", [(1, 1), (3, 3), (16, 16), (129, 1025)]) -def test_pool_ranges_nonstandard_capacities_preserve_correctness(batch, max_pools, single_query): - _check_pool_ranges(batch, max_pools, single_query) - - -def test_pool_compression_reuses_compilation_within_static_query_layouts(monkeypatch): - torch.manual_seed(1525) - tail = torch.zeros((130, 4, 256), device="cuda", dtype=torch.bfloat16) - ape = torch.zeros((4, 128), device="cuda") - hashes = {} - layout = None - misses = [] - original_run = _compress_pools.run - original_compile = JITFunction._do_compile - - def observe_compile(self, *args, **kwargs): - if self is _compress_pools: - misses.append(layout) - return original_compile(self, *args, **kwargs) - - def observe(*args, **kwargs): - compiled = original_run(*args, **kwargs) - hashes.setdefault(layout, set()).add(compiled.hash) - return compiled - - monkeypatch.setattr(_compress_pools, "run", observe) - monkeypatch.setattr(JITFunction, "_do_compile", observe_compile) - for batch, single_query in [(1, False), (3, True), (33, False), (1, True), (17, True), (1, True)]: - values = torch.randint(-16, 17, (batch, 4, 256), device="cuda").bfloat16() - values[:, :, 128:] = 0 - tail[:batch].zero_() - packed = torch.zeros((batch * 4, 1, 720), device="cuda", dtype=torch.uint8) - requests = torch.arange(batch, device="cuda", dtype=torch.int32) - ragged = torch.arange(batch * 4, device="cuda", dtype=torch.int32) - seq_lens = torch.full((batch,), 4, device="cuda", dtype=torch.int32) - if single_query: - tail[:batch, :3] = values[:, :3] - raw = values[:, 3].contiguous() - lengths = seq_lens - mem_index = requests * 4 + 3 - cu = torch.arange(batch + 1, device="cuda", dtype=torch.int32) - else: - raw = values.reshape(batch * 4, 256) - lengths = torch.arange(1, 5, device="cuda", dtype=torch.int32).repeat(batch) - mem_index = ragged - cu = torch.arange(batch + 1, device="cuda", dtype=torch.int32) * 4 - # Query mode and address strides stay constexpr for the faster kernel. - layout = (single_query, raw.stride(0)) - seen_layout = layout in hashes - misses_before = len(misses) - compress_pools(raw, tail, packed, ape, lengths, mem_index, requests, cu, seq_lens, 1 if single_query else 4) - if seen_layout: - assert len(misses) == misses_before - keys, scales = gather_pools(packed, ragged.view(batch, 4), requests, seq_lens, 1) - expected = values[:, :, :128].float().mean(1).bfloat16() - expected_key, expected_scale = hadamard_transform_quant_fp8(expected, scale=128 ** -0.5) - assert torch.equal(keys.view(torch.uint8), expected_key.view(torch.uint8)) - assert torch.equal(scales, expected_scale.flatten()) - assert torch.equal(tail[:batch], values) - assert len(hashes) == 3 - assert all(len(variants) == 1 for variants in hashes.values()) - - -def test_gather_pools_reuses_compilation_across_pool_counts(): - torch.manual_seed(1525) - packed = torch.zeros((16, 1, 720), device="cuda", dtype=torch.uint8) - source_keys = torch.randint(0, 256, (16, 128), device="cuda", dtype=torch.uint8) - source_scales = torch.rand(16, device="cuda") + 0.1 - packed[:, 0, :128] = source_keys - packed[:, 0, 128:132] = source_scales.view(torch.uint8).view(16, 4) - locations = torch.randint(0, 16, (6, 129), device="cuda", dtype=torch.int32) - table = torch.zeros((6, 129 * 4), device="cuda", dtype=torch.int32) - table[:, 3::4] = locations - requests = torch.tensor([2, 0, 5], device="cuda", dtype=torch.int32) - hashes = set() - for pools in [1, 3, 16, 33, 64, 129]: - lengths = torch.tensor([pools * 4, min(3, pools) * 4, 2], device="cuda", dtype=torch.int32) - keys = torch.empty((3, pools, 128), device="cuda", dtype=torch.float8_e4m3fn) - scales = torch.empty((3, pools), device="cuda") - compiled = _gather_pools[(triton.cdiv(pools, 32), 3)]( - packed, - table, - requests, - lengths, - keys, - scales, - packed.stride(0), - table.stride(0), - pools, - POOLS_PER_BLOCK=32, - num_warps=4, - ) - hashes.add(compiled.hash) - for row, (request, count) in enumerate([(2, pools), (0, min(3, pools)), (5, 0)]): - indices = locations[request, :count].long() - assert torch.equal(keys.view(torch.uint8)[row, :count], source_keys[indices]) - assert torch.equal(scales[row, :count], source_scales[indices]) - assert not keys.view(torch.uint8)[row, count:].any() - assert (scales[row, count:] == 1).all() - assert len(hashes) == 1 - - -def test_pool_compression_reuses_compilation_across_request_counts(monkeypatch): - torch.manual_seed(1525) - tail = torch.zeros((130, 4, 256), device="cuda", dtype=torch.bfloat16) - # Exact uniform weights isolate indexing from softmax/FP8 rounding differences. - ape = torch.zeros((4, 128), device="cuda") - hashes = set() - original_run = _compress_pools.run - - def observe(*args, **kwargs): - compiled = original_run(*args, **kwargs) - hashes.add(compiled.hash) - return compiled - - monkeypatch.setattr(_compress_pools, "run", observe) - for batch in [1, 3, 16, 17, 127, 128, 129]: - raw = torch.randint(-16, 17, (batch * 4, 256), device="cuda").bfloat16() - raw[:, 128:] = 0 - packed = torch.zeros((batch * 4, 1, 720), device="cuda", dtype=torch.uint8) - requests = torch.arange(batch, device="cuda", dtype=torch.int32) - lengths = torch.arange(1, 5, device="cuda", dtype=torch.int32).repeat(batch) - ragged = torch.arange(batch * 4, device="cuda", dtype=torch.int32) - cu = torch.arange(batch + 1, device="cuda", dtype=torch.int32) * 4 - seq_lens = torch.full((batch,), 4, device="cuda", dtype=torch.int32) - compress_pools( - raw, - tail, - packed, - ape, - lengths, - ragged, - requests, - cu, - seq_lens, - max_q_seq_len=4, - ) - keys, scales = gather_pools(packed, ragged.view(batch, 4), requests, seq_lens, 1) - values = raw.view(batch, 4, 256) - expected = (values[:, :, :128].float() * (values[:, :, 128:].float() + ape).softmax(1)).sum(1).bfloat16() - expected_key, expected_scale = hadamard_transform_quant_fp8(expected, scale=128 ** -0.5) - assert torch.equal(keys.view(torch.uint8), expected_key.view(torch.uint8)) - assert torch.equal(scales, expected_scale.flatten()) - assert torch.equal(tail[:batch], values) - assert len(hashes) == 1 diff --git a/unit_tests/models/glm5_next/test_layers.py b/unit_tests/models/glm5_next/test_layers.py deleted file mode 100644 index 1aca878b2d..0000000000 --- a/unit_tests/models/glm5_next/test_layers.py +++ /dev/null @@ -1,451 +0,0 @@ -import weakref -import dataclasses -import json -from types import SimpleNamespace - -import pytest -import torch -import torch.nn.functional as F - -from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.common.basemodel.attention.nsa.glm5_next import Glm5NextSparsePrefillState, Glm5NextSparseDecodeState -from lightllm.common.quantization import Quantcfg -from lightllm.models.glm5_next.layer_infer.transformer_layer_infer import Glm5NextTransformerLayerInfer -from lightllm.models.glm5_next.layer_weights.transformer_layer_weight import Glm5NextTransformerLayerWeight -from lightllm.models.glm5_next.model import Glm5NextTpPartModel -from lightllm.models.glm5_next_mtp.model import Glm5NextMTPModel -from lightllm.server.core.objs.start_args_type import StartArgs -from lightllm.utils.envs_utils import get_env_start_args - - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") - - -@pytest.fixture(autouse=True) -def runtime(monkeypatch): - monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(dataclasses.asdict(StartArgs()))) - for name, value in { - "GLOBAL_RANK": 0, - "GLOBAL_WORLD_SIZE": 1, - "DP_WORLD_SIZE": 1, - "CURRENT_RANK_IN_DP": 0, - "CURRENT_RANK_IN_NODE": 0, - "CURRENT_DEVICE_ID": 0, - }.items(): - monkeypatch.setenv("LIGHTLLM_" + name, str(value)) - get_env_start_args.cache_clear() - torch.manual_seed(53) - yield - get_env_start_args.cache_clear() - - -@pytest.fixture -def config(): - return { - "hidden_size": 128, - "n_embed": 128, - "intermediate_size": 128, - "moe_intermediate_size": 128, - "num_hidden_layers": 4, - "n_layer": 4, - "num_attention_heads": 4, - "num_key_value_heads": 4, - "vocab_size": 256, - "rms_norm_eps": 1e-5, - "q_lora_rank": 128, - "kv_lora_rank": 512, - "qk_nope_head_dim": 128, - "qk_rope_head_dim": 0, - "v_head_dim": 128, - "index_n_heads": 4, - "index_head_dim": 128, - "index_topk": 2048, - "index_kpool": 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, - "n_shared_experts": 1, - "first_k_dense_replace": 1, - "num_experts_per_tok": 2, - "norm_topk_prob": True, - "scoring_func": "sigmoid", - "n_group": 1, - "topk_group": 1, - "swiglu_limit": 10.0, - "num_nextn_predict_layers": 1, - } - - -def random_weight(*shape): - return torch.randn(*shape, device="cuda", dtype=torch.bfloat16) * 0.05 - - -def rmsnorm(x, eps): - return (x.float() * torch.rsqrt(x.float().square().mean(-1, keepdim=True) + eps)).to(x.dtype) - - -def test_kda_projections_need_no_sparse_attention_config(config): - for key in ("q_lora_rank", "kv_lora_rank", "qk_nope_head_dim", "qk_rope_head_dim", "v_head_dim"): - del config[key] - for key in list(config): - if key.startswith("index_"): - del config[key] - layer = Glm5NextTransformerLayerInfer(0, config) - weight = Glm5NextTransformerLayerWeight(0, torch.bfloat16, config, Quantcfg(config)) - projections = { - name: random_weight(size, 128) - for name, size in {"q": 512, "k": 512, "v": 512, "b": 4, "f_a": 128, "g_a": 128, "f_b": 512, "g_b": 512}.items() - } - weight.load_hf_weights( - {f"model.language_model.layers.0.self_attn.{name}_proj.weight": value for name, value in projections.items()} - ) - hidden = random_weight(3, 128) - actual = layer._kda_projections(hidden, SimpleNamespace(), weight) - expected = ( - torch.cat([F.linear(hidden, projections[name]) for name in ("q", "k", "v")], dim=-1), - F.linear(F.linear(hidden, projections["f_a"]), projections["f_b"]), - F.linear(hidden, projections["b"]), - F.linear(F.linear(hidden, projections["g_a"]), projections["g_b"]), - ) - for result, reference in zip(actual, expected): - torch.testing.assert_close(result, reference, atol=1e-3, rtol=0.01) - - -@pytest.mark.parametrize("rank", [0, 1]) -@pytest.mark.parametrize("is_prefill", [False, True]) -def test_mla_loading_preserves_replicated_qkv_and_sharded_bmm(config, monkeypatch, rank, is_prefill): - monkeypatch.setenv("LIGHTLLM_DP_WORLD_SIZE", "2") - monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_DP", str(rank)) - layer = Glm5NextTransformerLayerInfer(3, config) - weight = Glm5NextTransformerLayerWeight(3, torch.bfloat16, config, Quantcfg(config)) - q_a, kv_a = random_weight(128, 128), random_weight(512, 128) - q_b, kv_b, o = random_weight(512, 128), random_weight(1024, 512), random_weight(128, 512) - tensors = { - "q_a_proj.weight": q_a, - "kv_a_proj_with_mqa.weight": kv_a, - "q_b_proj.weight": q_b, - "kv_b_proj.weight": kv_b, - "o_proj.weight": o, - "q_a_layernorm.weight": torch.ones(128, device="cuda", dtype=torch.bfloat16), - "kv_a_layernorm.weight": torch.ones(512, device="cuda", dtype=torch.bfloat16), - } - weight.load_hf_weights( - {f"model.language_model.layers.3.self_attn.{name}": value for name, value in tensors.items()} - ) - hidden = random_weight(3, 128) - state = SimpleNamespace(is_prefill=is_prefill, need_dp_prefill_balance=False) - q, kv = layer._get_qkv(hidden, state, weight) - q_lora = rmsnorm(F.linear(hidden, q_a), config["rms_norm_eps"]) - expected_q = F.linear(q_lora, q_b).view(3, 4, 128)[:, rank * 2 : (rank + 1) * 2] - expected_kv = rmsnorm(F.linear(hidden, kv_a), config["rms_norm_eps"]).unsqueeze(1) - torch.testing.assert_close(q, expected_q, atol=0.015, rtol=0.015) - torch.testing.assert_close(kv, expected_kv, atol=0.015, rtol=0.015) - head_weights = kv_b.view(4, 256, 512)[rank * 2 : (rank + 1) * 2] - projected = weight.k_b_proj_.bmm(q.transpose(0, 1)) - torch.testing.assert_close(projected, torch.bmm(q.transpose(0, 1), head_weights[:, :128])) - latent_output = random_weight(3, 2, 512) - monkeypatch.setattr(layer, "_tpsp_reduce", lambda input, infer_state: input) - actual = layer._get_o(latent_output, state, weight) - values = torch.bmm(latent_output.transpose(0, 1), head_weights[:, 128:].transpose(1, 2)) - expected = F.linear(values.transpose(0, 1).reshape(3, 256), o[:, rank * 256 : (rank + 1) * 256]) - torch.testing.assert_close(actual, expected) - - -def test_sparse_layer_prefill_and_decode_match_dense_nope_attention(config): - if torch.cuda.get_device_capability()[0] != 9: - pytest.skip("GLM sparse attention kernels require Hopper") - layer = Glm5NextTransformerLayerInfer(3, config) - weight = Glm5NextTransformerLayerWeight(3, torch.bfloat16, config, Quantcfg(config)) - kv_b = random_weight(1024, 512) - weight.load_hf_weights({"model.layers.3.self_attn.kv_b_proj.weight": kv_b}) - q = random_weight(3, 4, 128) - # Use noncontiguous token slots and a padded sparse index table. - packed = random_weight(8, 1, 584) - kv = packed[:, :, :512] - slots = torch.tensor([7, 1, 5, 0, 3], device="cuda", dtype=torch.int32) - lengths = torch.tensor([2, 3, 5], device="cuda", dtype=torch.int32) - indices = torch.full((3, 128), -1, device="cuda", dtype=torch.int32) - for row, length in enumerate(lengths.tolist()): - indices[row, :length] = slots[:length] - layer.indexer = SimpleNamespace(_get_indices=lambda **kwargs: (indices, indices)) - state = InferStateInfo() - 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)) - ) - 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 - state.decode_att_state = Glm5NextSparseDecodeState( - _infer_state_weak_ref=weakref.ref(state), - nsa_cache_seqlens=lengths, - nsa_cu_seqlens_k_new=F.pad(lengths.cumsum(0, dtype=torch.int32), (1, 0)), - ) - projected = torch.bmm(q.transpose(0, 1), kv_b.view(4, 256, 512)[:, :128]).transpose(0, 1) - expected = [] - for row, length in enumerate(lengths.tolist()): - keys = kv[slots[:length].long(), 0].float() - scores = projected[row].float() @ keys.T * config["qk_nope_head_dim"] ** -0.5 - expected.append(scores.softmax(-1) @ keys) - expected = torch.stack(expected) - prefill = layer._context_attention_kernel(q, None, state, weight) - state.get_topk_indices_params = {"hidden_states": None, "q_lora": None} - decode = layer._token_attention_kernel(q, state, weight) - for output in (prefill, decode): - torch.testing.assert_close(output.float(), expected, atol=1e-3, rtol=0.015) - - -@pytest.mark.parametrize("mode", ["dense", "moe", "fused_moe"]) -def test_ffn_dispatch_preserves_clamp_and_shared_expert(config, monkeypatch, mode): - args = StartArgs(enable_fused_shared_experts=mode == "fused_moe") - monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(dataclasses.asdict(args))) - get_env_start_args.cache_clear() - index = 0 if mode == "dense" else 3 - layer = Glm5NextTransformerLayerInfer(index, config) - weight = Glm5NextTransformerLayerWeight(index, torch.bfloat16, config, Quantcfg(config)) - prefix = f"model.language_model.layers.{index}.mlp" - eye = torch.eye(128, device="cuda", dtype=torch.bfloat16) - tensors = {} - - def mlp_weights(name, gate_scale, up_scale): - tensors[f"{name}.gate_proj.weight"] = eye * gate_scale - tensors[f"{name}.up_proj.weight"] = eye * up_scale - tensors[f"{name}.down_proj.weight"] = eye - - if mode == "dense": - mlp_weights(prefix, 1.5, 0.75) - else: - mlp_weights(f"{prefix}.shared_experts", 1.5, 0.75) - for expert in range(4): - mlp_weights(f"{prefix}.experts.{expert}", 1 + expert / 4, 2 + expert / 8) - router = torch.zeros(4, 128, device="cuda", dtype=torch.float32) - router[:, :4] = torch.eye(4, device="cuda") * 0.2 - tensors[f"{prefix}.gate.weight"] = router - tensors[f"{prefix}.gate.e_score_correction_bias"] = torch.zeros(4, device="cuda") - weight.load_hf_weights(tensors) - x = torch.linspace(-25, 25, 3 * 128, device="cuda", dtype=torch.bfloat16).view(3, 128) - - def reference_mlp(gate_scale, up_scale): - gate = (x * gate_scale).float().clamp(max=10) - up = (x * up_scale).float().clamp(-10, 10) - return (F.silu(gate).bfloat16().float() * up).bfloat16() - - expected = reference_mlp(1.5, 0.75) - if mode != "dense": - scores = F.linear(x.float(), router).sigmoid() - probabilities, ids = scores.topk(2, dim=-1) - probabilities /= probabilities.sum(-1, keepdim=True) - routed = torch.zeros_like(x) - for expert in range(4): - coefficient = (probabilities * (ids == expert)).sum(-1, keepdim=True) - routed += (reference_mlp(1 + expert / 4, 2 + expert / 8).float() * coefficient).bfloat16() - expected += routed - actual = layer._ffn(x.clone(), SimpleNamespace(is_prefill=True), weight) - torch.testing.assert_close(actual, expected, atol=0.25, rtol=0.015) - - -@pytest.mark.parametrize("is_prefill", [False, True]) -def test_ep_ffn_uses_combined_output_and_replicated_shared_expert(config, monkeypatch, is_prefill): - args = StartArgs(tp=2, enable_ep_moe=True, enable_fused_shared_experts=True) - monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(dataclasses.asdict(args))) - monkeypatch.setenv("LIGHTLLM_DP_WORLD_SIZE", "2") - monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_DP", "1") - monkeypatch.setenv("LIGHTLLM_GLOBAL_WORLD_SIZE", "2") - monkeypatch.setenv("LIGHTLLM_GLOBAL_RANK", "1") - get_env_start_args.cache_clear() - layer = Glm5NextTransformerLayerInfer(3, config) - weight = Glm5NextTransformerLayerWeight(3, torch.bfloat16, config, Quantcfg(config)) - assert weight.num_fused_shared_experts == 0 - prefix = "model.language_model.layers.3.mlp.shared_experts" - eye = torch.eye(128, device="cuda", dtype=torch.bfloat16) - weight.load_hf_weights({f"{prefix}.{name}.weight": eye for name in ("gate_proj", "up_proj", "down_proj")}) - x = torch.linspace(-25, 25, 3 * 128, device="cuda", dtype=torch.bfloat16).view(3, 128) - original = x.clone() - monkeypatch.setattr(weight.moe_gate, "mm", lambda x: torch.zeros(3, 4, device="cuda")) - monkeypatch.setattr(layer, "_tpsp_reduce", lambda **kwargs: pytest.fail("EP output must not be TP-reduced again")) - - def combined_output(hidden_states, **kwargs): - assert kwargs["is_prefill"] is is_prefill - assert kwargs["alpha"] == 1.0 and kwargs["limit"] == 10.0 - assert kwargs["clamp_up_add_one"] is False - return torch.full_like(hidden_states, 17) - - monkeypatch.setattr(weight.experts, "experts", combined_output) - actual = layer._ffn(x, SimpleNamespace(is_prefill=is_prefill), weight) - shared = (F.silu(x.float().clamp(max=10)).bfloat16().float() * x.float().clamp(-10, 10)).bfloat16() - torch.testing.assert_close(actual, shared + 17, atol=0.25, rtol=0.01) - torch.testing.assert_close(x, original, atol=0, rtol=0) - - -def test_fp8_shared_expert_scales_and_bf16_kv_b_load_together(config, monkeypatch): - monkeypatch.setenv( - "LIGHTLLM_START_ARGS", json.dumps(dataclasses.asdict(StartArgs(enable_fused_shared_experts=True))) - ) - get_env_start_args.cache_clear() - config["quantization_config"] = {"quant_method": "fp8", "weight_block_size": [128, 128]} - weight = Glm5NextTransformerLayerWeight(3, torch.bfloat16, config, Quantcfg(config)) - prefix = "model.language_model.layers.3" - tensors = {} - for expert in range(5): - name = f"experts.{expert}" if expert < 4 else "shared_experts" - for projection in ("gate_proj", "up_proj", "down_proj"): - tensors[f"{prefix}.mlp.{name}.{projection}.weight"] = random_weight(128, 128).to(torch.float8_e4m3fn) - tensors[f"{prefix}.mlp.{name}.{projection}.weight_scale_inv"] = torch.ones(1, 1, device="cuda") - tensors[f"{prefix}.mlp.gate.e_score_correction_bias"] = torch.zeros(4, device="cuda") - kv_b = random_weight(1024, 512) - tensors[f"{prefix}.self_attn.kv_b_proj.weight"] = kv_b - # There is deliberately no kv_b scale: this matrix stays BF16 in the FP8 checkpoint. - weight.load_hf_weights(tensors) - assert weight.experts.verify_load() - assert weight.k_b_proj_.verify_load() and weight.v_b_proj_.verify_load() - q = random_weight(4, 2, 128) - actual = weight.k_b_proj_.bmm(q) - expected = torch.bmm(q, kv_b.view(4, 256, 512)[:, :128]) - torch.testing.assert_close(actual, expected) - - -@pytest.mark.parametrize("method", ["context_forward", "token_forward"]) -@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -def test_mtp_pre_layer_applies_main_norm_before_hidden_fusion(config, monkeypatch, method, dtype): - from lightllm.common.basemodel.layer_weights.meta_weights import RMSNormWeight - from lightllm.models.glm5_next_mtp.layer_infer.pre_layer_infer import Glm5NextMTPPreLayerInfer - from lightllm.models.qwen_vl.layer_infer.pre_layer_infer import LlamaMultimodalPreLayerInfer - - hidden_size, eps = config["hidden_size"], config["rms_norm_eps"] - norms = [RMSNormWeight(hidden_size, name, dtype) for name in ("main_norm", "enorm", "hnorm")] - for norm in norms: - norm.weight.copy_(torch.randn_like(norm.weight)) - embeddings = torch.randn(5, hidden_size, device="cuda", dtype=dtype) - hidden = torch.randn_like(embeddings) - projection = torch.randn(2 * hidden_size, hidden_size, device="cuda", dtype=dtype) * 0.1 - weight = SimpleNamespace( - main_norm_weight_=norms[0], - enorm_weight_=norms[1], - hnorm_weight_=norms[2], - eh_proj_weight_=SimpleNamespace(mm=lambda x: x @ projection), - ) - state = SimpleNamespace(mtp_draft_input_hiddens=hidden.clone()) - monkeypatch.setattr(LlamaMultimodalPreLayerInfer, method, lambda self, ids, state, weight: embeddings.clone()) - actual = getattr(Glm5NextMTPPreLayerInfer(config), method)(None, state, weight) - - def reference_norm(value, norm): - return F.rms_norm(value.float(), (hidden_size,), norm.weight.float(), eps).to(dtype) - - normalized_hidden = reference_norm(reference_norm(hidden, norms[0]), norms[2]) - expected = torch.cat((reference_norm(embeddings, norms[1]), normalized_hidden), dim=-1) @ projection - torch.testing.assert_close(actual, expected, atol=0.015 if dtype == torch.bfloat16 else 1e-5, rtol=0.015) - - -@pytest.mark.parametrize("model_class", [Glm5NextTpPartModel, Glm5NextMTPModel]) -def test_glm_weight_loading_requires_hf(config, model_class): - model = object.__new__(model_class) - model.config, model.tp_world_size_ = config, 1 - model.args, model.load_way = StartArgs(), "HF" - model._verify_params() - model.load_way = "DS" - with pytest.raises(AssertionError, match="only support HF format weights"): - model._verify_params() - - -@pytest.mark.parametrize( - "kv_type, data_type, kv_lora_rank", - [ - ("fp8kv_sph", torch.bfloat16, 512), - ("fp8kv_spt", torch.bfloat16, 512), - ("fp8kv_dsa", torch.float16, 512), - ("fp8kv_dsa", torch.bfloat16, 256), - ], -) -def test_glm_rejects_unsupported_kv_config(config, kv_type, data_type, kv_lora_rank): - model = object.__new__(Glm5NextTpPartModel) - model.config = dict(config, kv_lora_rank=kv_lora_rank) - model.tp_world_size_, model.data_type = 1, data_type - model.args, model.load_way = StartArgs(llm_kv_type=kv_type), "HF" - with pytest.raises(ValueError, match="GLM-5.3 Flash"): - model._verify_params() - - -@pytest.mark.parametrize("tp_world_size", [1, 2]) -def test_glm_rejects_tpsp(config, tp_world_size): - model = object.__new__(Glm5NextTpPartModel) - model.config, model.tp_world_size_ = config, tp_world_size - model.args, model.load_way = StartArgs(), "HF" - model._verify_params() - - model.args.enable_tpsp_mix_mode = True - with pytest.raises(AssertionError, match="GLM-5.3 Flash does not support TP/SP mixed mode"): - model._verify_params() - - -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(data_type="bfloat16", mtp_mode="vanilla_with_att", mtp_step=2), "HF" - main._verify_params() - main._init_some_value() - 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 = 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.tp_rank_, draft.data_type = 1, 0, torch.bfloat16 - draft.args = main.args - draft.load_way = "HF" - draft._init_config() - draft._verify_params() - draft.quant_cfg = Quantcfg(draft.config) - draft._init_weights() - draft._init_infer_layer() - draft._init_some_value() - 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"] - assert draft.layers_num == 1 and draft.head_dim_ == 512 - assert draft.layers_infer[0].layer_num_ == 4 + step - assert draft.trans_layers_weight[0].layer_num_ == 4 - 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 - drafts.append(draft) - - 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() - assert main.att_state_manager.main_kda_conv_bf16_ssm_fp32_state.conv_state.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 - draft._init_att_backend1() - 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.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}_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): - state = model.infer_state_class() - state.is_prefill = True - state.input_ids = torch.zeros(3, device="cuda", dtype=torch.int64) - state.b_ready_cache_len = torch.tensor([2], device="cuda", dtype=torch.int32) - state.b_seq_len = torch.tensor([5], device="cuda", dtype=torch.int32) - state.init_some_extra_state(model) - torch.testing.assert_close( - state.position_ids, torch.arange(2, 5, device="cuda", dtype=state.position_ids.dtype) - ) diff --git a/unit_tests/models/glm5_next/test_mtp.py b/unit_tests/models/glm5_next/test_mtp.py deleted file mode 100644 index 30c4fc22c4..0000000000 --- a/unit_tests/models/glm5_next/test_mtp.py +++ /dev/null @@ -1,308 +0,0 @@ -import weakref -import pytest -import torch - -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.common.basemodel.triton_kernel.linear_att.replayssm_fold import fold_prefix -from lightllm.models.glm5_next.triton_kernel.kpool import compress_pools - - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") - - -@pytest.mark.parametrize("mtp_step", [1, 2, 3, 5]) -@pytest.mark.parametrize("state_dtype", [torch.float32, torch.bfloat16]) -def test_kda_mtp_matches_single_token_decode_and_restores_accepted_state(mtp_step, state_dtype): - torch.manual_seed(53) - width, heads, dim = mtp_step + 1, 2, 128 - # Nonconsecutive requests plus a zero-length graph-padding sequence. - req_ids = torch.tensor([3, 1, 5], device="cuda", dtype=torch.int32) - slots = req_ids[:, None] * width + torch.arange(width, device="cuda", dtype=torch.int32) - states = torch.randn(6 * width, heads, dim, dim, device="cuda", dtype=state_dtype) * 0.1 - accepted = torch.tensor([width, 1, 1], device="cuda", dtype=torch.int32) - for lengths in ([width, 1, 0], [2, width, 0], [1, 1, 0]): - cu = torch.tensor( - [0, lengths[0], sum(lengths), sum(lengths)], - device="cuda", - dtype=torch.int32, - ) - tokens = sum(lengths) - q, k, v, gate = [torch.randn(1, tokens, heads, dim, device="cuda", dtype=torch.bfloat16) for _ in range(4)] - beta = torch.randn(1, tokens, heads, device="cuda", dtype=torch.bfloat16) - a = torch.randn(heads, device="cuda") - bias = torch.randn(heads * dim, device="cuda") - reference = states.clone() - before = states.clone() - expected = torch.empty_like(v) - for request, length in enumerate(lengths): - if not length: - continue - start = sum(lengths[:request]) - current = reference[slots[request, accepted[request] - 1]].clone().unsqueeze(0) - for offset in range(length): - token = start + offset - out, _ = fused_recurrent_kda( - q[:, token : token + 1], - k[:, token : token + 1], - v[:, token : token + 1], - gate[:, token : token + 1].reshape(1, 1, -1), - beta[:, token : token + 1], - a, - bias, - current, - torch.zeros(1, device="cuda", dtype=torch.int32), - ) - expected[:, token : token + 1] = out - reference[slots[request, offset]] = current[0] - actual, _ = fused_recurrent_kda( - q, - k, - v, - gate.reshape(1, tokens, -1), - beta, - a, - bias, - states, - slots, - cu_seqlens=cu, - num_accepted_tokens=accepted, - ) - torch.testing.assert_close(actual, expected, atol=2e-3, rtol=1e-2) - torch.testing.assert_close(states, reference, atol=2e-3, rtol=1e-2) - assert torch.equal(states[slots[2]], before[slots[2]]) - # The next verify must begin at the accepted candidate, not the last written slot. - accepted = torch.tensor([min(2, lengths[0]), 1, 1], device="cuda", dtype=torch.int32) - - -@pytest.mark.parametrize("mtp_step", [1, 2, 3, 5]) -@pytest.mark.parametrize("prefix", [1, 2, 3, 4, 7, 15]) -def test_kpool_verify_and_draft_rewind_match_full_prefill(mtp_step, prefix): - torch.manual_seed(53) - width, capacity = mtp_step + 1, 128 - raw = torch.randn(capacity, 256, device="cuda", dtype=torch.bfloat16) - ape = torch.randn(4, 128, device="cuda", dtype=torch.bfloat16) - ring = torch.zeros(2, 4 + mtp_step, 256, device="cuda", dtype=torch.bfloat16) - packed = torch.zeros(capacity, 132, device="cuda", dtype=torch.uint8) - ragged = torch.arange(capacity, device="cuda", dtype=torch.int32) - - def prefill(values, destination, tail): - count = values.shape[0] - compress_pools( - values, - tail, - destination, - ape, - torch.arange(1, count + 1, device="cuda", dtype=torch.int32), - ragged[:count], - torch.zeros(1, device="cuda", dtype=torch.int32), - torch.tensor([0, count], device="cuda", dtype=torch.int32), - torch.tensor([count], device="cuda", dtype=torch.int32), - count, - ) - - prefill(raw[:prefix], packed, ring) - # Alternate expanded verification, rewound single-token drafting, and re-verification. - 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[start : start + count], - ring, - packed, - ape, - lengths, - ragged[start : start + count], - torch.zeros(count, device="cuda", dtype=torch.int32), - torch.arange(count + 1, device="cuda", dtype=torch.int32), - lengths, - 1, - b_mtp_index=torch.arange(count, device="cuda", dtype=torch.int32), - ) - expected = torch.zeros_like(packed) - prefill( - raw[: start + count], - expected, - torch.zeros(2, 4, 256, device="cuda", dtype=torch.bfloat16), - ) - closing = torch.arange(start, start + count, device="cuda") - closing = closing[(closing + 1) % 4 == 0] - assert torch.equal(packed[closing], expected[closing]) - - -@pytest.mark.parametrize("dynamic", [False, True]) -@pytest.mark.parametrize("role", ["main", "draft0"]) -@pytest.mark.parametrize("mode,mtp_step", [("native", 2), ("replay", 2), ("native", 0), ("replay", 0)]) -def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str, mode, mtp_step): - from types import SimpleNamespace - - from lightllm.common.basemodel.attention.base_att import AttControl - from lightllm.common.basemodel.attention.linear.kda import KDADecodeAttState - from lightllm.common.basemodel.attention.linear.replay_kda import ( - ReplayKDADecodeAttState, - ) - from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ( - ReplayKDAState, - ) - from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import ( - causal_conv1d_update, - ) - - torch.manual_seed(53) - heads, dim, width = 2, 128, mtp_step + 1 - hidden = heads * dim - lengths = [2, 1, 3] if dynamic and mtp_step else [width] * 3 - requests = [2, 0, 1] - request_rows = [req for req, length in zip(requests, lengths) for _ in range(length)] - offsets = [offset for length in lengths for offset in range(length)] - if dynamic: - request_rows += [3, 3] - offsets += [0, 0] - tokens = len(request_rows) - conv = torch.randn(4, 3 * hidden, 3 + mtp_step, device="cuda", dtype=torch.bfloat16) - ssm = torch.randn(4 * width, heads, dim, dim, device="cuda") * 0.01 - old_conv, expected_ssm = conv.clone(), ssm.clone() - accepted_offsets = torch.tensor([1, 2, 0, 0] if mtp_step else [0, 0, 0, 0], dtype=torch.int32, device="cuda") - mixed = torch.randn(tokens, 3 * hidden, device="cuda", dtype=torch.bfloat16) - gate = torch.randn(tokens, hidden, device="cuda", dtype=torch.bfloat16) - beta = torch.randn(tokens, heads, device="cuda", dtype=torch.bfloat16) - conv_weight = torch.randn(3 * hidden, 4, device="cuda", dtype=torch.bfloat16) * 0.1 - a, bias = torch.randn(heads, device="cuda"), torch.randn(hidden, device="cuda") - expected = [] - row = 0 - zero = torch.zeros(1, dtype=torch.int32, device="cuda") - for req, length in zip(requests, lengths): - offset = int(accepted_offsets[req]) - history = old_conv[req : req + 1, :, offset : offset + 3].contiguous() - state = expected_ssm[req * width + offset].clone().unsqueeze(0) - for token in range(length): - convolved = causal_conv1d_update( - mixed[row : row + 1].clone(), - history, - conv_weight, - bias=None, - activation="silu", - conv_state_indices=zero, - ) - q, k, v = [part.reshape(1, 1, heads, dim) for part in convolved.split(hidden, -1)] - output, _ = fused_recurrent_kda( - q, - k, - v, - gate[row : row + 1].reshape(1, 1, hidden), - beta[row : row + 1].reshape(1, 1, heads), - a, - bias, - state, - zero, - ) - expected.append(output.reshape(1, heads, dim)) - expected_ssm[req * width + token] = state[0] - row += 1 - backend = SimpleNamespace( - mtp_step=mtp_step, - tp_num_heads=heads, - head_dim=dim, - tp_hidden_size=hidden, - lower_bound=-5.0, - uses_dynamic_spec_verify_layout=lambda: dynamic, - split_qkv=lambda x: x.split(hidden, -1), - ) - 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, - ) - selected = None - for model_role in ("main", "draft0"): - state_cls = ReplayKDAState if mode == "replay" else KDAState - linear = state_cls(config, request_capacity=4, mtp_step=mtp_step, device="cuda") - if mtp_step: - linear.mtp_state_index.copy_(accepted_offsets if model_role == role else torch.zeros_like(accepted_offsets)) - if model_role == role: - linear.conv_state[0].copy_(conv) - if mode == "native": - linear.ssm_state[0].copy_(ssm) - else: - indexes = torch.arange(4, device="cuda") * width + accepted_offsets - linear.ssm_state[0].copy_(ssm[indexes]) - selected = linear - 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_state_manager = resources - infer.att_state_role = role - weight = SimpleNamespace( - get_merged_kda_conv_weight=lambda: conv_weight, - linear_A_log=SimpleNamespace(weight=a), - linear_dt_bias=SimpleNamespace(weight=bias), - ) - decode_cls = ReplayKDADecodeAttState if mode == "replay" else KDADecodeAttState - state = decode_cls(backend=backend, _infer_state_weak_ref=weakref.ref(infer)) - state.init_state() - actual = state.decode_att( - None, - None, - None, - AttControl( - linear_att_decode=True, - linear_att_decode_dict=dict( - layer_weight=weight, - layer_num=0, - mixed_qkv=mixed, - raw_gate=gate, - raw_beta=beta, - ), - ), - ) - torch.testing.assert_close( - actual.reshape(tokens, heads, dim)[:row], - torch.cat(expected), - atol=2e-3, - rtol=1e-2, - ) - if mode == "native": - checked_slots = len(requests) if dynamic and not mtp_step else expected_ssm.shape[0] - torch.testing.assert_close( - selected.ssm_state[0, :checked_slots], - expected_ssm[:checked_slots], - atol=2e-3, - rtol=1e-2, - ) - else: - for req, length in zip(requests, lengths): - saved = selected.ssm_state - clone = saved.clone() - selected.ssm_state = clone - try: - fold_prefix(selected, torch.tensor([req], dtype=torch.int32, device="cuda")) - finally: - selected.ssm_state = saved - torch.testing.assert_close( - clone[0, req], - expected_ssm[req * width + length - 1], - atol=2e-3, - rtol=1e-2, - ) - if dynamic and mtp_step: - assert torch.equal(selected.conv_state[0, 3], old_conv[3]) diff --git a/unit_tests/models/glm5_next/test_mtp_prefill_cache.py b/unit_tests/models/glm5_next/test_mtp_prefill_cache.py deleted file mode 100644 index 5b60fece98..0000000000 --- a/unit_tests/models/glm5_next/test_mtp_prefill_cache.py +++ /dev/null @@ -1,278 +0,0 @@ -from types import SimpleNamespace -from unittest.mock import Mock - -import pytest -import torch - -from lightllm.common.attention_state import AttentionStateManager -from lightllm.common.attention_state.bounded_att import WindowCompressorState, WindowCompressorStateConfig -from lightllm.common.attention_state.full_att.packed_mla.fp8_e8m0_packed_mla_state import ( - FP8E8M0PackedMLAConfig, - FP8E8M0PackedMLAState, -) -from lightllm.common.attention_state.full_att.packed_mla.packed_mla_state import ( - PackedMLAAttentionConfig, - PackedMLAFullAttentionState, - PackedMLALogicalConfig, -) -from lightllm.common.basemodel.batch_objs import ModelInput, PostLayerOutput -from lightllm.common.basemodel.hidden_collector import FinalHiddenCollector -from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.common.basemodel.triton_kernel.norm.rmsnorm import rmsnorm_forward -from lightllm.models.glm5_next.indexer import Glm5NextNsaInfer -from lightllm.models.glm5_next_mtp.layer_infer.post_layer_infer import Glm5NextMTPPostLayerInfer -from lightllm.models.glm5_next_mtp.layer_infer.transformer_layer_infer import Glm5NextMTPTransformerLayerInfer -from lightllm.models.glm5_next_mtp.model import Glm5NextMTPModel -from lightllm.models.llama.layer_infer.post_layer_infer import LlamaPostLayerInfer -from lightllm.server.core.objs.start_args_type import StartArgs - - -def make_model(): - model = Glm5NextMTPModel.__new__(Glm5NextMTPModel) - model.args = StartArgs() - model.args.mtp_mode = "eagle_with_att" - model.args.mtp_step = 1 - model.args.llm_kv_type = "fp8kv_dsa" - model.tp_world_size_ = model.layers_num = 1 - model.prefill_graph = None - model.layers_infer = [SimpleNamespace(is_linear_attention_layer=False, use_mhc=False)] - model.trans_layers_weight = [object()] - return model - - -@pytest.mark.parametrize( - "field,value", - [ - ("mtp_mode", "eagle_no_att"), - ], -) -def test_unsupported_execution_uses_full_forward(field, value): - model = make_model() - assert model._supports_prefill_kv_only() - setattr(model.args, field, value) - assert not model._supports_prefill_kv_only() - - -@pytest.mark.parametrize("field,value", [("mtp_mode", "eagle_no_att")]) -def test_unsupported_layout_uses_full_forward(field, value): - model = make_model() - setattr(model.args, field, value) - assert not model._supports_prefill_kv_only() - - -def test_eagle_multistep_and_vanilla_multistep_capabilities(): - model = make_model() - model.args.mtp_step = 3 - assert model._supports_prefill_kv_only() - model.args.mtp_mode = "vanilla_with_att" - assert not model._supports_prefill_kv_only() - - -@pytest.mark.parametrize("kv_type", ["fp8kv_dsa", "bf16"]) -@pytest.mark.parametrize("tp_world_size", [1, 4]) -def test_cache_only_capability_does_not_depend_on_kv_format_or_tp(kv_type, tp_world_size): - model = make_model() - model.args.llm_kv_type = kv_type - model.tp_world_size_ = tp_world_size - assert model._supports_prefill_kv_only() - - -@pytest.mark.parametrize("tokens", [0, 1, 8192]) -@pytest.mark.parametrize("mode", ["eagle_with_att", "vanilla_with_att"]) -def test_forward_uses_common_prefill(monkeypatch, tokens, mode): - from lightllm.common.basemodel import basemodel - - model = make_model() - model.args.mtp_mode = mode - model.mtp_previous_draft_models = [] - model.req_manager = SimpleNamespace(HOLD_REQUEST_ID=3) - model.att_state_manager = object() - model.hidden_collector_prototype = FinalHiddenCollector() - model.return_all_prompt_logics = False - model.prefill_att_backend = model.prefill_att_backend1 = None - model._select_mem_indexes = lambda model_input: torch.zeros_like(model_input.input_ids) - batch = int(tokens > 0) - model_input = ModelInput( - is_prefill=True, - batch_size=batch, - total_token_num=tokens, - max_q_seq_len=tokens, - max_kv_seq_len=tokens, - max_cache_len=0, - input_ids=torch.zeros(tokens, dtype=torch.int64), - b_req_idx=torch.zeros(batch, dtype=torch.int32), - b_mtp_index=torch.zeros(batch, dtype=torch.int32), - b_seq_len=torch.full((batch,), tokens, dtype=torch.int32), - b_ready_cache_len=torch.zeros(batch, dtype=torch.int32), - b_prefill_start_loc=torch.zeros(batch, dtype=torch.int32), - b_is_decode_req=torch.zeros(batch, dtype=torch.bool), - b_prefill_has_output_cpu=[True] * batch, - multimodal_params=[{"images": [], "audios": []} for _ in range(batch)], - mtp_draft_input_hiddens=torch.ones(tokens, 2), - ) - model.pre_infer = SimpleNamespace( - context_forward=Mock(side_effect=lambda ids, *_: torch.zeros(ids.shape[0], 2)), - _tpsp_sp_split=lambda input, **_: input, - ) - model.post_infer = Glm5NextMTPPostLayerInfer.__new__(Glm5NextMTPPostLayerInfer) - model.post_infer.tp_world_size_ = 1 - model.post_infer._tpsp_allgather = Mock(side_effect=lambda input, **_: input) - post_forward = Mock(side_effect=lambda hidden, state, _: PostLayerOutput(torch.ones(state.batch_size, 3))) - monkeypatch.setattr(LlamaPostLayerInfer, "token_forward", post_forward) - model.pre_post_weight = object() - layer = Glm5NextMTPTransformerLayerInfer.__new__(Glm5NextMTPTransformerLayerInfer) - layer.is_linear_attention_layer = layer.use_mhc = False - layer.layer_num_, layer.num_hidden_layers = 45, 1 - layer.embed_dim_ = 2 - layer._att_norm = layer._ffn_norm = lambda hidden, *_: hidden - layer.context_attention_forward = Mock(side_effect=lambda hidden, *_: torch.zeros_like(hidden)) - layer._ffn = Mock(side_effect=lambda hidden, *_: torch.zeros_like(hidden)) - layer._project_sparse_qkv_a = lambda hidden, *_: (hidden, None, hidden) - layer._normalize_sparse_kv = lambda kv, _: kv - layer.indexer = SimpleNamespace(write_k_pools=Mock()) - writer = Mock() - layer._post_cache_kv = writer - model.layers_infer = [layer] - allocator = SimpleNamespace(cache_env_in=Mock(), cache_env_out=Mock()) - monkeypatch.setattr(basemodel, "g_cache_manager", allocator) - monkeypatch.setattr(basemodel.dist_group_manager, "get_group", Mock()) - monkeypatch.setattr(basemodel.dist_group_manager, "clear_deepep_buffer", Mock()) - monkeypatch.setattr(ModelInput, "to_cuda", lambda self: self.check_input()) - monkeypatch.setattr(InferStateInfo, "init_some_extra_state", Mock()) - event = Mock() - monkeypatch.setattr(torch.cuda, "Event", Mock(return_value=event)) - output = model.forward(model_input) - assert output.prefill_mem_indexes_ready_event is event - assert output.logits.numel() == 0 - assert output.mtp_collector.spec_hidden is None - ids, infer, _ = model.pre_infer.context_forward.call_args.args - assert ids.shape == (max(1, tokens),) - assert infer.b_req_idx.tolist() == ([0] if tokens else [model.req_manager.HOLD_REQUEST_ID]) - if mode == "eagle_with_att": - assert infer.mtp_draft_input_hiddens is None - writer.assert_called_once() - layer.context_attention_forward.assert_not_called() - layer._ffn.assert_not_called() - post_forward.assert_not_called() - placeholder = model.post_infer._tpsp_allgather.call_args.kwargs["input"] - assert placeholder.shape == (0, layer.embed_dim_) - assert placeholder.untyped_storage().nbytes() == 0 - layer.indexer.write_k_pools.assert_called_once() - else: - writer.assert_called_once() - layer.context_attention_forward.assert_not_called() - layer._ffn.assert_not_called() - post_forward.assert_not_called() - allocator.cache_env_in.assert_called_once() - allocator.cache_env_out.assert_called_once() - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("kv_type", ["fp8kv_dsa", "bf16"]) -@pytest.mark.parametrize( - "requests", - [((3, 0, 1),), ((0, 0, 8192),), ((0, 131071, 129), (2, 1048573, 3)), ((1, 3, 1), (0, 5, 7))], -) -def test_cache_only_matches_full_prefill_cache_writes(monkeypatch, kv_type, requests): - from lightllm.models.glm5_next import indexer as indexer_module - - torch.manual_seed(53) - layer_id, tokens = 45, sum(q for _, _, q in requests) - layer = Glm5NextMTPTransformerLayerInfer.__new__(Glm5NextMTPTransformerLayerInfer) - layer.layer_num_, layer.num_hidden_layers = layer_id, 1 - layer.use_mhc = layer.is_linear_attention_layer = False - layer.tp_world_size_ = layer.tp_q_head_num_ = 1 - layer.embed_dim_ = layer.q_lora_rank = 64 - layer.kv_lora_rank = layer.qk_nope_head_dim = 512 - layer.eps_ = 1e-5 - layer.indexer = Glm5NextNsaInfer( - layer_id, dict(index_topk=2048, index_n_heads=1, index_head_dim=128, rms_norm_eps=1e-5, index_kpool=4), 1 - ) - - def projection(rows, cols): - weight = torch.randn(rows, cols, device="cuda", dtype=torch.bfloat16) * 0.1 - return SimpleNamespace(mm=Mock(side_effect=lambda x: x @ weight)) - - def norm(input, eps, **_): - return rmsnorm_forward(input, None, eps) - - weights = SimpleNamespace( - att_norm_weight_=norm, - qkv_a_proj_with_mqa_=projection(64, 576), - q_a_layernorm_=SimpleNamespace(weight=torch.ones(64, device="cuda", dtype=torch.bfloat16)), - kv_a_layernorm_=SimpleNamespace(weight=torch.ones(512, device="cuda", dtype=torch.bfloat16)), - q_b_proj_=projection(64, 512), - wk_proj_=projection(64, 128), - k_norm_=norm, - index_kpool_compress_gate=projection(64, 128), - index_kpool_compress_ape=SimpleNamespace(weight=torch.randn(4, 128, device="cuda", dtype=torch.bfloat16)), - ) - resources = AttentionStateManager(request_capacity=4) - stores = [] - for role in ("main", "draft0"): - if kv_type == "fp8kv_dsa": - full = FP8E8M0PackedMLAState(FP8E8M0PackedMLAConfig(logical_layer_groups=((layer_id,),))) - else: - full = PackedMLAFullAttentionState( - PackedMLAAttentionConfig( - PackedMLALogicalConfig( - logical_layer_groups=((layer_id,),), - latent_dim=512, - format="packed_bf16", - indexer_dim=128, - ) - ), - device="cuda", - ) - full.allocate_storage(2 * tokens + 2) - full.storage.fill_(0x5A) - tail = WindowCompressorState( - WindowCompressorStateConfig((layer_id,), 4, 4, 128, "indexer", torch.bfloat16), 4, mtp_step=1 - ) - tail.buffer.normal_() - resources.register_state(role, full) - resources.register_state(role, tail) - stores.extend((full.storage, tail.buffer)) - resources.finish_state_registration() - initial = [s.clone() for s in stores] - q_lengths = [q for _, _, q in requests] - lengths = torch.cat([torch.arange(past + 1, past + q + 1, device="cuda") for _, past, q in requests]).int() - infer = SimpleNamespace( - is_prefill=True, - is_draft_model=True, - prefill_kv_only=True, - need_dp_prefill_balance=False, - att_state_role="draft0", - att_state_manager=resources, - mem_index=2 * torch.arange(tokens, device="cuda", dtype=torch.int32) + 1, - b_req_idx=torch.tensor([r for r, _, _ in requests], device="cuda", dtype=torch.int32), - b_seq_len=torch.tensor([p + q for _, p, q in requests], device="cuda", dtype=torch.int32), - b1_cu_q_seq_len=torch.tensor([0, *q_lengths], device="cuda", dtype=torch.int32).cumsum(0).int(), - max_q_seq_len=max(q_lengths), - max_kv_seq_len=2048, - prefill_att_state=SimpleNamespace(q_kv_seq_len=lengths, q_req_idx=None), - req_manager=SimpleNamespace(req_to_token_indexs=None), - ) - # Top-k outputs are discarded here; run the normal projections and cache writers. - monkeypatch.setattr(indexer_module, "expand_topk", Mock(return_value=(None, None))) - hidden = torch.randn(tokens, 64, device="cuda", dtype=torch.bfloat16) - _, kv = layer._get_qkv(layer._att_norm(hidden, infer, weights), infer, weights) - layer._post_cache_kv(kv, infer, weights) - layer.indexer._get_indices( - **infer.get_topk_indices_params, infer_state=infer, att_state=infer.prefill_att_state, layer_weight=weights - ) - expected = [s.clone() for s in stores] - assert not torch.equal(initial[2], expected[2]) - for store, before in zip(stores, initial): - store.copy_(before) - weights.q_b_proj_.mm.reset_mock() - placeholder = layer.context_forward(hidden, infer, weights) - assert placeholder.shape == (0, layer.embed_dim_) - assert placeholder.untyped_storage().nbytes() == 0 - weights.q_b_proj_.mm.assert_not_called() - for actual, reference in zip(stores, expected): - assert torch.equal(actual, reference) - for actual, before in zip(stores[:2], initial[:2]): - assert torch.equal(actual, before) - assert torch.equal(stores[2][:, ::2], initial[2][:, ::2]) - assert torch.equal(stores[3][:, 3], initial[3][:, 3]) diff --git a/unit_tests/models/glm5_next/test_pd_cache.py b/unit_tests/models/glm5_next/test_pd_cache.py deleted file mode 100644 index 4d15d6bf1c..0000000000 --- a/unit_tests/models/glm5_next/test_pd_cache.py +++ /dev/null @@ -1,443 +0,0 @@ -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 -from types import SimpleNamespace - -import pytest -import torch - -from lightllm.common.attention_state.allocators import FullAttKVAllocator -from unit_tests.common.att_resources.model_helpers import create_full_attention_roles - -from lightllm.common.attention_state.full_att import ( - PackedMLAFullAttentionState, - 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 -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"), - pytest.mark.usefixtures("kv_allocator_environment"), -] - - -@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 = request.param - args = StartArgs( - data_type="bfloat16", - pd_kv_page_size=2048, - linear_att_hash_page_size=4, - linear_att_page_block_num=2, - mtp_step=mtp_step, - mtp_mode="eagle_with_att" if mtp_step else None, - ) - set_unique_server_name(args) - get_env_start_args.cache_clear() - set_env_start_args(dataclasses.asdict(args)) - - def create(tp, target_mtp_step=None, mode="native"): - 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_head_dim=512, - global_linear_k_heads=8, - global_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, - full_attention_interval=4, - all_layer_num=8, - ) - 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_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, - ) - full.allocate_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, - ssm_state_mode=mode, - 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, - ) - mems.append( - SimpleNamespace( - 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_compressor_indexer_w4_s4_bf16_state, - tails=window_states(resources), - resources=resources, - ) - ) - return mems - - return create - - -@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]) -@pytest.mark.parametrize("source_mode", ["native", "replay"]) -@pytest.mark.parametrize("dest_mode", ["native", "replay"]) -def test_pd_kv_and_runtime_state_roundtrip( - make_mems, - prefill_tp, - decode_tp, - remainder, - decode_mtp_step, - source_mode, - dest_mode, -): - source = make_mems(prefill_tp, mode=source_mode) - dest = make_mems(decode_tp, target_mtp_step=decode_mtp_step, mode=dest_mode) - dst_ssm_stride = 1 if dest_mode == "replay" else dest[0].linear.mtp_step + 1 - mtp_size = get_env_start_args().mtp_step + 1 - 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() - dst_indexes = torch.randperm(32)[:length].tolist() - 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( - 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 - for state in mem.full_states: - 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)) - src_slot = src_req if source_mode == "replay" else src_req * (mem.linear.mtp_step + 1) - mem.linear.ssm_state[:, src_slot].copy_(ssm[:, heads]) - if source_mode == "replay": - backend = mem.linear - backend.raw_keys.zero_() - backend.raw_values.zero_() - backend.gates.fill_(-0.5) - backend.betas.zero_() - backend.cursors[src_req] = 1 - offset = 0 - for state in mem.tails: - end = offset + len(state.config.logical_layer_ids) - state.buffer[:, src_req].copy_(tail[offset:end]) - offset = end - - untouched = [] - original_tails = [] - for mem in dest: - for state in mem.full_states: - state.storage.zero_() - mem.linear.conv_state.normal_() - mem.linear.ssm_state.normal_() - if dest_mode == "replay": - mem.linear.cursors[dst_req] = 3 - if mem.linear.mtp_enabled: - mem.linear.mtp_state_index[dst_req] = decode_mtp_step - for state in mem.tails: - 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_ssm_stride), - *((state.buffer, 1) for state in mem.tails), - ) - ] - ) - - 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_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, - 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"), - task_kind=kind, - req_idx=src_req, - state_token_len=length, - write_to_pd_page=True, - ) - dest_buffer.copy_(source_buffer) - copy_page( - 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"), - task_kind=kind, - req_idx=dst_req, - state_token_len=length, - write_to_pd_page=False, - ) - torch.cuda.synchronize() - - for rank, mem in enumerate(dest): - 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 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), - ) - 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_ssm_stride], ssm[:, heads]) - 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]] - assert torch.equal(restored_tail, expected_tail) - if dest_mode == "replay": - assert mem.linear.cursors[dst_req] == 0 - if mem.linear.mtp_enabled: - assert mem.linear.mtp_state_index[dst_req] == 0 - for (state, stride), expected in zip( - ( - (mem.linear.conv_state, 1), - (mem.linear.ssm_state, dst_ssm_stride), - *((state.buffer, 1) for state in mem.tails), - ), - untouched[rank], - ): - indexes = torch.tensor([0, 1, 3], device="cuda") * stride - 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. - count = 4 - remainder - 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 = [] - for mem, req_idx, indexes in ((source[0], src_req, src_indexes), (dest[0], dst_req, dst_indexes)): - new_indexes = [i for i in range(32) if i not in indexes][:count] - compress_pools( - k_and_gate=raw, - indexer_tail=mem.tail.buffer[0], - indexer_k=mem.full.get_layer_view(3).indexer_k, - index_kpool_compress_ape=ape, - q_kv_seq_len=tensor(range(length + 1, length + count + 1)), - mem_index=tensor(new_indexes), - b_req_idx=tensor([req_idx]), - b1_cu_q_seq_len=tensor([0, count]), - b_seq_len=tensor([length + count]), - max_q_seq_len=count, - ) - 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, 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() - ) - runtime_bytes = max( - 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_operator().get_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)]) -@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.bounded_att import WindowCompressorStateConfig, create_window_compressor_state - from lightllm.common.attention_state.types import StateTransferLayout - - torch.manual_seed(53) - config = WindowCompressorStateConfig( - logical_layer_ids=(0,), - window_size=4, - compress_stride=4, - candidate_dim=128, - target="indexer", - dtype=torch.bfloat16, - ) - source = create_window_compressor_state( - config=config, - request_capacity=2, - mtp_step=source_step, - device="cuda", - ) - target = create_window_compressor_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( - k_and_gate=values, - indexer_tail=tail, - indexer_k=destination, - index_kpool_compress_ape=ape, - q_kv_seq_len=torch.arange( - start=1, - end=count + 1, - device="cuda", - dtype=torch.int32, - ), - mem_index=ragged[:count], - b_req_idx=torch.zeros(1, device="cuda", dtype=torch.int32), - b1_cu_q_seq_len=torch.tensor([0, count], device="cuda", dtype=torch.int32), - b_seq_len=torch.tensor([count], device="cuda", dtype=torch.int32), - max_q_seq_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( - 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_page( - req_idx=0, - pd_page_buffers=pages, - pd_page_index=0, - state_token_len=prefix, - layout=layout, - ) - target.buffer.fill_(17) - target.read_from_pd_page( - pd_page_buffers=pages, - req_idx=0, - pd_page_index=0, - state_token_len=prefix, - layout=layout, - ) - 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)]: - 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( - k_and_gate=raw[start : start + count], - indexer_tail=target.get_layer_view(0), - indexer_k=packed, - index_kpool_compress_ape=ape, - q_kv_seq_len=lengths, - mem_index=ragged[start : start + count], - b_req_idx=torch.zeros(count, device="cuda", dtype=torch.int32), - b1_cu_q_seq_len=torch.arange(count + 1, device="cuda", dtype=torch.int32), - b_seq_len=lengths, - max_q_seq_len=1, - b_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/glm5_next/test_prefill_workspace.py b/unit_tests/models/glm5_next/test_prefill_workspace.py deleted file mode 100644 index 0d3a85d965..0000000000 --- a/unit_tests/models/glm5_next/test_prefill_workspace.py +++ /dev/null @@ -1,116 +0,0 @@ -from types import SimpleNamespace - -import pytest -import torch - -from lightllm.common.basemodel.basemodel import TpPartBaseModel -from lightllm.common.basemodel.attention.nsa import fp8_glm5_next -from lightllm.models.glm5_next.model import Glm5NextTpPartModel -from lightllm.models.glm5_next_mtp.model import Glm5NextMTPModel -from unit_tests.models.glm5_next.test_fp8_kv import _PrefillCase, _make_workspace_backend, _capture_cuda_graph - - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") - - -@pytest.fixture(autouse=True) -def current_cuda_device(monkeypatch): - monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", str(torch.cuda.current_device())) - - -@pytest.mark.parametrize("overlap", [False, True]) -@pytest.mark.parametrize("run_mode", ["normal", "decode"]) -def test_prefill_workspace_is_reserved_once_and_reused_at_runtime(monkeypatch, overlap, run_mode): - model = SimpleNamespace( - args=SimpleNamespace(enable_prefill_microbatch_overlap=overlap, run_mode=run_mode), - batch_max_tokens=8, - max_req_num=2, - max_seq_length=13, - config={"index_topk": 2048}, - att_state_role=f"workspace_test_{run_mode}_{overlap}", - ) - backend_class = fp8_glm5_next.Fp8Glm5NextSparseAttBackend - reserved = backend_class.allocate_prefill_workspaces(model) - pointers = [{key: value.data_ptr() for key, value in buffers.items()} for buffers in reserved] - backend = object.__new__(backend_class) - backend.__init__(model) - assert [ - {key: value.data_ptr() for key, value in buffers.items()} for buffers in backend._prefill_workspaces - ] == pointers - assert len(reserved) == 1 + int(overlap) - capacity = 8 if run_mode == "decode" else 34 - assert reserved[0]["gathered_kv"].shape == (capacity, 1, 512) - assert reserved[0]["chunk_topk_indices"].shape == (8, 2176) - - def forbidden(*args, **kwargs): - raise AssertionError("Runtime must reuse the workspace reserved before KV capacity planning") - - monkeypatch.setattr(backend, "allocate_prefill_workspaces", forbidden) - monkeypatch.setattr(backend, "get_gpu_workspace_buffer", forbidden) - for slot in range(1 + int(overlap)): - for q_tokens in (1, 8, 3): - topk = torch.empty((q_tokens, 2176), dtype=torch.int32, device="cuda") - views = backend.get_prefill_workspace(SimpleNamespace(microbatch_index=slot), topk) - assert {key: value.data_ptr() for key, value in views.items()} == pointers[slot] - torch.testing.assert_close( - views["cu_seqlens_q"], torch.arange(q_tokens + 1, dtype=torch.int32, device="cuda") - ) - - -def test_prefill_microbatches_use_independent_scratch_during_graph_replay(monkeypatch): - monkeypatch.setattr(fp8_glm5_next, "KV_GATHER_CHUNK_SIZE", 7) - backend = _make_workspace_backend(8, overlap=True) - first = _PrefillCase([15], [3], 8, backend=backend) - second = _PrefillCase([17], [3], 8, backend=backend) - second.infer_state.microbatch_index = 1 - first.check(first.run()) - second.check(second.run()) - pointers = [{key: value.data_ptr() for key, value in buffers.items()} for buffers in backend._prefill_workspaces] - assert all(pointers[0][key] != pointers[1][key] for key in pointers[0]) - - def forbidden(*args, **kwargs): - raise AssertionError("Attention must not allocate a new prefill workspace") - - monkeypatch.setattr(backend, "allocate_prefill_workspaces", forbidden) - first_graph, first_output = _capture_cuda_graph(first.run) - second_graph, second_output = _capture_cuda_graph(second.run) - first_stream, second_stream = torch.cuda.Stream(), torch.cuda.Stream() - first.mutate() - second.mutate() - for _ in range(3): - first_stream.wait_stream(torch.cuda.current_stream()) - second_stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(first_stream): - first_graph.replay() - with torch.cuda.stream(second_stream): - second_graph.replay() - torch.cuda.current_stream().wait_stream(first_stream) - torch.cuda.current_stream().wait_stream(second_stream) - first.check(first_output) - second.check(second_output) - assert [ - {key: value.data_ptr() for key, value in buffers.items()} for buffers in backend._prefill_workspaces - ] == pointers - - -@pytest.mark.parametrize("model_class", [Glm5NextTpPartModel, Glm5NextMTPModel]) -@pytest.mark.parametrize("kv_type", ["None", "fp8kv_dsa"]) -def test_glm_reserves_prefill_workspace_during_resource_preparation(monkeypatch, model_class, kv_type): - model = object.__new__(model_class) - model.args = SimpleNamespace(llm_kv_type=kv_type) - model.mtp_previous_draft_models = [] - model.att_state_manager = SimpleNamespace(get_full_attention_states=lambda role: [object()]) - calls = [] - - def prepare(self): - assert self is model - calls.append("prepare_states") - - def reserve(cls, prepared_model): - assert prepared_model is model - calls.append("reserve_workspace") - - monkeypatch.setattr(TpPartBaseModel, "prepare_attention_states", prepare) - monkeypatch.setattr(fp8_glm5_next.Fp8Glm5NextSparseAttBackend, "allocate_prefill_workspaces", classmethod(reserve)) - model.prepare_attention_states() - assert calls == (["prepare_states", "reserve_workspace"] if kv_type == "fp8kv_dsa" else ["prepare_states"]) diff --git a/unit_tests/models/glm5_next/test_token_major_mla_bmm.py b/unit_tests/models/glm5_next/test_token_major_mla_bmm.py deleted file mode 100644 index fa9e727f87..0000000000 --- a/unit_tests/models/glm5_next/test_token_major_mla_bmm.py +++ /dev/null @@ -1,92 +0,0 @@ -from types import SimpleNamespace -from unittest.mock import Mock - -import pytest -import torch - -from lightllm.common.basemodel.layer_weights.meta_weights.mm_weight.mm_weight import BMMWeightTpl -from lightllm.models.glm5_next.layer_infer.transformer_layer_infer import Glm5NextTransformerLayerInfer - - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") - - -def make_projection(layer, heads: int, inner: int, cols: int, dtype: torch.dtype): - projection = BMMWeightTpl.__new__(BMMWeightTpl) - projection.weight = torch.randn(heads, inner, cols, device="cuda", dtype=dtype) * 0.1 - layer.tp_q_head_num_ = heads - if inner == 256: - layer.kv_lora_rank = cols - layer.layer_num_ = 0 - layer.softmax_scale = inner ** -0.5 - layer.indexer = SimpleNamespace(_get_indices=lambda **kwargs: (None, None)) - state = SimpleNamespace( - att_state_role="main", - att_state_manager=Mock(), - prefill_att_state=SimpleNamespace(prefill_att=lambda **kwargs: kwargs["q"]), - ) - weight = SimpleNamespace(k_b_proj_=projection) - - def run(x): - state.get_topk_indices_params = {"hidden_states": None, "q_lora": None} - return layer._context_attention_kernel(x, None, state, weight) - - else: - layer.v_head_dim = cols - layer._tpsp_reduce = lambda input, infer_state: input - state = SimpleNamespace(is_prefill=False, need_dp_prefill_balance=False) - weight = SimpleNamespace(v_b_proj_=projection, o_weight_=SimpleNamespace(mm=lambda x: x)) - - def run(x): - return layer._get_o(x, state, weight).view(x.shape[0], heads, cols) - - return projection, run - - -@pytest.mark.parametrize("rows", [1, 128, 1024, 1659, 4096, 8192]) -@pytest.mark.parametrize("heads", [16, 64]) -@pytest.mark.parametrize("inner,cols", [(256, 512), (512, 256)]) -@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) -def test_token_major_projection_matches_reference_and_keeps_guard_rows(rows, heads, inner, cols, dtype): - torch.manual_seed(272) - layer = Glm5NextTransformerLayerInfer.__new__(Glm5NextTransformerLayerInfer) - projection, run = make_projection(layer, heads, inner, cols, dtype) - x = torch.randn(rows, heads, inner, device="cuda", dtype=dtype) - expected = projection.bmm(x.transpose(0, 1)).transpose(0, 1).contiguous() - slab = torch.full((rows + 2, heads, cols), 113.0, device="cuda", dtype=dtype) - output = slab[1:-1] - layer.alloc_tensor = Mock(return_value=output) - actual = run(x) - assert actual.data_ptr() == output.data_ptr() and actual.is_contiguous() - assert torch.equal(actual, expected) - assert (slab[0] == 113).all() and (slab[-1] == 113).all() - - -@pytest.mark.parametrize("inner,cols", [(256, 512), (512, 256)]) -def test_token_major_projection_accepts_head_padding(inner, cols): - # TP4 FlashMLA crops a 64-head output to 16 heads without copying it. - layer = Glm5NextTransformerLayerInfer.__new__(Glm5NextTransformerLayerInfer) - projection, run = make_projection(layer, 16, inner, cols, torch.bfloat16) - x = torch.randn(128, 64, inner, device="cuda", dtype=torch.bfloat16)[:, :16] - assert not x.is_contiguous() - expected = projection.bmm(x.transpose(0, 1)).transpose(0, 1) - actual = run(x) - assert actual.is_contiguous() - assert torch.equal(actual, expected) - - -@pytest.mark.parametrize("inner,cols", [(256, 512), (512, 256)]) -def test_token_major_projection_replays_new_inputs_in_cuda_graph(inner, cols): - layer = Glm5NextTransformerLayerInfer.__new__(Glm5NextTransformerLayerInfer) - projection, run = make_projection(layer, 16, inner, cols, torch.bfloat16) - x = torch.randn(128, 16, inner, device="cuda", dtype=torch.bfloat16) - for _ in range(3): - run(x) - torch.cuda.synchronize() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - actual = run(x) - x.normal_() - graph.replay() - expected = projection.bmm(x.transpose(0, 1)).transpose(0, 1) - assert torch.equal(actual, expected) diff --git a/unit_tests/models/glm5_next/test_visual.py b/unit_tests/models/glm5_next/test_visual.py deleted file mode 100644 index 6a25a8d4ea..0000000000 --- a/unit_tests/models/glm5_next/test_visual.py +++ /dev/null @@ -1,216 +0,0 @@ -import dataclasses -import json -from types import SimpleNamespace - -import numpy as np -import pytest -import torch -from PIL import Image - -from lightllm.models.glm5_next.tokenizer import Glm5NextTokenizer -from lightllm.models.glm5_next.vision_process import Glm5NextImageProcessor, smart_resize -from lightllm.server.core.objs.start_args_type import StartArgs -from lightllm.server.multimodal_params import ImageItem, MultimodalParams -from lightllm.utils.envs_utils import set_env_start_args - - -@pytest.mark.parametrize("width,height,canvas", [(7, 11, (168, 112)), (173, 89, (112, 196)), (448, 448, (448, 448))]) -def test_image_canvas_and_token_expansion(tmp_path, width, height, canvas): - assert smart_resize(height, width) == canvas - (tmp_path / "processor_config.json").write_text(json.dumps({"image_processor": {}})) - tokenizer = Glm5NextTokenizer( - None, {"image_start_token_id": 100, "image_end_token_id": 101, "image_token_id": 102}, str(tmp_path) - ) - image = ImageItem(type="base64", data="") - image.image_w, image.image_h = width, height - image.token_id = 1000 - image.token_num = tokenizer.get_image_token_length(image) - assert image.token_num == canvas[0] * canvas[1] // 28 ** 2 - multi = MultimodalParams() - multi.images = [image] - ids = tokenizer.encode([1, 100, 102, 101, 2], multi) - assert ids == [1, 100, *range(1000, 1000 + image.token_num), 101, 2] - assert image.start_idx == 2 - assert image.grid_thwd[-1] == 0 - assert tokenizer.encode([1, 2, 3]) == [1, 2, 3] - assert tokenizer.encode([100, 102, 101]) == [100, 101] - - -def test_image_patch_order_temporal_repeat_and_black_padding(): - pixels = np.zeros((29, 31, 3), dtype=np.uint8) - pixels[:, :, 0] = np.arange(31)[None, :] - pixels[:, :, 1] = np.arange(29)[:, None] - pixels[:, :, 2] = 127 - processor = Glm5NextImageProcessor(min_image_tokens=1, do_rescale=False, do_normalize=False) - patches, grid = processor._preprocess_bydevice(Image.fromarray(pixels), device="cpu") - assert grid.tolist() == [[1, 4, 4]] - patches = patches.reshape(16, 3, 2, 14, 14) - assert torch.equal(patches[:, :, 0], patches[:, :, 1]) - assert torch.equal(patches[0, :, 0], torch.from_numpy(pixels[:14, :14]).permute(2, 0, 1)) - assert torch.equal(patches[1, :, 0], torch.from_numpy(pixels[:14, 14:28]).permute(2, 0, 1)) - assert torch.equal(patches[2, :, 0], torch.from_numpy(pixels[14:28, :14]).permute(2, 0, 1)) - assert not patches[-1].any() - - -def test_vision_attention_uses_triton_rope_with_fa3_backend(monkeypatch): - import lightllm.models.glm5_next.glm5_next_visual as glm_visual - import lightllm.server.visualserver as visualserver - from lightllm.common.basemodel.attention_vit.fa3.fp import Fa3VitAttBackend - from lightllm.server.visualserver import set_vit_att_backend - - rope_inputs = [] - norm_inputs = [] - attn_inputs = [] - - def fake_qk_norm(qkv, q_weight, k_weight, eps): - norm_inputs.append((qkv, q_weight, k_weight, eps)) - return qkv[:, 0].contiguous(), qkv[:, 1].contiguous() - - def fake_rope(x, cos, sin): - rope_inputs.append((x, cos, sin)) - return x - - def fake_fa3(q, k, v, o, cu_seqlens, max_seqlen): - attn_inputs.append((q, k, v, cu_seqlens, max_seqlen)) - o.copy_(q) - return o - - monkeypatch.setattr(glm_visual, "qk_rms_norm", fake_qk_norm) - monkeypatch.setattr(glm_visual, "apply_rotary_pos_emb_triton", fake_rope) - monkeypatch.setattr(Fa3VitAttBackend, "_vit_att_fwd", staticmethod(fake_fa3)) - monkeypatch.setattr(visualserver, "VIT_ATTN_BACKEND", visualserver.VIT_ATTN_BACKEND) - set_vit_att_backend("fa3") - - attention = glm_visual.Glm5NextVisionAttention(hidden_size=16, num_heads=2, eps=1e-5, bias=True) - x = torch.randn(3, 16) - cos = torch.randn(3, 4) - sin = torch.randn(3, 4) - cu_seqlens = torch.tensor([0, 3], dtype=torch.int32) - result = attention(x, cu_seqlens, 3, cos, sin) - - assert result.shape == x.shape - assert len(norm_inputs) == 1 - assert norm_inputs[0][0].shape == (3, 3, 2, 8) - assert [(q.shape, cos_.shape, sin_.shape) for q, cos_, sin_ in rope_inputs] == [ - ((3, 2, 8), (3, 4), (3, 4)), - ((3, 2, 8), (3, 4), (3, 4)), - ] - assert len(attn_inputs) == 1 - assert attn_inputs[0][0].shape == attn_inputs[0][1].shape == attn_inputs[0][2].shape == (3, 2, 8) - - -def test_vision_rms_norm_preserves_glm_rounding(monkeypatch): - import lightllm.models.glm5_next.glm5_next_visual as glm_visual - - norm_calls = [] - - def fake_rms_norm(x, weight, eps, round_norm_before_weight): - norm_calls.append((x, weight, eps, round_norm_before_weight)) - return x - - monkeypatch.setattr(glm_visual, "rms_norm", fake_rms_norm) - norm = glm_visual.Glm5NextVisionRMSNorm(hidden_size=8, eps=1e-5) - x = torch.randn(2, 8) - assert norm(x) is x - assert len(norm_calls) == 1 - assert norm_calls[0][0] is x - assert norm_calls[0][1] is norm.weight - assert norm_calls[0][2:] == (1e-5, True) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_triton_glm_vision_norms_match_reference(): - from lightllm.models.vit.triton_kernel.rms_norm_vit import qk_rms_norm, rms_norm - - def glm_rms_norm(x, weight, eps): - normalized = x.float() * torch.rsqrt(x.float().square().mean(-1, keepdim=True) + eps) - return normalized.to(x.dtype) * weight - - torch.manual_seed(53) - eps = 1e-5 - qkv = torch.randn(7, 3, 16, 64, dtype=torch.bfloat16, device="cuda") - q_weight = torch.randn(64, dtype=torch.bfloat16, device="cuda") - k_weight = torch.randn(64, dtype=torch.bfloat16, device="cuda") - q, k = qk_rms_norm(qkv, q_weight, k_weight, eps) - - torch.testing.assert_close(q, glm_rms_norm(qkv[:, 0], q_weight, eps), atol=1e-2, rtol=1e-2) - torch.testing.assert_close(k, glm_rms_norm(qkv[:, 1], k_weight, eps), atol=1e-2, rtol=1e-2) - - x = torch.randn(7, 1024, dtype=torch.bfloat16, device="cuda") - weight = torch.randn(1024, dtype=torch.bfloat16, device="cuda") - torch.testing.assert_close( - rms_norm(x, weight, eps, round_norm_before_weight=True), - glm_rms_norm(x, weight, eps), - atol=1e-2, - rtol=1e-2, - ) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_vision_batch_preserves_independent_image_attention(): - from lightllm.models.glm5_next.glm5_next_visual import Glm5NextVisionTransformer - from lightllm.server.visualserver import set_vit_att_backend - - torch.manual_seed(53) - set_vit_att_backend("sdpa") - model = ( - Glm5NextVisionTransformer( - {"data_type": "float32"}, - hidden_size=128, - out_hidden_size=256, - depth=2, - intermediate_size=256, - projection_intermediate_size=512, - num_heads=2, - ) - .cuda() - .eval() - ) - grid = torch.tensor([[1, 4, 6], [1, 6, 4]]) - x = torch.randn(48, 3 * 2 * 14 ** 2, device="cuda") - with torch.inference_mode(): - batch = model(x, grid) - individual = torch.cat([model(x[:24], grid[:1]), model(x[24:], grid[1:])]) - assert batch.shape == (12, 256) - torch.testing.assert_close(batch, individual, atol=2e-5, rtol=2e-5) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_target_and_mtp_prefill_embed_images_before_hidden_fusion(monkeypatch): - from lightllm.models.glm5_next.layer_infer.pre_layer_infer import Glm5NextPreLayerInfer - from lightllm.models.glm5_next_mtp.layer_infer.pre_layer_infer import Glm5NextMTPPreLayerInfer - from lightllm.server.router.model_infer.infer_batch import g_infer_context - - set_env_start_args(dataclasses.asdict(StartArgs(mtp_mode="eagle_with_att", mtp_step=2))) - monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_DP", "0") - monkeypatch.setenv("LIGHTLLM_DP_WORLD_SIZE", "1") - hidden = 128 - config = {"hidden_size": hidden, "rms_norm_eps": 1e-5, "hc_mult": 4} - image_cache = torch.randn(3, 1, hidden, pin_memory=True) - monkeypatch.setattr(g_infer_context, "cpu_embed_cache_client", SimpleNamespace(cpu_embed_cache_tensor=image_cache)) - text = torch.randn(8, hidden, device="cuda") - ids = torch.tensor([1, 1000, 1001, 2], device="cuda") - metadata = {"token_id": 1000, "token_num": 2, "start_index_in_embed_cache": 1} - state = SimpleNamespace(multimodal_params=[{"images": [metadata], "audios": []}]) - weight = SimpleNamespace(wte_weight_=SimpleNamespace(weight=text, tp_vocab_start_id=0, tp_vocab_end_id=8)) - expected = torch.cat((text[1:2], image_cache[1:, 0].cuda(), text[2:3])) - expanded = Glm5NextPreLayerInfer(config).context_forward(ids, state, weight) - torch.testing.assert_close(expanded.reshape(4, 4, hidden), expected[:, None].expand(-1, 4, -1)) - - def norm(input, eps, out): - out.copy_(input * torch.rsqrt(input.square().mean(-1, keepdim=True) + eps)) - return out - - state.mtp_draft_input_hiddens = torch.randn(4, hidden, device="cuda") - main_norm_weight = torch.randn(hidden, device="cuda") - old_hidden = state.mtp_draft_input_hiddens.clone() - projection = torch.randn(2 * hidden, hidden, device="cuda") * 0.1 - weight.enorm_weight_ = weight.hnorm_weight_ = norm - weight.main_norm_weight_ = lambda input, eps, out: out.copy_( - torch.nn.functional.rms_norm(input, (hidden,), main_norm_weight, eps) - ) - weight.eh_proj_weight_ = SimpleNamespace(mm=lambda x: x @ projection) - actual = Glm5NextMTPPreLayerInfer(config).context_forward(ids, state, weight) - old_hidden = torch.nn.functional.rms_norm(old_hidden, (hidden,), main_norm_weight, 1e-5) - normalized = [x * torch.rsqrt(x.square().mean(-1, keepdim=True) + 1e-5) for x in (expected, old_hidden)] - torch.testing.assert_close(actual, torch.cat(normalized, -1) @ projection) diff --git a/unit_tests/models/phi3/test_rotary_emb.py b/unit_tests/models/phi3/test_rotary_emb.py deleted file mode 100644 index 3127f110ad..0000000000 --- a/unit_tests/models/phi3/test_rotary_emb.py +++ /dev/null @@ -1,50 +0,0 @@ -import pytest -import torch - -from lightllm.models.phi3.triton_kernel.rotary_emb import rotary_emb_fwd - - -def ref_rotary(x, cos, sin, head_dim): - rot = head_dim // 2 - x = x.float() - x1, x2 = x[..., :rot], x[..., rot:head_dim] - c, s = cos.float()[:, None, :], sin.float()[:, None, :] - out = x.clone() - out[..., :rot] = x1 * c - x2 * s - out[..., rot:head_dim] = x1 * s + x2 * c - return out - - -@pytest.mark.parametrize( - "shape", - [ - # (total_len, q_heads, k_heads, head_dim); 96 is Phi-3-mini's head_dim, - # whose rot_dim 48 is not a power of two, so BLOCK_DMODEL has padding - # lanes that every load must mask off - (33, 32, 32, 96), - (128, 32, 8, 80), - (64, 16, 16, 128), - ], -) -@pytest.mark.parametrize("partial_rotary_factor", [1.0, 0.5]) -def test_rotary_matches_reference(shape, partial_rotary_factor): - torch.manual_seed(0) - device = "cuda" - T, HQ, HK, D = shape - rot_dim = int(D * partial_rotary_factor) // 2 - - q = torch.randn(T, HQ, D, device=device, dtype=torch.float16) - k = torch.randn(T, HK, D, device=device, dtype=torch.float16) - pos = torch.arange(T, device=device, dtype=torch.float32) - inv = 1.0 / (10000 ** (torch.arange(0, rot_dim, device=device, dtype=torch.float32) / rot_dim)) - angles = pos[:, None] * inv[None, :] - cos = angles.cos().to(torch.float16).contiguous() - sin = angles.sin().to(torch.float16).contiguous() - - q_ref = ref_rotary(q, cos, sin, int(D * partial_rotary_factor)) - k_ref = ref_rotary(k, cos, sin, int(D * partial_rotary_factor)) - - rotary_emb_fwd(q, k, cos, sin, partial_rotary_factor=partial_rotary_factor) - - assert torch.allclose(q.float(), q_ref, atol=5e-3, rtol=1e-2) - assert torch.allclose(k.float(), k_ref, atol=5e-3, rtol=1e-2) diff --git a/unit_tests/models/qwen2_vl/test_get_mrope_position_ids.py b/unit_tests/models/qwen2_vl/test_get_mrope_position_ids.py deleted file mode 100644 index 7226b05239..0000000000 --- a/unit_tests/models/qwen2_vl/test_get_mrope_position_ids.py +++ /dev/null @@ -1,48 +0,0 @@ -import torch -import pytest - -from lightllm.models.qwen2_vl.triton_kernel.get_mrope_position_ids import get_mrope_position_triton - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="需要 CUDA 支持") -def test_get_mrope_position_triton(): - """ - 测试 get_mrope_position_triton 函数的正确性 - """ - b_image_start_idx = torch.tensor([0, 0, 4], dtype=torch.int32, device="cuda") - b_image_thwd = torch.tensor([[1, 2, 2, -2], [1, 2, 2, -2], [1, 2, 2, -2]], dtype=torch.int32, device="cuda") - b_image_nums = torch.tensor([1, 2], dtype=torch.int32, device="cuda") - b_image_start_num = torch.tensor([0, 1], dtype=torch.int32, device="cuda") - b_image_len = torch.tensor([4, 4, 4], dtype=torch.int32, device="cuda") - position_ids = ( - torch.tensor([0, 1, 2, 3, 4, 5, 6, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12], dtype=torch.int32, device="cuda") - .unsqueeze(0) - .expand(3, -1) - .contiguous() - ) - b_ready_cache_len = torch.tensor([0, 0], dtype=torch.int32, device="cuda") - b_q_seq_len = torch.tensor([7, 13], dtype=torch.int32, device="cuda") - b_start_loc = torch.tensor([0, 7], dtype=torch.int32, device="cuda") - get_mrope_position_triton( - b_image_start_idx, - b_image_thwd, - b_image_nums, - b_image_start_num, - b_image_len, - position_ids, - b_ready_cache_len, - b_q_seq_len, - b_start_loc, - ) - # 预期的输出结果 - expected_output = torch.tensor( - [ - [0, 0, 0, 0, 2, 3, 4, 0, 0, 0, 0, 2, 2, 2, 2, 4, 5, 6, 7, 8], - [0, 0, 1, 1, 2, 3, 4, 0, 0, 1, 1, 2, 2, 3, 3, 4, 5, 6, 7, 8], - [0, 1, 0, 1, 2, 3, 4, 0, 1, 0, 1, 2, 3, 2, 3, 4, 5, 6, 7, 8], - ], - dtype=torch.int32, - device="cuda", - ) - - assert torch.equal(position_ids, expected_output), "position_ids 输出与预期结果不一致" diff --git a/unit_tests/models/qwen2_vl/test_infer_struct.py b/unit_tests/models/qwen2_vl/test_infer_struct.py deleted file mode 100644 index b239e64067..0000000000 --- a/unit_tests/models/qwen2_vl/test_infer_struct.py +++ /dev/null @@ -1,68 +0,0 @@ -from types import SimpleNamespace - -import pytest -import torch - -from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.models.qwen2_vl.infer_struct import Qwen2VLInferStateInfo - - -def _patch_base_position_ids(monkeypatch): - def init_positions(self, model): - self.position_ids = torch.tensor([9, 19], dtype=torch.int32) - self.b_q_seq_len = torch.ones(2, dtype=torch.int32) - - monkeypatch.setattr(InferStateInfo, "init_some_extra_state", init_positions) - - -def _make_model(): - return SimpleNamespace( - config={"rope_scaling": {}}, - _cos_cached=torch.arange(32, dtype=torch.float32).view(32, 1), - _sin_cached=torch.arange(32, dtype=torch.float32).view(32, 1), - ) - - -def test_normal_prompt_prefill_builds_multimodal_positions(monkeypatch): - _patch_base_position_ids(monkeypatch) - expected_position_ids = torch.tensor([[1, 2], [3, 4], [5, 6]], dtype=torch.int32) - monkeypatch.setattr( - Qwen2VLInferStateInfo, - "get_mrope_position", - lambda self, multimodal_params: expected_position_ids, - ) - - infer_state = Qwen2VLInferStateInfo() - infer_state.is_prefill = True - infer_state.b_position_delta = None - infer_state.multimodal_params = [{"images": [], "audios": []}] * 2 - - infer_state.init_some_extra_state(_make_model()) - - assert torch.equal(infer_state.position_ids, expected_position_ids) - - -def test_normal_decode_applies_position_delta(monkeypatch): - _patch_base_position_ids(monkeypatch) - - infer_state = Qwen2VLInferStateInfo() - infer_state.is_prefill = False - infer_state.b_position_delta = torch.tensor([3, 5], dtype=torch.int32) - infer_state.multimodal_params = [{"images": [], "audios": []}] * 2 - - infer_state.init_some_extra_state(_make_model()) - - expected_position_ids = torch.tensor([[12, 24]] * 3, dtype=torch.int32) - assert torch.equal(infer_state.position_ids, expected_position_ids) - - -def test_prefill_rejects_position_delta(monkeypatch): - _patch_base_position_ids(monkeypatch) - - infer_state = Qwen2VLInferStateInfo() - infer_state.is_prefill = True - infer_state.b_position_delta = torch.tensor([3, 5], dtype=torch.int32) - infer_state.multimodal_params = [{"images": [], "audios": []}] * 2 - - with pytest.raises(AssertionError, match="prefill must not provide b_position_delta"): - infer_state.init_some_extra_state(_make_model()) diff --git a/unit_tests/models/qwen2_vl/test_mrope.py b/unit_tests/models/qwen2_vl/test_mrope.py deleted file mode 100644 index e98989179a..0000000000 --- a/unit_tests/models/qwen2_vl/test_mrope.py +++ /dev/null @@ -1,66 +0,0 @@ -import torch -import pytest - -# Import the Triton kernel function under test. Adjust the import path as needed. -from lightllm.models.qwen2_vl.triton_kernel.mrope import mrope_triton_fused - -# Reference Python implementation for multimodal rotary positional embeddings - - -def rotate_half(x): - x1 = x[..., : x.shape[-1] // 2] - x2 = x[..., x.shape[-1] // 2 :] - return torch.cat((-x2, x1), dim=-1) - - -def apply_multimodal_rotary_pos_emb(q, k, cos, sin, mrope_section, unsqueeze_dim=1): - mrope_section = mrope_section * 2 - cos = torch.cat([m[i % 3] for i, m in enumerate(cos.split(mrope_section, dim=-1))], dim=-1).unsqueeze(unsqueeze_dim) - sin = torch.cat([m[i % 3] for i, m in enumerate(sin.split(mrope_section, dim=-1))], dim=-1).unsqueeze(unsqueeze_dim) - - q_embed = (q * cos) + (rotate_half(q) * sin) - k_embed = (k * cos) + (rotate_half(k) * sin) - - return q_embed, k_embed - - -@pytest.mark.parametrize( - "B,H_q,H_k,L,D,mrope_section", - [ - (1, 4, 2, 32, 16, [4, 2, 2]), - ], -) -def test_mrope_triton_correctness(B, H_q, H_k, L, D, mrope_section): - """ - Test that the Triton kernel matches the reference PyTorch implementation. - """ - torch.manual_seed(0) - device = "cuda" - HALF = D // 2 - - q = torch.rand((B, H_q, L, D), dtype=torch.float32, device=device) - k = torch.rand((B, H_k, L, D), dtype=torch.float32, device=device) - - cos_half = torch.rand((1, L, HALF), dtype=torch.float32, device=device).expand(3, -1, -1) - sin_half = torch.rand((1, L, HALF), dtype=torch.float32, device=device).expand(3, -1, -1) - - cos_full = torch.cat([cos_half, cos_half], dim=-1) - sin_full = torch.cat([sin_half, sin_half], dim=-1) - - cos_ref = cos_full.unsqueeze(1) - sin_ref = sin_full.unsqueeze(1) - - ref_q, ref_k = apply_multimodal_rotary_pos_emb(q, k, cos_ref, sin_ref, mrope_section, unsqueeze_dim=1) - - q = q.transpose(1, 2).contiguous().view(L, H_q, D) - k = k.transpose(1, 2).contiguous().view(L, H_k, D) - mrope_section = torch.tensor(mrope_section, dtype=torch.int32, device="cuda") - mrope_triton_fused(q, k, cos_half, sin_half, mrope_section, is_interleaved=False) - q = q.transpose(0, 1).contiguous().view(B, H_q, L, D) - k = k.transpose(0, 1).contiguous().view(B, H_k, L, D) - assert torch.allclose(q, ref_q, rtol=1e-3, atol=1e-3) - assert torch.allclose(k, ref_k, rtol=1e-3, atol=1e-3) - - -if __name__ == "__main__": - pytest.main() diff --git a/unit_tests/models/qwen2_vl/test_rotary_pos_emb.py b/unit_tests/models/qwen2_vl/test_rotary_pos_emb.py deleted file mode 100644 index 34e41578a6..0000000000 --- a/unit_tests/models/qwen2_vl/test_rotary_pos_emb.py +++ /dev/null @@ -1,55 +0,0 @@ -import math -import torch -import pytest - -from lightllm.models.qwen2_vl.triton_kernel.rotary_pos_emb import apply_rotary_pos_emb_triton - - -def rotate_half(x): - """Rotates half the hidden dims of the input.""" - x1 = x[..., : x.shape[-1] // 2] - x2 = x[..., x.shape[-1] // 2 :] - return torch.cat((-x2, x1), dim=-1) - - -def apply_rotary_pos_emb_vision(tensor: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: - orig_dtype = tensor.dtype - tensor = tensor.float() - cos = cos.unsqueeze(1).repeat(1, 1, 2).unsqueeze(0).float() - sin = sin.unsqueeze(1).repeat(1, 1, 2).unsqueeze(0).float() - output = (tensor * cos) + (rotate_half(tensor) * sin) - output = output.to(orig_dtype) - return output - - -@pytest.mark.parametrize( - "shape", - [ - (1296, 64, 80), - (1024, 2, 192), - (1024, 1, 256), - (1024, 3, 160), - ], -) -def test_triton_matches_reference(shape): - L, H, D = shape - assert D % 2 == 0 - - torch.manual_seed(0) - - freqs = torch.randn(L, D // 2, device="cuda", dtype=torch.bfloat16) - cos = freqs.cos() - sin = freqs.sin() - - tensor = torch.randn(L, H, D, device="cuda", dtype=torch.bfloat16) - - ref = apply_rotary_pos_emb_vision(tensor.unsqueeze(0), cos, sin).squeeze(0) - out = apply_rotary_pos_emb_triton(tensor, cos, sin) - - assert out.dtype == tensor.dtype, "输出 dtype 应与输入一致" - assert out.shape == tensor.shape, "输出形状应与输入一致" - assert torch.allclose(out, ref, rtol=1e-2, atol=1e-2), "Triton 与参考实现不一致" - - -if __name__ == "__main__": - pytest.main() diff --git a/unit_tests/models/qwen3-vl/test_deepstack_emb.py b/unit_tests/models/qwen3-vl/test_deepstack_emb.py deleted file mode 100644 index f629a16352..0000000000 --- a/unit_tests/models/qwen3-vl/test_deepstack_emb.py +++ /dev/null @@ -1,82 +0,0 @@ -import torch -import triton -import triton.language as tl -from lightllm.models.qwen3_vl.infer_struct import Qwen3VLInferStateInfo -from lightllm.models.qwen3_vl.triton_kernel.deepstack_multimodal_emb import add_deepstack_embs - - -def test_deepstack_same_image_twice(): - device = "cuda" - - # 1. 构造 input_ids，包含两段相同的 image token 范围 [100, 101, 102] - input_ids = torch.tensor( - [1, 100, 101, 102, 2, 100, 101, 102, 3], - device=device, - dtype=torch.long, - ) - seq_len = input_ids.shape[0] - - hidden_size = 4 - token_len = 3 # 每张图 3 个 token - - # 2. 构造初始 embedding，全 0，方便看增量 - input_embeddings = torch.zeros(seq_len, hidden_size, device=device, dtype=torch.float32) - - # 3. 构造 deepstack_embs（这一层的 deepstack） - # 只有一张图片，所以 deepstack_embs 形状是 [token_len, hidden_size] - # 每一行是 [1,1,1,1], [2,2,2,2], [3,3,3,3] - deepstack_embs = torch.tensor( - [ - [1.0, 1.0, 1.0, 1.0], # 对应 token_id = 100 - [2.0, 2.0, 2.0, 2.0], # 对应 token_id = 101 - [3.0, 3.0, 3.0, 3.0], # 对应 token_id = 102 - ], - device=device, - dtype=torch.float32, - ) - - # 4. image 相关索引信息（与 multimodal_emb 一致的语义） - img_start_token_ids = torch.tensor([100], device=device, dtype=torch.long) # 只有一个 image handle，从 100 开始 - img_token_lens = torch.tensor([token_len], device=device, dtype=torch.long) - img_start_locs = torch.tensor([0], device=device, dtype=torch.long) # deepstack_embs 从第 0 行开始是这张图的 - - # 5. 保存一份原始 embedding，方便求差 - before = input_embeddings.clone() - - # 6. 调用 Triton 算子 - add_deepstack_embs( - out=input_embeddings, - input_ids=input_ids, - deepstack_embs=deepstack_embs, - img_token_lens=img_token_lens, - img_start_token_ids=img_start_token_ids, - img_start_locs_in_cache=img_start_locs, - ) - - # 7. 看看相同图片两段上的增量 - delta = input_embeddings - before - - print("input_ids:", input_ids) - print("delta:\n", delta) - - # 第一次 image：位置 1,2,3 - print("first image span delta:\n", delta[1:4]) - # 第二次 image：位置 5,6,7 - print("second image span delta:\n", delta[5:8]) - - # 8. 断言它们和预期一致 - expected = deepstack_embs # [3, 4] - - assert torch.allclose(delta[1:4], expected), "first image span does not match expected deepstack" - assert torch.allclose(delta[5:8], expected), "second image span does not match expected deepstack" - - # 其他位置应该仍然是 0 - assert torch.all(delta[0] == 0) - assert torch.all(delta[4] == 0) - assert torch.all(delta[8] == 0) - - print("OK: same image appears twice, both spans get deepstack added correctly.") - - -if __name__ == "__main__": - test_deepstack_same_image_twice() diff --git a/unit_tests/models/qwen3next/test_shared_expert_gate.py b/unit_tests/models/qwen3next/test_shared_expert_gate.py deleted file mode 100644 index 44d46e4471..0000000000 --- a/unit_tests/models/qwen3next/test_shared_expert_gate.py +++ /dev/null @@ -1,22 +0,0 @@ -import pytest -import torch - -from lightllm.models.qwen3next.triton_kernel.shared_expert_gate import sigmoid_mul_ - - -@pytest.mark.parametrize("gate_width", [1, 6144]) -def test_sigmoid_mul_matches_unfused_dtype_rounding(gate_width): - if not torch.cuda.is_available(): - pytest.skip("CUDA is required") - - torch.manual_seed(2026) - x = torch.randn((17, 6144), device="cuda", dtype=torch.bfloat16) * 2 - gate = torch.randn((17, gate_width), device="cuda", dtype=torch.bfloat16) * 3 - - expected = x.clone() - expected.mul_(gate.clone().sigmoid_()) - - actual = x.clone() - sigmoid_mul_(actual, gate) - - torch.testing.assert_close(actual, expected, rtol=0, atol=0) diff --git a/unit_tests/models/sensenova_loop/test_gdn_state.py b/unit_tests/models/sensenova_loop/test_gdn_state.py deleted file mode 100644 index c25576a19a..0000000000 --- a/unit_tests/models/sensenova_loop/test_gdn_state.py +++ /dev/null @@ -1,151 +0,0 @@ -from types import SimpleNamespace - -import pytest -import torch - -from lightllm.common.attention_state import AttentionStateManager -from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState -from lightllm.common.basemodel.attention.base_att import AttControl -from lightllm.common.basemodel.attention.linear.triton import TritonLinearAttBackend -from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.models.sensenova_loop.config import get_loop_layer_indices, get_loop_layer_types -from lightllm.server.core.objs.start_args_type import StartArgs - - -@pytest.mark.parametrize("device", ["cpu", "cuda"]) -@pytest.mark.parametrize("mtp_step", [0, 2]) -@pytest.mark.parametrize( - "start,end,count,logical_layer,storage_layer", - [(4, 28, 1, 4, 3), (4, 28, 3, 28, 21), (1, 6, 2, 13, 10)], - ids=["no-loop", "aligned-loop", "unaligned-loop"], -) -def test_gdn_decode_keeps_logical_layer_state( - monkeypatch, device, mtp_step, start, end, count, logical_layer, storage_layer -): - if device == "cuda" and not torch.cuda.is_available(): - pytest.skip("CUDA required") - torch.manual_seed(1585) - settings = dict(start=start, end=end, eval_count=count, use_gate=False) - config = dict( - num_hidden_layers=32, - full_attention_interval=4, - layer_types=["full_attention" if (i + 1) % 4 == 0 else "linear_attention" for i in range(32)], - linear_num_key_heads=1, - linear_num_value_heads=2, - linear_key_head_dim=32, - linear_value_head_dim=32, - linear_conv_kernel_dim=4, - hidden_act="silu", - ) - indices = get_loop_layer_indices(settings, 32) - kinds = get_loop_layer_types(settings, config) - layers = tuple(i for i, kind in enumerate(kinds) if kind == "linear_attention") - args = StartArgs(mtp_step=mtp_step, linear_att_ssm_data_type="float32") - monkeypatch.setattr("lightllm.common.basemodel.attention.linear.gdn.get_env_start_args", lambda: args) - monkeypatch.setattr("lightllm.common.basemodel.attention.base_att.get_env_start_args", lambda: args) - backend = object.__new__(TritonLinearAttBackend) - backend.__init__( - SimpleNamespace( - config=config, - tp_world_size_=1, - is_mtp_draft_model=False, - mtp_manager=SimpleNamespace(get_decode_draft_step=lambda _: mtp_step), - ) - ) - dtype = torch.bfloat16 if device == "cuda" else torch.float32 - layout = ConvSSMStateConfig(layers, 1, 2, 32, 32, 3, dtype, torch.float32) - resource = GDNState(layout, request_capacity=3, mtp_step=mtp_step, device=device) - resources = AttentionStateManager(request_capacity=3) - resources.register_state("main", resource) - resources.finish_state_registration() - conv_view, ssm_view = resource.get_layer_view(logical_layer_id=logical_layer) - assert conv_view.data_ptr() == resource.conv_state[storage_layer].data_ptr() - assert ssm_view.data_ptr() == resource.ssm_state[storage_layer].data_ptr() - physical_weights = [ - SimpleNamespace( - layer_num_=i, - linear_conv1d=SimpleNamespace( - mm_param=SimpleNamespace(weight=torch.randn(layout.conv_channels, 4, device=device, dtype=dtype) * 0.1), - bias=None, - ), - linear_A_log=SimpleNamespace(weight=torch.zeros(2, device=device)), - linear_dt_bias=SimpleNamespace(weight=torch.zeros(2, device=device)), - ) - for i in range(32) - ] - execution_weights = [physical_weights[i] for i in indices] - weight = execution_weights[logical_layer] - assert weight is physical_weights[indices[logical_layer]] - - if device == "cpu": - - def conv_pack(mixed, z, a, b, conv, *args): - conv[1].add_(1) - return (*backend._rearrange_mixed_qkv(mixed, decode=True), z, a, b) - - def conv_mtp(mixed, conv, *args, **kwargs): - conv[1].add_(1) - return mixed - - def recurrent(*, initial_state, v, **kwargs): - initial_state[mtp_step + 1 : 2 * (mtp_step + 1)].add_(1) - return v, None - - monkeypatch.setattr("lightllm.common.basemodel.attention.linear.gdn.conv_pack_gdn_decode_inputs", conv_pack) - monkeypatch.setattr( - "lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_mtp.causal_conv1d_update", conv_mtp - ) - monkeypatch.setattr( - "lightllm.common.basemodel.attention.linear.gdn.fused_recurrent_gated_delta_rule", recurrent - ) - monkeypatch.setattr( - "lightllm.common.basemodel.attention.linear.gdn.mtp_fused_recurrent_gated_delta_rule", - recurrent, - ) - - infer = InferStateInfo() - infer.att_state_manager, infer.att_state_role = resources, "main" - infer.b_req_idx = torch.tensor([1], device=device, dtype=torch.int32) - control = dict(layer_num=logical_layer, layer_weight=weight) - mixed_dim = backend.tp_qkvz_dim + backend.tp_ba_dim - if device == "cuda": - infer.b1_cu_q_seq_len = torch.tensor([0, 5], device=device, dtype=torch.int32) - infer.b_ready_cache_len = torch.zeros(1, device=device, dtype=torch.int32) - prefill = backend.create_att_prefill_state(infer) - prefill.init_state() - prefill.prefill_att( - None, - None, - None, - AttControl( - linear_att_prefill=True, - linear_att_prefill_dict=dict( - **control, mixed_qkvzba=torch.randn(5, mixed_dim, device=device, dtype=dtype) - ), - ), - ) - - infer.batch_size = mtp_step + 1 - infer.b_req_idx = infer.b_req_idx.repeat(infer.batch_size) - decode = backend.create_att_decode_state(infer) - untouched = [i for i in range(len(layers)) if i != storage_layer] - for _ in range(2): - before_conv, before_ssm = resource.conv_state.clone(), resource.ssm_state.clone() - decode.init_state() - output, _ = decode.decode_att( - None, - None, - None, - AttControl( - linear_att_decode=True, - linear_att_decode_dict=dict( - **control, mixed_qkvzba=torch.randn(infer.batch_size, mixed_dim, device=device, dtype=dtype) - ), - ), - ) - assert torch.isfinite(output).all() - assert not torch.equal(resource.conv_state[storage_layer], before_conv[storage_layer]) - assert not torch.equal(resource.ssm_state[storage_layer], before_ssm[storage_layer]) - assert torch.equal(resource.conv_state[untouched], before_conv[untouched]) - assert torch.equal(resource.ssm_state[untouched], before_ssm[untouched]) - assert weight.layer_num_ == indices[logical_layer] diff --git a/unit_tests/models/sensenova_loop/test_loop.py b/unit_tests/models/sensenova_loop/test_loop.py deleted file mode 100644 index 7caaa09f3f..0000000000 --- a/unit_tests/models/sensenova_loop/test_loop.py +++ /dev/null @@ -1,376 +0,0 @@ -import dataclasses -import json -from copy import deepcopy -from types import SimpleNamespace - -import pytest -import torch - -from lightllm.models.registry import ModelRegistry -from lightllm.models.sensenova_loop.config import get_loop_layer_indices, get_loop_layer_types -from lightllm.models.qwen3next.config import Qwen3NextResourceConfig -from lightllm.server.core.objs.start_args_type import StartArgs - - -def _config(layers=32, start=4, end=28, count=3): - return { - "model_type": "sensenova_loop", - "loop_settings": {"start": start, "end": end, "eval_count": count, "use_gate": False}, - "text_config": { - "num_hidden_layers": layers, - "full_attention_interval": 4, - "layer_types": ["full_attention" if (i + 1) % 4 == 0 else "linear_attention" for i in range(layers)], - "num_key_value_heads": 4, - "head_dim": 256, - "linear_num_key_heads": 16, - "linear_num_value_heads": 32, - "linear_key_head_dim": 128, - "linear_value_head_dim": 128, - "linear_conv_kernel_dim": 4, - }, - } - - -@pytest.mark.parametrize("start,end,count", [(4, 28, 3), (4, 28, 2), (1, 6, 2)]) -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.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() - 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) - full = linear = 0 - slots = {} - for execution, (physical, kind) in enumerate(zip(indices, kinds)): - if kind == "full_attention": - slot = full_layers.index(execution) - assert slot == full - full += 1 - else: - slot = linear_layers.index(execution) - assert slot == linear - linear += 1 - assert slot not in slots.setdefault(physical, set()) - slots[physical].add(slot) - 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): - assert (cache.all_layer_num, full, linear) == (56, 14, 42) - - -def test_physical_weights_are_created_and_loaded_once(monkeypatch): - from lightllm.models.sensenova_loop.model import SenseNovaLoopTpPartModel - from lightllm.models.qwen3_5.model import Qwen3_5TpPartModel - - model = object.__new__(SenseNovaLoopTpPartModel) - model.loop_layer_indices = get_loop_layer_indices(_config()["loop_settings"], 32) - model.num_physical_layers = 32 - model.config = {"n_layer": 80} - model.data_type, model.quant_cfg = torch.bfloat16, None - loaded = [] - - def create_weight(index, dtype, network_config, quant_cfg): - assert network_config is model.config - return object() - - def load_weights(self, weights): - loaded.extend(self.trans_layers_weight) - - model.pre_and_post_weight_class = lambda dtype, network_config: object() - model.transformer_weight_class = create_weight - monkeypatch.setattr(Qwen3_5TpPartModel, "load_weights", load_weights) - model._init_weights() - model.load_weights({}) - assert len(loaded) == len(set(loaded)) == 32 - assert len(model.trans_layers_weight) == 80 - for execution, physical in enumerate(model.loop_layer_indices): - assert model.trans_layers_weight[execution] is loaded[physical] - assert model.config["n_layer"] == 80 - - -@pytest.mark.parametrize("streams,count", [(2, 2), (4, 3)]) -def test_mhc_registry_and_embedded_mtp_reuse_existing_qwen_draft(tmp_path, streams, count): - from lightllm.models.draft_registry import get_draft_model_class - from lightllm.models.sensenova_loop.model import SenseNovaLoopTpPartModel - from lightllm.models.qwen3_5_mtp.model import Qwen3_5MTPModel - - config = _config(count=count) - config["mhc"] = dict(start=4, end=28, streams=streams, iterations=20) - config["text_config"].update(num_attention_heads=16, hidden_size=4096, vocab_size=128) - config["text_config"]["layer_types"] = ["linear_attention"] * 3 + ["full_attention"] - config["text_config"]["layer_types"] *= 8 - (tmp_path / "config.json").write_text(json.dumps(config)) - registration = ModelRegistry.get_model_config(config) - assert registration.model_class.endswith(":SenseNovaLoopTpPartModel") - assert registration.is_multimodal - assert get_draft_model_class(config, "eagle_with_att") is Qwen3_5MTPModel - - target = object.__new__(SenseNovaLoopTpPartModel) - target.weight_dir_, target.finetune_config, target.tp_world_size_ = str(tmp_path), None, 4 - target.args = SimpleNamespace(**dataclasses.asdict(StartArgs())) - target._init_config() - assert target.config["n_layer"] == 32 + 24 * (count - 1) - assert target.num_physical_layers == 32 - - draft = object.__new__(Qwen3_5MTPModel) - 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"] == ["full_attention"] - - -def test_loop_tokenizer_loads_without_hf_loop_autoconfig(tmp_path): - from tokenizers import Tokenizer, models - from transformers import PreTrainedTokenizerFast, Qwen2VLImageProcessorFast - from lightllm.server.tokenizer import get_tokenizer - - tokenizer = PreTrainedTokenizerFast( - tokenizer_object=Tokenizer(models.WordLevel({"hello": 0, "[UNK]": 1}, unk_token="[UNK]")), - unk_token="[UNK]", - ) - tokenizer.save_pretrained(tmp_path) - Qwen2VLImageProcessorFast().save_pretrained(tmp_path) - config = _config() - config.update( - architectures=["SenseNovaLoopForConditionalGeneration"], - vision_config={}, - vision_start_token_id=2, - vision_end_token_id=3, - image_token_id=4, - ) - (tmp_path / "config.json").write_text(json.dumps(config)) - wrapped = get_tokenizer(str(tmp_path)) - assert wrapped.encode("hello") == [0] - assert isinstance(wrapped.image_processor, Qwen2VLImageProcessorFast) - - -@pytest.mark.parametrize("custom_layout", [False, True]) -def test_model_and_worker_use_the_same_loop_cache_layout(monkeypatch, custom_layout): - from lightllm.models.sensenova_loop.model import SenseNovaLoopTpPartModel - - config = _config(start=1, end=6, count=2) - 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", 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, {})) - 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, - ) - (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 - - -def test_plain_loop_registry_does_not_enable_mhc(): - registration = ModelRegistry.get_model_config(_config()) - assert registration.model_class == "lightllm.models.sensenova_loop.model:SenseNovaLoopTpPartModel" - assert registration.is_multimodal - - -def test_awq_config_applies_to_main_model_and_excludes_float_mtp(): - from lightllm.models.sensenova_loop.model import SenseNovaLoopTpPartModel - from lightllm.models.qwen3_5_mtp.model import Qwen3_5MTPModel - - config = _config() - quantization = {"quant_method": "awq", "bits": 4, "group_size": 128, "modules_to_not_convert": ["mtp"]} - config["quantization_config"] = quantization - original = deepcopy(config) - assert SenseNovaLoopTpPartModel.normalize_model_config(config)["quantization_config"] == quantization - assert "quantization_config" not in Qwen3_5MTPModel.normalize_model_config(config) - assert config == original - - config["text_config"]["quantization_config"] = None - assert SenseNovaLoopTpPartModel.normalize_model_config(config)["quantization_config"] is None - - -@pytest.mark.parametrize("suffix,factor", [("qweight", 8), ("qzeros", 8), ("scales", 1)]) -def test_awq_fused_projections_split_on_output_axis(suffix, factor): - from lightllm.models.qwen3_5.layer_weights.transformer_layer_weight import Qwen35TransformerLayerWeight - - method = SimpleNamespace(method_name="awq_marlin", pack_factor=8) - layer = object.__new__(Qwen35TransformerLayerWeight) - layer.layer_num_ = 0 - layer.linear_num_k_heads = 2 - layer.linear_num_v_heads = 4 - layer.linear_k_head_dim = layer.linear_v_head_dim = 16 - layer.linear_in_proj = SimpleNamespace(quant_method=method) - prefix = "model.layers.0.linear_attn.in_proj_" - qkv = torch.arange(256 // factor).reshape(2, 128 // factor) - fused_name = f"{prefix}qkv.{suffix}" - weights = {fused_name: qkv} - layer._split_linear_in_proj_qkv(weights) - q_end, k_end = 32 // factor, 64 // factor - assert fused_name not in weights - assert torch.equal(weights[f"{prefix}q.{suffix}"], qkv[:, :q_end]) - assert torch.equal(weights[f"{prefix}k.{suffix}"], qkv[:, q_end:k_end]) - assert torch.equal(weights[f"{prefix}v.{suffix}"], qkv[:, k_end:]) - - layer.is_linear_attention_layer = False - layer.q_head_num_, layer.head_dim = 2, 16 - layer.qkvo_gate_proj = SimpleNamespace(quant_method=method) - layer._q_weight_name = "model.layers.0.self_attn.q_proj.weight" - layer._o_gate_weight_name = "model.layers.0.self_attn.o_gate_proj.weight" - q_name = f"model.layers.0.self_attn.q_proj.{suffix}" - gate_name = f"model.layers.0.self_attn.o_gate_proj.{suffix}" - head_width = layer.head_dim // factor - q_gate = torch.arange(8 * head_width).reshape(2, 4 * head_width) - q1, gate1, q2, gate2 = q_gate.split(head_width, dim=1) - weights = {q_name: q_gate} - layer._split_q_with_gate(weights) - assert torch.equal(weights[q_name], torch.cat((q1, q2), dim=1)) - assert torch.equal(weights[gate_name], torch.cat((gate1, gate2), dim=1)) - layer._split_q_with_gate(weights) - assert torch.equal(weights[q_name], torch.cat((q1, q2), dim=1)) - assert torch.equal(weights[gate_name], torch.cat((gate1, gate2), dim=1)) - - -def test_loop_binds_split_gdn_projection_only_for_weights_with_separate_gates(monkeypatch): - from lightllm.models.qwen3_5.model import Qwen3_5TpPartModel - from lightllm.models.qwen3_5.layer_infer.transformer_layer_infer import Qwen35TransformerLayerInfer - from lightllm.models.qwen3next.layer_infer.transformer_layer_infer import Qwen3NextTransformerLayerInfer - from lightllm.models.sensenova_loop.model import SenseNovaLoopTpPartModel - from lightllm.models.sensenova_loop.layer_infer.transformer_layer_infer import SenseNovaLoopTransformerLayerInfer - - qkvz = torch.arange(8).reshape(2, 4) - gates = torch.arange(8, 12).reshape(2, 2) - fused = torch.cat((qkvz, gates), dim=-1) - fused_weight = SimpleNamespace(linear_in_proj=SimpleNamespace(mm=lambda _: fused)) - split_weight = SimpleNamespace( - linear_in_proj=SimpleNamespace(mm=lambda _: qkvz), - linear_gate_proj=SimpleNamespace(mm=lambda _: gates), - ) - model = object.__new__(SenseNovaLoopTpPartModel) - model.trans_layers_weight = [fused_weight, split_weight, split_weight] - layers = [object.__new__(SenseNovaLoopTransformerLayerInfer) for _ in model.trans_layers_weight] - for layer in layers: - layer.embed_dim_ = 4 - monkeypatch.setattr(Qwen3_5TpPartModel, "_init_infer_layer", lambda self, _: setattr(self, "layers_infer", layers)) - - model._init_infer_layer() - assert Qwen35TransformerLayerInfer._linear_in_proj is Qwen3NextTransformerLayerInfer._linear_in_proj - assert layers[0]._linear_in_proj.__func__ is Qwen3NextTransformerLayerInfer._linear_in_proj - for layer in layers[1:]: - assert layer._linear_in_proj.__func__ is SenseNovaLoopTransformerLayerInfer._linear_in_proj_with_gates - for layer, weight in zip(layers, model.trans_layers_weight): - assert torch.equal(layer._linear_in_proj(torch.zeros(1, 2, 4), weight), fused) - - -def test_weight_load_failure_restores_execution_weights(monkeypatch): - from lightllm.models.sensenova_loop.model import SenseNovaLoopTpPartModel - from lightllm.models.qwen3_5.model import Qwen3_5TpPartModel - - model = object.__new__(SenseNovaLoopTpPartModel) - model.config = {"n_layer": 80} - model.physical_layer_weights = [object()] - model.trans_layers_weight = [object()] - execution_config, execution_weights = model.config, model.trans_layers_weight - - def fail(*args): - raise RuntimeError("weight load failed") - - monkeypatch.setattr(Qwen3_5TpPartModel, "load_weights", fail) - with pytest.raises(RuntimeError, match="weight load failed"): - model.load_weights({}) - assert model.config is execution_config - assert model.trans_layers_weight is execution_weights - - -@pytest.fixture -def layer_runtime(monkeypatch): - from lightllm.utils.envs_utils import get_env_start_args - - monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(dataclasses.asdict(StartArgs()))) - for key, value in { - "DP_WORLD_SIZE": 1, - "CURRENT_RANK_IN_DP": 0, - "CURRENT_DEVICE_ID": 0, - "CURRENT_RANK_IN_NODE": 0, - }.items(): - monkeypatch.setenv("LIGHTLLM_" + key, str(value)) - get_env_start_args.cache_clear() - yield - get_env_start_args.cache_clear() - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("mhc_enabled", [False, True]) -def test_loop_layer_types_are_local_to_sensenova(layer_runtime, tmp_path, mhc_enabled): - from lightllm.common.quantization import Quantcfg - from lightllm.models.qwen3_5.layer_infer.transformer_layer_infer import Qwen35TransformerLayerInfer - from lightllm.models.qwen3_5.layer_weights.transformer_layer_weight import Qwen35TransformerLayerWeight - from lightllm.models.qwen3_5_mtp.model import Qwen3_5MTPModel - from lightllm.models.sensenova_loop.layer_infer.transformer_layer_infer import SenseNovaLoopTransformerLayerInfer - from lightllm.models.sensenova_loop.layer_weights.transformer_layer_weight import ( - SenseNovaLoopTransformerLayerWeight, - ) - - config = _config(layers=8, start=1, end=6, count=2) - text = config["text_config"] - text.update( - n_layer=8, - hidden_size=128, - intermediate_size=256, - num_attention_heads=1, - num_key_value_heads=1, - head_dim=128, - linear_num_key_heads=1, - linear_num_value_heads=1, - rms_norm_eps=1e-6, - vocab_size=128, - layer_types=["full_attention", "linear_attention"] * 4, - ) - if mhc_enabled: - text["mhc"] = dict(start=1, end=6, streams=2, iterations=20) - indices = get_loop_layer_indices(config["loop_settings"], 8) - execution_config = { - **text, - "loop_physical_indices": indices, - "layer_types": get_loop_layer_types(config["loop_settings"], text), - } - weights = [ - SenseNovaLoopTransformerLayerWeight(i, torch.bfloat16, execution_config, Quantcfg(execution_config)) - for i in range(8) - ] - for execution, physical in enumerate(indices): - layer = SenseNovaLoopTransformerLayerInfer(execution, execution_config) - expected_linear = text["layer_types"][physical] == "linear_attention" - assert layer.is_linear_attention_layer == weights[physical].is_linear_attention_layer == expected_linear - assert hasattr(weights[physical], "linear_in_proj") == expected_linear - assert hasattr(weights[physical], "qkvo_gate_proj") != expected_linear - - # The base model still uses its original periodic layout, even with layer_types present. - base_layer = Qwen35TransformerLayerInfer(0, text) - base_weight = Qwen35TransformerLayerWeight(0, torch.bfloat16, text, Quantcfg(text)) - assert base_layer.is_linear_attention_layer and base_weight.is_linear_attention_layer - - (tmp_path / "config.json").write_text(json.dumps(config)) - draft = object.__new__(Qwen3_5MTPModel) - draft.weight_dir_, draft.finetune_config, draft.tp_world_size_ = str(tmp_path), None, 1 - draft._init_config() - for global_index in (len(indices), len(indices) + 1): - assert not Qwen35TransformerLayerInfer(global_index, draft.config).is_linear_attention_layer diff --git a/unit_tests/models/sensenova_loop/test_loop_mhc.py b/unit_tests/models/sensenova_loop/test_loop_mhc.py deleted file mode 100644 index 34d4fe9696..0000000000 --- a/unit_tests/models/sensenova_loop/test_loop_mhc.py +++ /dev/null @@ -1,107 +0,0 @@ -import pytest -import torch -import torch.nn.functional as F - -from lightllm.models.sensenova_loop.triton_kernel.mhc import hc_contract, hc_expand, hc_post, hc_pre_norm -from lightllm.models.sensenova_loop.triton_kernel.mhc.pre_norm import has_deepgemm_mhc -from lightllm.common.basemodel.triton_kernel.norm.rmsnorm import rmsnorm_forward - - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") - - -def _route_reference(x, weight, bias, streams, iterations): - # Training normalizes the residual before the FP32 router projection. - flat = x.float() - logits = F.linear(flat * torch.rsqrt(flat.square().mean(-1, keepdim=True) + 1e-6), weight, bias) - pre = logits[:, :streams].sigmoid() - post = 2 * logits[:, streams : 2 * streams].sigmoid() - logits = logits[:, 2 * streams :].reshape(-1, streams, streams) - for _ in range(iterations): - logits = logits - torch.logsumexp(logits, dim=-1, keepdim=True) - logits = logits - torch.logsumexp(logits, dim=-2, keepdim=True) - mix = logits.exp() - branch = (pre.to(x.dtype).unsqueeze(1) @ x.view(-1, streams, x.shape[1] // streams)).squeeze(1) - return branch, mix, post - - -@pytest.mark.parametrize("tokens", [1, 19, 257]) -@pytest.mark.parametrize("streams", [2, 4]) -def test_loop_mhc_matches_training_routing_and_rounding(tokens, streams): - torch.manual_seed(35) - hidden = 4096 - x = torch.randn(tokens, streams * hidden, device="cuda", dtype=torch.bfloat16) - weight = torch.randn(2 * streams + streams * streams, streams * hidden, device="cuda") * 0.008 - bias = torch.randn(weight.shape[0], device="cuda") - norm = torch.randn(hidden, device="cuda", dtype=torch.bfloat16) - expected = _route_reference(x, weight, bias, streams, 20) - actual = hc_pre_norm( - x, - weight, - bias, - norm, - streams, - 1e-6, - 1e-6, - 20, - use_deepgemm=False, - ) - torch.testing.assert_close(actual[0], rmsnorm_forward(expected[0], norm, 1e-6), atol=0.02, rtol=0.02) - torch.testing.assert_close(actual[1], expected[1], atol=2e-6, rtol=2e-5) - torch.testing.assert_close(actual[2], expected[2], atol=2e-6, rtol=2e-5) - layer_out = torch.randn(tokens, hidden, device="cuda", dtype=torch.bfloat16) - # Qwen uses output,input order; both matmul and the post product round to bf16. - expected_out = expected[1].to(x.dtype) @ x.view(tokens, streams, hidden) - expected_out += expected[2].to(x.dtype).unsqueeze(-1) * layer_out.unsqueeze(1) - actual_out = hc_post(layer_out, x, actual[1].transpose(1, 2), actual[2], streams) - torch.testing.assert_close(actual_out.view_as(expected_out), expected_out, atol=0.02, rtol=0.01) - - -@pytest.mark.skipif(not has_deepgemm_mhc(), reason="DeepGEMM mHC projection unavailable") -def test_loop_deepgemm_projection_matches_torch(): - x = torch.randn(19, 16384, device="cuda", dtype=torch.bfloat16) - weight = torch.randn(24, 16384, device="cuda") * 0.005 - bias = torch.randn(24, device="cuda") - norm = torch.randn(4096, device="cuda", dtype=torch.bfloat16) - actual = hc_pre_norm(x, weight, bias, norm, 4, 1e-6, 1e-6, 20) - expected = hc_pre_norm(x, weight, bias, norm, 4, 1e-6, 1e-6, 20, use_deepgemm=False) - for a, b in zip(actual, expected): - torch.testing.assert_close(a, b, atol=0.02, rtol=0.02) - - -@pytest.mark.parametrize("streams", [2, 4]) -def test_loop_mhc_decode_can_capture_and_replay_cuda_graph(streams): - x = torch.randn(3, streams * 4096, device="cuda", dtype=torch.bfloat16) - weight = torch.randn(2 * streams + streams * streams, streams * 4096, device="cuda") * 0.005 - bias = torch.randn(weight.shape[0], device="cuda") - norm = torch.ones(4096, device="cuda", dtype=torch.bfloat16) - - def forward(): - branch, mix, post = hc_pre_norm( - x, - weight, - bias, - norm, - streams, - 1e-6, - 1e-6, - 20, - use_deepgemm=False, - ) - return hc_post(branch, x, mix.transpose(1, 2), post, streams) - - stream = torch.cuda.Stream() - stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(stream): - for _ in range(3): - forward() - torch.cuda.current_stream().wait_stream(stream) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - output = forward() - x.normal_() - expected = forward() - graph.replay() - torch.testing.assert_close(output, expected, atol=0, rtol=0) - collapsed = torch.randn(3, 4096, device="cuda", dtype=torch.bfloat16) - assert torch.equal(hc_contract(hc_expand(collapsed, streams), streams), collapsed) diff --git a/unit_tests/models/test_multimodal_tokenizer_positions.py b/unit_tests/models/test_multimodal_tokenizer_positions.py deleted file mode 100644 index 9425e02587..0000000000 --- a/unit_tests/models/test_multimodal_tokenizer_positions.py +++ /dev/null @@ -1,98 +0,0 @@ -from types import SimpleNamespace -from unittest.mock import Mock - -import pytest - -from lightllm.models.gemma3.model import Gemma3Tokenizer -from lightllm.models.internvl.model import InternvlTokenizer -from lightllm.models.llava.model import LlavaTokenizer -from lightllm.models.qwen3_omni_moe_thinker.model import QWen3OmniTokenizer -from lightllm.models.qwen_vl.model import QWenVLTokenizer -from lightllm.models.tarsier2.model import Tarsier2Tokenizer -from lightllm.server.multimodal_params import MultimodalParams -from lightllm.utils.multimodal_chunk_utils import get_mm_input_spans - - -def make_params(images=(), audios=()): - params = MultimodalParams( - images=[{"type": "url", "data": "unused"} for _ in images], - audios=[{"type": "url", "data": "unused"} for _ in audios], - ) - for items, specs in ((params.images, images), (params.audios, audios)): - for item, (token_id, token_num) in zip(items, specs): - item.token_id, item.token_num = token_id, token_num - return params - - -def make_tokenizer(tokenizer_class, origin_ids): - tokenizer = tokenizer_class.__new__(tokenizer_class) - tokenizer.tokenizer = Mock() - tokenizer.tokenizer.encode.return_value = origin_ids - tokenizer.image_start_id, tokenizer.image_end_id, tokenizer.image_token_id = 10, 11, 12 - tokenizer.audio_start_id, tokenizer.audio_end_id, tokenizer.audio_token_id = 20, 21, 22 - return tokenizer - - -@pytest.mark.parametrize( - "tokenizer_class", [LlavaTokenizer, Gemma3Tokenizer, QWenVLTokenizer, Tarsier2Tokenizer, InternvlTokenizer] -) -def test_images_record_positions_even_when_embeddings_are_reused(tokenizer_class): - tokenizer = make_tokenizer(tokenizer_class, [1, 10, 11, 2, 10, 11, 3]) - tokenizer.tokenizer.side_effect = [SimpleNamespace(input_ids=[token]) for token in (1, 2, 3)] - tokenizer.tokenizer.tokenizer = tokenizer.tokenizer - tokenizer.image_token = tokenizer.boi_token = "" - tokenizer.image_length = tokenizer.mm_tokens_per_image = 2 - tokenizer.boi_token_index, tokenizer.eoi_token_index = 10, 11 - tokenizer.image_start_tag, tokenizer.image_end_tag = "", "" - tokenizer.skip_start = False - params = make_params(images=[(100, 2), (100, 2)]) - - prompt = "abc" if tokenizer_class is QWenVLTokenizer else "abc" - kwargs = {"add_special_tokens": False} if tokenizer_class is InternvlTokenizer else {} - input_ids = tokenizer.encode(prompt, params, **kwargs) - - if tokenizer_class is LlavaTokenizer: - assert input_ids == [1, 100, 101, 2, 100, 101, 3] - expected_spans = [(1, 3), (4, 6)] - else: - assert input_ids == [1, 10, 100, 101, 11, 2, 10, 100, 101, 11, 3] - expected_spans = [(2, 4), (7, 9)] - assert get_mm_input_spans(input_ids, params.to_dict()) == expected_spans - - -@pytest.mark.parametrize("tokenizer_class", [InternvlTokenizer, QWen3OmniTokenizer]) -@pytest.mark.parametrize( - "origin_ids,image_starts,audio_starts,expected_ids", - [ - ( - [1, 20, 21, 2, 10, 11, 3, 20, 21, 4, 10, 11, 5], - [7, 19], - [2, 13], - [1, 20, 200, 201, 21, 2, 10, 100, 101, 102, 11, 3, 20, 300, 301, 302, 21, 4, 10, 100, 101, 102, 11, 5], - ), - ( - [1, 10, 11, 2, 20, 21, 3, 10, 11, 4, 20, 21, 5], - [2, 13], - [8, 19], - [1, 10, 100, 101, 102, 11, 2, 20, 200, 201, 21, 3, 10, 100, 101, 102, 11, 4, 20, 300, 301, 302, 21, 5], - ), - ], -) -def test_positions_refer_to_final_input_after_both_modalities_expand( - tokenizer_class, origin_ids, image_starts, audio_starts, expected_ids -): - tokenizer = make_tokenizer(tokenizer_class, origin_ids) - params = make_params(images=[(100, 3), (100, 3)], audios=[(200, 2), (300, 3)]) - - input_ids = tokenizer.encode("mixed image/audio input", params, add_special_tokens=False) - - assert input_ids == expected_ids - assert [image.start_idx for image in params.images] == image_starts - assert [audio.start_idx for audio in params.audios] == audio_starts - spans = get_mm_input_spans(input_ids, params.to_dict()) - assert len(spans) == 4 - for item in params.images + params.audios: - assert (item.start_idx, item.start_idx + item.token_num) in spans - assert input_ids[item.start_idx : item.start_idx + item.token_num] == list( - range(item.token_id, item.token_id + item.token_num) - ) diff --git a/unit_tests/models/test_neo_registry.py b/unit_tests/models/test_neo_registry.py deleted file mode 100644 index ab3efee352..0000000000 --- a/unit_tests/models/test_neo_registry.py +++ /dev/null @@ -1,30 +0,0 @@ -import sys -from types import ModuleType - -import pytest - -from lightllm import models -from lightllm.models.registry import ModelRegistry - - -@pytest.mark.parametrize( - "llm_type,module_name,class_name", - [ - ("qwen3", "lightllm.models.neo_chat.model", "NeoTpPartModel"), - ("qwen3_moe", "lightllm.models.neo_chat_moe.model", "NeoTpMOEPartModel"), - ], -) -def test_neo_models_resolve_through_lazy_registry(monkeypatch, llm_type, module_name, class_name): - config = {"model_type": "neo_chat", "llm_config": {"model_type": llm_type}} - registration = ModelRegistry.get_model_config(config) - assert registration.is_multimodal - assert registration.model_class == f"{module_name}:{class_name}" - - implementation = ModuleType(module_name) - model_class = type(class_name, (), {}) - setattr(implementation, class_name, model_class) - monkeypatch.setitem(sys.modules, module_name, implementation) - - assert models.get_model_class(config) is model_class - assert getattr(models, class_name) is model_class - assert models.get_model_class(config) is model_class diff --git a/unit_tests/models/test_post_layer_output.py b/unit_tests/models/test_post_layer_output.py index e020c6709b..0e229938a4 100644 --- a/unit_tests/models/test_post_layer_output.py +++ b/unit_tests/models/test_post_layer_output.py @@ -78,72 +78,6 @@ def candidates(*args): assert prompt.logits_token_ids is None -def test_vocab_parallel_topk_returns_candidates_on_single_rank(monkeypatch) -> None: - _mock_vocab_parallel_pack(monkeypatch) - head = llama_post.LlamaPostLayerInfer.__new__(llama_post.LlamaPostLayerInfer) - head.tp_world_size_ = 1 - head.alloc_tensor = lambda shape, dtype, **kwargs: torch.empty(shape, dtype=dtype) - local_logits = torch.tensor( - [ - [1.0, 8.0], - [7.0, 2.0], - [3.0, 6.0], - [5.0, 4.0], - ] - ) - head._project_local_logits = lambda *args: local_logits - weight = SimpleNamespace(lm_head_weight_=SimpleNamespace(vocab_size=4, tp_vocab_start_id=0)) - - logits, token_ids = head._vocab_parallel_topk(None, 2, weight, SimpleNamespace(dist_group=None), 2) - - expected_logits, expected_token_ids = torch.topk(local_logits, k=2, dim=0, sorted=False) - assert logits.dtype == torch.float32 and token_ids.dtype == torch.int64 - assert token_ids.is_contiguous() - torch.testing.assert_close(logits, expected_logits.permute(1, 0)) - torch.testing.assert_close(token_ids, expected_token_ids.permute(1, 0)) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") -@pytest.mark.parametrize("vocab_size, token_num, top_k", [(257, 15, 512), (257, 65, 64), (248320, 64, 64)]) -def test_single_rank_vocab_topk_graph_preserves_output_contract(vocab_size, token_num, top_k) -> None: - head = llama_post.LlamaPostLayerInfer.__new__(llama_post.LlamaPostLayerInfer) - head.tp_world_size_ = 1 - local_logits = torch.randn((vocab_size, token_num), dtype=torch.bfloat16, device="cuda") - head._project_local_logits = lambda *args: local_logits - weight = SimpleNamespace(lm_head_weight_=SimpleNamespace(vocab_size=vocab_size, tp_vocab_start_id=0)) - state = SimpleNamespace(dist_group=None) - - for _ in range(3): - head._vocab_parallel_topk(None, token_num, weight, state, top_k) - torch.cuda.synchronize() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - values, token_ids = head._vocab_parallel_topk(None, token_num, weight, state, top_k) - - local_logits.normal_() - graph.replay() - torch.cuda.synchronize() - candidate_count = min(top_k, vocab_size) - expected_values = torch.topk(local_logits.T, k=candidate_count, dim=1, sorted=True).values.float() - assert values.shape == token_ids.shape == (token_num, candidate_count) - assert values.dtype == torch.float32 and token_ids.dtype == torch.int64 - assert token_ids.is_contiguous() - assert torch.all((token_ids >= 0) & (token_ids < vocab_size)) - assert torch.all(token_ids.sort(dim=1).values.diff(dim=1) > 0) - torch.testing.assert_close(values, local_logits.T.gather(1, token_ids).float(), rtol=0, atol=0) - torch.testing.assert_close(values.sort(dim=1, descending=True).values, expected_values, rtol=0, atol=0) - - -def test_vocab_parallel_topk_requires_enough_local_vocabulary() -> None: - head = llama_post.LlamaPostLayerInfer.__new__(llama_post.LlamaPostLayerInfer) - head.tp_world_size_ = 2 - head._project_local_logits = lambda *args: torch.ones(2, 1) - weight = SimpleNamespace(lm_head_weight_=SimpleNamespace(vocab_size=6, tp_vocab_start_id=0)) - - with pytest.raises(AssertionError, match="local vocabulary size 2 must be at least the candidate count 3"): - head._vocab_parallel_topk(None, 1, weight, SimpleNamespace(dist_group=None), 3) - - def test_vocab_parallel_topk_merges_rank_candidates(monkeypatch) -> None: _mock_vocab_parallel_pack(monkeypatch) head = llama_post.LlamaPostLayerInfer.__new__(llama_post.LlamaPostLayerInfer) @@ -210,32 +144,6 @@ def test_vocab_parallel_topk_pack_round_trip_preserves_token_id_bits(transposed) torch.testing.assert_close(token_ids, expected_token_ids) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") -@pytest.mark.parametrize("shape", [(129, 15), (129, 16), (129, 17), (62080, 256)]) -@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) -def test_local_vocab_topk_preserves_values_and_token_ids(shape, dtype) -> None: - logits = torch.randn(shape, dtype=dtype, device="cuda") - values, token_ids = local_vocab_topk(logits, top_k=64) - expected_values = torch.topk(logits, k=64, dim=0, sorted=True).values.T - - assert values.shape == token_ids.shape == (shape[1], 64) - assert values.dtype == dtype and token_ids.dtype == torch.int64 - assert torch.all((token_ids >= 0) & (token_ids < shape[0])) - assert torch.all(token_ids.sort(dim=1).values.diff(dim=1) > 0) - torch.testing.assert_close(values, logits.T.gather(1, token_ids), rtol=0, atol=0) - # Equal-valued boundary candidates can have different valid IDs/order. - torch.testing.assert_close(values.sort(dim=1, descending=True).values, expected_values, rtol=0, atol=0) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") -def test_local_vocab_topk_accepts_noncontiguous_logits() -> None: - logits = torch.randn((65, 129), dtype=torch.bfloat16, device="cuda").T - values, token_ids = local_vocab_topk(logits, top_k=8) - expected_values = torch.topk(logits, k=8, dim=0, sorted=True).values.T - torch.testing.assert_close(values, logits.T.gather(1, token_ids), rtol=0, atol=0) - torch.testing.assert_close(values.sort(dim=1, descending=True).values, expected_values, rtol=0, atol=0) - - @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") @pytest.mark.parametrize("top_k", [8, 64, 512]) def test_local_vocab_topk_graph_replay_reads_updated_logits(top_k) -> None: @@ -286,17 +194,6 @@ def test_prefill_prompt_logits_do_not_replace_candidate_output(draft) -> None: assert state.prompt_logics is prompt -def test_model_output_and_unpadding_preserve_candidate_mapping() -> None: - model = TpPartBaseModel.__new__(TpPartBaseModel) - logits = torch.randn(3, 2) - ids = torch.tensor([[10, 20], [30, 40], [50, 60]]) - state = SimpleNamespace(hidden_collector=NoopHiddenCollector(), prompt_logics=None) - output = model._create_model_output(PostLayerOutput(logits, ids), state) - unpadded = model._create_unpad_decode_model_output(output, 2) - torch.testing.assert_close(unpadded.logits, logits[:2]) - torch.testing.assert_close(unpadded.logits_token_ids, ids[:2]) - - def test_gemma_softcap_preserves_candidate_ids(monkeypatch) -> None: ids = torch.tensor([[10, 20]]) logits = torch.tensor([[1.0, 4.0]]) @@ -307,12 +204,3 @@ def test_gemma_softcap_preserves_candidate_ids(monkeypatch) -> None: output = head.token_forward(None, state, None) torch.testing.assert_close(output.logits, torch.tanh(logits / 2) * 2) assert output.logits_token_ids is ids - - -def test_overlap_keeps_each_microbatch_mapping() -> None: - head = llama_post.LlamaPostLayerInfer.__new__(llama_post.LlamaPostLayerInfer) - outputs = [PostLayerOutput(torch.ones(1, 2), torch.tensor([[10, 20]])), PostLayerOutput(torch.ones(1, 3))] - head.token_forward = lambda hidden, state, layer_weight: outputs[state] - output0, output1 = head.overlap_tpsp_token_forward(None, None, 0, 1, None) - assert output0 is outputs[0] - assert output1 is outputs[1] diff --git a/unit_tests/models/test_qwen3_dspark_model_output.py b/unit_tests/models/test_qwen3_dspark_model_output.py deleted file mode 100644 index 8a1187a7f5..0000000000 --- a/unit_tests/models/test_qwen3_dspark_model_output.py +++ /dev/null @@ -1,517 +0,0 @@ -from types import SimpleNamespace - -import pytest -import torch - -from lightllm.common.basemodel import batch_objs -from lightllm.common.basemodel.batch_objs import ModelMtpOutputCollector, ModelOutput -from lightllm.models.qwen3_5_dflash.model import Qwen3_5DFlashModel -from lightllm.models.qwen3_5_dspark.model import Qwen3_5DSparkModel -from lightllm.models.qwen3_dflash.layer_infer import transformer_layer_infer as qwen3_dflash_layer_infer -from lightllm.models.qwen3_dflash.layer_infer.transformer_layer_infer import Qwen3DFlashTransformerLayerInfer -from lightllm.models.qwen3_dflash.model import Qwen3DFlashModel -from lightllm.models.qwen3_dspark.layer_weights import pre_and_post_layer_weight as dspark_pre_post_weight -from lightllm.models.qwen3_dspark.layer_infer.post_layer_infer import Qwen3DSparkPostLayerInfer -from lightllm.models.qwen3_dspark.model import Qwen3DSparkModel -from lightllm.models.llama.model import LlamaTpPartModel - - -@pytest.mark.parametrize("mtp_step", [1, 5, 7]) -@pytest.mark.parametrize("model_class", [Qwen3DFlashModel, Qwen3DSparkModel, Qwen3_5DSparkModel]) -def test_parallel_block_runtime_width_follows_mtp_step(monkeypatch, mtp_step, model_class): - monkeypatch.setattr(LlamaTpPartModel, "_verify_params", lambda self: None) - model = model_class.__new__(model_class) - model.args = SimpleNamespace(mtp_mode="dspark", mtp_step=mtp_step) - model.config = {"block_size": 7} - model.enable_tpsp_mix_mode = False - - model._verify_params() - - assert model.config["block_size"] == mtp_step - - -@pytest.mark.parametrize("model_class", [Qwen3DFlashModel, Qwen3DSparkModel, Qwen3_5DSparkModel]) -def test_parallel_block_rejects_mtp_step_above_checkpoint_capacity(monkeypatch, model_class): - monkeypatch.setattr(LlamaTpPartModel, "_verify_params", lambda self: None) - model = model_class.__new__(model_class) - model.args = SimpleNamespace(mtp_mode="dflash", mtp_step=8) - model.config = {"block_size": 7} - model.enable_tpsp_mix_mode = False - - with pytest.raises(AssertionError): - model._verify_params() - - -@pytest.mark.parametrize("model_class", [Qwen3DSparkModel, Qwen3_5DSparkModel]) -@pytest.mark.parametrize("mtp_step", [1, 6]) -def test_dspark_without_block_size_uses_mtp_step(model_class, mtp_step): - model = model_class.__new__(model_class) - model.args = SimpleNamespace(mtp_step=mtp_step) - model.config = {"num_attention_heads": 40, "num_key_value_heads": 8} - model.load_way = "HF" - model.tp_world_size_ = 4 - model.enable_tpsp_mix_mode = False - - model._verify_params() - - assert model.config["block_size"] == mtp_step - - -@pytest.mark.parametrize( - "model_class", - [Qwen3DFlashModel, Qwen3_5DFlashModel, Qwen3DSparkModel, Qwen3_5DSparkModel], -) -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_state_manager = object() - model.is_mtp_draft_model = True - model.mtp_previous_draft_models = [] - model.pre_post_weight = object() - - target_hiddens = torch.arange(6, dtype=torch.float32).view(2, 3) - mem_indexes = torch.tensor([7, 11]) - model._select_mem_indexes = lambda _: mem_indexes - observed_states = [] - - class PreInfer: - def context_forward(self, input_ids, infer_state, layer_weight): - assert input_ids is None - assert layer_weight is model.pre_post_weight - assert infer_state.mtp_draft_input_hiddens is target_hiddens - observed_states.append(infer_state) - return target_hiddens + 1 - - class TransformerLayer: - def __init__(self, increment): - self.increment = increment - - def context_forward(self, hidden, infer_state, layer_weight): - assert infer_state is observed_states[0] - assert layer_weight == self.increment - return hidden + self.increment - - model.pre_infer = PreInfer() - model.layers_infer = [TransformerLayer(2), TransformerLayer(3)] - model.trans_layers_weight = [2, 3] - model_input = SimpleNamespace( - batch_size=2, - b_seq_len=torch.tensor([3, 5]), - mtp_draft_input_hiddens=target_hiddens, - ) - - output = model._decode(model_input) - - assert isinstance(output, ModelOutput) - assert output.logits.shape == (2, 1) - infer_state = observed_states[0] - assert not hasattr(infer_state, "full_att_state") - 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]]) - assert torch.equal(infer_state.position_sin, model._sin_cached[[2, 4]]) - - -@pytest.mark.parametrize("model_class", [Qwen3DSparkModel, Qwen3_5DSparkModel]) -def test_dspark_decode_unpad_uses_common_output_and_slices_mtp_fields(model_class): - model = model_class.__new__(model_class) - output = ModelOutput( - logits=torch.arange(48).view(12, 4), - mtp_collector=ModelMtpOutputCollector( - spec_hidden=torch.arange(36).view(12, 3), - confidence_logits=torch.arange(12).view(3, 4), - draft_token_ids=torch.arange(12), - ), - ) - - unpadded = model._create_unpad_decode_model_output(output, origin_batch_size=8) - - assert isinstance(unpadded, ModelOutput) - assert unpadded.logits.shape == (8, 4) - assert unpadded.mtp_collector.spec_hidden.shape == (8, 3) - assert unpadded.mtp_collector.confidence_logits.shape == (2, 4) - assert unpadded.mtp_collector.draft_token_ids.shape == (8,) - assert output.logits.shape == (12, 4) - assert output.mtp_collector.confidence_logits.shape == (3, 4) - assert output.mtp_collector.draft_token_ids.shape == (12,) - - -def test_common_output_no_ref_conversion_includes_dspark_fields(monkeypatch): - monkeypatch.setattr(batch_objs, "tensor_to_no_ref_tensor", torch.clone) - output = ModelOutput( - logits=torch.ones((2, 4)), - mtp_collector=ModelMtpOutputCollector( - spec_hidden=torch.ones((2, 3)), - confidence_logits=torch.ones((1, 2)), - draft_token_ids=torch.ones((2,), dtype=torch.int64), - ), - ) - original_ptrs = ( - output.logits.data_ptr(), - output.mtp_collector.spec_hidden.data_ptr(), - output.mtp_collector.confidence_logits.data_ptr(), - output.mtp_collector.draft_token_ids.data_ptr(), - ) - - output.to_no_ref_tensor() - - converted_ptrs = ( - output.logits.data_ptr(), - output.mtp_collector.spec_hidden.data_ptr(), - output.mtp_collector.confidence_logits.data_ptr(), - output.mtp_collector.draft_token_ids.data_ptr(), - ) - assert all(converted != original for converted, original in zip(converted_ptrs, original_ptrs)) - - -def test_dspark_post_layer_publishes_head_results_through_collector(): - post_infer = Qwen3DSparkPostLayerInfer.__new__(Qwen3DSparkPostLayerInfer) - post_infer.block_size_ = 2 - post_infer.markov_rank_ = 2 - post_infer._slice_get_last_input = lambda input_embeddings, infer_state: ( - input_embeddings, - 4, - ) - post_infer._norm = lambda hidden, infer_state, layer_weight: hidden - post_infer._sample_markov = lambda *args, **kwargs: torch.tensor([[1, 2], [3, 4]]) - confidence_logits = torch.arange(4, dtype=torch.float32).view(2, 2) - post_infer.predict_confidence_logits = lambda *args, **kwargs: confidence_logits - - class RecordingCollector: - def add_mtp_outputs(self, **kwargs): - self.outputs = kwargs - - class LMHead: - def __call__(self, input, alloc_func): - return torch.empty((8, input.shape[1])) - - collector = RecordingCollector() - infer_state = SimpleNamespace( - is_prefill=False, - input_ids=torch.tensor([10, 0, 20, 0]), - hidden_collector=collector, - ) - layer_weight = SimpleNamespace(lm_head_weight_=LMHead()) - - output = post_infer.token_forward( - input_embdings=torch.randn(4, 3), - infer_state=infer_state, - layer_weight=layer_weight, - ) - - assert output.logits.shape == (4, 1) - assert output.logits_token_ids is None - assert torch.equal(collector.outputs["draft_token_ids"], torch.tensor([1, 2, 3, 4])) - assert collector.outputs["confidence_logits"] is confidence_logits - - -def test_dspark_non_markov_head_does_not_treat_candidate_columns_as_token_ids(): - post_infer = Qwen3DSparkPostLayerInfer.__new__(Qwen3DSparkPostLayerInfer) - post_infer.block_size_ = 2 - post_infer.markov_rank_ = 0 - post_infer._slice_get_last_input = lambda input_embeddings, infer_state: (input_embeddings, 4) - candidate_output = batch_objs.PostLayerOutput( - logits=torch.randn(4, 2), - logits_token_ids=torch.tensor([[10, 20], [30, 40], [50, 60], [70, 80]]), - ) - post_infer._draft_lm_head_and_gather = lambda *args: candidate_output - observed = {} - - def predict_confidence_logits(*args, sampled_tokens, **kwargs): - observed["sampled_tokens"] = sampled_tokens - return None - - post_infer.predict_confidence_logits = predict_confidence_logits - - class RecordingCollector: - def add_mtp_outputs(self, **kwargs): - self.outputs = kwargs - - collector = RecordingCollector() - infer_state = SimpleNamespace( - is_prefill=False, - input_ids=torch.tensor([10, 0, 20, 0]), - hidden_collector=collector, - ) - - output = post_infer.token_forward(torch.randn(4, 3), infer_state, object()) - - assert output is candidate_output - assert observed["sampled_tokens"] is None - assert collector.outputs["draft_token_ids"] is None - - -@pytest.mark.parametrize( - "head_config", - [ - {"enable_confidence_head": True}, - {"markov_rank": 4, "markov_head_type": "gated"}, - {"markov_rank": 4, "markov_head_type": "rnn"}, - ], -) -def test_biased_dspark_heads_do_not_inherit_model_quantization(monkeypatch, head_config): - captured_kwargs = [] - - def init_base_weight(self, data_type, network_config, quant_cfg): - self.data_type_ = data_type - self.quant_cfg = quant_cfg - - class RecordingROWMMWeight: - def __init__(self, **kwargs): - captured_kwargs.append(kwargs) - - class StubWeight: - def __init__(self, **kwargs): - pass - - monkeypatch.setattr( - dspark_pre_post_weight.Qwen3DFlashPreAndPostLayerWeight, - "__init__", - init_base_weight, - ) - monkeypatch.setattr(dspark_pre_post_weight, "EmbeddingWeight", StubWeight) - monkeypatch.setattr(dspark_pre_post_weight, "LMHeadWeight", StubWeight) - monkeypatch.setattr(dspark_pre_post_weight, "ROWMMWeight", RecordingROWMMWeight) - - quant_method = object() - quant_cfg = SimpleNamespace(get_quant_method=lambda *_: quant_method) - network_config = { - "hidden_size": 16, - "vocab_size": 32, - **head_config, - } - dspark_pre_post_weight.Qwen3DSparkPreAndPostLayerWeight( - data_type=torch.bfloat16, - network_config=network_config, - quant_cfg=quant_cfg, - ) - - assert captured_kwargs - assert all(kwargs["quant_method"] is None for kwargs in captured_kwargs) - - -def test_fixed_dspark_does_not_require_confidence_head(monkeypatch): - monkeypatch.setattr(Qwen3DFlashModel, "_verify_params", lambda self: None) - model = Qwen3DSparkModel.__new__(Qwen3DSparkModel) - model.args = SimpleNamespace(mtp_step=6) - model.config = {"enable_confidence_head": False} - - model._verify_params() - - -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): - rotary_factors = [] - monkeypatch.setattr(qwen3_dflash_layer_infer, "qk_rmsnorm_forward", lambda *args, **kwargs: None) - monkeypatch.setattr( - qwen3_dflash_layer_infer, - "rotary_emb_fwd", - lambda *args, **kwargs: rotary_factors.append(kwargs["partial_rotary_factor"]), - ) - - class Projection: - def __init__(self, width): - self.width = width - - def mm(self, input, **kwargs): - return input.new_zeros((input.shape[0], self.width)) - - class QKNorm: - k_weight = object() - - def __call__(self, *args, **kwargs): - return None - - layer = Qwen3DFlashTransformerLayerInfer.__new__(Qwen3DFlashTransformerLayerInfer) - layer.tp_q_head_num_ = 1 - layer.tp_k_head_num_ = 1 - layer.tp_v_head_num_ = 1 - layer.head_dim_ = 8 - layer.eps_ = 1e-6 - layer.partial_rotary_factor = 0.25 - layer._post_cache_kv = lambda *args, **kwargs: None - infer_state = SimpleNamespace(position_cos=object(), position_sin=object()) - layer_weight = SimpleNamespace( - q_proj=Projection(8), - kv_proj=Projection(16), - qk_norm_weight_=QKNorm(), - ) - inputs = torch.zeros((2, 4)) - - layer.context_forward(inputs, infer_state, layer_weight) - layer._get_qkv(inputs, infer_state, layer_weight) - - assert rotary_factors == [0.25, 0.25] - - -def _cuda_empty(*size, **kwargs): - kwargs.setdefault("device", "cuda") - return torch.empty(*size, **kwargs) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_vanilla_markov_local_sampling_matches_full_logits(): - post_infer = Qwen3DSparkPostLayerInfer.__new__(Qwen3DSparkPostLayerInfer) - post_infer.block_size_ = 3 - post_infer.markov_rank_ = 2 - post_infer.markov_head_type_ = "vanilla" - post_infer.tp_world_size_ = 1 - - vocab_size = 5 - request_count = 2 - local_logits = torch.randn(vocab_size, request_count * post_infer.block_size_, device="cuda") - - class MarkovEmbedding: - def __init__(self, weight): - self.weight = weight - - def __call__(self, input_ids, alloc_func): - return torch.nn.functional.embedding(input_ids, self.weight) - - class MarkovLMHead: - def __init__(self, weight): - self.weight = weight - self.tp_vocab_start_id = 0 - - def __call__(self, input, alloc_func): - return self.weight @ input - - markov_w1 = torch.randn(vocab_size, post_infer.markov_rank_, device="cuda") - markov_w2 = torch.randn(vocab_size, post_infer.markov_rank_, device="cuda") - layer_weight = SimpleNamespace( - markov_w1_weight_=MarkovEmbedding(markov_w1), - markov_w2_weight_=MarkovLMHead(markov_w2), - ) - anchor_token_ids = torch.tensor([1, 3], device="cuda") - block_hidden = torch.empty(request_count, post_infer.block_size_, 0, device="cuda") - post_infer.alloc_tensor = _cuda_empty - - sampled_tokens = post_infer._sample_markov( - local_logits=local_logits, - block_hidden=block_hidden, - infer_state=SimpleNamespace(), - anchor_token_ids=anchor_token_ids, - layer_weight=layer_weight, - ) - - base_logits = local_logits.T.reshape(request_count, post_infer.block_size_, vocab_size) - prev_token_ids = anchor_token_ids - expected_tokens = [] - for step_idx in range(post_infer.block_size_): - markov_bias = torch.nn.functional.linear(markov_w1[prev_token_ids], markov_w2) - prev_token_ids = torch.argmax(base_logits[:, step_idx] + markov_bias, dim=-1) - expected_tokens.append(prev_token_ids) - expected_tokens = torch.stack(expected_tokens, dim=1) - - torch.testing.assert_close(sampled_tokens, expected_tokens) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_vanilla_markov_tp4_sampling_matches_full_logits(monkeypatch): - torch.manual_seed(0) - post_infer = Qwen3DSparkPostLayerInfer.__new__(Qwen3DSparkPostLayerInfer) - post_infer.block_size_ = 3 - post_infer.markov_rank_ = 4 - post_infer.markov_head_type_ = "vanilla" - post_infer.tp_world_size_ = 4 - post_infer.alloc_tensor = _cuda_empty - - vocab_size = 11 - request_count = 2 - split_indexes = torch.linspace(0, vocab_size, post_infer.tp_world_size_ + 1, dtype=torch.int64) - tp_rank = 2 - local_start = int(split_indexes[tp_rank]) - local_end = int(split_indexes[tp_rank + 1]) - - base_logits = torch.randn(request_count, post_infer.block_size_, vocab_size, device="cuda") - markov_w1 = torch.randn(vocab_size, post_infer.markov_rank_, device="cuda") - markov_w2 = torch.randn(vocab_size, post_infer.markov_rank_, device="cuda") - anchor_token_ids = torch.tensor([1, 7], device="cuda") - - class MarkovEmbedding: - weight = markov_w1 - - class MarkovLMHead: - weight = markov_w2[local_start:local_end] - tp_vocab_start_id = local_start - - layer_weight = SimpleNamespace( - markov_w1_weight_=MarkovEmbedding(), - markov_w2_weight_=MarkovLMHead(), - ) - local_logits = base_logits[:, :, local_start:local_end].reshape(-1, local_end - local_start).T.contiguous() - - expected_tokens = [] - prev_token_ids = anchor_token_ids - for step_idx in range(post_infer.block_size_): - scores = base_logits[:, step_idx] + torch.nn.functional.linear(markov_w1[prev_token_ids], markov_w2) - prev_token_ids = torch.argmax(scores, dim=-1) - expected_tokens.append(prev_token_ids) - expected_tokens = torch.stack(expected_tokens, dim=1) - - step = 0 - - def gather_tp_winners(output, local_winners, group, async_op): - nonlocal step - prev_tokens = anchor_token_ids if step == 0 else expected_tokens[:, step - 1] - scores = base_logits[:, step] + torch.nn.functional.linear(markov_w1[prev_tokens], markov_w2) - winners = [] - for rank in range(post_infer.tp_world_size_): - start = int(split_indexes[rank]) - end = int(split_indexes[rank + 1]) - values, indexes = scores[:, start:end].max(dim=-1) - winners.append(torch.stack((values, (indexes + start).float()), dim=-1)) - torch.testing.assert_close(local_winners, winners[tp_rank]) - output.copy_(torch.stack(winners).reshape_as(output)) - step += 1 - - monkeypatch.setattr( - "lightllm.models.qwen3_dspark.layer_infer.post_layer_infer.all_gather_into_tensor", - gather_tp_winners, - ) - - sampled_tokens = post_infer._sample_markov( - local_logits=local_logits, - block_hidden=torch.empty(request_count, post_infer.block_size_, 0, device="cuda"), - infer_state=SimpleNamespace(dist_group=None), - anchor_token_ids=anchor_token_ids, - layer_weight=layer_weight, - ) - - torch.testing.assert_close(sampled_tokens, expected_tokens) diff --git a/unit_tests/models/test_qwen3_eagle_model_input.py b/unit_tests/models/test_qwen3_eagle_model_input.py deleted file mode 100644 index 1395153351..0000000000 --- a/unit_tests/models/test_qwen3_eagle_model_input.py +++ /dev/null @@ -1,100 +0,0 @@ -from types import SimpleNamespace - -import pytest -import torch - -from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaTransformerLayerInfer -from lightllm.models.llama.model import LlamaTpPartModel -from lightllm.models.qwen3_eagle.layer_infer.pre_layer_infer import Qwen3EaglePreLayerInfer -from lightllm.models.qwen3_eagle.layer_infer.transformer_layer_infer import Qwen3EagleTransformerLayerInfer -from lightllm.models.qwen3_eagle.model import Qwen3EagleModel - - -def test_qwen3_eagle_uses_configured_attention_head_dim(monkeypatch): - def init_llama_layer(self, layer_num, network_config): - self.head_dim_ = network_config["hidden_size"] // network_config["num_attention_heads"] - - monkeypatch.setattr(LlamaTransformerLayerInfer, "__init__", init_llama_layer) - layer = Qwen3EagleTransformerLayerInfer( - layer_num=0, - network_config={ - "hidden_size": 2560, - "num_attention_heads": 32, - "head_dim": 128, - }, - ) - - assert layer.head_dim_ == 128 - - -def test_qwen3_eagle_projects_concatenated_target_hiddens_before_inferstate(monkeypatch): - parent_calls = [] - monkeypatch.setattr( - LlamaTpPartModel, - "_create_inferstate", - lambda self, model_input, microbatch_index=0: parent_calls.append((model_input, microbatch_index)) - or SimpleNamespace(model_input=model_input), - ) - projected_hiddens = torch.empty((3, 4)) - projection_inputs = [] - model = Qwen3EagleModel.__new__(Qwen3EagleModel) - model.config = {"hidden_size": 4} - model.pre_post_weight = SimpleNamespace( - fc_weight_=SimpleNamespace( - mm=lambda hiddens: projection_inputs.append(hiddens) or projected_hiddens, - ) - ) - target_hiddens = torch.empty((3, 12)) - model_input = SimpleNamespace( - input_ids=torch.arange(3), - mtp_draft_input_hiddens=target_hiddens, - ) - - infer_state = model._create_inferstate(model_input=model_input, microbatch_index=1) - - assert projection_inputs == [target_hiddens] - assert len(parent_calls) == 1 - normalized_input, microbatch_index = parent_calls[0] - assert normalized_input is not model_input - assert normalized_input.mtp_draft_input_hiddens is projected_hiddens - assert model_input.mtp_draft_input_hiddens is target_hiddens - assert microbatch_index == 1 - assert infer_state.model_input is normalized_input - - -def test_qwen3_eagle_keeps_recursive_draft_hiddens_without_projection(monkeypatch): - monkeypatch.setattr( - LlamaTpPartModel, - "_create_inferstate", - lambda self, model_input, microbatch_index=0: model_input, - ) - model = Qwen3EagleModel.__new__(Qwen3EagleModel) - model.config = {"hidden_size": 4} - model.pre_post_weight = SimpleNamespace( - fc_weight_=SimpleNamespace( - mm=lambda _: pytest.fail("fixed-width recursive hidden must not be projected"), - ) - ) - model_input = SimpleNamespace( - input_ids=torch.arange(3), - mtp_draft_input_hiddens=torch.empty((3, 4)), - ) - - normalized_input = model._create_inferstate(model_input=model_input) - - assert normalized_input is model_input - - -def test_qwen3_eagle_pre_layer_rejects_non_normalized_hidden_width(): - pre_layer = Qwen3EaglePreLayerInfer.__new__(Qwen3EaglePreLayerInfer) - pre_layer.hidden_size_ = 4 - infer_state = SimpleNamespace(mtp_draft_input_hiddens=torch.empty((2, 12))) - - with pytest.raises(AssertionError): - pre_layer.prepare_spec_draft_hiddens(infer_state) - - normalized_hiddens = torch.empty((2, 4)) - infer_state.mtp_draft_input_hiddens = normalized_hiddens - pre_layer.prepare_spec_draft_hiddens(infer_state) - - assert infer_state.eagle_draft_hidden_states is normalized_hiddens diff --git a/unit_tests/models/test_qwen_vl_tokenizer.py b/unit_tests/models/test_qwen_vl_tokenizer.py deleted file mode 100644 index b1d8d87fb0..0000000000 --- a/unit_tests/models/test_qwen_vl_tokenizer.py +++ /dev/null @@ -1,69 +0,0 @@ -from types import SimpleNamespace -from unittest.mock import Mock - -import pytest - -from lightllm.models.qwen2_vl.model import QWen2VLTokenizer -from lightllm.models.qwen3_vl.model import QWen3VLTokenizer -from lightllm.utils.error_utils import InvalidRequestError - - -@pytest.fixture(params=[QWen2VLTokenizer, QWen3VLTokenizer]) -def tokenizer(request): - instance = request.param.__new__(request.param) - instance.tokenizer = Mock() - instance.image_start_id = 10 - instance.image_end_id = 11 - instance.image_token_id = 12 - return instance - - -def test_image_expansion(tokenizer): - images = [SimpleNamespace(token_id=100, token_num=2), SimpleNamespace(token_id=200, token_num=3)] - prompt = [1, 10, 12, 11, 2, 10, 11, 3] - assert tokenizer.encode(prompt, SimpleNamespace(images=images)) == [ - 1, - 10, - 100, - 101, - 11, - 2, - 10, - 200, - 201, - 202, - 11, - 3, - ] - assert [image.start_idx for image in images] == [2, 7] - assert prompt == [1, 10, 12, 11, 2, 10, 11, 3] - assert tokenizer.encode(prompt, None) == [1, 10, 11, 2, 10, 11, 3] - - -@pytest.mark.parametrize("prompt,image_count", [([10, 12, 11], 0), ([10, 11, 10, 11], 1), ([1], 1)]) -def test_image_count_mismatch(tokenizer, prompt, image_count): - images = [SimpleNamespace(token_id=100, token_num=2) for _ in range(image_count)] - with pytest.raises(InvalidRequestError, match="invalid image tag num"): - tokenizer.encode(prompt, SimpleNamespace(images=images)) - assert all(not hasattr(image, "start_idx") for image in images) - - -@pytest.mark.parametrize("prompt", [[10], [10, 12], [10, 1, 11], [10, 10, 11]]) -@pytest.mark.parametrize("params", [None, SimpleNamespace(images=[])]) -def test_invalid_image_sequence(tokenizer, prompt, params): - with pytest.raises(InvalidRequestError, match="image"): - tokenizer.encode(prompt, params) - - -def test_text_and_encode_options(tokenizer): - tokenizer.tokenizer.encode.return_value = [1, 12, 2] - assert tokenizer.encode("text", SimpleNamespace(images=[]), add_special_tokens=False) == [1, 12, 2] - tokenizer.tokenizer.encode.assert_called_once_with("text", add_special_tokens=False) - - -def test_unpaired_image_end_tokens(tokenizer): - assert tokenizer.encode([1, 11, 2], SimpleNamespace(images=[])) == [1, 11, 2] - images = [SimpleNamespace(token_id=100, token_num=2), SimpleNamespace(token_id=200, token_num=1)] - prompt = [11, 1, 10, 12, 11, 11, 2, 10, 11, 11] - assert tokenizer.encode(prompt, SimpleNamespace(images=images)) == [11, 1, 10, 100, 101, 11, 11, 2, 10, 200, 11, 11] - assert tokenizer.encode(prompt, None) == [11, 1, 10, 11, 11, 2, 10, 11, 11] diff --git a/unit_tests/models/test_yarn_rotary.py b/unit_tests/models/test_yarn_rotary.py deleted file mode 100644 index ef791c20d1..0000000000 --- a/unit_tests/models/test_yarn_rotary.py +++ /dev/null @@ -1,142 +0,0 @@ -import json - -import pytest -import torch -from transformers.modeling_rope_utils import _compute_yarn_parameters -from transformers.models.qwen3_5.configuration_qwen3_5 import Qwen3_5TextConfig - -from lightllm.models.qwen2_vl.triton_kernel.mrope import mrope_triton_fused -from lightllm.models.qwen3_5.model import Qwen3_5TpPartModel - - -@pytest.fixture -def qwen35_model(tmp_path): - # Qwen3.5-27B with the documented 4x YaRN extension. - text_config = { - "hidden_size": 5120, - "head_dim": 256, - "num_attention_heads": 24, - "num_key_value_heads": 4, - "num_hidden_layers": 64, - "max_position_embeddings": 262144, - "rope_parameters": { - "mrope_interleaved": True, - "mrope_section": [11, 11, 10], - "rope_type": "yarn", - "rope_theta": 10000000, - "partial_rotary_factor": 0.25, - "factor": 4.0, - "original_max_position_embeddings": 262144, - }, - } - (tmp_path / "config.json").write_text(json.dumps({"text_config": text_config, "vision_config": {}})) - model = Qwen3_5TpPartModel.__new__(Qwen3_5TpPartModel) - model.weight_dir_ = str(tmp_path) - model.finetune_config = None - model.tp_world_size_ = 1 - model._init_config() - model.head_dim_ = model.config["head_dim"] - model.max_seq_length = 262148 - model.data_type = torch.float32 - return model - - -@pytest.mark.parametrize("type_key", ["rope_type", "type"]) -@pytest.mark.parametrize("rope_type,expected", [("default", "default"), ("mrope", "default"), ("yarn", "yarn")]) -def test_mrope_respects_explicit_scaling_type(qwen35_model, monkeypatch, type_key, rope_type, expected): - rope_scaling = qwen35_model.config["rope_scaling"] - rope_scaling.pop("rope_type") - rope_scaling[type_key] = rope_type - selected = [] - monkeypatch.setattr(qwen35_model, "_init_to_get_rotary", lambda: selected.append("default")) - monkeypatch.setattr(qwen35_model, "_init_to_get_yarn_rotary", lambda: selected.append("yarn")) - - qwen35_model._init_custom() - - assert selected == [expected] - - -def _reference_yarn_cache(model, position_ids): - config = Qwen3_5TextConfig( - hidden_size=model.config["hidden_size"], - head_dim=model.head_dim_, - num_attention_heads=model.config["num_attention_heads"], - max_position_embeddings=model.config["max_position_embeddings"], - rope_parameters=model.config["rope_parameters"].copy(), - partial_rotary_factor=model.config["partial_rotary_factor"], - ) - inv_freq, attention_factor = _compute_yarn_parameters(config, position_ids.device) - freqs = position_ids.float().unsqueeze(-1) * inv_freq - return (freqs.cos() * attention_factor).to(model.data_type), (freqs.sin() * attention_factor).to(model.data_type) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for rotary caches") -@pytest.mark.parametrize("data_type", [torch.float32, torch.bfloat16]) -@pytest.mark.parametrize("partial_rotary_factor", [0.25, 1.0]) -@pytest.mark.parametrize("factor", [1.0, 4.0]) -def test_yarn_cache_matches_transformers(qwen35_model, data_type, partial_rotary_factor, factor): - model = qwen35_model - model.data_type = data_type - model.config["partial_rotary_factor"] = partial_rotary_factor - model.config["rope_scaling"]["partial_rotary_factor"] = partial_rotary_factor - model.config["rope_scaling"]["factor"] = factor - - model._init_custom() - - half_rotary_dim = int(model.head_dim_ * partial_rotary_factor) // 2 - assert model._cos_cached.shape == (model.max_seq_length, half_rotary_dim) - assert model._sin_cached.shape == model._cos_cached.shape - # Include both sides of the original context boundary. - positions = torch.tensor([0, 1, 127, 8191, 262143, 262144, 262147], device="cuda") - expected_cos, expected_sin = _reference_yarn_cache(model, positions) - torch.testing.assert_close(model._cos_cached[positions], expected_cos, rtol=1e-5, atol=1e-6) - torch.testing.assert_close(model._sin_cached[positions], expected_sin, rtol=1e-5, atol=1e-6) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for the MRoPE kernel") -@pytest.mark.parametrize("data_type", [torch.float32, torch.bfloat16]) -def test_qwen35_yarn_interleaved_mrope_matches_reference(qwen35_model, data_type): - model = qwen35_model - model.data_type = data_type - model._init_custom() - position_ids = torch.tensor( - [[0, 1, 127, 8191, 262143, 262147], [0, 7, 83, 8189, 262144, 262145], [0, 5, 61, 8190, 262145, 262146]], - device="cuda", - ) - cos, sin = _reference_yarn_cache(model, position_ids) - half_rotary_dim = cos.shape[-1] - rotary_dim = 2 * half_rotary_dim - channels = torch.arange(half_rotary_dim, device="cuda") - # [11, 11, 10] assigns the 32 frequencies to repeating T/H/W axes. - cos = cos[channels % 3, :, channels].T.unsqueeze(1).float() - sin = sin[channels % 3, :, channels].T.unsqueeze(1).float() - cos = torch.cat((cos, cos), dim=-1) - sin = torch.cat((sin, sin), dim=-1) - - def rotate_reference(x): - result = x.clone() - rotary = x[..., :rotary_dim].float() - rotated_half = torch.cat((-rotary[..., half_rotary_dim:], rotary[..., :half_rotary_dim]), dim=-1) - result[..., :rotary_dim] = (rotary * cos + rotated_half * sin).to(data_type) - return result - - torch.manual_seed(0) - q = torch.randn((position_ids.shape[1], 24, 256), dtype=data_type, device="cuda") - k = torch.randn((position_ids.shape[1], 4, 256), dtype=data_type, device="cuda") - expected_q, expected_k = rotate_reference(q), rotate_reference(k) - - mrope_triton_fused( - q, - k, - model._cos_cached[position_ids], - model._sin_cached[position_ids], - torch.tensor(model.config["rope_scaling"]["mrope_section"], dtype=torch.int32, device="cuda"), - is_interleaved=True, - partial_rotary_factor=model.config["partial_rotary_factor"], - ) - - tolerance = 1e-5 if data_type == torch.float32 else 2e-2 - torch.testing.assert_close(q, expected_q, rtol=tolerance, atol=tolerance) - torch.testing.assert_close(k, expected_k, rtol=tolerance, atol=tolerance) - torch.testing.assert_close(q[..., rotary_dim:], expected_q[..., rotary_dim:], rtol=0, atol=0) - torch.testing.assert_close(k[..., rotary_dim:], expected_k[..., rotary_dim:], rtol=0, atol=0) diff --git a/unit_tests/models/vit/test_flash_attention_forward.py b/unit_tests/models/vit/test_flash_attention_forward.py deleted file mode 100644 index f736ec4cd6..0000000000 --- a/unit_tests/models/vit/test_flash_attention_forward.py +++ /dev/null @@ -1,139 +0,0 @@ -import torch -import math -import time -import pytest -from lightllm.models.vit.triton_kernel.flashattention_nopad import flash_attention_fwd -from lightllm.models.vit.triton_kernel import flashattention_nopad -from lightllm.common.basemodel.attention_vit.fa3.fp import Fa3VitAttBackend - - -_fa3_mark = pytest.mark.skipif( - not flashattention_nopad._flash_attn_v3_available - or not torch.cuda.is_available() - or torch.cuda.get_device_capability()[0] < 9, - reason="FA3 requires sgl_kernel and a Hopper GPU", -) - - -def test_fa3_failure_computes_fallback_on_same_invocation(monkeypatch): - calls = [] - q = torch.zeros(2, 1, 4) - out = torch.full_like(q, -123) - cu = torch.tensor([0, 2], dtype=torch.int32) - - def fail_fa3(q, k, v, o, cu_seqlens, max_seqlen): - calls.append("fa3") - o[0].fill_(99) # A backend may fail after partially writing output. - raise RuntimeError("shape-dependent FA3 failure") - - def fallback(q, k, v, o, cu_seqlens, max_seqlen): - calls.append("triton") - o.fill_(7) - - monkeypatch.setattr(flashattention_nopad, "_flash_attn_v3_available", True) - monkeypatch.setattr(flashattention_nopad, "is_hopper", lambda: True) - monkeypatch.setattr(flashattention_nopad, "flash_attention_v3_fwd", fail_fa3, raising=False) - monkeypatch.setattr(flashattention_nopad, "_flash_attention_triton_fwd", fallback) - flash_attention_fwd(q, q, q, out, cu, 2) - assert calls == ["fa3", "triton"] - assert torch.equal(out, torch.full_like(q, 7)) - assert not flashattention_nopad._flash_attn_v3_available - - # Later calls should use the disabled-backend path without retrying FA3. - out.fill_(-123) - flash_attention_fwd(q, q, q, out, cu, 2) - assert calls == ["fa3", "triton", "triton"] - assert torch.equal(out, torch.full_like(q, 7)) - - -def test_fa3_success_does_not_run_fallback(monkeypatch): - q = torch.zeros(2, 1, 4) - out = torch.full_like(q, -123) - cu = torch.tensor([0, 2], dtype=torch.int32) - - def fa3(q, k, v, o, cu_seqlens, max_seqlen): - o.fill_(9) - - monkeypatch.setattr(flashattention_nopad, "_flash_attn_v3_available", True) - monkeypatch.setattr(flashattention_nopad, "is_hopper", lambda: True) - monkeypatch.setattr(flashattention_nopad, "flash_attention_v3_fwd", fa3, raising=False) - monkeypatch.setattr( - flashattention_nopad, - "_flash_attention_triton_fwd", - lambda *args: pytest.fail("a successful FA3 call must not run fallback"), - ) - flash_attention_fwd(q, q, q, out, cu, 2) - assert torch.equal(out, torch.full_like(q, 9)) - assert flashattention_nopad._flash_attn_v3_available - - -def reference_attention_varlen(q, k, v, cu): - """ - q, k, v : (total_len, n_head, D) - cu_seqlen : prefix sums (batch+1,) - """ - total, n_head, d = q.shape - out = torch.empty_like(q) - scale = 1.0 / math.sqrt(d) - - for b in range(cu.numel() - 1): - s, e = cu[b].item(), cu[b + 1].item() - q_b, k_b, v_b = q[s:e], k[s:e], v[s:e] # (seq, head, D) - - q_hsd = q_b.permute(1, 0, 2) # (head, seq, D) - k_hds = k_b.permute(1, 2, 0) # (head, D, seq) - v_hsd = v_b.permute(1, 0, 2) # (head, seq, D) - - scores = torch.matmul(q_hsd, k_hds) * scale # (head, seq, seq) - probs = torch.softmax(scores.float(), dim=-1) - - out_hsd = torch.matmul(probs, v_hsd.float()) # (head, seq, D) - out[s:e] = out_hsd.permute(1, 0, 2).to(q.dtype) # back to (seq, head, D) - - return out - - -@pytest.mark.parametrize("dtype,atol", [(torch.float16, 1e-2), (torch.bfloat16, 2e-2)]) -@pytest.mark.parametrize( - "attention_fn", - [ - flash_attention_fwd, - pytest.param( - getattr(flashattention_nopad, "flash_attention_v3_fwd", None), - marks=_fa3_mark, - ), - pytest.param( - Fa3VitAttBackend._vit_att_fwd, - marks=_fa3_mark, - ), - ], -) -def test_varlen(dtype, atol, attention_fn, batch=4, heads=8, d=80, device="cuda:0"): - torch.manual_seed(0) - lengths = torch.randint(1, 257, (batch,)) - max_len = int(lengths.max().item()) - - cu = torch.zeros(batch + 1, dtype=torch.int32, device=device) - cu[1:] = torch.cumsum(lengths, 0) - tot = int(cu[-1]) - - q = torch.randn(tot, heads, d, dtype=dtype, device=device) - k = torch.randn_like(q) - v = torch.randn_like(q) - out_tri = torch.randn_like(q) - attention_fn(q, k, v, out_tri, cu, max_len) - a = time.time() - for _ in range(100): - attention_fn(q, k, v, out_tri, cu, max_len) - b = time.time() - print(f"flash_attention_fwd time: {(b - a) / 100 * 1000:.2f} ms") - out_ref = reference_attention_varlen(q, k, v, cu) - - max_err = (out_ref - out_tri).abs().max().item() - mean_err = (out_ref - out_tri).abs().mean().item() - print(f"{dtype}: max {max_err:.6f}, mean {mean_err:.6f}") - torch.testing.assert_close(out_tri, out_ref, atol=atol, rtol=0) - - -if __name__ == "__main__": - pytest.main() diff --git a/unit_tests/server/core/objs/test_out_token_circlequeue.py b/unit_tests/server/core/objs/test_out_token_circlequeue.py deleted file mode 100644 index 80b91b7b5d..0000000000 --- a/unit_tests/server/core/objs/test_out_token_circlequeue.py +++ /dev/null @@ -1,127 +0,0 @@ -import pytest - -from lightllm.server.core.objs.out_token_circlequeue import CircularQueue, LIGHTLLM_OUT_TOKEN_QUEUE_SIZE - - -def test_queue_initialization(): - queue = CircularQueue() - assert queue.head == 0 - assert queue.tail == 0 - assert len(queue) == 0 - assert queue.is_empty() is True - assert queue.is_full() is False - - -def test_queue_push_and_pop(): - queue = CircularQueue() - - # Push an item - queue.push("token1", 0, False, 1) - assert len(queue) == 1 - assert queue.is_empty() is False - assert queue.is_full() is False - - # Pop the item - token, src_index, special, out_token_count = queue.pop() - assert token == "token1" - assert src_index == 0 - assert len(queue) == 0 - assert special is False - assert out_token_count == 1 - - -def test_queue_push_multiple_items(): - queue = CircularQueue() - - # Push multiple items - queue.push("token1", 0, False, 1) - queue.push("token2", 1, False, 2) - queue.push("token3", 2, False, 3) - - assert len(queue) == 3 - assert queue.is_empty() is False - if len(queue) == LIGHTLLM_OUT_TOKEN_QUEUE_SIZE - 1: - assert queue.is_full() is True - - # Pop items - token, src_index, special, out_token_count = queue.pop() - assert token == "token1" - assert src_index == 0 - assert special is False - assert out_token_count == 1 - - token, src_index, special, out_token_count = queue.pop() - assert token == "token2" - assert src_index == 1 - assert special is False - assert out_token_count == 2 - - token, src_index, special, out_token_count = queue.pop() - assert token == "token3" - assert src_index == 2 - assert special is False - assert out_token_count == 3 - - assert len(queue) == 0 - - -def test_queue_full_condition(): - queue = CircularQueue() - - # Fill the queue - for i in range(LIGHTLLM_OUT_TOKEN_QUEUE_SIZE - 1): - queue.push(f"token{i}", i, False, 1) - - assert queue.is_full() is True - - # Try to push another item (should raise an exception) - with pytest.raises(Exception, match="Queue is full"): - queue.push("token_overflow", LIGHTLLM_OUT_TOKEN_QUEUE_SIZE, False, 1) - - -def test_queue_pop_empty_condition(): - queue = CircularQueue() - - # Attempt to pop from an empty queue (should raise an exception) - with pytest.raises(Exception, match="Queue is empty"): - queue.pop() - - -def test_queue_wrap_around(): - queue = CircularQueue() - - # Push enough items to wrap around - for i in range(LIGHTLLM_OUT_TOKEN_QUEUE_SIZE - 1): - queue.push(f"token{i}", i, False, i) - - # Pop all items - for i in range(LIGHTLLM_OUT_TOKEN_QUEUE_SIZE - 1): - token, src_index, _, _ = queue.pop() - assert token == f"token{i}" - assert src_index == i - - # Now the queue should be empty - assert queue.is_empty() is True - - # Push again - queue.push("token1", 0, False, 1) - queue.push("token2", 1, False, 2) - - assert len(queue) == 2 - assert queue.is_empty() is False - - # Pop the items - token, src_index, _, _ = queue.pop() - assert token == "token1" - assert src_index == 0 - - token, src_index, _, _ = queue.pop() - assert token == "token2" - assert src_index == 1 - - assert len(queue) == 0 - - -# Run with pytest -if __name__ == "__main__": - pytest.main() diff --git a/unit_tests/server/core/objs/test_req.py b/unit_tests/server/core/objs/test_req.py deleted file mode 100644 index 5966497ef1..0000000000 --- a/unit_tests/server/core/objs/test_req.py +++ /dev/null @@ -1,127 +0,0 @@ -import pytest -import easydict -from types import SimpleNamespace -from lightllm.server.core.objs.req import Req, ChunkedPrefillReq, SamplingParams -from lightllm.server.core.objs.token_metadata import ReqFinalTokenMetadata -from lightllm.utils import shm_utils -from lightllm.utils.envs_utils import get_env_start_args, set_env_start_args - - -@pytest.fixture(scope="module", autouse=True) -def setup_module_env(): - with pytest.MonkeyPatch.context() as monkeypatch: - monkeypatch.setattr(shm_utils, "get_unique_server_name", lambda: "test_req_service_0") - monkeypatch.setenv("LIGHTLLM_START_ARGS", "{}") - set_env_start_args( - easydict.EasyDict( - { - "mtp_step": 0, - "llm_prefill_att_backend": ["None"], - "llm_decode_att_backend": ["None"], - "cpu_cache_token_page_size": 256, - "enable_cpu_cache": False, - "model_dir": "", - "page_size": 4, - } - ) - ) - get_env_start_args.cache_clear() - yield - get_env_start_args.cache_clear() - - -@pytest.fixture -def req(): - req_instance = Req() - req_instance.init(1, [1, 2, 3], {"max_new_tokens": 1}, None, chunked_prefill_size=128) - return req_instance - - -def test_req_init(req): - assert req.request_id == 1 - assert req.input_len == 3 - - -def test_req_init_validates_dictionary_params(): - with pytest.raises(ValueError, match="max_new_tokens must be at least 1"): - Req().init(1, [1], {"max_new_tokens": 0}, None) - - -def test_create_prompt_ids_shm_array(req): - assert hasattr(req, "shm_prompt_ids") - - -def test_detach_shm_arrays(req): - prompt_ids = req.shm_prompt_ids - logprobs = req.shm_logprobs - - req.detach_shm_arrays() - - assert prompt_ids.shm is None - assert logprobs.shm is None - assert not hasattr(req, "shm_prompt_ids") - assert not hasattr(req, "shm_logprobs") - - -def test_get_used_tokens(req): - req.shm_cur_kv_len = 5 - assert req.get_used_tokens() == 5 - - -def test_final_token_metadata_read_returns_actual_prompt_tokens(req): - req.sample_params.prompt_logprobs = 0 - req.shm_logprobs.arr["logprob"][1] = -0.5 - req.shm_logprobs.arr["logprob"][2] = -1.25 - req.shm_logprobs.arr["rank"][1] = 315 - req.shm_logprobs.arr["rank"][2] = 4 - - metadata = ReqFinalTokenMetadata(req).read() - - assert metadata["prompt_token_ids"] == [1, 2, 3] - assert metadata["prompt_logprobs"] == [ - None, - {2: {"logprob": -0.5, "rank": 315, "decoded_token": None}}, - {3: {"logprob": -1.25, "rank": 4, "decoded_token": None}}, - ] - - -def test_chunked_req_get_tuple_tokens_adds_page_and_async_reserve(): - req = SimpleNamespace( - input_len=10, - shm_cur_output_len=0, - shm_cur_kv_len=0, - sample_params=SimpleNamespace(ignore_eos=True, max_new_tokens=5), - ) - - assert ChunkedPrefillReq.get_tuple_tokens(req, False, 10) == (11, 26) - - -def test_finish_status(req): - req.finish_status.set_status(req.finish_status.FINISHED_STOP) - assert req.finish_status.is_finished() - assert not req.finish_status.is_error_finished() - assert req.finish_status.get_finish_reason() == "stop" - - req.finish_status.set_status(req.finish_status.FINISHED_LENGTH) - assert req.finish_status.is_finished() - assert not req.finish_status.is_error_finished() - - req.finish_status.set_status(req.finish_status.FINISHED_ABORTED) - assert req.finish_status.is_finished() - assert req.finish_status.is_error_finished() - - req.finish_status.set_status(req.finish_status.FINISHED_ERROR) - assert req.finish_status.is_finished() - assert req.finish_status.is_finished_error() - assert req.finish_status.is_error_finished() - assert req.finish_status.get_finish_reason() == "error" - - req.finish_status.set_status(req.finish_status.FINISHED_PD_DECODE_CAPACITY) - assert req.finish_status.is_finished() - assert req.finish_status.is_finished_pd_decode_capacity() - assert not req.finish_status.is_error_finished() - assert req.finish_status.get_finish_reason() == "length" - - -if __name__ == "__main__": - pytest.main() diff --git a/unit_tests/server/core/objs/test_sampling_params.py b/unit_tests/server/core/objs/test_sampling_params.py index b78608bf55..f8027b0fec 100644 --- a/unit_tests/server/core/objs/test_sampling_params.py +++ b/unit_tests/server/core/objs/test_sampling_params.py @@ -37,48 +37,6 @@ }""" -@pytest.mark.parametrize( - "sequence, expected", - [ - ([1, 2, 3], [1, 2, 3]), - ([1] * (STOP_SEQUENCE_MAX_LENGTH), [1] * STOP_SEQUENCE_MAX_LENGTH), - ], -) -def test_stop_sequence_initialization(sequence, expected): - seq = StopSequence() - seq.initialize(sequence) - assert seq.size == len(expected) - assert seq.to_list() == expected - - -def test_stop_sequence_initialization_too_many(): - seq = StopSequence() - with pytest.raises(AssertionError): - seq.initialize([1] * (STOP_SEQUENCE_MAX_LENGTH + 1)) - - -@pytest.mark.parametrize( - "stop_sequences, expected_ids, expected_strings, expected_original", - [ - (["stop1", "stop2"], [], ["stop1", "stop2"], ["stop1", "stop2"]), - ("END", [], ["END"], ["END"]), - ([[1, 2], [3, 4]], [[1, 2], [3, 4]], [], [[1, 2], [3, 4]]), - (["END", [1, 2], "结束"], [[1, 2]], ["END", "结束"], ["END", [1, 2], "结束"]), - (["", [], "END", [1]], [[1]], ["END"], ["END", [1]]), - ([], [], [], []), - (None, [], [], []), - ], -) -def test_stop_sequence_groups_initialization(stop_sequences, expected_ids, expected_strings, expected_original): - groups = StopSequenceGroups() - groups.initialize(stop_sequences) - assert groups.size == len(expected_original) - assert groups.to_token_ids_list() == expected_ids - assert groups.to_stop_str_list() == expected_strings - assert groups.has_stop_strs() is bool(expected_strings) - assert groups.to_origin_list() == expected_original - - @pytest.mark.parametrize("mode", [StopStrMatchMode.ALL_TEXT, StopStrMatchMode.CONTENT_ONLY]) def test_string_stops_are_not_encoded_and_round_trip_with_token_ids(mode): tokenizer = Mock() @@ -101,53 +59,6 @@ def test_string_stops_are_not_encoded_and_round_trip_with_token_ids(mode): tokenizer.encode.assert_not_called() -def test_regular_constraint_initialization(): - constraint = RegularConstraint() - constraint.initialize("[a-zA-Z]+") - assert constraint.length == len("[a-zA-Z]+") - assert constraint.to_str() == "[a-zA-Z]+" - - with pytest.raises(AssertionError): - constraint.initialize("a" * (REGULAR_CONSTRAINT_MAX_LENGTH + 1)) - - -def test_guided_grammar_initialization(): - grammar = GuidedGrammar() - grammar.initialize(grammar_str, None) - assert grammar.to_str() == grammar_str - - with pytest.raises(AssertionError): - grammar.initialize("a" * (GRAMMAR_CONSTRAINT_MAX_LENGTH + 1), None) - - -def test_guided_json_schema_initialization(): - schema = GuidedJsonSchema() - schema.initialize(schema_str, None) - assert schema.to_str() == schema_str - - with pytest.raises(AssertionError): - schema.initialize("a" * (JSON_SCHEMA_MAX_LENGTH + 1), None) - - -def test_exponential_decay_length_penalty_initialization(): - penalty = ExponentialDecayLengthPenalty() - penalty.initialize((5, 1.5)) - assert penalty.to_tuple() == (5, 1.5) - - with pytest.raises(AssertionError): - penalty.initialize((5, 0.5)) - - -def test_node_uuid_initialization(): - node_id = 12345678901234567890 - node_uuid = NodeUUId() - node_uuid.initialize(node_id) - - assert node_uuid.node_id_high == (node_id >> 64) & 0xFFFFFFFFFFFFFFFF - assert node_uuid.node_id_low == node_id & 0xFFFFFFFFFFFFFFFF - assert node_uuid.get() == node_id - - def test_sampling_params_initialization(): params = SamplingParams() pd_master_node_id = 12345678901234567890 @@ -238,15 +149,6 @@ def test_stop_matching_policy_round_trip(kwargs, expected): assert py_restored.stop_str_match_mode.get_mode() == expected -@pytest.mark.parametrize("mode", [-1, 0, 3, 2 ** 32, "CONTENT_ONLY", None, 0.0]) -def test_invalid_stop_str_match_mode(mode): - params = SamplingParams() - with pytest.raises((ValueError, TypeError)): - params.init(None, stop_str_match_mode=mode) - with pytest.raises((ValueError, TypeError)): - PySamplingParams(stop_str_match_mode=mode) - - def test_stop_str_match_mode_updates_shared_buffer(): params = SamplingParams() params.init(None, include_stop_str_in_output=True, _initial_reasoning_state=1) @@ -264,12 +166,6 @@ def test_stop_str_match_mode_updates_shared_buffer(): assert reader.stop_str_match_mode.is_content_only() -def test_infer_high_priority_is_internal_and_ignores_external_value(): - params = SamplingParams() - params.init(None, infer_high_priority=-3) - assert params.infer_high_priority == 0 - - # Mock tokenizer for testing class MockTokenizer: def encode(self, text, add_special_tokens=False): diff --git a/unit_tests/server/core/objs/test_shm_array.py b/unit_tests/server/core/objs/test_shm_array.py deleted file mode 100644 index af02c7de0b..0000000000 --- a/unit_tests/server/core/objs/test_shm_array.py +++ /dev/null @@ -1,74 +0,0 @@ -# test_shm_array.py - -import pytest -import numpy as np -from multiprocessing import shared_memory -from lightllm.server.core.objs.shm_array import ShmArray # Replace 'your_module' with the actual module name -from lightllm.utils import shm_utils - - -@pytest.fixture(scope="module", autouse=True) -def service_name(): - monkeypatch = pytest.MonkeyPatch() - monkeypatch.setattr(shm_utils, "get_unique_server_name", lambda: "test_shm_array_service_0") - yield - monkeypatch.undo() - - -@pytest.fixture(scope="module") -def shm_array(): - name = "test_shm_array" - shape = (10, 10) - dtype = np.float32 - shm = ShmArray(name, shape, dtype) - shm.create_shm() - yield shm - shm.close_shm() - - -def test_create_shm(shm_array): - assert shm_array.shm is not None - assert shm_array.arr.shape == shm_array.shape - assert shm_array.arr.dtype == shm_array.dtype - - -def test_link_shm(shm_array): - # Create a new instance to link the existing shm - shm_linked = ShmArray(shm_array.name, shm_array.shape, shm_array.dtype) - shm_linked.link_shm() - - assert shm_linked.shm is not None - assert shm_linked.arr.shape == shm_array.shape - assert shm_linked.arr.dtype == shm_array.dtype - - -def test_shm_size_mismatch(shm_array): - # Create a new shm with a different size - shape_mismatch = (5, 5) - shm_mismatch = ShmArray(shm_array.name, shape_mismatch, shm_array.dtype) - - with pytest.raises(Exception): - shm_mismatch.link_shm() - - -def test_close_shm(shm_array): - shm_array.close_shm() - assert shm_array.shm is None - - -def test_recreate_shm_after_close(shm_array): - shm_array.close_shm() - - # Attempting to create a new shm after close - shm_new = ShmArray(shm_array.name, shm_array.shape, shm_array.dtype) - shm_new.create_shm() - - assert shm_new.shm is not None - assert shm_new.arr.shape == shm_array.shape - assert shm_new.arr.dtype == shm_array.dtype - - shm_new.close_shm() - - -if __name__ == "__main__": - pytest.main() diff --git a/unit_tests/server/core/objs/test_shm_req_manager.py b/unit_tests/server/core/objs/test_shm_req_manager.py index c51e3e6b95..3b59f7d7ca 100644 --- a/unit_tests/server/core/objs/test_shm_req_manager.py +++ b/unit_tests/server/core/objs/test_shm_req_manager.py @@ -1,7 +1,6 @@ import os import asyncio import pytest -import time import uuid import numpy as np from unittest.mock import MagicMock @@ -96,18 +95,6 @@ def test_alloc_req_index_no_available(shm_req_manager): assert index is None # No indices should be available -def test_performance_alloc_release(shm_req_manager): - start_time = time.time() - - for _ in range(100): - index = shm_req_manager.alloc_req_index() - if index is not None: - shm_req_manager.release_req_index(index) - - duration = time.time() - start_time - print(f"Time taken for 100 alloc/release cycles: {duration:.6f} seconds") - - @pytest.fixture def batch_req_manager(monkeypatch): service_name = f"test_batch_req_manager_{uuid.uuid4().hex}" diff --git a/unit_tests/server/core/objs/test_stop_sequence_output.py b/unit_tests/server/core/objs/test_stop_sequence_output.py index 4e0d514c63..d935550dc7 100644 --- a/unit_tests/server/core/objs/test_stop_sequence_output.py +++ b/unit_tests/server/core/objs/test_stop_sequence_output.py @@ -210,170 +210,6 @@ async def stop_request(req_id): assert not buffer.ready_token_list -@pytest.mark.parametrize("include", [False, True]) -def test_only_possible_stop_prefix_delays_whole_token(include): - buffer = StopStrMatchBuffer(_http_manager(), _params(include=include), []) - assert buffer.state is buffer.State.SCANNING - _handle_output(buffer, *_output(1, "hello", 1)) - ready_tokens = _drain_ready_tokens(buffer) - assert ready_tokens == [("hello", 1)] - assert buffer.state is buffer.State.SCANNING - _handle_output(buffer, *_output(2, "worldE", 2)) - ready_tokens = _drain_ready_tokens(buffer) - assert not ready_tokens - assert buffer.state is buffer.State.WAITING - _handle_output(buffer, *_output(3, "X", 3)) - ready_tokens = _drain_ready_tokens(buffer) - assert ready_tokens == [("worldE", 2), ("X", 3)] - assert buffer.state is buffer.State.SCANNING - - -@pytest.mark.parametrize("include", [False, True]) -def test_releasing_one_token_preserves_next_token_before_stop(include): - buffer = StopStrMatchBuffer(_http_manager(), _params(include=include), []) - _handle_output(buffer, *_output(1, "aE", "first")) - assert not buffer.ready_token_list - assert buffer.state is buffer.State.WAITING - _handle_output(buffer, *_output(2, "bE", "second")) - released_tokens = _drain_ready_tokens(buffer) - assert released_tokens == [("aE", "first")] - assert buffer.state is buffer.State.WAITING - _handle_output(buffer, *_output(3, "NDextra", "third")) - stopped_tokens = _drain_ready_tokens(buffer) - assert buffer.state is buffer.State.STOPPED - assert stopped_tokens == [("bE" if include else "b", "second"), ("ND" if include else "", "third")] - assert released_tokens == [("aE", "first")] - - -@pytest.mark.parametrize("include", [False, True]) -def test_regex_partial_match_keeps_longest_tail_and_releases_safe_tokens(include): - manager = _http_manager() - manager.stop_request = AsyncMock() - buffer = StopStrMatchBuffer(manager, _params(["abcd", "bcx"], include), []) - _handle_output(buffer, *_output(1, "safe", "first")) - assert _drain_ready_tokens(buffer) == [("safe", "first")] - _handle_output(buffer, *_output(2, "a", "second")) - assert not buffer.ready_token_list - _handle_output(buffer, *_output(3, "bc", "third")) - assert not buffer.ready_token_list - assert buffer._pending_text == "abc" - assert buffer.state is buffer.State.WAITING - _handle_output(buffer, *_output(4, "xab", "fourth")) - # 后一个 stop 已完整命中，即使文本尾部仍是前一个 stop 的前缀，也应立即结束。 - assert _drain_ready_tokens(buffer) == [ - ("a", "second"), - ("bc" if include else "", "third"), - ("x" if include else "", "fourth"), - ] - assert buffer.state is buffer.State.STOPPED - manager.stop_request.assert_awaited_once_with(80) - - -@pytest.mark.parametrize("include", [False, True]) -@pytest.mark.parametrize("mode", [StopStrMatchMode.ALL_TEXT, StopStrMatchMode.CONTENT_ONLY]) -def test_unicode_stop_matching_does_not_normalize_text(include, mode): - params = _params("é", include) - params.stop_str_match_mode.set_mode(mode) - manager = _http_manager() - manager.stop_request = AsyncMock() - buffer = StopStrMatchBuffer(manager, params, []) - _handle_output(buffer, *_output(1, "e", "first")) - assert _drain_ready_tokens(buffer) == [("e", "first")] - _handle_output(buffer, *_output(2, "\u0301", "last", finish=FinishStatus.FINISHED_LENGTH)) - assert buffer.ready_token_list[-1][3].is_finished_length() - assert _drain_ready_tokens(buffer) == [("\u0301", "last")] - assert buffer.state is buffer.State.STOPPED - manager.stop_request.assert_not_awaited() - - -def test_none_token_resets_text_matching_without_simulated_flag(): - buffer = StopStrMatchBuffer(_http_manager(), _params(), []) - _handle_output(buffer, *_output(10, "helloE", "first")) - first_tokens = _drain_ready_tokens(buffer) - assert buffer.state is buffer.State.WAITING - _handle_output(buffer, *_output(None, "N", "marker")) - marker_tokens = _drain_ready_tokens(buffer) - assert buffer.state is buffer.State.SCANNING - _handle_output(buffer, *_output(11, "D", "last")) - last_tokens = _drain_ready_tokens(buffer) - assert first_tokens + marker_tokens + last_tokens == [("helloE", "first"), ("N", "marker"), ("D", "last")] - _handle_output(buffer, *_output(10, "E", "new first")) - assert buffer.state is buffer.State.WAITING - _handle_output(buffer, *_output(11, "ND", "new last")) - assert buffer.ready_token_list[-1][3].is_stopped() - - -@pytest.mark.parametrize( - "token_id,finish", - [(None, FinishStatus.NO_FINISH), (1001, FinishStatus.FINISHED_ERROR), (1001, FinishStatus.FINISHED_ABORTED)], -) -def test_non_real_token_does_not_leave_reasoning(token_id, finish): - buffer = _stop_buffer() - _handle_output(buffer, *_output(token_id, "", "marker", finish=finish)) - marker_tokens = _drain_ready_tokens(buffer) - assert marker_tokens == [("", "marker")] - assert buffer._reasoning_state.in_reasoning - if finish == FinishStatus.NO_FINISH: - _handle_output(buffer, *_output(10, "END", "thought")) - thought_tokens = _drain_ready_tokens(buffer) - assert thought_tokens == [("END", "thought")] - else: - assert buffer.state is buffer.State.STOPPED - - -@pytest.mark.parametrize("include", [False, True]) -def test_unmatched_tail_flushes(include): - buffer = StopStrMatchBuffer(_http_manager(), _params(include=include), []) - _handle_output(buffer, *_output(1, "helloE", 1)) - assert not buffer.ready_token_list - assert buffer.state is buffer.State.WAITING - _handle_output(buffer, *_output(2, "N", 2, finish=FinishStatus.FINISHED_LENGTH)) - ready_tokens = _drain_ready_tokens(buffer) - assert ready_tokens == [("helloE", 1), ("N", 2)] - assert buffer.state is buffer.State.STOPPED - - -def test_no_string_stops_emit_immediately(): - buffer = StopStrMatchBuffer(_http_manager(), _params([[10, 11]]), []) - _handle_output(buffer, *_output(1, "hello", 1)) - ready_tokens = _drain_ready_tokens(buffer) - assert ready_tokens == [("hello", 1)] - - -@pytest.mark.parametrize("include", [False, True]) -@pytest.mark.parametrize("alphabet", ["ab结束", "ab结束🙂.*[]\\\n"]) -def test_string_stop_matches_reference_for_random_token_boundaries(include, alphabet): - async def run(): - rng = random.Random(20261011) - for _ in range(300): - stops = ["".join(rng.choices(alphabet, k=rng.randint(1, 5))) for _ in range(rng.randint(1, 4))] - chunks = ["".join(rng.choices(alphabet, k=rng.randint(0, 4))) for _ in range(20)] - manager = _http_manager() - manager.stop_request = AsyncMock() - buffer = StopStrMatchBuffer(manager, _params(stops, include), []) - text = "" - expected_text = None - expected_count = len(chunks) - outputs = [] - for index, chunk in enumerate(chunks): - text += chunk - if expected_text is None: - hits = [(text.find(stop), -len(stop), stop) for stop in stops if stop in text] - if hits: - position, _, stop = min(hits) - expected_text = text[: position + (len(stop) if include else 0)] - expected_count = index + 1 - finish = FinishStatus.FINISHED_LENGTH if index == len(chunks) - 1 else FinishStatus.NO_FINISH - await buffer.handle(*_output(index, chunk, index, finish)) - outputs.extend(_drain_ready_tokens(buffer)) - assert "".join(chunk for chunk, _ in outputs) == (text if expected_text is None else expected_text) - assert [payload for _, payload in outputs] == list(range(expected_count)) - assert buffer.state is buffer.State.STOPPED - assert manager.stop_request.await_count == int(expected_text is not None and expected_count < len(chunks)) - - asyncio.run(run()) - - @pytest.mark.parametrize("mode", [StopStrMatchMode.ALL_TEXT, StopStrMatchMode.CONTENT_ONLY]) def test_token_id_stop_is_wrapped_without_text_buffering_even_in_reasoning(mode): params = _params([[10, 11]]) @@ -402,205 +238,6 @@ async def run(): manager.stop_request.assert_awaited_once_with(80) -@pytest.mark.parametrize( - "stop,chunks,expected_text", - [ - ("END", [(1, "E"), (2, "E"), (3, "ND")], "E"), - ], -) -def test_waiting_tracks_overlapping_stop_prefixes(stop, chunks, expected_text): - buffer = StopStrMatchBuffer(_http_manager(), _params(stop), []) - outputs = [] - for index, (token_id, text) in enumerate(chunks): - _handle_output(buffer, *_output(token_id, text, index)) - ready_tokens = _drain_ready_tokens(buffer) - outputs.extend(ready_tokens) - if index < len(chunks) - 1: - assert buffer.state is buffer.State.WAITING - else: - assert buffer.state is buffer.State.STOPPED - assert buffer.state is buffer.State.STOPPED - assert "".join(text for text, _ in outputs) == expected_text - assert [payload for _, payload in outputs] == list(range(len(chunks))) - - -@pytest.mark.parametrize("mode", [StopStrMatchMode.ALL_TEXT, StopStrMatchMode.CONTENT_ONLY]) -def test_token_id_stop_does_not_enter_text_waiting(mode): - params = _params([[10, 11]]) - params.stop_str_match_mode.set_mode(mode) - buffer = StopStrMatchBuffer(_http_manager(), params, []) - _handle_output(buffer, *_output(10, "first", 1)) - first_tokens = _drain_ready_tokens(buffer) - assert buffer.state is buffer.State.SCANNING - _handle_output(buffer, *_output(11, "last", 2, finish=FinishStatus.FINISHED_LENGTH)) - last_tokens = _drain_ready_tokens(buffer) - assert buffer.state is buffer.State.STOPPED - assert first_tokens + last_tokens == [("first", 1), ("last", 2)] - - -@pytest.mark.parametrize("stop_in_reasoning,stop_at", [(False, 3), (True, 0)]) -@pytest.mark.parametrize("hidden_delimiter", [False, True]) -def test_reasoning_stop_matches_only_selected_scope(stop_in_reasoning, stop_at, hidden_delimiter): - buffer = _stop_buffer(stop_in_reasoning) - chunks = [(10, "END"), (1001, "" if hidden_delimiter else ""), (11, "answerE"), (12, "NDextra")] - results = [] - for index, (token_id, text) in enumerate(chunks): - _handle_output(buffer, *_output(token_id, text, None)) - ready_tokens = _drain_ready_tokens(buffer) - results.extend(ready_tokens) - assert (buffer.state is buffer.State.STOPPED) is (index == stop_at) - if buffer.state is buffer.State.STOPPED: - break - assert "".join(text for text, _ in results) == ( - "END" + ("" if hidden_delimiter else "") + "answer" if stop_at == 3 else "" - ) - - -@pytest.mark.parametrize("include", [False, True]) -def test_reasoning_eos_preserves_ignored_stop_text(include): - buffer = _stop_buffer(include=include) - _handle_output(buffer, *_output(10, "reasoning END", None)) - ready_tokens = _drain_ready_tokens(buffer) - assert ready_tokens == [("reasoning END", None)] - - -@pytest.mark.parametrize("include", [False, True]) -def test_scoped_match_filters_content_and_keeps_reasoning(include): - buffer = _stop_buffer(include=include) - tokens = [] - for token_id, text in [(10, "END"), (1001, ""), (11, "answerENDextra")]: - _handle_output(buffer, *_output(token_id, text, None)) - tokens.extend(_drain_ready_tokens(buffer)) - assert "".join(text for text, _ in tokens) == ("ENDanswer" + ("END" if include else "")) - - -def test_stop_cannot_cross_reasoning_content_boundary(): - buffer = _stop_buffer() - tokens = [] - for token_id, text in [(1, "E"), (1001, ""), (2, "ND")]: - _handle_output(buffer, *_output(token_id, text, None)) - assert buffer.state is buffer.State.SCANNING - tokens.extend(_drain_ready_tokens(buffer)) - assert "".join(text for text, _ in tokens) == "END" - - -def test_entering_reasoning_releases_and_resets_pending_string_stop(): - buffer = _stop_buffer(force=0) - _handle_output(buffer, *_output(10, "helloE", "first")) - first_tokens = _drain_ready_tokens(buffer) - assert buffer.state is buffer.State.WAITING - _handle_output(buffer, *_output(1000, "", "start")) - start_tokens = _drain_ready_tokens(buffer) - assert buffer.state is buffer.State.SCANNING - assert first_tokens + start_tokens == [("helloE", "first"), ("", "start")] - _handle_output(buffer, *_output(11, "ND", "thought")) - thought_tokens = _drain_ready_tokens(buffer) - assert thought_tokens == [("ND", "thought")] - _handle_output(buffer, *_output(1001, "", "end")) - assert not buffer._reasoning_state.in_reasoning - _handle_output(buffer, *_output(10, "E", "answer first")) - _handle_output(buffer, *_output(11, "ND", "answer last")) - assert buffer.ready_token_list[-1][3].is_stopped() - - -def test_text_buffer_does_not_match_token_id_stops_in_any_phase(): - buffer = _stop_buffer(stop=[[10, 11]]) - for token_id in [10, 11, 1001, 10, 11]: - _handle_output(buffer, *_output(token_id, "", None)) - assert buffer.state is buffer.State.SCANNING - assert all(not output[3].is_finished() for output in buffer.ready_token_list) - - -def test_pd_first_token_updates_reasoning_state(): - buffer = _stop_buffer() - _handle_output(buffer, *_output(1001, "", None)) - assert not buffer._reasoning_state.in_reasoning - _handle_output(buffer, *_output(10, "END", None)) - assert buffer.ready_token_list[-1][3].is_stopped() - - -def test_prompt_start_marker_prevents_stop_inside_reasoning(): - manager = _http_manager() - manager.args.reasoning_parser = "qwen3" - params = _params(force=0) - params.stop_str_match_mode.set_mode(StopStrMatchMode.CONTENT_ONLY) - buffer = StopStrMatchBuffer(manager, params, [1, 1000]) - _handle_output(buffer, *_output(10, "END", "thought")) - thought_tokens = _drain_ready_tokens(buffer) - assert thought_tokens == [("END", "thought")] - _handle_output(buffer, *_output(1001, "", "end")) - assert _drain_ready_tokens(buffer) == [("", "end")] - _handle_output(buffer, *_output(11, "answerEND", "answer")) - answer_tokens = _drain_ready_tokens(buffer) - assert buffer.state is buffer.State.STOPPED - assert answer_tokens == [("answer", "answer")] - - -def test_end_marker_outside_reasoning_can_match_stop(): - buffer = _stop_buffer(stop="", force=0) - payload = object() - _handle_output(buffer, *_output(1001, "", payload)) - ready_tokens = _drain_ready_tokens(buffer) - assert buffer.state is buffer.State.STOPPED - assert ready_tokens == [("", payload)] - - -def test_pd_capacity_marker_preserves_reasoning_state(): - buffer = _stop_buffer() - _handle_output(buffer, *_output(None, "", None)) - # 模拟 token 的 不结束思考，后续 END 仍属于 reasoning。 - _handle_output(buffer, *_output(10, "END", None)) - assert buffer._reasoning_state.in_reasoning - assert buffer.state is buffer.State.SCANNING - assert _drain_ready_tokens(buffer) == [("", None), ("END", None)] - - -@pytest.mark.parametrize( - "finish", - [ - FinishStatus.FINISHED_LENGTH, - FinishStatus.FINISHED_STOP, - FinishStatus.FINISHED_ERROR, - FinishStatus.FINISHED_ABORTED, - FinishStatus.FINISHED_PD_DECODE_CAPACITY, - ], -) -def test_detokenization_forwards_raw_tokens_before_release(monkeypatch, finish): - tokens = list(range(10, 30)) - params = _params() - params.stop_str_match_mode.set_mode(StopStrMatchMode.ALL_TEXT) - decode = DecodeReq(_req(tokens, finish, params)) - simulated = finish in ( - FinishStatus.FINISHED_ERROR, - FinishStatus.FINISHED_ABORTED, - FinishStatus.FINISHED_PD_DECODE_CAPACITY, - ) - chunks = {tokens[-2]: "helloE", tokens[-1]: "" if simulated else "NDextra"} - monkeypatch.setattr( - "lightllm.server.detokenization.manager.decode_token", - lambda tokenizer, req, token_id, eos: chunks.get(token_id, ""), - ) - manager = object.__new__(DeTokenizationManager) - manager.req_id_to_out = {80: decode} - manager.tokenizer = _ReasoningTokenizer() - manager.eos_id = [] - manager.all_special_ids = set() - manager.pub_to_httpserver = SimpleNamespace(send_pyobj=lambda *a, **kw: None) - manager.shm_req_manager = SimpleNamespace(put_back_req_obj=lambda req: None) - results = [] - for _ in range(100): - manager.gen_token_out() - if decode.out_queue_is_full() and decode.need_detoken(): - assert not decode.req.can_released_mark - while not decode.req.out_tokens_queue.is_empty(): - results.append(decode.req.out_tokens_queue.pop()) - if decode.req.can_released_mark: - break - assert decode.req.can_released_mark - assert "".join(text for text, *_ in results) == ("helloE" if simulated else "helloENDextra") - assert [count for *_, count in results] == list(range(1, len(tokens) + 1)) - - @pytest.mark.parametrize("include", [False, True]) @pytest.mark.parametrize("empty_tokens", [2, 35]) def test_http_stop_drains_detokenization_before_release(monkeypatch, include, empty_tokens): @@ -647,565 +284,6 @@ async def raw_results(): asyncio.run(run()) -def test_http_uses_request_finish_index(): - async def run(): - manager = _http_manager() - manager.pd_mode = SimpleNamespace(is_normal=lambda: False) - req = _req([10, 11, 12], FinishStatus.FINISHED_ABORTED) - req.out_tokens_queue.push("hello", 2, False, 1) - first = await _read_http_tokens(manager, req) - req.out_tokens_queue.push("", 4, False, 3) - last = await _read_http_tokens(manager, req) - assert first[0][1] == "hello" - assert not first[0][3].is_finished() - assert last[0][2]["id"] is None - assert last[0][3].status == FinishStatus.FINISHED_ABORTED - req.merge_final_token_metadata.assert_awaited_once() - - asyncio.run(run()) - - -@pytest.mark.parametrize("include", [False, True]) -def test_http_stop_aborts_only_matching_choice_and_drains_late_tokens(include): - async def run(): - manager = _http_manager() - params = _params(include=include) - params.n = params.best_of = 2 - req = _req([10, 11, 12, 13], params=params) - sibling = _req([50, 51], FinishStatus.FINISHED_LENGTH, params, request_id=81) - drained = [] - - async def raw_results(): - for target, index, text in [ - (req, 0, "helloE"), - (sibling, 0, "otherE"), - (req, 1, "NDextra"), - (req, 2, "late"), - (sibling, 1, "X"), - (req, 3, ""), - ]: - if target is req and index >= 2: - assert req.is_aborted - assert not sibling.is_aborted - if target is req and index == 3: - req.finish_status.set_status(FinishStatus.FINISHED_ABORTED) - req.finish_token_index = 5 - req.can_released_mark = True - target.out_tokens_queue.push(text, 2 + index, False, index + 1) - drained.append((target.request_id, index)) - for result in await _read_http_tokens(manager, target): - yield result - - results = [ - result async for result in _match_http_stream(manager, raw_results(), params, [1, 2], [req, sibling]) - ] - assert len(drained) == 6 - assert req.out_tokens_queue.is_empty() - assert "".join(r[1] for r in results if r[0] == 80) == "hello" + ("END" if include else "") - assert "".join(r[1] for r in results if r[0] == 81) == "otherEX" - assert [r[2]["id"] for r in results if r[0] == 80] == [10, 11] - assert {r[0]: r[3].get_finish_reason() for r in results if r[3].is_finished()} == {80: "stop", 81: "length"} - req.merge_final_token_metadata.assert_awaited_once() - sibling.merge_final_token_metadata.assert_awaited_once() - - asyncio.run(run()) - - -@pytest.mark.parametrize( - "finish", - [ - FinishStatus.FINISHED_STOP, - FinishStatus.FINISHED_LENGTH, - FinishStatus.FINISHED_ERROR, - FinishStatus.FINISHED_ABORTED, - ], -) -def test_http_finish_flushes_unmatched_prefix_and_preserves_final_metadata(finish): - async def run(): - manager = _http_manager() - params = _params() - req = _req([10, 11], finish, params) - req.merge_final_token_metadata.side_effect = lambda metadata, *args, **kwargs: metadata.update( - prompt_token_ids=[1, 2] - ) - - async def raw_results(): - for index, text in enumerate(["helloE", ""]): - req.out_tokens_queue.push(text, 2 + index, False, index + 1) - for result in await _read_http_tokens(manager, req): - yield result - - results = [result async for result in _match_http_stream(manager, raw_results(), params, [1, 2], [req])] - assert "".join(r[1] for r in results) == "helloE" - assert results[-1][2]["prompt_token_ids"] == [1, 2] - assert results[-1][3].status == finish - assert results[-1][2]["id"] == (None if finish == FinishStatus.FINISHED_ABORTED else 11) - assert not req.is_aborted - - asyncio.run(run()) - - -@pytest.mark.parametrize("stop_in_reasoning", [False, True]) -@pytest.mark.parametrize("token_stop", [False, True]) -def test_http_scopes_stop_matching_and_preserves_hidden_reasoning_marker(stop_in_reasoning, token_stop): - async def run(): - manager = _http_manager() - manager.args.reasoning_parser = "qwen3" - params = _params([[10, 11]] if token_stop else "END", force=1) - params.stop_str_match_mode.set_mode( - StopStrMatchMode.ALL_TEXT if stop_in_reasoning else StopStrMatchMode.CONTENT_ONLY - ) - tokens = [10, 11, 1001, 10, 11] - req = _req(tokens, params=params) - infer_req = infer_batch.InferReq.__new__(infer_batch.InferReq) - infer_req.stop_sequences = params.stop_sequences.to_token_ids_list() - infer_req.shm_req = req - infer_req.sampling_param = SimpleNamespace(shm_param=params) - infer_req.finish_status = req.finish_status - - async def raw_results(): - for index, text in enumerate(["E", "ND", "", "answerE", "NDextra"]): - infer_req.update_finish_status([], output_len=index + 1) - if req.finish_status.is_finished(): - req.finish_token_index = req.input_len + index - req.out_tokens_queue.push(text, 2 + index, False, index + 1) - for result in await _read_http_tokens(manager, req): - yield result - if req.finish_status.is_finished(): - break - - results = [result async for result in _match_http_stream(manager, raw_results(), params, [1, 2], [req])] - expected = "END" if token_stop else ("" if stop_in_reasoning else "ENDanswer") - assert "".join(r[1] for r in results) == expected - assert [r[2]["id"] for r in results] == tokens[: 2 if token_stop or stop_in_reasoning else 5] - assert results[-1][3].is_stopped() - - asyncio.run(run()) - - -def test_stop_groups_can_separate_explicit_token_ids_from_strings(): - params = _params(["END", [10, 11]]) - assert params.stop_sequences.to_token_ids_list() == [[10, 11]] - assert params.stop_sequences.to_stop_str_list() == ["END"] - assert params.to_dict()["stop_sequences"] == ["END", [10, 11]] - - -@pytest.mark.parametrize("mode", [StopStrMatchMode.ALL_TEXT, StopStrMatchMode.CONTENT_ONLY]) -def test_string_stop_is_matched_only_after_detokenization(mode): - params = _params() - params.stop_str_match_mode.set_mode(mode) - tokens = [ord(c) for c in "END"] - shm_req = _req(tokens, params=params) - req = infer_batch.InferReq.__new__(infer_batch.InferReq) - req.stop_sequences = params.stop_sequences.to_token_ids_list() - req.shm_req = shm_req - req.sampling_param = SimpleNamespace(shm_param=params) - req.finish_status = FinishStatus() - req.update_finish_status([], output_len=len(tokens)) - assert not req.finish_status.is_finished() - - decode = DecodeReq(shm_req) - assert decode.need_detoken() - assert not decode.can_set_release_mark() - - -@pytest.mark.parametrize("mode", [StopStrMatchMode.ALL_TEXT, StopStrMatchMode.CONTENT_ONLY]) -@pytest.mark.parametrize("reasoning", [False, True]) -@pytest.mark.parametrize( - "stop_ids, output_ids, matched", - [ - ([10], [10], True), - ([10, 11], [10, 11], True), - ([10, 11], [10], False), - ([10, 11], [10, 12, 11], False), - ([10, 10, 11], [10, 10, 10, 11], True), - ], -) -def test_inference_token_stop_is_independent_of_string_mode_and_reasoning( - mode, reasoning, stop_ids, output_ids, matched -): - req = infer_batch.InferReq.__new__(infer_batch.InferReq) - params = _params(["END", stop_ids], force=int(reasoning)) - params.stop_str_match_mode.set_mode(mode) - req.stop_sequences = params.stop_sequences.to_token_ids_list() - req.shm_req = SimpleNamespace(input_len=1, shm_prompt_ids=SimpleNamespace(arr=np.array([1, *output_ids]))) - req.sampling_param = SimpleNamespace(shm_param=params) - req.finish_status = FinishStatus() - req.update_finish_status([], output_len=len(output_ids)) - assert req.finish_status.is_stopped() is matched - - # EOS and the output length limit remain independent of user stop sequences. - req.finish_status = FinishStatus() - req.update_finish_status([output_ids[-1]], output_len=len(output_ids)) - assert req.finish_status.is_stopped() is True - req.finish_status = FinishStatus() - req.stop_sequences = [] - req.sampling_param.shm_param.max_new_tokens = 1 - req.update_finish_status([], output_len=1) - assert req.finish_status.is_finished_length() is True - - -@pytest.mark.parametrize("include", [False, True]) -@pytest.mark.parametrize("last_status", [FinishStatus.NO_FINISH, FinishStatus.FINISHED_LENGTH]) -@pytest.mark.parametrize( - "stops, output_ids, match_index", - [ - ([[10]], [10, 20], 0), - ([[10, 11]], [10, 11, 20], 1), - ([[10, 10, 11]], [10, 10, 10, 11, 20], 3), - ([[10, 11], [11, 12]], [11, 12, 20], 1), - ([[10, 11]], [10, 12, 11], None), - ], -) -def test_token_id_buffer_preserves_tokens_and_only_corrects_matched_finish( - include, last_status, stops, output_ids, match_index -): - async def run(): - manager = _http_manager() - manager.stop_request = AsyncMock() - buffer = StopTokenIdsMatchBuffer(manager, _params(stops, include)) - results = [] - raw_outputs = [] - for index, token_id in enumerate(output_ids): - status = last_status if index == match_index or index == len(output_ids) - 1 else FinishStatus.NO_FINISH - output = _output(token_id, "" if token_id == 10 else str(token_id), index, status) - raw_outputs.append(output) - await buffer.handle(*output) - while buffer.ready_token_list: - results.append(buffer.ready_token_list.popleft()) - expected = raw_outputs if match_index is None else raw_outputs[: match_index + 1] - assert [r[:3] for r in results] == [r[:3] for r in expected] - assert all(result[2] is raw[2] for result, raw in zip(results, expected)) - if match_index is None: - assert results[-1][3].status == last_status - manager.stop_request.assert_not_awaited() - else: - assert results[-1][3].is_stopped() - assert raw_outputs[match_index][3].status == last_status - if last_status == FinishStatus.NO_FINISH: - manager.stop_request.assert_awaited_once_with(80) - else: - manager.stop_request.assert_not_awaited() - - asyncio.run(run()) - - -@pytest.mark.parametrize("last_status", [FinishStatus.FINISHED_ERROR, FinishStatus.FINISHED_LENGTH]) -def test_token_id_stream_preserves_error_and_simulated_final_output(last_status): - manager = _http_manager() - manager.stop_request = AsyncMock() - params = _params([[10, 11]]) - - async def raw_results(): - yield _output(10, "first", 0) - yield _output(11 if last_status == FinishStatus.FINISHED_ERROR else None, "final", 1, last_status) - yield _output(11, "late", 2) - - async def run(): - return [result async for result in match_stop_token_ids_stream(manager, raw_results(), params)] - - results = asyncio.run(run()) - assert [r[1] for r in results] == ["first", "final"] - assert results[-1][3].status == last_status - manager.stop_request.assert_not_awaited() - - -def test_token_id_stream_keeps_interleaved_choices_independent(): - manager = _http_manager() - manager.stop_request = AsyncMock() - params = _params([[10, 11]]) - params.n = params.best_of = 2 - - async def raw_results(): - for req_id, token_id in [(80, 10), (81, 11), (80, 11), (80, 20), (81, 10), (81, 11), (81, 20)]: - yield req_id, str(token_id), {"id": token_id}, FinishStatus() - - async def run(): - return [result async for result in match_stop_token_ids_stream(manager, raw_results(), params)] - - results = asyncio.run(run()) - assert [r[2]["id"] for r in results if r[0] == 80] == [10, 11] - assert [r[2]["id"] for r in results if r[0] == 81] == [11, 10, 11] - assert [r[0] for r in results if r[3].is_stopped()] == [80, 81] - assert [call.args for call in manager.stop_request.await_args_list] == [(80,), (81,)] - - -def test_token_id_stream_does_not_match_prompt_ids(): - manager = _http_manager() - manager.stop_request = AsyncMock() - params = _params([[10, 11]]) - - async def raw_results(): - yield 80, "answer", {"id": 11, "prompt_token_ids": [10]}, FinishStatus(FinishStatus.FINISHED_LENGTH) - - async def run(): - return [result async for result in wrap_stop_match_stream(manager, raw_results(), params, [10])] - - results = asyncio.run(run()) - assert results[0][3].is_finished_length() - manager.stop_request.assert_not_awaited() - - -def test_token_id_matching_runs_before_string_matching(): - manager = _http_manager() - manager.stop_request = AsyncMock() - params = _params(["END", [10, 11]]) - - async def raw_results(): - yield _output(10, "helloE", 0) - yield _output(11, "N", 1) - yield _output(12, "Dextra", 2, FinishStatus.FINISHED_LENGTH) - - async def run(): - # prompt 的 ID 即使是 stop 前缀，也不进入 token ID 匹配历史。 - return [result async for result in wrap_stop_match_stream(manager, raw_results(), params, [10])] - - results = asyncio.run(run()) - assert "".join(r[1] for r in results) == "helloEN" - assert [r[2]["id"] for r in results] == [10, 11] - assert results[-1][3].is_stopped() - manager.stop_request.assert_awaited_once_with(80) - - -@pytest.mark.parametrize("with_stop", [None, "END", [[10, 11]], ["END", [10, 11]]]) -def test_stop_stream_close_closes_raw_generator(with_stop): - params = _params(stop=with_stop) - closed = [] - - async def raw_results(): - try: - yield 80, "hello", {"id": 10}, FinishStatus() - pytest.fail("Closing the wrapper must not consume another raw token") - finally: - closed.append(True) - - async def run(): - manager = SimpleNamespace( - args=SimpleNamespace(enable_rl=False, enable_return_routed_experts=False, reasoning_parser=None), - tokenizer=_ReasoningTokenizer(), - stop_request=AsyncMock(), - ) - raw_generator = raw_results() - output = wrap_stop_match_stream(manager, raw_generator, params, "prompt") - assert (output is raw_generator) is (not with_stop) - result = await output.__anext__() - assert result[1] == "hello" - await output.aclose() - - asyncio.run(run()) - assert closed == [True] - - -@pytest.mark.parametrize("node_finished", [False, True]) -def test_stop_stream_calls_manager_without_mutating_raw_finish(node_finished): - params = _params() - raw_finish = FinishStatus(FinishStatus.FINISHED_LENGTH if node_finished else FinishStatus.NO_FINISH) - manager = _http_manager() - manager.stop_request = AsyncMock() - - async def raw_results(): - yield 80, "helloENDextra", {"id": 10}, raw_finish - yield 80, "late", {"id": 11}, FinishStatus(FinishStatus.FINISHED_LENGTH) - - async def run(): - return [result async for result in match_stop_str_stream(manager, raw_results(), params, "prompt")] - - results = asyncio.run(run()) - assert results[0][1] == "hello" - assert results[0][3].is_stopped() - assert raw_finish.status == (FinishStatus.FINISHED_LENGTH if node_finished else FinishStatus.NO_FINISH) - if node_finished: - manager.stop_request.assert_not_awaited() - else: - manager.stop_request.assert_awaited_once_with(80) - - -@pytest.mark.parametrize("include", [False, True]) -def test_http_stream_discards_late_tokens_in_same_full_queue_batch(include): - async def run(): - manager = _http_manager() - params = _params(include=include) - req = _req(list(range(10, 20)), params=params) - index = 0 - while not req.out_tokens_queue.is_full(): - chunks = ["helloE", "N", "Dextra"] - req.out_tokens_queue.push(chunks[index] if index < len(chunks) else "late", 2 + index, False, index + 1) - index += 1 - - async def raw_results(): - for result in await _read_http_tokens(manager, req): - yield result - - results = [result async for result in _match_http_stream(manager, raw_results(), params, [1, 2], [req])] - assert req.out_tokens_queue.is_empty() - assert req.is_aborted - assert "".join(r[1] for r in results) == "hello" + ("END" if include else "") - assert [r[2]["id"] for r in results] == [10, 11, 12] - assert results[-1][3].is_stopped() - - asyncio.run(run()) - - -def test_stop_stream_uses_prompt_tail_for_all_choices_and_ignores_metadata(): - params = _params(force=0) - params.n = params.best_of = 2 - params.stop_str_match_mode.set_mode(StopStrMatchMode.CONTENT_ONLY) - manager = _http_manager() - manager.args.reasoning_parser = "qwen3" - manager.stop_request = AsyncMock() - manager.tokenizer.encode = Mock(wraps=manager.tokenizer.encode) - - async def raw_results(): - yield 80, "END", {"id": 10, "prompt_token_ids": [1, 1001]}, FinishStatus() - yield 81, "END", {"id": 10}, FinishStatus() - for req_id in (80, 81): - yield req_id, "", {"id": 1001}, FinishStatus() - yield req_id, "answerENDextra", {"id": 11}, FinishStatus() - - async def run(): - return [result async for result in match_stop_str_stream(manager, raw_results(), params, [1] * 4096 + [1000])] - - results = asyncio.run(run()) - for req_id in (80, 81): - assert "".join(result[1] for result in results if result[0] == req_id) == "ENDanswer" - assert [call.args for call in manager.stop_request.await_args_list] == [(80,), (81,)] - assert [call.args for call in manager.tokenizer.encode.call_args_list] == [("",), ("",)] * 2 - assert results[0][2]["prompt_token_ids"] == [1, 1001] - - -@pytest.mark.parametrize("prompt", ["prompt", "x" * 4096 + "prompt", [1, 2], [1] * 4096 + [2]]) -@pytest.mark.parametrize("prompt_metadata", [{}, {"prompt_token_ids": None}, {"prompt_token_ids": [1000]}]) -@pytest.mark.parametrize( - "mode, parser", - [ - (StopStrMatchMode.ALL_TEXT, "qwen3"), - (StopStrMatchMode.CONTENT_ONLY, None), - (StopStrMatchMode.CONTENT_ONLY, "qwen3"), - ], -) -def test_stop_stream_encodes_prompt_tail_independent_of_mode_and_metadata(prompt, prompt_metadata, mode, parser): - params = _params() - params.stop_str_match_mode.set_mode(mode) - manager = _http_manager() - manager.args.reasoning_parser = parser - manager.stop_request = AsyncMock() - manager.tokenizer.encode = Mock(wraps=manager.tokenizer.encode) - - async def raw_results(): - yield 80, "helloENDextra", {"id": 10, **prompt_metadata}, FinishStatus() - - async def run(): - return [result async for result in match_stop_str_stream(manager, raw_results(), params, prompt)] - - results = asyncio.run(run()) - assert results[0][1] == "hello" - assert results[0][3].is_stopped() - prompt_calls = [ - call for call in manager.tokenizer.encode.call_args_list if call.args[0] not in ("", "") - ] - assert len(prompt_calls) == int(isinstance(prompt, str)) - if isinstance(prompt, str): - assert prompt_calls[0].args == (prompt[-2048:],) - assert prompt_calls[0].kwargs == {"add_special_tokens": False} - - -def test_stop_stream_keeps_reasoning_state_independent_for_interleaved_choices(): - params = _params(force=0) - params.n = params.best_of = 2 - params.stop_str_match_mode.set_mode(StopStrMatchMode.CONTENT_ONLY) - manager = _http_manager() - manager.args.reasoning_parser = "qwen3" - manager.stop_request = AsyncMock() - - async def raw_results(): - for req_id, token_id, text in [ - (80, 1000, ""), - (81, 10, "END"), - (80, 10, "END"), - (81, 11, "late"), - (80, 1001, ""), - (80, 12, "answerENDextra"), - ]: - yield req_id, text, {"id": token_id}, FinishStatus() - - async def run(): - return [result async for result in match_stop_str_stream(manager, raw_results(), params, [])] - - results = asyncio.run(run()) - assert "".join(r[1] for r in results if r[0] == 80) == "ENDanswer" - assert "".join(r[1] for r in results if r[0] == 81) == "" - assert [call.args for call in manager.stop_request.await_args_list] == [(81,), (80,)] - assert params._initial_reasoning_state == 0 - - -@pytest.mark.parametrize("first_finished", [False, True]) -@pytest.mark.parametrize( - "late_status", [FinishStatus.NO_FINISH, FinishStatus.FINISHED_LENGTH, FinishStatus.FINISHED_ERROR] -) -def test_stopped_choice_ignores_late_output_before_and_after_sibling_starts(first_finished, late_status): - params = _params() - params.n = params.best_of = 2 - manager = _http_manager() - manager.stop_request = AsyncMock() - closed = [] - exhausted = [] - - async def raw_results(): - try: - finish = FinishStatus(FinishStatus.FINISHED_LENGTH if first_finished else FinishStatus.NO_FINISH) - yield 80, "first" if first_finished else "firstENDextra", {"id": 10}, finish - yield 80, "lateEND", {"id": 11}, FinishStatus(late_status) - yield 81, "secondE", {"id": 20}, FinishStatus() - yield 80, "laterEND", {"id": 12}, FinishStatus(late_status) - yield 81, "NDextra", {"id": 21}, FinishStatus(FinishStatus.FINISHED_LENGTH) - yield 80, "lastEND", {"id": 13}, FinishStatus(late_status) - yield 81, "lastEND", {"id": 22}, FinishStatus(late_status) - exhausted.append(True) - finally: - closed.append(True) - - async def run(): - return [result async for result in match_stop_str_stream(manager, raw_results(), params, [])] - - results = asyncio.run(run()) - assert [result[0] for result in results] == [80, 81, 81] - assert "".join(result[1] for result in results) == "firstsecond" - assert results[0][3].status == (FinishStatus.FINISHED_LENGTH if first_finished else FinishStatus.FINISHED_STOP) - assert results[-1][3].is_stopped() - if first_finished: - manager.stop_request.assert_not_awaited() - else: - manager.stop_request.assert_awaited_once_with(80) - assert closed == [True] - assert exhausted == [True] - - -def test_application_stop_keeps_hit_metadata_and_discards_backend_finish(): - params = _params() - manager = _http_manager() - manager.stop_request = AsyncMock() - - async def raw_results(): - yield 80, "helloE", {"id": 10, "logprob": -0.1, "count_output_tokens": 1}, FinishStatus() - yield 80, "NDextra", {"id": 11, "logprob": -0.2, "count_output_tokens": 2}, FinishStatus() - yield 80, "late", {"id": 12, "logprob": -0.3, "count_output_tokens": 3}, FinishStatus( - FinishStatus.FINISHED_LENGTH - ) - - async def run(): - return [result async for result in match_stop_str_stream(manager, raw_results(), params, [])] - - results = asyncio.run(run()) - assert "".join(r[1] for r in results) == "hello" - assert results[-1][3].is_stopped() - assert results[-1][2] == { - "id": 11, - "logprob": -0.2, - "count_output_tokens": 2, - } - manager.stop_request.assert_awaited_once_with(80) - - @pytest.mark.parametrize("include", [False, True]) @pytest.mark.parametrize("finished", [False, True]) def test_ready_queue_preserves_unconsumed_tokens_after_stop(include, finished): @@ -1232,32 +310,3 @@ def test_ready_queue_preserves_unconsumed_tokens_after_stop(include, finished): ("N" if finished else "ND" if include else "", payloads[3]), ] assert not ready_queue - - -@pytest.mark.parametrize("matched,finished", [(False, True), (True, False), (True, True)]) -def test_stopped_buffer_ignores_later_tokens(matched, finished): - buffer = _stop_buffer(stop_in_reasoning=True) - finish = FinishStatus.FINISHED_LENGTH if finished else FinishStatus.NO_FINISH - _handle_output(buffer, *_output(10, "helloEND" if matched else "helloE", "first", finish=finish)) - first_finish = buffer.ready_token_list[-1][3] - first_tokens = _drain_ready_tokens(buffer) - assert first_finish.is_stopped() is matched - assert "".join(text for text, _ in first_tokens) == ("hello" if matched else "helloE") - assert buffer.state is buffer.State.STOPPED - for finished in [False, True]: - finish = FinishStatus.FINISHED_LENGTH if finished else FinishStatus.NO_FINISH - _handle_output(buffer, *_output(11, "laterEND", "later", finish=finish)) - ready_tokens = _drain_ready_tokens(buffer) - assert not ready_tokens - assert "".join(text for text, _ in first_tokens) == ("hello" if matched else "helloE") - - -def test_simulated_finish_releases_prefix_and_preserves_payloads(): - buffer = _stop_buffer(stop_in_reasoning=True) - first_payload, finish_payload = object(), object() - _handle_output(buffer, *_output(10, "helloE", first_payload)) - assert not buffer.ready_token_list - _handle_output(buffer, *_output(None, "", finish_payload, finish=FinishStatus.FINISHED_LENGTH)) - ready_tokens = _drain_ready_tokens(buffer) - assert ready_tokens == [("helloE", first_payload), ("", finish_payload)] - assert buffer.state is buffer.State.STOPPED diff --git a/unit_tests/server/core/objs/test_token_ids_validation.py b/unit_tests/server/core/objs/test_token_ids_validation.py deleted file mode 100644 index d848137c6b..0000000000 --- a/unit_tests/server/core/objs/test_token_ids_validation.py +++ /dev/null @@ -1,109 +0,0 @@ -import pytest -from lightllm.server.core.objs.sampling_params import ( - StopSequence, - InvalidTokenIds, - SamplingParams, - _check_and_store_int_token_ids, - STOP_SEQUENCE_MAX_LENGTH, - INVALID_TOKEN_IDS_MAX_LENGTH, -) - - -def test_invalid_token_ids_accepts_valid_ints(): - invalid_ids = InvalidTokenIds() - invalid_ids.initialize([4, 5, 6]) - assert invalid_ids.size == 3 - assert invalid_ids.to_list() == [4, 5, 6] - - -@pytest.mark.parametrize("bad_ids", [[4, "5"], [4, 5.0]]) -def test_invalid_token_ids_rejects_non_int(bad_ids): - invalid_ids = InvalidTokenIds() - with pytest.raises(AssertionError): - invalid_ids.initialize(bad_ids) - - -def test_invalid_token_ids_rejects_too_many(): - invalid_ids = InvalidTokenIds() - with pytest.raises(AssertionError): - invalid_ids.initialize([1] * (INVALID_TOKEN_IDS_MAX_LENGTH + 1)) - - -def test_invalid_token_ids_rejects_negative(): - # A negative id survives the GPU-side vocab-size filter and is then used as a - # pointer offset into the logits tensor by the apply_invalid_token triton kernel. - invalid_ids = InvalidTokenIds() - with pytest.raises(AssertionError, match="invalid token ids"): - invalid_ids.initialize([-5]) - - -@pytest.mark.parametrize("bad_ids", [[2 ** 31], [4, 2 ** 40]]) -def test_invalid_token_ids_rejects_int32_overflow(bad_ids): - # An id >= 2**31 would wrap to a different value in the c_int buffer - # (2**31 -> -2147483648), silently banning the wrong token. - invalid_ids = InvalidTokenIds() - with pytest.raises(AssertionError, match="invalid token ids"): - invalid_ids.initialize(bad_ids) - - -def test_invalid_token_ids_accepts_int32_boundary_ids(): - invalid_ids = InvalidTokenIds() - invalid_ids.initialize([0, 2 ** 31 - 1]) - assert invalid_ids.size == 2 - assert invalid_ids.to_list() == [0, 2 ** 31 - 1] - - -def test_stop_sequence_rejects_non_int(): - seq = StopSequence() - with pytest.raises(AssertionError): - seq.initialize([1, "2"]) - - -@pytest.mark.parametrize("bad_ids", [[-7], [2 ** 31 + 5]]) -def test_stop_sequence_rejects_out_of_range_ids(bad_ids): - seq = StopSequence() - with pytest.raises(AssertionError, match="stop token ids"): - seq.initialize(bad_ids) - - -def test_stop_sequence_accepts_int32_boundary_ids(): - seq = StopSequence() - seq.initialize([0, 2 ** 31 - 1]) - assert seq.to_list() == [0, 2 ** 31 - 1] - - -def test_check_and_store_int_token_ids_returns_size_and_writes_buffer(): - import ctypes - - buf = (ctypes.c_int * 8)() - size = _check_and_store_int_token_ids(buf, [7, 8, 9], 8, "test ids") - assert size == 3 - assert list(buf[:size]) == [7, 8, 9] - - -def test_check_and_store_int_token_ids_rejects_overflow(): - import ctypes - - buf = (ctypes.c_int * 2)() - with pytest.raises(AssertionError): - _check_and_store_int_token_ids(buf, [1, 2, 3], 2, "test ids") - - -@pytest.mark.parametrize("bad_key", ["-3", str(2 ** 31)]) -def test_logit_bias_out_of_range_keys_rejected(bad_key): - # logit_bias keys are int()-converted and stored as invalid_token_ids by - # SamplingParams.init, so out-of-range keys must fail the same range check - # on the real production path, not only when initialize() is called directly. - params = SamplingParams() - with pytest.raises(AssertionError, match="invalid token ids"): - params.init(None, logit_bias={bad_key: 0.5}) - - -def test_logit_bias_boundary_keys_round_trip(): - params = SamplingParams() - params.init(None, logit_bias={"0": 1.0, str(2 ** 31 - 1): 0.5}) - assert params.invalid_token_ids.to_list() == [0, 2 ** 31 - 1] - - -if __name__ == "__main__": - pytest.main([__file__, "-v"]) diff --git a/unit_tests/server/embed_cache/test_multimodal_token_id_range.py b/unit_tests/server/embed_cache/test_multimodal_token_id_range.py deleted file mode 100644 index b47a1f912f..0000000000 --- a/unit_tests/server/embed_cache/test_multimodal_token_id_range.py +++ /dev/null @@ -1,61 +0,0 @@ -from types import SimpleNamespace - -from lightllm.server.embed_cache.impl import naive_memory_cache - - -def test_multimodal_token_id_range_selection_uses_time_seed(monkeypatch): - seed_values = [] - selected_index = 1234 - - monkeypatch.setattr(naive_memory_cache.time, "time", lambda: 123.5) - monkeypatch.setattr(naive_memory_cache.random, "seed", seed_values.append) - monkeypatch.setattr(naive_memory_cache.random, "randint", lambda start, end: selected_index) - - cache = naive_memory_cache.InMemoryCache.__new__(naive_memory_cache.InMemoryCache) - range_index = cache._set_prefill_node_random_token_id_range() - range_size = ((2 ** 63 - 1) - 100000000) // 10000 - - assert seed_values == [123.5] - assert range_index == selected_index - assert cache.token_id_range_start == 100000000 + selected_index * range_size - assert cache.token_id_range_end == cache.token_id_range_start + range_size - - -def test_cache_manager_selects_and_logs_probabilistic_range(monkeypatch): - node_uuid = 0x12345678123456781234567812345678 - expected_index = 4321 - range_size = ((2 ** 63 - 1) - 100000000) // 10000 - expected_start = 100000000 + expected_index * range_size - expected_end = expected_start + range_size - warning_messages = [] - - def set_expected_range(cache): - cache.token_id_range_start = expected_start - cache.token_id_range_end = expected_end - return expected_index - - monkeypatch.setattr(naive_memory_cache, "CpuEmbedCacheClient", lambda **kwargs: object()) - monkeypatch.setattr( - naive_memory_cache.InMemoryCache, - "_set_prefill_node_random_token_id_range", - set_expected_range, - ) - monkeypatch.setattr(naive_memory_cache.logger, "warning", warning_messages.append) - - cache = naive_memory_cache.InMemoryCache( - SimpleNamespace( - cache_capacity=200, - config_server_host=None, - config_server_port=None, - pd_node_id=node_uuid, - run_mode="prefill", - ) - ) - cache._check_and_set_new_id_range(1) - - assert cache.token_id_range_start == expected_start - assert cache.token_id_range_end == expected_end - assert len(warning_messages) == 1 - assert f"slot={expected_index}/9999" in warning_messages[0] - assert "does not guarantee global uniqueness" in warning_messages[0] - assert "0.0100% probability" in warning_messages[0] diff --git a/unit_tests/server/grammar_helpers.py b/unit_tests/server/grammar_helpers.py deleted file mode 100644 index a27c025fb1..0000000000 --- a/unit_tests/server/grammar_helpers.py +++ /dev/null @@ -1,154 +0,0 @@ -"""Shared byte-vocabulary fixtures and CPU references for grammar tests.""" - -import asyncio -import json -from itertools import count -from types import SimpleNamespace -from unittest.mock import patch - -import pytest -import torch - -from lightllm.server.core.objs import ConstraintOutParams, FinishStatus, StartArgs -from lightllm.server.router.model_infer.mode_backend.generic_post_process import ( - _get_post_sample_tensors, - g_pin_mem_manager, -) -from lightllm.utils.envs_utils import get_env_start_args - - -@pytest.fixture(autouse=True) -def reset_args_cache(): - get_env_start_args.cache_clear() - yield - get_env_start_args.cache_clear() - - -def configure(monkeypatch, **kwargs): - monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(vars(StartArgs(**kwargs)))) - get_env_start_args.cache_clear() - - -req_indices = count() - - -class SamplingRequest(SimpleNamespace): - def get_cur_total_len(self): - return self.shm_req.input_len + self.cur_output_len - - -def make_req(guided_reasoning_end=(), **constraints): - constraint_out_params = ConstraintOutParams() - constraint_out_params.initialize( - tokenizer=None, - regular_constraint=constraints.pop("regular_constraint", None) or "", - guided_grammar=constraints.pop("guided_grammar", None) or "", - guided_json=constraints.pop("guided_json", None) or "", - guided_reasoning_end=list(guided_reasoning_end), - ) - params = dict( - invalid_token_ids=[], - shm_param=SimpleNamespace( - constraint_out_params=constraint_out_params, - exponential_decay_length_penalty=SimpleNamespace(to_tuple=lambda: (0, 1.0)), - min_new_tokens=0, - temperature=1.0, - top_p=1.0, - top_k=1, - ), - ) - params.update(constraints) - return SamplingRequest( - req_idx=next(req_indices), - vocab_size=257, - generator=None, - sampling_param=SimpleNamespace(**params), - output_constraint=None, - finish_status=FinishStatus(), - cur_output_len=0, - shm_req=SimpleNamespace( - input_len=1, shm_prompt_ids=SimpleNamespace(arr=[ord("P")]), get_compiled_grammar=lambda: None - ), - ) - - -def init_request(compiler, req): - params = req.sampling_param.shm_param.constraint_out_params - for kind, value in ( - ("grammar", params.guided_grammar), - ("regex", params.regular_constraint), - ("json", params.guided_json), - ): - if value.length: - payload = asyncio.run(compiler.compile(kind, value.to_str())) - req.shm_req.get_compiled_grammar = lambda: payload - req.output_constraint = compiler.output_constraint_state_factory.create_state( - req.shm_req, req.sampling_param - ) - return - - -def make_mask_buffers(reqs, verify_width=1): - capacity = max((req.req_idx for req in reqs), default=0) + 1 - return SimpleNamespace( - vocab_size=257, - req_to_next_token_ids=torch.zeros((capacity, verify_width), dtype=torch.int64), - req_to_next_token_ids_cpu=torch.zeros((capacity, verify_width), dtype=torch.int64), - req_to_bitmask=torch.empty((capacity, verify_width, 9), dtype=torch.int32), - req_to_bitmask_enabled=torch.zeros(capacity, dtype=torch.bool), - ) - - -def prepare_sampling_tensors(reqs, manager, has_output=None): - # Exercise production sampling preparation with CPU parameter tensors; - # CUDA integration tests use the real pinned-memory manager and transfers. - def cpu_tensor(key, data, dtype): - tensor = torch.tensor(data, dtype=dtype) - return SimpleNamespace(cuda=lambda non_blocking: tensor) - - with patch.object(g_pin_mem_manager, "gen_from_list", side_effect=cpu_tensor): - return _get_post_sample_tensors(reqs, manager, has_output) - - -def build_mask(reqs, has_output=None, draft_input_ids=None): - manager = make_mask_buffers(reqs, len(reqs)) - if draft_input_ids is not None: - for index, req in enumerate(reqs): - if index == 0 or reqs[index - 1] is not req: - request_start = index - manager.req_to_next_token_ids_cpu[req.req_idx, index - request_start] = draft_input_ids[index] - return manager if prepare_sampling_tensors(reqs, manager, has_output)[-1] else None - - -def apply_masks(reqs, logits, manager, b_mtp_index=None): - if manager is None: - return - if b_mtp_index is None: - positions = [] - for index, req in enumerate(reqs): - if index == 0 or reqs[index - 1] is not req: - request_start = index - positions.append(index - request_start) - b_mtp_index = torch.tensor(positions, dtype=torch.int32, device=logits.device) - # Independent CPU reference; CUDA integration tests use the production kernel. - token_ids = torch.arange(manager.vocab_size) - for index, req in enumerate(reqs): - if manager.req_to_bitmask_enabled[req.req_idx]: - bitmask = manager.req_to_bitmask[req.req_idx, b_mtp_index[index]] - allowed = ((bitmask[token_ids // 32] >> (token_ids % 32)) & 1).bool() - logits[index, : manager.vocab_size].masked_fill_(~allowed, float("-inf")) - logits[:, manager.vocab_size :].fill_(float("-inf")) - - -def allowed(reqs, has_output=None): - mask = build_mask(reqs, has_output) - logits = torch.zeros(len(reqs), 257) - apply_masks(reqs, logits, mask) - return torch.isfinite(logits) - - -def commit_token(req, token): - req.shm_req.shm_prompt_ids.arr.append(token) - req.cur_output_len += 1 - req.output_constraint.commit(token) - return req.output_constraint.is_terminated() diff --git a/unit_tests/server/httpserver/test_grammar.py b/unit_tests/server/httpserver/test_grammar.py deleted file mode 100644 index 916c0f665c..0000000000 --- a/unit_tests/server/httpserver/test_grammar.py +++ /dev/null @@ -1,340 +0,0 @@ -import asyncio -import json -import pickle -from threading import Event -from types import SimpleNamespace -from unittest.mock import AsyncMock, MagicMock, patch - -import httpx -import pytest -import xgrammar as xgr - -from lightllm.server.core.objs import FinishStatus -from lightllm.server.core.objs.sampling_params import SamplingParams -from unit_tests.server.grammar_helpers import configure, reset_args_cache - - -@pytest.fixture(autouse=True) -def start_args(monkeypatch): - configure(monkeypatch) - - -def params(compiler=None, **kwargs): - result = SamplingParams() - result.init(None, **kwargs) - asyncio.run(result.verify_cpu_async(compiler)) - return result - - -@pytest.fixture -def http_client(compiler, monkeypatch): - from lightllm.server import api_http, api_openai - from lightllm.server.api_lightllm import lightllm_generate, lightllm_generate_stream - - configure(monkeypatch, run_mode="pd_master", disable_delay_response_start=True) - generate = MagicMock(side_effect=AssertionError("Invalid request reached generation")) - manager = SimpleNamespace( - tokenizer=None, - output_grammar_compiler=compiler, - generate=generate, - args=SimpleNamespace(enable_return_routed_experts=False), - ) - monkeypatch.setattr(api_http.g_objs, "httpserver_manager", manager) - monkeypatch.setattr(api_http.g_objs, "g_generate_func", lightllm_generate) - monkeypatch.setattr(api_http.g_objs, "g_generate_stream_func", lightllm_generate_stream) - monkeypatch.setattr(api_http.g_objs, "metric_client", SimpleNamespace(counter_inc=lambda *args: None)) - monkeypatch.setattr(api_openai, "build_prompt", AsyncMock(return_value="test")) - monkeypatch.setattr(api_openai, "_is_force_thinking_mode", lambda request: False) - yield httpx.AsyncClient(transport=httpx.ASGITransport(app=api_http.app), base_url="http://test") - generate.assert_not_called() - - -def assert_bad_request(client, path, body): - async def send(): - async with client: - return await client.post(path, json=body) - - response = asyncio.run(send()) - assert response.status_code == 400, response.text - assert response.json()["error"]["type"] == "BadRequestError" - return response.text - - -@pytest.mark.parametrize("failure", ["syntax", "exception", "timeout"]) -@pytest.mark.parametrize("early_headers", [False, True]) -def test_compilation_failure_returns_400_before_generation(compiler, http_client, monkeypatch, failure, early_headers): - configure(monkeypatch, disable_delay_response_start=early_headers) - if failure == "exception": - monkeypatch.setattr(compiler, "_compile", MagicMock(side_effect=RuntimeError("injected compiler error"))) - elif failure == "timeout": - compiler.timeout = 0 - grammar = "[" if failure == "syntax" else "ab" - body = {"inputs": "test", "parameters": {"regular_constraint": grammar}} - error = assert_bad_request(http_client, "/generate_stream", body) - assert ("timed out" if failure == "timeout" else "Failed to compile") in error - - -@pytest.mark.parametrize("path", ["/generate", "/generate_stream"]) -@pytest.mark.parametrize("early_headers", [False, True]) -@pytest.mark.parametrize("invalid_params", [{"max_new_tokens": 0}, {"do_sample": True, "top_p": 0}]) -def test_invalid_sampling_params_return_400(http_client, monkeypatch, path, early_headers, invalid_params): - configure(monkeypatch, disable_delay_response_start=early_headers) - assert_bad_request(http_client, path, {"inputs": "test", "parameters": invalid_params}) - - -@pytest.mark.parametrize("path", ["/v1/chat/completions", "/v1/completions"]) -@pytest.mark.parametrize("stream", [False, True]) -@pytest.mark.parametrize("failure", ["syntax", "exception", "timeout"]) -@pytest.mark.parametrize("early_headers", [False, True]) -def test_openai_grammar_failures_return_400(compiler, http_client, monkeypatch, path, stream, failure, early_headers): - configure(monkeypatch, disable_delay_response_start=early_headers) - if failure == "exception": - monkeypatch.setattr(compiler, "_compile", MagicMock(side_effect=RuntimeError("injected compiler error"))) - elif failure == "timeout": - compiler.timeout = 0 - schema = {"type": 42} if failure == "syntax" else {"type": "object"} - body = { - "model": "test", - "stream": stream, - "response_format": {"type": "json_schema", "json_schema": {"name": "test", "schema": schema}}, - } - body.update({"messages": []} if "chat" in path else {"prompt": "test"}) - assert_bad_request(http_client, path, body) - - -@pytest.mark.parametrize( - "constraint,text", - [ - ({"guided_grammar": 'root ::= "ab"'}, b"ab"), - ({"regular_constraint": "ab"}, b"ab"), - ({"guided_json": {"const": "ok"}}, b'"ok"'), - ], -) -def test_copied_and_pickled_params_get_artifact_from_compiler_cache(compiler, monkeypatch, constraint, text): - sampling_params = params(compiler=compiler, **constraint) - # Choices and PD continuation transport scalar parameters, then get the - # immutable artifact from the compiler's completed-result cache. - copied = pickle.loads(pickle.dumps(SamplingParams.from_buffer_copy(sampling_params))) - with monkeypatch.context() as patcher: - patcher.setattr(compiler._executor, "submit", MagicMock(side_effect=AssertionError("Cache miss"))) - patcher.setattr(xgr, "GrammarCompiler", MagicMock(side_effect=AssertionError("Backend recompiled"))) - payload = asyncio.run(copied.get_compiled_grammar(compiler)) - assert payload is asyncio.run(sampling_params.get_compiled_grammar(compiler)) - grammar = compiler.output_constraint_state_factory.get_grammar(payload) - matcher = xgr.GrammarMatcher(grammar) - assert all(matcher.accept_token(token) for token in text) - assert matcher.accept_token(256) and matcher.is_terminated() - assert "compiled_grammar" not in sampling_params.to_dict() - assert not hasattr(copied, "compiled_grammar") - copied.max_new_tokens = 2 - assert sampling_params.max_new_tokens != copied.max_new_tokens - - -def test_completed_artifact_is_returned_without_executor_work(compiler, monkeypatch): - async def run(): - payload = await compiler.compile("regex", "ab") - monkeypatch.setattr( - compiler._executor, "submit", MagicMock(side_effect=AssertionError("Cache hit reached executor")) - ) - assert await compiler.compile("regex", "ab") is payload - - asyncio.run(run()) - - -def test_completed_artifact_cache_evicts_least_recently_used(compiler, monkeypatch): - compiler._cache_size = 2 - compile_grammar = MagicMock(wraps=compiler._compile) - monkeypatch.setattr(compiler, "_compile", compile_grammar) - - async def run(): - for grammar in ("a", "b", "a", "c", "a", "b"): - await compiler.compile("regex", grammar) - - asyncio.run(run()) - assert [call.args[1] for call in compile_grammar.call_args_list] == ["a", "b", "c", "b"] - - -def test_failed_compilation_is_retried_before_caching_success(compiler, monkeypatch): - compile_grammar = MagicMock(side_effect=[RuntimeError("temporary failure"), b"compiled artifact"]) - monkeypatch.setattr(compiler, "_compile", compile_grammar) - - async def run(): - with pytest.raises(ValueError, match="temporary failure"): - await compiler.compile("regex", "ab") - payload = await compiler.compile("regex", "ab") - assert payload == b"compiled artifact" - assert await compiler.compile("regex", "ab") is payload - - asyncio.run(run()) - assert compile_grammar.call_count == 2 - - -def test_no_grammar_does_not_use_compiler_and_none_mode_rejects_constraints(): - compiler = MagicMock(side_effect=AssertionError("Ordinary request compiled")) - params(compiler=compiler) - compiler.compile.assert_not_called() - with pytest.raises(ValueError, match="output_constraint_mode xgrammar"): - params(regular_constraint="ab") - - -def test_getter_uses_current_constraints_after_reinitialization(compiler): - sampling_params = SamplingParams() - - async def run(): - sampling_params.init(None, regular_constraint="ab") - await sampling_params.verify_cpu_async(compiler) - first = await sampling_params.get_compiled_grammar(compiler) - sampling_params.init(None, regular_constraint="cd") - assert await sampling_params.get_compiled_grammar(compiler) != first - sampling_params.init(None) - assert await sampling_params.get_compiled_grammar(None) is None - - asyncio.run(run()) - - -@pytest.mark.parametrize("stop_waiter", ["cancel", "timeout"]) -def test_cold_requests_compile_independently_and_cache_only_success(compiler, monkeypatch, stop_waiter): - release = Event() - compile_grammar = compiler._compile - - async def run(): - loop = asyncio.get_running_loop() - started, finished = asyncio.Event(), asyncio.Event() - - def slow_first_compile(*args): - loop.call_soon_threadsafe(started.set) - try: - assert release.wait(5) - return b"late artifact" - finally: - loop.call_soon_threadsafe(finished.set) - - monkeypatch.setattr(compiler, "_compile", slow_first_compile) - compiler.timeout = 0.02 if stop_waiter == "timeout" else 5 - first = asyncio.create_task(compiler.compile("regex", "ab")) - try: - await asyncio.wait_for(started.wait(), timeout=5) - # The same grammar can finish on another worker while the first job - # is blocked. Neither a cache hit nor shared pending work is involved. - second_compile = MagicMock(wraps=compile_grammar) - monkeypatch.setattr(compiler, "_compile", second_compile) - compiler.timeout = 5 - payload = await compiler.compile("regex", "ab") - second_compile.assert_called_once_with("regex", "ab") - - if stop_waiter == "cancel": - first.cancel() - with pytest.raises(asyncio.CancelledError): - await first - else: - with pytest.raises(ValueError, match="timed out"): - await first - finally: - release.set() - - await asyncio.wait_for(finished.wait(), timeout=5) - # The abandoned worker's eventual result must not replace a cached success. - assert await compiler.compile("regex", "ab") is payload - - try: - asyncio.run(run()) - finally: - release.set() - - -@pytest.mark.parametrize("explicit_eos", [None, [255]]) -def test_pd_master_compiler_resolves_eos_before_model_loading(monkeypatch, tmp_path, explicit_eos): - from lightllm.server.core.objs import StartArgs - from lightllm.server.httpserver.grammar import create_output_grammar_compiler - - (tmp_path / "config.json").write_text('{"model_type":"llama","vocab_size":257,"eos_token_id":256}') - args = StartArgs(run_mode="pd_master", model_dir=str(tmp_path), eos_id=explicit_eos) - stop_ids = explicit_eos if explicit_eos is not None else [256] - info = xgr.TokenizerInfo([bytes([i]) for i in range(256)] + [b""], stop_token_ids=stop_ids) - tokenizer = object() - with patch.object(xgr.TokenizerInfo, "from_huggingface", return_value=info) as create_info: - compiler = create_output_grammar_compiler(args, tokenizer) - try: - create_info.assert_called_once_with(tokenizer, vocab_size=257, stop_token_ids=stop_ids) - payload = asyncio.run(compiler.compile("regex", "a")) - matcher = xgr.GrammarMatcher(xgr.CompiledGrammar.deserialize_json(payload.decode(), info)) - assert matcher.accept_token(ord("a")) - assert matcher.accept_token(stop_ids[0]) and matcher.is_terminated() - finally: - compiler.shutdown() - - -def test_dictionary_schema_is_serialized_during_initialization(): - params = SamplingParams() - params.init(None, guided_json={"type": "object"}) - assert json.loads(params.constraint_out_params.guided_json.to_str()) == {"type": "object"} - - -@pytest.mark.parametrize( - "invalid_params", [{"max_new_tokens": 0}, {"do_sample": True, "top_p": 0}, {"best_of": 2, "n": 1}] -) -def test_parameter_validation_precedes_grammar_compilation(invalid_params): - params = SamplingParams() - compiler = SimpleNamespace(compile=AsyncMock()) - with pytest.raises(ValueError): - params.init(None, regular_constraint="ab", **invalid_params) - asyncio.run(params.verify_cpu_async(compiler)) - compiler.compile.assert_not_called() - assert not hasattr(params, "compiled_grammar") - - -@pytest.mark.parametrize( - "constraint", [{"guided_grammar": "bad grammar"}, {"guided_json": "not json"}, {"regular_constraint": "["}] -) -def test_invalid_grammar_is_rejected_by_output_constraint_validation(compiler, constraint): - params = SamplingParams() - params.init(None, **constraint) - with pytest.raises(ValueError): - asyncio.run(params.verify_cpu_async(compiler)) - assert not hasattr(params, "compiled_grammar") - - -@pytest.mark.parametrize("status", [FinishStatus.FINISHED_ERROR, FinishStatus.FINISHED_STOP]) -def test_http_preserves_finish_status_and_inflight_token(status): - import numpy as np - - from lightllm.server.core.objs import FinishStatus, Req - from lightllm.server.httpserver.manager import HttpServerManager, ReqStatus - from lightllm.server.pd_io_struct import NodeRole - - req = Req() - req.request_id = req.group_req_id = 1 - req.input_len = 1 - req.shm_prompt_ids = SimpleNamespace(arr=np.array([ord("P"), ord("!")])) - req.shm_logprobs = SimpleNamespace(arr=np.array([(0.0,), (-0.2,)], dtype=[("logprob", np.float32)])) - req.finish_status.set_status(status) - req.finish_token_index = 1 - req.out_tokens_queue.push("!", 1, False, 1) - req.get_output_logprobs_metadata = MagicMock(return_value={"token": "!"}) - req.merge_final_token_metadata = AsyncMock() - - async def collect(): - manager = HttpServerManager.__new__(HttpServerManager) - manager.args = SimpleNamespace(use_reward_model=False, enable_return_routed_experts=False) - manager.tokenizer = None - manager.is_multinode_tp_slave = False - manager.pd_mode = NodeRole.NORMAL - wakeups = asyncio.Queue() - wakeups.put_nowait(None) - manager.zmq_recv_socket = SimpleNamespace(recv_pyobj=wakeups.get) - manager.recycle_resource_loop = AsyncMock() - status = ReqStatus(1, None, [req], 0) - manager.req_id_to_out_inf = {1: status} - task = asyncio.create_task(manager.handle_loop()) - try: - await asyncio.wait_for(status.event.wait(), timeout=2) - return status.out_token_info_list[0] - finally: - task.cancel() - await asyncio.gather(task, return_exceptions=True) - - _, text, metadata, finish = asyncio.run(collect()) - assert text == "!" and finish.status == status - assert metadata["id"] == ord("!") and metadata["logprob"] == pytest.approx(-0.2) - assert metadata["logprobs"] == {"token": "!"} - req.merge_final_token_metadata.assert_awaited_once() diff --git a/unit_tests/server/httpserver/test_pd_continuation.py b/unit_tests/server/httpserver/test_pd_continuation.py deleted file mode 100644 index bdf18c3fb3..0000000000 --- a/unit_tests/server/httpserver/test_pd_continuation.py +++ /dev/null @@ -1,93 +0,0 @@ -import asyncio -import pickle -from types import SimpleNamespace -from unittest.mock import AsyncMock, Mock - -import pytest -from tokenizers import Tokenizer, models, processors - -from lightllm.models.glm5_next.tokenizer import Glm5NextTokenizer -from lightllm.models.qwen3_5.model import QWen3_5Tokenizer -from lightllm.server.core.objs import SamplingParams -from lightllm.server.httpserver.manager import HttpServerManager -from lightllm.server.httpserver_for_pd_master.manager import HttpServerManagerForPDMaster -from lightllm.server.pd_io_struct import NodeRole - - -@pytest.mark.parametrize("multimodal", [False, True]) -@pytest.mark.parametrize("add_special_tokens", [False, True]) -def test_master_counts_prompt_tokens_and_prefill_preserves_output_ids(multimodal, add_special_tokens): - tokenizer = Tokenizer(models.BPE({"P": 0, "a": 1, "Pa": 2, "": 3}, [("P", "a")])) - tokenizer.post_processor = processors.TemplateProcessing(single=" $A", special_tokens=[("", 3)]) - master = object.__new__(HttpServerManagerForPDMaster) - master.tokenizer = SimpleNamespace( - encode=Mock(side_effect=lambda text, _, **kwargs: tokenizer.encode(text, **kwargs).ids) - ) - params = SamplingParams() - params.init(None, add_special_tokens=add_special_tokens) - resources = SimpleNamespace(images=[], audios=[], verify_resource_limits=lambda: None) - token_count = master.tokens("P", resources, params) - prompt_ids = ([3] if add_special_tokens else []) + [0] - assert token_count == len(prompt_ids) - master.tokenizer.encode.assert_called_once_with("P", None, add_special_tokens=add_special_tokens) - - manager = object.__new__(HttpServerManager) - manager.enable_multimodal = multimodal - manager.pd_mode = NodeRole.P - manager.vocab_size = 4 - manager._alloc_multimodal_resources = AsyncMock() - # P expands media placeholders in the supplied IDs without text tokenization. - manager.tokenizer = SimpleNamespace(encode=Mock(side_effect=lambda ids, *_args, **_kwargs: [2 ** 32] + ids)) - params.constraint_out_params.pd_previous_output_len = 1 - continuation_ids = pickle.loads(pickle.dumps(prompt_ids + [1])) - result = asyncio.run(manager._encode(continuation_ids, resources, params)) - - assert tokenizer.encode("Pa", add_special_tokens=False).ids == [2] - assert result == ([2 ** 32] if multimodal else []) + prompt_ids + [1] - if multimodal: - manager.tokenizer.encode.assert_called_once_with( - continuation_ids, resources, add_special_tokens=add_special_tokens - ) - manager._alloc_multimodal_resources.assert_awaited_once_with(resources, params) - else: - manager.tokenizer.encode.assert_not_called() - manager._alloc_multimodal_resources.assert_not_awaited() - - -@pytest.mark.parametrize("tokenizer_class", [QWen3_5Tokenizer, Glm5NextTokenizer], ids=["qwen3.5", "glm5.3-flash"]) -@pytest.mark.parametrize("image_count", [0, 1, 2]) -def test_multimodal_continuation_expands_images_without_retokenizing(tokenizer_class, image_count): - tokenizer = object.__new__(tokenizer_class) - tokenizer.tokenizer = SimpleNamespace(encode=Mock(return_value=[10, 12, 11] * image_count + [1])) - tokenizer.image_start_id, tokenizer.image_end_id, tokenizer.image_token_id = 10, 11, 12 - tokenizer.get_image_token_length = lambda image: image.token_num - images = [SimpleNamespace(token_id=1000, token_num=2) for _ in range(image_count)] - resources = SimpleNamespace(images=images, audios=[], verify_resource_limits=lambda: None) - params = SamplingParams() - params.init(None) - master = object.__new__(HttpServerManagerForPDMaster) - master.tokenizer = tokenizer - master.args = SimpleNamespace(max_image_token_count=16) - token_count = master.tokens("prompt", resources, params) - prompt_ids = [10, 11] * image_count + [1] - - manager = object.__new__(HttpServerManager) - manager.enable_multimodal = True - manager.pd_mode = NodeRole.P - manager.tokenizer = tokenizer - manager.args = master.args - manager._alloc_multimodal_resources = AsyncMock() - history = [2, 3] - params.constraint_out_params.pd_previous_output_len = len(history) - encoded = asyncio.run(manager._encode(prompt_ids + history, resources, params)) - assert encoded == [10, 1000, 1001, 11] * image_count + [1] + history - assert token_count + len(history) == len(encoded) - assert [image.start_idx for image in images] == [4 * index + 1 for index in range(image_count)] - - # D receives P's expanded IDs and must preserve them, including the history suffix. - manager.pd_mode = NodeRole.D - assert asyncio.run(manager._encode(encoded, resources, params)) == encoded - tokenizer.tokenizer.encode.assert_called_once_with("prompt", add_special_tokens=params.add_special_tokens) - manager._alloc_multimodal_resources.assert_awaited_once_with(resources, params) - assert master.tokens("prompt", resources, params) == token_count - assert manager.tokens("prompt", resources, params) == token_count diff --git a/unit_tests/server/httpserver/test_pd_generate_error.py b/unit_tests/server/httpserver/test_pd_generate_error.py deleted file mode 100644 index acb8abc714..0000000000 --- a/unit_tests/server/httpserver/test_pd_generate_error.py +++ /dev/null @@ -1,506 +0,0 @@ -import asyncio -import pickle -from contextlib import suppress -from types import SimpleNamespace -from unittest.mock import AsyncMock, MagicMock - -import pytest - -from lightllm.server.core.objs import FinishStatus, SamplingParams -from lightllm.server.httpserver.pd_loop import _pd_process_generate -from lightllm.server.httpserver_for_pd_master.manager import ( - HttpServerManagerForPDMaster, - ReqStatus, -) -from lightllm.server.multimodal_params import MultimodalParams -from lightllm.server.pd_io_struct import ObjType -from lightllm.utils.error_utils import InvalidRequestError, PDPrefillNodeStopGenToken, ServerBusyError - - -class _FailingManager: - args = SimpleNamespace(run_mode="prefill") - - async def generate(self, **_kwargs): - raise RuntimeError("prefill failed") - yield - - -class _CancelledManager: - args = SimpleNamespace(run_mode="prefill") - - async def generate(self, **_kwargs): - raise asyncio.CancelledError() - yield - - -class _SuccessfulManager: - args = SimpleNamespace(run_mode="prefill") - - async def generate(self, **_kwargs): - yield 123, "token", {}, FinishStatus(FinishStatus.FINISHED_STOP) - - -class _StopPrefillManager: - args = SimpleNamespace(run_mode="prefill") - - async def generate(self, **_kwargs): - raise PDPrefillNodeStopGenToken(group_request_id=123) - yield - - -class _FatalGenerateError(BaseException): - pass - - -class _FatalManager: - args = SimpleNamespace(run_mode="decode") - - async def generate(self, **_kwargs): - raise _FatalGenerateError("fatal generate failure") - yield - - -class _BusyManager: - args = SimpleNamespace(run_mode="decode") - - async def generate(self, **_kwargs): - raise ServerBusyError("decode node could not allocate a shm_req object within 20 seconds") - yield - - -def test_pd_node_reports_generate_error_to_master(): - async def run(): - sampling_params = SamplingParams() - sampling_params.group_request_id = 123 - websocket = AsyncMock() - - await _pd_process_generate( - manager=_FailingManager(), - prompt="prompt", - sampling_params=sampling_params, - multimodal_params=MagicMock(), - forwarding_queue=MagicMock(), - pd_upload_websocket=websocket, - pd_event=asyncio.Event(), - ) - - websocket.send.assert_awaited_once() - obj = pickle.loads(websocket.send.await_args.args[0]) - assert obj == ( - ObjType.PD_UPLOAD_GENERATE_ERROR, - 123, - "RuntimeError: prefill failed", - ) - - asyncio.run(run()) - - -def test_pd_node_reports_local_request_rejection_to_master(): - async def run(): - sampling_params = SamplingParams() - sampling_params.group_request_id = 123 - websocket = AsyncMock() - - await _pd_process_generate( - manager=_BusyManager(), - prompt="prompt", - sampling_params=sampling_params, - multimodal_params=MagicMock(), - forwarding_queue=MagicMock(), - pd_upload_websocket=websocket, - pd_event=asyncio.Event(), - ) - - websocket.send.assert_awaited_once() - obj = pickle.loads(websocket.send.await_args.args[0]) - assert obj == ( - ObjType.PD_UPLOAD_SERVER_BUSY, - 123, - "decode node could not allocate a shm_req object within 20 seconds", - ) - - asyncio.run(run()) - - -def test_pd_node_cancellation_finishes_without_reporting_generate_error(): - async def run(): - sampling_params = SamplingParams() - sampling_params.group_request_id = 123 - websocket = AsyncMock() - - await _pd_process_generate( - manager=_CancelledManager(), - prompt="prompt", - sampling_params=sampling_params, - multimodal_params=MagicMock(), - forwarding_queue=MagicMock(), - pd_upload_websocket=websocket, - pd_event=asyncio.Event(), - ) - - websocket.send.assert_not_awaited() - - asyncio.run(run()) - - -def test_pd_node_success_forwards_token_without_error_report(): - async def run(): - sampling_params = SamplingParams() - sampling_params.group_request_id = 123 - websocket = AsyncMock() - forwarding_queue = AsyncMock() - - await _pd_process_generate( - manager=_SuccessfulManager(), - prompt="prompt", - sampling_params=sampling_params, - multimodal_params=MagicMock(), - forwarding_queue=forwarding_queue, - pd_upload_websocket=websocket, - pd_event=asyncio.Event(), - ) - - forwarding_queue.put.assert_awaited_once() - forwarded = forwarding_queue.put.await_args.args[0] - assert forwarded[0] == 123 - assert forwarded[1] == "token" - assert forwarded[2]["node_mode"] == "prefill" - assert forwarded[3].is_finished() - websocket.send.assert_not_awaited() - - asyncio.run(run()) - - -def test_pd_prefill_full_cache_stop_is_not_reported_as_error(): - async def run(): - sampling_params = SamplingParams() - sampling_params.group_request_id = 123 - websocket = AsyncMock() - - await _pd_process_generate( - manager=_StopPrefillManager(), - prompt="prompt", - sampling_params=sampling_params, - multimodal_params=MagicMock(), - forwarding_queue=AsyncMock(), - pd_upload_websocket=websocket, - pd_event=asyncio.Event(), - ) - - websocket.send.assert_not_awaited() - - asyncio.run(run()) - - -def test_pd_node_reports_non_exception_base_exception(): - async def run(): - sampling_params = SamplingParams() - sampling_params.group_request_id = 123 - websocket = AsyncMock() - - await _pd_process_generate( - manager=_FatalManager(), - prompt="prompt", - sampling_params=sampling_params, - multimodal_params=MagicMock(), - forwarding_queue=AsyncMock(), - pd_upload_websocket=websocket, - pd_event=asyncio.Event(), - ) - - obj = pickle.loads(websocket.send.await_args.args[0]) - assert obj == ( - ObjType.PD_UPLOAD_GENERATE_ERROR, - 123, - "_FatalGenerateError: fatal generate failure", - ) - - asyncio.run(run()) - - -def test_pd_node_error_report_send_failure_is_contained(): - async def run(): - sampling_params = SamplingParams() - sampling_params.group_request_id = 123 - websocket = AsyncMock() - websocket.send.side_effect = ConnectionError("websocket closed") - - await _pd_process_generate( - manager=_FailingManager(), - prompt="prompt", - sampling_params=sampling_params, - multimodal_params=MagicMock(), - forwarding_queue=AsyncMock(), - pd_upload_websocket=websocket, - pd_event=asyncio.Event(), - ) - - websocket.send.assert_awaited_once() - - asyncio.run(run()) - - -def test_pd_master_generate_error_marks_request_and_wakes_all_waiters(): - async def run(): - manager = HttpServerManagerForPDMaster.__new__(HttpServerManagerForPDMaster) - manager.args = SimpleNamespace(config_server_host=None) - manager.pd_manager = MagicMock() - manager.timer_log = AsyncMock() - manager.infos_queues = None - - p_node = SimpleNamespace(websocket=SimpleNamespace(send_bytes=AsyncMock())) - d_node = SimpleNamespace(websocket=SimpleNamespace(send_bytes=AsyncMock())) - req_status = ReqStatus(123, p_node, d_node) - manager.req_id_to_out_inf = {123: req_status} - - handle_task = asyncio.create_task(manager.handle_loop()) - try: - while manager.infos_queues is None: - await asyncio.sleep(0) - await manager.put_to_handle_queue((ObjType.PD_UPLOAD_GENERATE_ERROR, 123, "RuntimeError: prefill failed")) - await asyncio.wait_for(req_status.event.wait(), timeout=1) - - assert req_status.error_info == "RuntimeError: prefill failed" - assert req_status.prefill_prompt_ids_event.is_set() - assert req_status.up_status_event.is_set() - assert manager.req_id_to_out_inf[123] is req_status - p_node.websocket.send_bytes.assert_not_awaited() - d_node.websocket.send_bytes.assert_not_awaited() - - with pytest.raises( - RuntimeError, - match="PD node generate failed: RuntimeError: prefill failed", - ): - req_status.raise_if_error() - finally: - handle_task.cancel() - with suppress(asyncio.CancelledError): - await handle_task - - asyncio.run(run()) - - -@pytest.mark.parametrize( - "terminal_status, output_index", - [ - (FinishStatus.FINISHED_ERROR, 1), - (FinishStatus.FINISHED_ABORTED, 1), - (FinishStatus.FINISHED_LENGTH, 2), - (FinishStatus.FINISHED_STOP, 2), - ], -) -def test_pd_master_first_token_dedup_preserves_terminal_status(terminal_status, output_index): - async def run(): - manager = HttpServerManagerForPDMaster.__new__(HttpServerManagerForPDMaster) - manager.args = SimpleNamespace(pd_node_id=1) - manager.req_id_to_out_inf = {} - p_node = SimpleNamespace(websocket=SimpleNamespace(send_bytes=AsyncMock())) - d_node = SimpleNamespace(websocket=SimpleNamespace(send_bytes=AsyncMock())) - sampling_params = SamplingParams() - sampling_params.group_request_id = 123 - sampling_params.max_new_tokens = 4 - - async def handshake(event, _request, **kwargs): - if kwargs["stage"] == "prefill": - event.prompt_ids = [1, 2] - else: - event.upkv_status = SimpleNamespace(pd_kv_trans_params=pickle.dumps(SimpleNamespace(ready_kv_len=0))) - status = manager.req_id_to_out_inf[123] - status.out_token_info_list.extend( - [ - ( - 123, - "first", - {"count_output_tokens": 1, "node_mode": "prefill"}, - FinishStatus(FinishStatus.FINISHED_LENGTH), - ), - ( - 123, - "duplicate", - {"count_output_tokens": 1, "node_mode": "decode"}, - FinishStatus(FinishStatus.NO_FINISH), - ), - ( - 123, - "terminal", - {"count_output_tokens": output_index, "node_mode": "decode"}, - FinishStatus(terminal_status), - ), - ] - ) - status.event.set() - - manager._wait_for_event_or_disconnect = AsyncMock(side_effect=handshake) - stream = manager.fetch_pd_stream( - p_node, - d_node, - "prompt", - sampling_params, - MultimodalParams(), - SimpleNamespace(is_disconnected=AsyncMock(return_value=False)), - ) - try: - first = await asyncio.wait_for(anext(stream), timeout=1) - assert first[1] == "first" - assert not first[3].is_finished() - terminal = await asyncio.wait_for(anext(stream), timeout=1) - assert terminal[1] == "terminal" - assert terminal[3].status == terminal_status - finally: - await stream.aclose() - - asyncio.run(run()) - - -def test_pd_master_request_rejection_becomes_server_busy_error(): - async def run(): - manager = HttpServerManagerForPDMaster.__new__(HttpServerManagerForPDMaster) - manager.args = SimpleNamespace(config_server_host=None) - manager.pd_manager = MagicMock() - manager.timer_log = AsyncMock() - manager.infos_queues = None - - req_status = ReqStatus(123, MagicMock(), MagicMock()) - manager.req_id_to_out_inf = {123: req_status} - - handle_task = asyncio.create_task(manager.handle_loop()) - try: - while manager.infos_queues is None: - await asyncio.sleep(0) - await manager.put_to_handle_queue( - ( - ObjType.PD_UPLOAD_SERVER_BUSY, - 123, - "prefill node could not allocate a shm_req object within 20 seconds", - ) - ) - await asyncio.wait_for(req_status.event.wait(), timeout=1) - - assert req_status.is_server_busy is True - with pytest.raises(ServerBusyError, match="prefill node could not allocate a shm_req object"): - req_status.raise_if_error() - finally: - handle_task.cancel() - with suppress(asyncio.CancelledError): - await handle_task - - asyncio.run(run()) - - -@pytest.mark.parametrize("event_name", ["prefill_prompt_ids_event", "up_status_event"]) -def test_pd_master_generate_error_wakes_resource_wait(event_name): - async def run(): - manager = HttpServerManagerForPDMaster.__new__(HttpServerManagerForPDMaster) - req_status = ReqStatus(123, MagicMock(), MagicMock()) - request = SimpleNamespace(is_disconnected=AsyncMock(return_value=False)) - wait_task = asyncio.create_task( - manager._wait_for_event_or_disconnect( - getattr(req_status, event_name), - request, - timeout=60, - group_request_id=123, - stage="test", - ) - ) - - await req_status.set_error("RuntimeError: node failed") - await asyncio.wait_for(wait_task, timeout=1) - - with pytest.raises( - RuntimeError, - match="PD node generate failed: RuntimeError: node failed", - ): - req_status.raise_if_error() - - asyncio.run(run()) - - -def test_pd_master_abort_removes_request_even_when_node_notifications_fail(): - async def run(): - manager = HttpServerManagerForPDMaster.__new__(HttpServerManagerForPDMaster) - manager.req_id_to_out_inf = {} - p_node = SimpleNamespace(websocket=SimpleNamespace(send_bytes=AsyncMock(side_effect=ConnectionError("p down")))) - d_node = SimpleNamespace(websocket=SimpleNamespace(send_bytes=AsyncMock(side_effect=ConnectionError("d down")))) - manager.req_id_to_out_inf[123] = ReqStatus(123, p_node, d_node) - - await manager.abort(123) - - assert 123 not in manager.req_id_to_out_inf - - asyncio.run(run()) - - -def test_pd_master_abort_uses_explicit_nodes_when_request_status_is_missing(): - async def run(): - manager = HttpServerManagerForPDMaster.__new__(HttpServerManagerForPDMaster) - manager.req_id_to_out_inf = {} - p_node = SimpleNamespace(websocket=SimpleNamespace(send_bytes=AsyncMock())) - d_node = SimpleNamespace(websocket=SimpleNamespace(send_bytes=AsyncMock())) - - await manager.abort(123, p_node=p_node, d_node=d_node) - - p_node.websocket.send_bytes.assert_awaited_once_with(pickle.dumps((ObjType.ABORT, 123))) - d_node.websocket.send_bytes.assert_awaited_once_with(pickle.dumps((ObjType.ABORT, 123))) - - asyncio.run(run()) - - -def test_pd_invalid_request_error_round_trip(): - class InvalidManager: - args = SimpleNamespace(run_mode="prefill") - - async def generate(self, **_kwargs): - raise InvalidRequestError("invalid image tag num: 0 images vs 1 tags") - yield - - async def run(): - sampling_params = SamplingParams() - sampling_params.group_request_id = 123 - websocket = AsyncMock() - await _pd_process_generate( - manager=InvalidManager(), - prompt="prompt", - sampling_params=sampling_params, - multimodal_params=MagicMock(), - forwarding_queue=AsyncMock(), - pd_upload_websocket=websocket, - pd_event=asyncio.Event(), - ) - kind, req_id, error_info = pickle.loads(websocket.send.await_args.args[0]) - assert kind == ObjType.PD_UPLOAD_GENERATE_ERROR - assert req_id == 123 - status = ReqStatus(req_id, MagicMock(), MagicMock()) - await status.set_error(error_info) - with pytest.raises(InvalidRequestError, match="invalid image tag num: 0 images vs 1 tags"): - status.raise_if_error() - - asyncio.run(run()) - - -def test_pd_stream_validation_returns_bad_request(monkeypatch): - from fastapi import FastAPI - from fastapi.testclient import TestClient - from lightllm.server import api_stream_obj - from lightllm.server.api_http import g_objs, invalid_request_exception_handler - from lightllm.server.api_openai import _safe_stream_wrapper - - monkeypatch.setattr( - api_stream_obj, "get_env_start_args", lambda: SimpleNamespace(disable_delay_response_start=False) - ) - monkeypatch.setattr(g_objs, "metric_client", MagicMock()) - app = FastAPI() - app.add_exception_handler(InvalidRequestError, invalid_request_exception_handler) - - async def generate(): - status = ReqStatus(123, MagicMock(), MagicMock()) - await status.set_error("InvalidRequestError: invalid image tag num: 0 images vs 1 tags") - status.raise_if_error() - yield - - @app.get("/") - async def stream(): - return api_stream_obj.CustomStreamingResponse(_safe_stream_wrapper(generate()), media_type="text/event-stream") - - response = TestClient(app).get("/") - assert response.status_code == 400 - assert "invalid image tag num" in response.text diff --git a/unit_tests/server/httpserver/test_pd_kv_status_channel.py b/unit_tests/server/httpserver/test_pd_kv_status_channel.py deleted file mode 100644 index 9f00bbbf5f..0000000000 --- a/unit_tests/server/httpserver/test_pd_kv_status_channel.py +++ /dev/null @@ -1,109 +0,0 @@ -"""KV status routing through the master owners discovered by P/D nodes.""" - -import asyncio -import pickle -from contextlib import asynccontextmanager -from types import SimpleNamespace -from unittest.mock import AsyncMock, MagicMock - -from lightllm.server.pd_io_struct import PD_Master_Obj, PDUpKVStatus -from lightllm.server.router.model_infer.mode_backend.pd.decode_node_impl import up_status - - -def _manager(): - manager = up_status.UpStatusManager.__new__(up_status.UpStatusManager) - manager.args = SimpleNamespace(pd_node_id=20) - manager.id_to_handle_queue = {} - return manager - - -def _websocket(send): - return SimpleNamespace( - send=send, - recv=AsyncMock(), - transport=SimpleNamespace(get_extra_info=lambda _key: MagicMock()), - ) - - -def test_status_is_sent_directly_without_control_frames_or_reply(monkeypatch): - async def run(): - manager = _manager() - status = PDUpKVStatus(8, 100, b"decode metadata") - manager.id_to_handle_queue[100] = asyncio.Queue() - manager.id_to_handle_queue[100].put_nowait(status) - sent = [] - status_sent = asyncio.Event() - - async def send(data): - sent.append(pickle.loads(data)) - status_sent.set() - - websocket = _websocket(send) - - @asynccontextmanager - async def connect(uri): - assert uri == "ws://master:8100/kv_move_status" - yield websocket - - monkeypatch.setattr(up_status.websockets, "connect", connect) - task = asyncio.create_task(manager.up_kv_status_task(PD_Master_Obj(100, "master:8100"))) - try: - await asyncio.wait_for(status_sent.wait(), timeout=1) - assert sent == [status] - websocket.recv.assert_not_called() - finally: - task.cancel() - await asyncio.gather(task, return_exceptions=True) - - asyncio.run(run()) - - -def test_topology_refresh_connects_all_new_workers_and_retires_previous_ids(monkeypatch): - from lightllm.server.httpserver import pd_loop - - async def run(): - manager = _manager() - responses = asyncio.Queue() - started = asyncio.Queue() - retired = set() - children = [] - original_sleep = asyncio.sleep - - async def discover(_args): - return await responses.get() - - async def dispatch(): - children.append(asyncio.current_task()) - await asyncio.Event().wait() - - async def connect(master): - children.append(asyncio.current_task()) - started.put_nowait(master.node_id) - try: - await asyncio.Event().wait() - finally: - retired.add(master.node_id) - - async def skip_poll_interval(seconds): - assert seconds == 30, "Topology refresh unexpectedly failed" - await original_sleep(0) - - monkeypatch.setattr(pd_loop, "_get_pd_master_objs", discover) - monkeypatch.setattr(manager, "dispatch_task_loop", dispatch) - monkeypatch.setattr(manager, "up_kv_status_task", connect) - monkeypatch.setattr(up_status.asyncio, "sleep", skip_poll_interval) - task = asyncio.create_task(manager.task_loop()) - try: - responses.put_nowait({node_id: PD_Master_Obj(node_id, f"master:{node_id}") for node_id in (100, 101)}) - assert {await asyncio.wait_for(started.get(), timeout=1) for _ in range(2)} == {100, 101} - responses.put_nowait({node_id: PD_Master_Obj(node_id, f"master:{node_id}") for node_id in (200, 201)}) - assert {await asyncio.wait_for(started.get(), timeout=1) for _ in range(2)} == {200, 201} - assert retired == {100, 101} - assert set(manager.id_to_handle_queue) == {200, 201} - finally: - task.cancel() - for child in children: - child.cancel() - await asyncio.gather(task, *children, return_exceptions=True) - - asyncio.run(run()) diff --git a/unit_tests/server/httpserver/test_pd_master_cached_tokens.py b/unit_tests/server/httpserver/test_pd_master_cached_tokens.py deleted file mode 100644 index f71534f9be..0000000000 --- a/unit_tests/server/httpserver/test_pd_master_cached_tokens.py +++ /dev/null @@ -1,769 +0,0 @@ -import asyncio -import copy -import pickle -from types import SimpleNamespace -from unittest.mock import AsyncMock - -import pytest -import numpy as np - -from lightllm.server.core.objs import StopStrMatchMode, FinishStatus, SamplingParams -from lightllm.server.router.model_infer.infer_batch import InferReq -from lightllm.server.httpserver_for_pd_master.manager import HttpServerManagerForPDMaster, PDManager -from lightllm.server.httpserver_for_pd_master.pd_selector import PDSelectionExtraInfo -from lightllm.server.pd_io_struct import PD_Client_Obj - -pytestmark = pytest.mark.usefixtures("single_http_worker_start_args") - - -def _make_manager(monkeypatch): - monkeypatch.setattr( - "lightllm.server.httpserver.manager.HttpServerManager._check_and_repair_length", - classmethod(lambda cls, *a, **k: asyncio.sleep(0)), - ) - monkeypatch.setattr(SamplingParams, "from_buffer_copy", classmethod(lambda cls, other: copy.copy(other))) - mgr = object.__new__(HttpServerManagerForPDMaster) - mgr.active_req_ids = set() - mgr.req_id_to_out_inf = {} - mgr.args = SimpleNamespace( - disable_pd_master_decode_capacity_limit=True, - reasoning_parser=None, - enable_rl=False, - enable_return_routed_experts=False, - select_p_d_node_strategy="cache_aware", - pd_master_mode="elastic", - ) - mgr.enable_pd_node_self_request_limit = True - mgr.pd_node_resource_wait_timeout_seconds = -1 - mgr.pd_node_continuation_resource_wait_timeout_seconds = 60 - mgr.pd_node_busy_retry_timeout_seconds = 120 - mgr.pd_cache_high_priority_max_age_seconds = 60 - mgr.pd_cache_high_priority_min_prompt_tokens = 8192 - mgr.disable_pd_cache_high_priority = False - mgr.running_request_count = 0 - counter = [0] - - def gen_id(): - counter[0] += 1 - return counter[0] - - mgr.metric_client = SimpleNamespace(counter_inc=lambda *a, **k: None, histogram_observe=lambda *a, **k: None) - mgr.tokens = lambda prompt, *a, **k: 10 + len(prompt) - len("hello") - mgr.tokenizer = SimpleNamespace(encode=lambda text, _multimodal=None, **kwargs: [1, 2]) - mgr._log_req_header = lambda *a, **k: asyncio.sleep(0) - mgr.recorded_cache_hit_rates = [] - mgr.inserted_prompt_caches = [] - mgr.segment_inputs = [] - mgr.pd_manager = PDManager(mgr.args) - mgr.pd_node_stats = mgr.pd_manager.selector.pd_node_stats - mgr.id_gen = SimpleNamespace(generate_id=gen_id) - original_record_hit_rate = mgr.pd_manager.selector.record_prompt_cache_hit_rate - prompt_cache_tree = mgr.pd_manager.selector.policy.prompt_cache_tree - original_insert_cache = prompt_cache_tree.insert_sampled - - def record_hit_rate(hit_rate): - mgr.recorded_cache_hit_rates.append(hit_rate) - original_record_hit_rate(hit_rate) - - async def insert_cache(prompt_key, connection_key, sample_stride): - mgr.inserted_prompt_caches.append((prompt_key, connection_key)) - await original_insert_cache(prompt_key, connection_key, sample_stride) - - mgr.pd_manager.selector.record_prompt_cache_hit_rate = record_hit_rate - prompt_cache_tree.insert_sampled = insert_cache - p_node = PD_Client_Obj( - node_id=100, - client_ip_port="prefill:8000", - mode="prefill", - start_args={}, - websocket=SimpleNamespace(send_bytes=AsyncMock()), - ) - d_node = PD_Client_Obj( - node_id=200, - client_ip_port="decode:8000", - mode="decode", - start_args={}, - websocket=SimpleNamespace(send_bytes=AsyncMock()), - ) - mgr.pd_manager.prefill_nodes = [p_node] - mgr.pd_manager.decode_nodes = [d_node] - mgr.pd_manager.url_to_pd_nodes = {node.connection_key: node for node in [p_node, d_node]} - mgr.pd_manager.selector.update_nodes([p_node], [d_node]) - mgr.select_p_d_node = lambda *a, **k: asyncio.sleep(0, result=(p_node, d_node, PDSelectionExtraInfo())) - mgr.remove_req = lambda *a, **k: asyncio.sleep(0) - return mgr - - -def _collect( - mgr, - sampling_params, - monkeypatch, - segments, - first_node_mode="prefill", - needs_prefill_first_token=True, - first_token_text="", - prompt="hello", -): - segment_iter = iter(segments) - - async def fake_wait(p_node, d_node, start_time, prompt_ids, sp, multimodal_params, request): - mgr.segment_inputs.append( - pickle.loads(pickle.dumps((prompt_ids, sp.constraint_out_params.pd_previous_output_len))) - ) - sub_req_id = sp.group_request_id - token_count, final_status = next(segment_iter) - hit = sampling_params.max_new_tokens * 10 - for token_index in range(1, token_count + 1): - finish_status = FinishStatus() - if token_index == token_count: - finish_status = FinishStatus(final_status) - metadata = { - "id": ord("x"), - "prompt_tokens": 100, - "prompt_cache_len": hit if token_index == 1 else 0, - "count_output_tokens": token_index, - "node_mode": first_node_mode if token_index == 1 else "decode", - } - if token_index == 1 and needs_prefill_first_token is not None: - metadata["needs_prefill_first_token"] = needs_prefill_first_token - yield sub_req_id, first_token_text if token_index == 1 else "x", metadata, finish_status - - monkeypatch.setattr(mgr, "_wait_to_token_package", fake_wait) - - async def run(): - out = [] - async for sub_id, out_str, metadata, finish in mgr.generate( - prompt=prompt, - sampling_params=sampling_params, - multimodal_params=SimpleNamespace(images=[], audios=[], verify_and_preload=lambda req: asyncio.sleep(0)), - request=None, - ): - assert "needs_prefill_first_token" not in metadata - out.append(metadata.get("prompt_cache_len", -1)) - return out - - return asyncio.run(run()) - - -def test_single_block_prefill_hit_persists_past_decode_zeros(monkeypatch): - mgr = _make_manager(monkeypatch) - sp = SamplingParams() - sp.init(None) - sp.n = 1 - sp.max_new_tokens = 3 - sp.best_of = 1 - sp.group_request_id = 0 - cached = _collect(mgr, sp, monkeypatch, segments=[(3, FinishStatus.FINISHED_STOP)]) - assert cached and all(c == 30 for c in cached), cached - assert mgr.recorded_cache_hit_rates == [pytest.approx(0.3)] - assert len(mgr.inserted_prompt_caches) == 1 - assert mgr.inserted_prompt_caches[0][0] == "h" - - -@pytest.mark.parametrize("mode", [StopStrMatchMode.ALL_TEXT, StopStrMatchMode.CONTENT_ONLY]) -@pytest.mark.parametrize("hidden_prefix", [False, True]) -@pytest.mark.parametrize("token_chunks", [[[10], [11, 12]], [[10], [11], [12, 13]]]) -def test_pd_token_id_stop_matches_across_capacity_segments(monkeypatch, mode, hidden_prefix, token_chunks): - async def run(): - manager = _make_manager(monkeypatch) - params = SamplingParams() - stop_ids = list(range(10, 10 + len(token_chunks))) - params.init(None, stop_sequences=[stop_ids], stop_str_match_mode=mode, max_new_tokens=6) - segments = [] - - async def raw_segment(p_node, d_node, start_time, prompt, sp, multimodal, request): - index = len(segments) - segments.append(sp.group_request_id) - req = InferReq.__new__(InferReq) - req.stop_sequences = sp.stop_sequences.to_token_ids_list() - req.sampling_param = SimpleNamespace(shm_param=sp) - tokens = token_chunks[index] - req.shm_req = SimpleNamespace(input_len=1, shm_prompt_ids=SimpleNamespace(arr=np.array([1, *tokens]))) - for output_len, token_id in enumerate(tokens, 1): - req.finish_status = FinishStatus() - req.update_finish_status([], output_len) - if index == len(token_chunks) - 1 and output_len == len(tokens) and not req.finish_status.is_finished(): - req.finish_status = FinishStatus(FinishStatus.FINISHED_LENGTH) - text = "" if token_id == 10 and hidden_prefix else str(token_id) - yield sp.group_request_id, text, {"id": token_id, "prompt_tokens": 2}, req.finish_status - if req.finish_status.is_finished(): - return - yield sp.group_request_id, "", {"id": None}, FinishStatus(FinishStatus.FINISHED_PD_DECODE_CAPACITY) - - manager._wait_to_token_package = raw_segment - # 分段内仍由真实 InferReq 判断；跨分段匹配由对外生成流的包装层完成。 - from lightllm.server.httpserver.stop_seq_match import wrap_stop_match_stream - - raw_generator = manager._generate_one_attempt("prompt", params, None, None, 0, 80, 2) - results = [result async for result in wrap_stop_match_stream(manager, raw_generator, params, "prompt")] - assert [result[2]["id"] for result in results] == stop_ids - assert results[-1][3].is_stopped() - assert len(segments) == len(token_chunks) - - asyncio.run(asyncio.wait_for(run(), timeout=2)) - - -def test_dynamic_split_keeps_first_segment_hit(monkeypatch): - mgr = _make_manager(monkeypatch) - sp = SamplingParams() - sp.init(None) - sp.n = 1 - sp.max_new_tokens = 5 - sp.best_of = 1 - sp.group_request_id = 0 - cached = _collect( - mgr, - sp, - monkeypatch, - segments=[ - (3, FinishStatus.FINISHED_PD_DECODE_CAPACITY), - (2, FinishStatus.FINISHED_PD_DECODE_CAPACITY), - (1, FinishStatus.FINISHED_STOP), - ], - ) - assert cached == [50, 50, 50, 50], cached - assert mgr.recorded_cache_hit_rates == [pytest.approx(0.5)] - assert len(mgr.inserted_prompt_caches) == 1 - assert mgr.inserted_prompt_caches[0][0] == "h" - # 无约束请求不需要计算语法状态重放长度。 - assert mgr.segment_inputs == [ - ("hello", 0), - ("hellox", 0), - ("hellox", 0), - ] - - -@pytest.mark.parametrize("add_special_tokens", [False, True]) -@pytest.mark.parametrize("prompt,expected_lengths", [("P", (0, 1)), ("PP", (0, 0))]) -@pytest.mark.parametrize("constraint_enabled", [False, True]) -def test_continuation_uses_reencoded_length_difference( - monkeypatch, add_special_tokens, prompt, expected_lengths, constraint_enabled -): - from tokenizers import Tokenizer, models, processors - from unittest.mock import Mock, call - - tokenizer = Tokenizer(models.BPE({"P": 0, "x": 1, "Px": 2, "": 3, "PPx": 4}, [("P", "x"), ("P", "Px")])) - tokenizer.post_processor = processors.TemplateProcessing(single=" $A", special_tokens=[("", 3)]) - manager = _make_manager(monkeypatch) - del manager.tokens - manager.tokenizer = SimpleNamespace( - encode=Mock(side_effect=lambda text, _multimodal=None, **kwargs: tokenizer.encode(text, **kwargs).ids) - ) - params = SamplingParams() - params.init( - None, - max_new_tokens=4, - add_special_tokens=add_special_tokens, - regular_constraint=".*" if constraint_enabled else "", - ) - _collect( - manager, - params, - monkeypatch, - segments=[ - (2, FinishStatus.FINISHED_PD_DECODE_CAPACITY), - (2, FinishStatus.FINISHED_PD_DECODE_CAPACITY), - (2, FinishStatus.FINISHED_STOP), - ], - first_token_text="x", - prompt=prompt, - ) - # 约束输出按拼接文本的编码长度差重放；普通请求跳过后续编码。 - assert manager.segment_inputs == [ - (prompt, 0), - (prompt + "x", expected_lengths[0] if constraint_enabled else 0), - (prompt + "xx", expected_lengths[1] if constraint_enabled else 0), - ] - expected_encoded_prompts = (prompt, prompt + "x", prompt + "xx") if constraint_enabled else (prompt,) - assert manager.tokenizer.encode.call_args_list == [ - call(text, None, add_special_tokens=add_special_tokens) for text in expected_encoded_prompts - ] - - -def test_error_result_records_hit_rate_without_inserting_prompt_cache(monkeypatch): - mgr = _make_manager(monkeypatch) - sampling_params = SamplingParams() - sampling_params.init(None) - sampling_params.n = 1 - sampling_params.best_of = 1 - sampling_params.max_new_tokens = 1 - - async def failed_wait(_p_node, _d_node, _start_time, _prompt_ids, sp, *_args): - yield ( - sp.group_request_id, - "", - { - "id": None, - "prompt_tokens": 100, - "prompt_cache_len": 20, - "count_output_tokens": 0, - "needs_prefill_first_token": True, - }, - FinishStatus(FinishStatus.FINISHED_ERROR), - ) - - monkeypatch.setattr(mgr, "_wait_to_token_package", failed_wait) - - async def run(): - async for _ in mgr.generate( - prompt="hello", - sampling_params=sampling_params, - multimodal_params=SimpleNamespace( - images=[], - audios=[], - verify_and_preload=lambda req: asyncio.sleep(0), - ), - request=None, - ): - pass - - asyncio.run(run()) - - assert mgr.recorded_cache_hit_rates == [pytest.approx(0.2)] - assert mgr.inserted_prompt_caches == [] - - -def test_pd_raw_stop_does_not_use_configured_reasoning_parser(monkeypatch): - manager = _make_manager(monkeypatch) - manager.args.reasoning_parser = "qwen3" - manager.tokenizer = SimpleNamespace(encode=lambda text, _multimodal=None, **kwargs: [ord(c) for c in text]) - params = SamplingParams() - params.init( - manager.tokenizer, - stop_sequences="END", - max_new_tokens=10, - stop_str_match_mode=StopStrMatchMode.ALL_TEXT, - _initial_reasoning_state=1, - ) - - async def fake_wait(p_node, d_node, start_time, prompt, sp, multimodal_params, request): - assert sp.stop_str_match_mode.get_mode() == (StopStrMatchMode.ALL_TEXT) - metadata = {"prompt_tokens": 10, "prompt_cache_len": 0, "count_output_tokens": 1, "id": 10} - yield sp.group_request_id, "reasonENDextra", metadata, FinishStatus(FinishStatus.FINISHED_LENGTH) - - monkeypatch.setattr(manager, "_wait_to_token_package", fake_wait) - - async def run(): - return [ - result - async for result in manager.generate( - "prompt", - params, - SimpleNamespace(images=[], audios=[], verify_and_preload=lambda req: asyncio.sleep(0)), - None, - ) - ] - - results = asyncio.run(run()) - assert "".join(result[1] for result in results) == "reason" - assert results[-1][3].is_stopped() - - -@pytest.mark.parametrize("in_reasoning", [False, True]) -def test_pd_continuation_uses_local_reasoning_state(monkeypatch, in_reasoning): - - manager = _make_manager(monkeypatch) - manager.args.reasoning_parser = "qwen3" - manager.tokenizer = SimpleNamespace( - encode=lambda text, _multimodal=None, **kwargs: {"": [1000], "": [1001]}.get( - text, [ord(c) for c in text] - ) - ) - params = SamplingParams() - params.init( - manager.tokenizer, - stop_sequences="END", - max_new_tokens=3, - _initial_reasoning_state=1, - stop_str_match_mode=StopStrMatchMode.CONTENT_ONLY, - ) - segment_modes = [] - - async def fake_wait(p_node, d_node, start_time, prompt, sp, multimodal_params, request): - segment_modes.append(sp._initial_reasoning_state) - assert sp.stop_str_match_mode.get_mode() == StopStrMatchMode.CONTENT_ONLY - metadata = {"prompt_tokens": 10, "prompt_cache_len": 0, "count_output_tokens": 1} - if len(segment_modes) == 1: - yield sp.group_request_id, "reason", {**metadata, "id": 10}, FinishStatus() - if not in_reasoning: - yield sp.group_request_id, "", {**metadata, "id": 1001}, FinishStatus() - # A simulated delimiter at the capacity boundary must not change state. - yield ( - sp.group_request_id, - "", - {**metadata, "id": 1001}, - FinishStatus(FinishStatus.FINISHED_PD_DECODE_CAPACITY), - ) - else: - yield sp.group_request_id, "END", {**metadata, "id": 11}, FinishStatus(FinishStatus.FINISHED_STOP) - - monkeypatch.setattr(manager, "_wait_to_token_package", fake_wait) - - async def run(): - return [ - result - async for result in manager.generate( - "hello", - params, - SimpleNamespace(images=[], audios=[], verify_and_preload=lambda req: asyncio.sleep(0)), - None, - ) - ] - - results = asyncio.run(run()) - assert segment_modes == [1, 1] # reasoning 状态只保存在应用侧，不再回写节点参数。 - assert "".join(result[1] for result in results) == "reason" + ("END" if in_reasoning else "") - - -@pytest.mark.parametrize("node_mode", ["decode", None]) -def test_skipped_prefill_does_not_insert_prompt_cache(monkeypatch, node_mode): - mgr = _make_manager(monkeypatch) - sampling_params = SamplingParams() - sampling_params.init(None) - sampling_params.n = sampling_params.best_of = 1 - sampling_params.max_new_tokens = 3 - cached = _collect( - mgr, - sampling_params, - monkeypatch, - segments=[(3, FinishStatus.FINISHED_STOP)], - first_node_mode=node_mode, - needs_prefill_first_token=False, - ) - assert cached == [30, 30, 30] - assert mgr.recorded_cache_hit_rates == [pytest.approx(0.3)] - assert mgr.inserted_prompt_caches == [] - - -@pytest.mark.parametrize("needs_prefill_first_token", [True, None]) -def test_completed_prefill_inserts_cache_when_first_output_is_decode(monkeypatch, needs_prefill_first_token): - mgr = _make_manager(monkeypatch) - sampling_params = SamplingParams() - sampling_params.init(None) - sampling_params.n = sampling_params.best_of = 1 - sampling_params.max_new_tokens = 3 - cached = _collect( - mgr, - sampling_params, - monkeypatch, - segments=[(3, FinishStatus.FINISHED_STOP)], - first_node_mode="decode", - needs_prefill_first_token=needs_prefill_first_token, - ) - assert cached == [30, 30, 30] - assert len(mgr.inserted_prompt_caches) == 1 - assert mgr.inserted_prompt_caches[0][0] == "h" - - -@pytest.mark.parametrize("include", [False, True]) -@pytest.mark.parametrize("parser", [None, "qwen3"]) -@pytest.mark.parametrize( - "continuation,node_finish", - [ - ("NDextra", FinishStatus.NO_FINISH), - ("NDextra", FinishStatus.FINISHED_STOP), - ("NX", FinishStatus.FINISHED_LENGTH), - ("N", FinishStatus.FINISHED_LENGTH), - ], -) -def test_pd_stop_prefix_survives_capacity_split(monkeypatch, include, parser, continuation, node_finish): - manager = _make_manager(monkeypatch) - manager.args.reasoning_parser = parser - manager.tokenizer = SimpleNamespace( - encode=lambda text, _multimodal=None, **kwargs: {"": [1000], "": [1001]}.get( - text, [ord(c) for c in text] - ) - ) - manager.abort = AsyncMock() - stop_request = manager.stop_request - manager.stop_request = AsyncMock(side_effect=stop_request) - params = SamplingParams() - params.init( - manager.tokenizer, - stop_sequences="END", - include_stop_str_in_output=include, - _initial_reasoning_state=int(parser is not None), - stop_str_match_mode=StopStrMatchMode.ALL_TEXT if parser is None else StopStrMatchMode.CONTENT_ONLY, - max_new_tokens=10, - ) - segments = [] - - async def fake_wait(p_node, d_node, start_time, prompt, sp, multimodal_params, request): - segments.append(prompt) - metadata = {"prompt_tokens": 10, "prompt_cache_len": 0, "count_output_tokens": 1} - if len(segments) == 1: - if parser: - yield sp.group_request_id, "END", {**metadata, "id": 10}, FinishStatus() - yield sp.group_request_id, "", {**metadata, "id": 1001}, FinishStatus() - yield sp.group_request_id, "helloE", {**metadata, "id": 11}, FinishStatus() - yield sp.group_request_id, "", dict(metadata), FinishStatus(FinishStatus.FINISHED_PD_DECODE_CAPACITY) - else: - yield sp.group_request_id, continuation, {**metadata, "id": 12}, FinishStatus(node_finish) - if continuation == "NDextra" and node_finish == FinishStatus.NO_FINISH: - yield sp.group_request_id, "", {**metadata, "id": None}, FinishStatus(FinishStatus.FINISHED_ABORTED) - - monkeypatch.setattr(manager, "_wait_to_token_package", fake_wait) - - async def run(): - return [ - result - async for result in manager.generate( - "prompt", - params, - SimpleNamespace(images=[], audios=[], verify_and_preload=lambda req: asyncio.sleep(0)), - None, - ) - ] - - results = asyncio.run(run()) - assert len(segments) == 2 - assert segments == ["prompt", "prompt" + ("END" if parser else "") + "helloE"] - expected = "hello" + ("END" if include else "") if continuation == "NDextra" else "helloE" + continuation - assert "".join(result[1] for result in results) == ("END" if parser else "") + expected - assert [result[2]["id"] for result in results] == ([10, 1001] if parser else []) + [11, 12] - assert results[-1][3].get_finish_reason() == ("stop" if continuation == "NDextra" else "length") - assert all(not result[3].is_finished() for result in results[:-1]) - if continuation == "NDextra" and node_finish == FinishStatus.NO_FINISH: - manager.stop_request.assert_awaited_once() - else: - manager.stop_request.assert_not_awaited() - - -def test_pd_stop_request_marks_only_matching_choice(monkeypatch): - manager = _make_manager(monkeypatch) - manager.active_req_ids.update((80, 81)) - - async def run(): - await manager.stop_request(80) - assert manager.active_req_ids == {81} - await manager.stop_request(80) - await manager.stop_request(999) - assert manager.active_req_ids == {81} - - asyncio.run(run()) - - -def test_pd_abort_removes_output_channel_by_default(monkeypatch): - manager = _make_manager(monkeypatch) - p_node = SimpleNamespace(websocket=SimpleNamespace(send_bytes=AsyncMock())) - d_node = SimpleNamespace(websocket=SimpleNamespace(send_bytes=AsyncMock())) - manager.req_id_to_out_inf = {88: SimpleNamespace(p_node=p_node, d_node=d_node), 96: object()} - - asyncio.run(manager.abort(88)) - - assert set(manager.req_id_to_out_inf) == {96} - for node in (p_node, d_node): - node.websocket.send_bytes.assert_awaited_once() - - -def test_application_stop_discards_late_pd_output_and_closes_segment(monkeypatch): - manager = _make_manager(monkeypatch) - manager.tokenizer = SimpleNamespace(encode=lambda text, _multimodal=None, **kwargs: [ord(c) for c in text]) - params = SamplingParams() - params.init(manager.tokenizer, stop_sequences="END", max_new_tokens=10) - - async def run(): - segments = [] - closed = asyncio.Event() - stopped = asyncio.Event() - manager.abort = AsyncMock() - stop_request = manager.stop_request - - async def stop_and_notify(request_id): - await stop_request(request_id) - stopped.set() - - manager.stop_request = stop_and_notify - - async def raw_segment(p_node, d_node, start_time, prompt, sp, multimodal_params, request): - segments.append(prompt) - metadata = {"prompt_tokens": 10, "prompt_cache_len": 0, "count_output_tokens": 1, "id": 10} - try: - yield sp.group_request_id, "helloENDextra", metadata, FinishStatus() - await stopped.wait() - yield sp.group_request_id, "late", {**metadata, "id": 11}, FinishStatus() - finally: - closed.set() - - manager._wait_to_token_package = raw_segment - results = [ - result - async for result in manager.generate( - "prompt", - params, - SimpleNamespace(images=[], audios=[], verify_and_preload=lambda req: asyncio.sleep(0)), - None, - ) - ] - assert segments == ["prompt"] - assert results[-1][1] == "hello" - assert results[-1][3].is_stopped() - assert not manager.active_req_ids - await asyncio.wait_for(closed.wait(), timeout=1) - assert closed.is_set() - manager.abort.assert_awaited_once() - - asyncio.run(asyncio.wait_for(run(), timeout=2)) - - -@pytest.mark.parametrize( - "exit_mode, token_count", [("finish", 2), ("close", 1), ("error", 1), ("empty", 0), ("empty_prompt", 2)] -) -def test_pd_segment_summary_is_recorded_on_exit(monkeypatch, caplog, exit_mode, token_count): - import math - import time - from unittest.mock import MagicMock - - manager = _make_manager(monkeypatch) - manager.first_time_costs = MagicMock() - manager.per_token_costs = MagicMock() - manager.metric_client = MagicMock() - params = SamplingParams() - params.init(SimpleNamespace(encode=lambda *_args, **_kwargs: []), max_new_tokens=4, group_request_id=80) - request = SimpleNamespace(headers={"X-Request-Id": "summary-test"}, is_disconnected=lambda: asyncio.sleep(0, False)) - caplog.set_level("INFO", logger="lightllm.server.httpserver_for_pd_master.manager") - - async def raw_stream(*_args): - if exit_mode == "empty": - return - metadata = { - "prompt_tokens": 0 if exit_mode == "empty_prompt" else 10, - "prompt_cache_len": 0 if exit_mode == "empty_prompt" else 3, - } - yield 80, "first", dict(metadata), FinishStatus() - if exit_mode == "error": - raise RuntimeError("decode failed") - yield 80, "last", dict(metadata), FinishStatus(FinishStatus.FINISHED_LENGTH) - - manager.fetch_pd_stream = raw_stream - - async def run(): - generator = manager._wait_to_token_package(None, None, time.time(), "prompt", params, None, request) - if exit_mode == "close": - await generator.__anext__() - await generator.aclose() - elif exit_mode == "error": - with pytest.raises(RuntimeError, match="decode failed"): - async for _ in generator: - pass - else: - assert len([result async for result in generator]) == token_count - - asyncio.run(run()) - summaries = [record.getMessage() for record in caplog.records if "out_token_counter:" in record.getMessage()] - if token_count: - assert len(summaries) == 1 - assert "X-Request-Id:summary-test" in summaries[0] - assert f"out_token_counter:{token_count}" in summaries[0] - manager.first_time_costs.add.assert_called_once() - manager.per_token_costs.add.assert_called_once() - assert math.isfinite(manager.per_token_costs.add.call_args.args[0]) - generated = [ - args - for args, _ in manager.metric_client.histogram_observe.call_args_list - if args[0] == "lightllm_request_generated_tokens" - ] - assert generated == [("lightllm_request_generated_tokens", token_count)] - else: - assert not summaries - manager.per_token_costs.add.assert_not_called() - manager.metric_client.histogram_observe.assert_not_called() - - -@pytest.mark.parametrize("late_status", [FinishStatus.NO_FINISH, FinishStatus.FINISHED_PD_DECODE_CAPACITY]) -def test_pd_stopped_choice_records_segment_summary(monkeypatch, caplog, late_status): - import time - from unittest.mock import AsyncMock, MagicMock - - manager = _make_manager(monkeypatch) - manager.first_time_costs = MagicMock() - manager.per_token_costs = MagicMock() - manager.metric_client = MagicMock() - manager.abort = AsyncMock() - params = SamplingParams() - params.init(SimpleNamespace(encode=lambda *_args, **_kwargs: []), max_new_tokens=4) - request = SimpleNamespace(headers={}, is_disconnected=lambda: asyncio.sleep(0, False)) - caplog.set_level("INFO", logger="lightllm.server.httpserver_for_pd_master.manager") - - async def raw_stream(_p_node, _d_node, _prompt, sampling_params, *_args): - metadata = {"id": 1, "prompt_tokens": 10, "prompt_cache_len": 3} - yield sampling_params.group_request_id, "first", dict(metadata), FinishStatus() - yield sampling_params.group_request_id, "late", dict(metadata), FinishStatus(late_status) - - manager.fetch_pd_stream = raw_stream - - async def run(): - generator = manager._generate_one("prompt", params, None, request, time.time(), 80, 10) - assert (await generator.__anext__())[1] == "first" - await manager.stop_request(80) - # stop_request 只做标记，恢复生成时才停止当前 P/D 分段。 - manager.abort.assert_not_awaited() - with pytest.raises(StopAsyncIteration): - await generator.__anext__() - assert not manager.active_req_ids - manager.abort.assert_awaited_once() - summaries = [record.getMessage() for record in caplog.records if "out_token_counter:" in record.getMessage()] - assert len(summaries) == 1 - assert "lightllm_req_id:1" in summaries[0] - manager.metric_client.histogram_observe.assert_any_call("lightllm_request_generated_tokens", 2) - - asyncio.run(run()) - - -def test_application_pd_stop_leaves_sibling_choice_running(monkeypatch): - manager = _make_manager(monkeypatch) - manager.tokenizer = SimpleNamespace(encode=lambda text, _multimodal=None, **kwargs: [ord(c) for c in text]) - params = SamplingParams() - params.init(manager.tokenizer, stop_sequences="END", n=2, best_of=2, max_new_tokens=10) - - async def run(): - stopped = asyncio.Event() - cleaned = asyncio.Event() - segments = [] - manager.abort = AsyncMock(side_effect=lambda *args, **kwargs: stopped.set()) - original_remove_req = manager.remove_req - - async def remove_req(*args, **kwargs): - await original_remove_req(*args, **kwargs) - if not manager.active_req_ids: - cleaned.set() - - manager.remove_req = remove_req - - async def raw_segment(p_node, d_node, start_time, prompt, sp, multimodal_params, request): - index = len(segments) - segments.append(sp.group_request_id) - metadata = {"prompt_tokens": 10, "prompt_cache_len": 0, "count_output_tokens": 1, "id": 10} - if index == 0: - yield sp.group_request_id, "helloENDextra", dict(metadata), FinishStatus() - await asyncio.sleep(0.01) - yield sp.group_request_id, "late", dict(metadata), FinishStatus() - else: - yield sp.group_request_id, "otherE", dict(metadata), FinishStatus() - await stopped.wait() - # 保持 sibling 未结束，验证另一个 choice 停止后仍能继续生成。 - await asyncio.sleep(0.01) - yield sp.group_request_id, "X", {**metadata, "id": 11}, FinishStatus(FinishStatus.FINISHED_LENGTH) - - manager._wait_to_token_package = raw_segment - results = [ - result - async for result in manager.generate( - "prompt", - params, - SimpleNamespace(images=[], audios=[], verify_and_preload=lambda req: asyncio.sleep(0)), - None, - ) - ] - ids = sorted({result[0] for result in results}) - assert len(ids) == 2 - assert "".join(result[1] for result in results if result[0] == ids[0]) == "hello" - assert "".join(result[1] for result in results if result[0] == ids[1]) == "otherEX" - assert {result[0]: result[3].get_finish_reason() for result in results if result[3].is_finished()} == { - ids[0]: "stop", - ids[1]: "length", - } - await asyncio.wait_for(cleaned.wait(), timeout=1) - assert not manager.active_req_ids - - asyncio.run(asyncio.wait_for(run(), timeout=2)) diff --git a/unit_tests/server/httpserver/test_pd_master_discovery.py b/unit_tests/server/httpserver/test_pd_master_discovery.py deleted file mode 100644 index 25c4a49163..0000000000 --- a/unit_tests/server/httpserver/test_pd_master_discovery.py +++ /dev/null @@ -1,371 +0,0 @@ -import asyncio -import base64 -import pickle -from types import SimpleNamespace -from unittest.mock import AsyncMock - -import httpx -import pytest - -from lightllm.server.httpserver import pd_loop -from lightllm.server.pd_io_struct import PD_Master_Obj - - -def _registry_payload(masters): - return {"data": base64.b64encode(pickle.dumps(masters)).decode("ascii")} - - -def _install_http(monkeypatch, handler): - original_client = httpx.AsyncClient - options = [] - - def client(**kwargs): - options.append(kwargs) - return original_client(transport=httpx.MockTransport(handler), **kwargs) - - monkeypatch.setattr(pd_loop.httpx, "AsyncClient", client) - monkeypatch.setattr( - pd_loop, "get_shm_port_args", lambda: SimpleNamespace(pd_master_port=8000, config_server_port=9000) - ) - return options - - -@pytest.mark.parametrize("config_host,config_port", [(None, None), ("config", None), (None, 9000)]) -@pytest.mark.parametrize("worker_ports", [[8000], [8100, 8101]]) -def test_without_config_server_reads_worker_ports_before_connecting( - monkeypatch, config_host, config_port, worker_ports -): - requests = [] - - def handler(request): - requests.append(str(request.url)) - if request.url.path == "/get_server_info": - worker_port = request.url.port - return httpx.Response( - 200, - json={ - "run_mode": "pd_master", - "pd_master_multi_ports": worker_ports, - "pd_node_id": 100 + max(0, worker_port - 8100), - "host": "internal-master", - "port": worker_port, - "httpserver_workers": len(worker_ports), - }, - ) - raise AssertionError(f"unexpected request: {request.url}") - - options = _install_http(monkeypatch, handler) - args = SimpleNamespace(config_server_host=config_host, config_server_port=config_port, pd_master_ip="seed") - masters = asyncio.run(pd_loop._get_pd_master_objs(args)) - if len(worker_ports) > 1: - assert masters == { - 8100: PD_Master_Obj(8100, "seed:8100"), - 8101: PD_Master_Obj(8101, "seed:8101"), - } - assert requests == ["http://seed:8000/get_server_info"] - else: - assert masters == {100: PD_Master_Obj(100, "seed:8000")} - assert requests == ["http://seed:8000/get_server_info"] - assert options == [{"trust_env": False}] - - -@pytest.mark.parametrize("status", [404, 405, 500, 503]) -def test_mode_probe_http_errors_do_not_guess_a_single_master(monkeypatch, status): - requests = [] - - def handler(request): - requests.append(str(request.url)) - return httpx.Response(status) - - _install_http(monkeypatch, handler) - sleep = AsyncMock() - monkeypatch.setattr(pd_loop.asyncio, "sleep", sleep) - args = SimpleNamespace(config_server_host=None, config_server_port=None, pd_master_ip="seed") - - assert asyncio.run(pd_loop._get_pd_master_objs(args)) is None - assert requests == ["http://seed:8000/get_server_info"] - sleep.assert_awaited_once_with(10) - - -@pytest.mark.parametrize( - "info", - [ - {"run_mode": "prefill", "pd_master_multi_ports": [8000], "pd_node_id": 0}, - {"run_mode": "pd_master", "pd_node_id": 0}, - {"run_mode": "pd_master", "pd_master_multi_ports": [], "pd_node_id": 0}, - ], -) -def test_invalid_mode_responses_do_not_guess_a_single_master(monkeypatch, info): - requests = [] - - def handler(request): - requests.append(str(request.url)) - return httpx.Response(200, json=info) - - _install_http(monkeypatch, handler) - monkeypatch.setattr(pd_loop.asyncio, "sleep", AsyncMock()) - args = SimpleNamespace(config_server_host=None, config_server_port=None, pd_master_ip="seed") - - assert asyncio.run(pd_loop._get_pd_master_objs(args)) is None - assert requests == ["http://seed:8000/get_server_info"] - - -def test_discovered_multiprocess_master_does_not_query_worker_ports(monkeypatch): - requests = [] - - def handler(request): - requests.append(str(request.url)) - return httpx.Response( - 200, - json={ - "run_mode": "pd_master", - "httpserver_workers": 2, - "pd_master_multi_ports": [8100, 8101], - }, - ) - - _install_http(monkeypatch, handler) - args = SimpleNamespace(config_server_host=None, config_server_port=None, pd_master_ip="seed") - - assert asyncio.run(pd_loop._get_pd_master_objs(args)) == { - 8100: PD_Master_Obj(8100, "seed:8100"), - 8101: PD_Master_Obj(8101, "seed:8101"), - } - assert requests == ["http://seed:8000/get_server_info"] - - -def test_single_worker_discovery_keeps_the_reachable_external_endpoint(monkeypatch): - requests = [] - - def handler(request): - requests.append(str(request.url)) - return httpx.Response( - 200, - json={ - "run_mode": "pd_master", - "httpserver_workers": 1, - "pd_master_multi_ports": [8080], - "pd_node_id": 100, - "port": 8080, - }, - ) - - _install_http(monkeypatch, handler) - monkeypatch.setattr( - pd_loop, "get_shm_port_args", lambda: SimpleNamespace(pd_master_port=28080, config_server_port=9000) - ) - args = SimpleNamespace(config_server_host=None, config_server_port=None, pd_master_ip="master") - - assert asyncio.run(pd_loop._get_pd_master_objs(args)) == {100: PD_Master_Obj(100, "master:28080")} - assert requests == ["http://master:28080/get_server_info"] - - -def test_config_server_discovery_keeps_the_original_registry_protocol_and_uses_http(monkeypatch): - requests = [] - registered = {100: PD_Master_Obj(100, "master:8100"), 101: PD_Master_Obj(101, "master:8101")} - - def handler(request): - requests.append(str(request.url)) - return httpx.Response(200, json=_registry_payload(registered)) - - options = _install_http(monkeypatch, handler) - args = SimpleNamespace(config_server_host="config", config_server_port=9000) - - assert asyncio.run(pd_loop._get_pd_master_objs(args)) == registered - assert requests == ["http://config:9000/registered_objects"] - assert options == [{"trust_env": False}] - - -@pytest.mark.parametrize("status", [404, 405, 400, 401, 500, 503]) -def test_config_server_http_errors_do_not_fall_back_to_a_direct_master(monkeypatch, status): - _install_http(monkeypatch, lambda request: httpx.Response(status)) - args = SimpleNamespace(config_server_host="config", config_server_port=9000, pd_master_ip="seed") - - assert asyncio.run(pd_loop._get_pd_master_objs(args)) is None - - -@pytest.mark.parametrize("failure", ["network", "invalid_pickle", "base_exception"]) -def test_discovery_retries_only_regular_exceptions(monkeypatch, failure): - class DiscoveryFailure(BaseException): - pass - - def handler(request): - if failure == "network": - raise httpx.ConnectError("not ready", request=request) - if failure == "base_exception": - raise DiscoveryFailure("discovery failed") - return httpx.Response(200, json={"data": base64.b64encode(b"invalid-pickle").decode()}) - - _install_http(monkeypatch, handler) - sleep = AsyncMock() - monkeypatch.setattr(pd_loop.asyncio, "sleep", sleep) - args = SimpleNamespace(config_server_host="config", config_server_port=9000, pd_master_ip="seed") - - if failure == "base_exception": - with pytest.raises(DiscoveryFailure): - asyncio.run(pd_loop._get_pd_master_objs(args)) - sleep.assert_not_awaited() - else: - assert asyncio.run(pd_loop._get_pd_master_objs(args)) is None - sleep.assert_awaited_once_with(10) - - -def test_discovery_follows_config_registry_registration_and_removal(monkeypatch): - registered = {} - _install_http(monkeypatch, lambda request: httpx.Response(200, json=_registry_payload(registered))) - args = SimpleNamespace(config_server_host="config", config_server_port=9000) - - assert asyncio.run(pd_loop._get_pd_master_objs(args)) == {} - registered[100] = PD_Master_Obj(100, "master:8100") - registered[101] = PD_Master_Obj(101, "[2001:db8::1]:8101") - assert asyncio.run(pd_loop._get_pd_master_objs(args)) == registered - del registered[100] - assert asyncio.run(pd_loop._get_pd_master_objs(args)) == registered - registered.clear() - assert asyncio.run(pd_loop._get_pd_master_objs(args)) == {} - - -@pytest.mark.parametrize("host", ["::1", "[::1]", "2001:db8::1234"]) -def test_ipv6_direct_master_address_is_bracketed(monkeypatch, host): - _install_http( - monkeypatch, - lambda request: httpx.Response( - 200, - json={ - "run_mode": "pd_master", - "httpserver_workers": 1, - "pd_master_multi_ports": [8000], - "pd_node_id": 0, - "port": 8000, - }, - ), - ) - args = SimpleNamespace(config_server_host=None, config_server_port=None, pd_master_ip=host) - masters = asyncio.run(pd_loop._get_pd_master_objs(args)) - expected_host = host if host.startswith("[") else f"[{host}]" - - assert masters == {0: PD_Master_Obj(0, f"{expected_host}:8000")} - - -def test_ipv6_config_server_address_is_bracketed_for_http(monkeypatch): - requests = [] - - def handler(request): - requests.append(str(request.url)) - return httpx.Response(200, json=_registry_payload({})) - - _install_http(monkeypatch, handler) - args = SimpleNamespace(config_server_host="2001:db8::1", config_server_port=9000) - - assert asyncio.run(pd_loop._get_pd_master_objs(args)) == {} - assert requests == ["http://[2001:db8::1]:9000/registered_objects"] - - -def test_discovery_cancellation_closes_http_client_without_retry(monkeypatch): - async def run(): - started = asyncio.Event() - closed = asyncio.Event() - - class Client: - async def __aenter__(self): - return self - - async def __aexit__(self, *_exc): - closed.set() - - async def get(self, _url): - started.set() - await asyncio.Event().wait() - - monkeypatch.setattr(pd_loop.httpx, "AsyncClient", lambda **_kwargs: Client()) - monkeypatch.setattr(pd_loop, "get_shm_port_args", lambda: SimpleNamespace(config_server_port=9000)) - sleep = AsyncMock() - monkeypatch.setattr(pd_loop.asyncio, "sleep", sleep) - args = SimpleNamespace(config_server_host="config", config_server_port=9000) - task = asyncio.create_task(pd_loop._get_pd_master_objs(args)) - await started.wait() - task.cancel() - with pytest.raises(asyncio.CancelledError): - await task - sleep.assert_not_awaited() - assert closed.is_set() - - asyncio.run(run()) - - -def test_registry_changes_connect_all_owners_and_retire_removed_ids(monkeypatch): - async def run(): - manager = SimpleNamespace(args=SimpleNamespace(host="prefill")) - responses = asyncio.Queue() - discoveries = asyncio.Queue() - started = asyncio.Queue() - children = [] - live = set() - original_sleep = asyncio.sleep - - async def skip_poll_interval(seconds): - assert seconds == 30, "Topology refresh unexpectedly failed" - await original_sleep(0) - - async def discover(_args): - value = await responses.get() - discoveries.put_nowait(value) - return value - - async def connect(_manager, master, _rank): - children.append(asyncio.current_task()) - live.add(master.node_id) - started.put_nowait(master.node_id) - try: - await asyncio.Event().wait() - finally: - live.remove(master.node_id) - - monkeypatch.setattr(pd_loop.asyncio, "sleep", skip_poll_interval) - monkeypatch.setattr(pd_loop, "_get_pd_master_objs", discover) - monkeypatch.setattr(pd_loop, "_pd_handle_task", connect) - task = asyncio.create_task(pd_loop.pd_handle_loop(manager)) - try: - responses.put_nowait({node_id: PD_Master_Obj(node_id, f"master:{node_id}") for node_id in (100, 101)}) - assert {await asyncio.wait_for(started.get(), timeout=1) for _ in range(2)} == {100, 101} - await discoveries.get() - responses.put_nowait(None) - assert await discoveries.get() is None - await original_sleep(0) - assert live == {100, 101} - responses.put_nowait({node_id: PD_Master_Obj(node_id, f"master:{node_id}") for node_id in (101, 200)}) - assert await asyncio.wait_for(started.get(), timeout=1) == 200 - assert live == {101, 200} - await discoveries.get() - responses.put_nowait({}) - assert await discoveries.get() == {} - await original_sleep(0) - assert not live - finally: - task.cancel() - for child in children: - child.cancel() - await asyncio.gather(task, *children, return_exceptions=True) - - asyncio.run(run()) - - -@pytest.mark.parametrize("dp_index,load", [(None, 0.5), (0, 0.2), (1, 0.8)]) -def test_load_report_keeps_its_original_fields_and_formats_ipv6_keys(monkeypatch, dp_index, load): - from lightllm.server import api_http - - monkeypatch.setattr( - api_http, - "g_objs", - SimpleNamespace( - args=SimpleNamespace(dp=2, nnodes=1), - httpserver_manager=SimpleNamespace(host_ip="2001:db8::1"), - shared_token_load=SimpleNamespace(get_dynamic_max_load=lambda rank: [0.2, 0.8][rank]), - ), - ) - monkeypatch.setattr(pd_loop, "get_shm_port_args", lambda: SimpleNamespace(port=8100)) - - assert pd_loop._get_load_info(dp_index) == { - "total_token_usage_rate": load, - "connection_key": f"[2001:db8::1]:8100/dp{dp_index}", - } diff --git a/unit_tests/server/httpserver/test_pd_split_connections.py b/unit_tests/server/httpserver/test_pd_split_connections.py deleted file mode 100644 index e4fe76eafb..0000000000 --- a/unit_tests/server/httpserver/test_pd_split_connections.py +++ /dev/null @@ -1,294 +0,0 @@ -import asyncio -import json -import pickle -from types import SimpleNamespace -from unittest.mock import AsyncMock, MagicMock - -import pytest - -from lightllm.server.api_cli import make_argument_parser -from lightllm.server.core.objs import FinishStatus, SamplingParams, StartArgs -from lightllm.server.httpserver import pd_loop -from lightllm.server.httpserver.manager import HttpServerManager -from lightllm.server.httpserver_for_pd_master.manager import PDManager -from lightllm.server.pd_io_struct import ObjType, PD_Master_Obj - -pytestmark = pytest.mark.usefixtures("single_http_worker_start_args") - - -def _registration(args, dp_index=None, mode="prefill"): - info = { - "node_id": 1, - "client_ip_port": "p:8000" if mode == "prefill" else "d:8000", - "mode": mode, - "start_args": dict(vars(args), dp=8, nnodes=2), - } - if dp_index is not None: - info["dp_index"] = dp_index - return info - - -def test_split_connections_register_select_update_and_reconnect_independently(): - args = StartArgs(select_p_d_node_strategy="cache_aware", pd_master_mode="1p1d") - manager = PDManager(args) - sockets = [object() for _ in range(4)] - for rank, websocket in enumerate(sockets): - manager.register_pd(_registration(args, rank), websocket) - manager.register_pd(_registration(args, mode="decode"), object()) - assert len(manager.prefill_nodes) == 4 - assert manager.is_pd_nodes_ready() - assert {node.to_llm_url() for node in manager.prefill_nodes} == {"http://p:8000/pd_generate_stream"} - - prompt = "x" * 1025 - cached_node = manager.prefill_nodes[3] - asyncio.run(manager.selector.insert_prompt_cache(prompt, cached_node)) - selected, _, extra = asyncio.run(manager.selector.select_p_d_node(prompt, None, None)) - assert selected is cached_node - assert extra.estimated_cache_hit_rate == 1.0 - for rank in range(4): - manager.update_node_load_info({"connection_key": f"p:8000/dp{rank}", "total_token_usage_rate": rank / 10}) - assert [node.run_status.total_token_usage_rate for node in manager.prefill_nodes] == [0, 0.1, 0.2, 0.3] - replacement_socket = object() - replacement_info = dict(_registration(args, 3), connection_id="replacement") - manager.register_pd(replacement_info, replacement_socket) - replacement = manager.url_to_pd_nodes["p:8000/dp3"] - assert replacement.connection_key == cached_node.connection_key - manager.remove_pd(_registration(args, 3)) - assert manager.url_to_pd_nodes["p:8000/dp3"] is replacement - assert replacement in manager.prefill_nodes - assert replacement in manager.selector.prefill_nodes - manager.remove_pd(replacement_info) - assert len(manager.prefill_nodes) == 3 - assert "p:8000/dp3" not in manager.url_to_pd_nodes - - -@pytest.mark.parametrize("mode", ["prefill", "decode"]) -def test_unsplit_reconnection_ignores_old_connection_id_cleanup(mode): - args = StartArgs() - manager = PDManager(args) - info = dict(_registration(args, mode=mode), connection_id="old") - new_info = dict(info, connection_id="new") - old_socket, new_socket = object(), object() - manager.register_pd(info, old_socket) - manager.register_pd(new_info, new_socket) - key = f"{info['client_ip_port']}/dpNone" - replacement = manager.url_to_pd_nodes[key] - manager.remove_pd(info) - assert manager.url_to_pd_nodes[key] is replacement - assert replacement.websocket is new_socket - manager.remove_pd(new_info) - assert key not in manager.url_to_pd_nodes - - -@pytest.mark.parametrize( - "split, dp, nnodes, expected", - [(False, 8, 2, [None]), (True, 8, 2, [0, 1, 2, 3]), (True, 1, 1, [0]), (True, 1, 2, [0])], -) -def test_prefill_opens_local_dp_connections(monkeypatch, split, dp, nnodes, expected): - async def run(): - args = StartArgs(host="p", run_mode="prefill", dp=dp, nnodes=nnodes, use_dp_split_mode_connect_pd_master=split) - manager = HttpServerManager.__new__(HttpServerManager) - manager.args = args - manager.pd_mode = pd_loop.NodeRole.P - manager.is_multinode_tp_slave = False - manager.recycle_resource_loop = AsyncMock() - started = [] - - async def connect(_manager, dp_index=None): - started.append(dp_index) - - async def receive(): - while len(started) < len(expected): - await asyncio.sleep(0) - raise asyncio.CancelledError() - - manager.zmq_recv_socket = SimpleNamespace(recv_pyobj=receive) - timer = AsyncMock() - monkeypatch.setattr(pd_loop, "timer_log", timer) - monkeypatch.setattr(pd_loop, "pd_handle_loop", connect) - with pytest.raises(asyncio.CancelledError): - await manager.handle_loop() - assert started == expected - timer.assert_awaited_once_with(manager) - - asyncio.run(run()) - - -@pytest.mark.parametrize("dp_index", [None, 0, 3]) -def test_pd_handle_loop_passes_its_fixed_rank_to_each_master(monkeypatch, dp_index): - async def run(): - manager = SimpleNamespace(args=StartArgs(host="p", run_mode="prefill")) - started = [] - tasks = [] - masters = {rank: PD_Master_Obj(rank, f"master{rank}:8000") for rank in (1, 2)} - - async def connect(_manager, master, rank): - started.append((master.node_id, rank)) - tasks.append(asyncio.current_task()) - await asyncio.Event().wait() - - async def stop_after_connections(_seconds): - while len(started) < len(masters): - await original_sleep(0) - raise asyncio.CancelledError() - - original_sleep = asyncio.sleep - monkeypatch.setattr(pd_loop, "_get_pd_master_objs", AsyncMock(return_value=masters)) - monkeypatch.setattr(pd_loop, "_pd_handle_task", connect) - monkeypatch.setattr(pd_loop.asyncio, "sleep", stop_after_connections) - try: - with pytest.raises(asyncio.CancelledError): - await pd_loop.pd_handle_loop(manager, dp_index) - assert started == [(1, dp_index), (2, dp_index)] - finally: - for task in tasks: - task.cancel() - await asyncio.gather(*tasks, return_exceptions=True) - - asyncio.run(run()) - - -@pytest.mark.parametrize("dp_index", [None, 0, 3]) -def test_connection_registers_rank_binds_requests_and_reports_rank_load(monkeypatch, dp_index): - from lightllm.server import api_http - - async def run(): - args = StartArgs(run_mode="prefill", dp=8, nnodes=2) - manager = SimpleNamespace(args=args, host_ip="p", pd_mode=SimpleNamespace(value="prefill")) - monkeypatch.setattr( - api_http, - "g_objs", - SimpleNamespace( - args=args, - httpserver_manager=manager, - shared_token_load=SimpleNamespace(get_dynamic_max_load=lambda rank: rank / 10), - ), - ) - monkeypatch.setattr(pd_loop, "get_shm_port_args", lambda: SimpleNamespace(port=8000)) - params = SamplingParams() - params.init(None, max_new_tokens=1) - params.group_request_id = 123 - params.suggested_dp_index = -1 - dispatched = asyncio.Event() - payloads = [] - - async def generate(**kwargs): - req_params = kwargs["sampling_params"] - assert req_params.suggested_dp_index == (-1 if dp_index is None else dp_index) - assert kwargs["pd_upload_websocket"] is websocket - await kwargs["forwarding_queue"].put((123, "P", {}, FinishStatus())) - dispatched.set() - - requests = 0 - - async def receive(): - nonlocal requests - requests += 1 - if requests == 1: - return pickle.dumps((ObjType.REQ, ("prompt", params, None))) - await dispatched.wait() - while not any(isinstance(payload, tuple) and payload[0] == ObjType.TOKEN_PACKS for payload in payloads): - await asyncio.sleep(0) - raise asyncio.CancelledError() - - async def send(payload): - payloads.append(json.loads(payload) if isinstance(payload, str) else pickle.loads(payload)) - - websocket = SimpleNamespace( - send=send, recv=receive, transport=SimpleNamespace(get_extra_info=lambda _key: MagicMock()) - ) - - class Connection: - async def __aenter__(self): - return websocket - - async def __aexit__(self, *_args): - return False - - monkeypatch.setattr(pd_loop.websockets, "connect", lambda *_args, **_kwargs: Connection()) - monkeypatch.setattr(pd_loop, "_pd_process_generate", generate) - await asyncio.wait_for(pd_loop._pd_handle_task(manager, PD_Master_Obj(1, "master:8000"), dp_index), 2) - assert payloads[0].get("dp_index") == dp_index - assert "dp_index" in payloads[0] - for payload in payloads[1:]: - if payload[0] != ObjType.TOKEN_PACKS: - continue - load = payload[-1] - assert "dp_index" not in load - assert load["connection_key"] == f"p:8000/dp{dp_index}" - assert load["total_token_usage_rate"] == pytest.approx(0.15 if dp_index is None else dp_index / 10) - assert args.host != "p" # Registration must not mutate shared startup arguments. - - asyncio.run(run()) - - -@pytest.mark.parametrize("dp_index", [None, 0, 3]) -@pytest.mark.parametrize("with_source", [False, True], ids=["legacy", "connection-source"]) -def test_reported_load_key_matches_registered_connection_without_rewriting(monkeypatch, dp_index, with_source): - from fastapi import WebSocketDisconnect - from lightllm.server import api_http, api_http_pd - - args = StartArgs(dp=8, nnodes=2, pd_node_id=1) - manager = PDManager(args) - info = _registration(args, dp_index) - if with_source: - info["connection_id"] = f"load-source-{dp_index}" - monkeypatch.setattr( - api_http, - "g_objs", - SimpleNamespace( - args=args, - httpserver_manager=SimpleNamespace(host_ip="p"), - shared_token_load=SimpleNamespace(get_dynamic_max_load=lambda _rank: 0.7), - ), - ) - monkeypatch.setattr(pd_loop, "get_shm_port_args", lambda: SimpleNamespace(port=8000)) - load_info = pd_loop._get_load_info(dp_index) - message = (ObjType.TOKEN_PACKS, [], load_info) - messages = iter([pickle.dumps(message)]) - updated_nodes = [] - - async def receive(): - try: - return next(messages) - except StopIteration: - raise WebSocketDisconnect() - - async def handle(obj): - assert obj[2] == load_info - manager.update_node_load_info(obj[2]) - updated_nodes.append(manager.url_to_pd_nodes[obj[2]["connection_key"]]) - - websocket = SimpleNamespace( - accept=AsyncMock(), - client=("p", 8000), - receive_text=AsyncMock(return_value=json.dumps(info)), - receive_bytes=receive, - ) - monkeypatch.setattr( - api_http, - "g_objs", - SimpleNamespace( - httpserver_manager=SimpleNamespace( - register_pd=AsyncMock(side_effect=manager.register_pd), - remove_pd=AsyncMock(side_effect=manager.remove_pd), - put_to_handle_queue=handle, - ) - ), - ) - asyncio.run(api_http_pd.register_and_keep_alive(websocket)) - assert len(updated_nodes) == 1 - assert updated_nodes[0].dp_index == dp_index - assert updated_nodes[0].websocket is websocket - assert updated_nodes[0].run_status.total_token_usage_rate == 0.7 - - -def test_split_option_defaults_off_and_is_prefill_only(): - from lightllm.server.api_server import launch_server - - parser = make_argument_parser() - assert parser.parse_args([]).use_dp_split_mode_connect_pd_master is False - assert parser.parse_args(["--use_dp_split_mode_connect_pd_master"]).use_dp_split_mode_connect_pd_master is True - assert StartArgs().use_dp_split_mode_connect_pd_master is False - with pytest.raises(ValueError, match="only supports prefill"): - launch_server(StartArgs(run_mode="decode", use_dp_split_mode_connect_pd_master=True)) diff --git a/unit_tests/server/httpserver/test_pd_websocket_lifecycle.py b/unit_tests/server/httpserver/test_pd_websocket_lifecycle.py deleted file mode 100644 index 0c261b5dda..0000000000 --- a/unit_tests/server/httpserver/test_pd_websocket_lifecycle.py +++ /dev/null @@ -1,81 +0,0 @@ -"""Existing PD WebSocket registration and message-forwarding behavior.""" - -import asyncio -import json -import pickle -from types import SimpleNamespace -from unittest.mock import AsyncMock - -import pytest - -from lightllm.server import api_http, api_http_pd - - -def _websocket(**overrides): - values = dict( - accept=AsyncMock(), - close=AsyncMock(), - client=("pd", 18000), - receive_text=AsyncMock(return_value=json.dumps({"node_id": 1, "mode": "prefill"})), - receive_bytes=AsyncMock(), - send_bytes=AsyncMock(), - ) - values.update(overrides) - return SimpleNamespace(**values) - - -def _manager(monkeypatch, **overrides): - values = dict( - register_pd=AsyncMock(), - remove_pd=AsyncMock(), - put_to_handle_queue=AsyncMock(), - args=SimpleNamespace(pd_node_id=100), - update_req_status=AsyncMock(), - ) - values.update(overrides) - manager = SimpleNamespace(**values) - monkeypatch.setattr(api_http, "g_objs", SimpleNamespace(httpserver_manager=manager)) - return manager - - -def test_rejected_registration_does_not_remove_an_existing_connection(monkeypatch): - manager = _manager(monkeypatch, register_pd=AsyncMock(side_effect=ValueError("wrong model parameters"))) - websocket = _websocket() - with pytest.raises(ValueError, match="wrong model parameters"): - asyncio.run(api_http_pd.register_and_keep_alive(websocket)) - manager.remove_pd.assert_not_awaited() - - -def test_heartbeat_wait_is_bounded(monkeypatch): - async def run(): - original_wait_for = asyncio.wait_for - timeouts = [] - - async def wait_forever(): - await asyncio.Event().wait() - - async def short_heartbeat_wait(awaitable, timeout): - timeouts.append(timeout) - return await original_wait_for(awaitable, timeout=0.01) - - manager = _manager(monkeypatch) - monkeypatch.setattr(api_http_pd.asyncio, "wait_for", short_heartbeat_wait) - websocket = _websocket(receive_bytes=wait_forever) - await original_wait_for(api_http_pd.register_and_keep_alive(websocket), timeout=1) - assert timeouts == [30] - assert websocket.close.await_args.kwargs["code"] == 1011 - manager.remove_pd.assert_awaited_once() - - asyncio.run(run()) - - -def test_kv_status_is_forwarded_without_registration_or_reply(monkeypatch): - from fastapi import WebSocketDisconnect - from lightllm.server.pd_io_struct import PDUpKVStatus - - status = PDUpKVStatus(8, 100, b"metadata") - manager = _manager(monkeypatch) - websocket = _websocket(receive_bytes=AsyncMock(side_effect=[pickle.dumps(status), WebSocketDisconnect()])) - asyncio.run(api_http_pd.kv_move_status(websocket)) - manager.update_req_status.assert_awaited_once_with(status) - websocket.send_bytes.assert_not_awaited() diff --git a/unit_tests/server/httpserver/test_running_request_lifecycle.py b/unit_tests/server/httpserver/test_running_request_lifecycle.py deleted file mode 100644 index 0487e45281..0000000000 --- a/unit_tests/server/httpserver/test_running_request_lifecycle.py +++ /dev/null @@ -1,565 +0,0 @@ -import asyncio -import pickle -from types import SimpleNamespace -from unittest.mock import AsyncMock, MagicMock, patch - -import pytest - -from lightllm.server.core.objs import StopStrMatchMode, FinishStatus, SamplingParams -from lightllm.server.httpserver.manager import HttpServerManager, ReqStatus -from lightllm.server.pd_io_struct import NodeRole, ObjType -from lightllm.utils.error_utils import ClientDisconnected, PDPrefillNodeStopGenToken, ServerBusyError - - -class _ValueMark: - def __init__(self, value=0): - self.value = value - - def get_value(self): - return self.value - - def set_value(self, value): - self.value = value - - -def _make_manager(mode: NodeRole): - manager = HttpServerManager.__new__(HttpServerManager) - manager.args = SimpleNamespace( - run_mode=mode.value, - running_max_req_size=2, - reasoning_parser=None, - enable_rl=False, - enable_return_routed_experts=False, - ) - manager.pd_mode = mode - manager.is_multinode_tp_slave = False - manager.alloc_req_id = MagicMock(return_value=123) - manager.is_multinode_tp_master = False - manager.rl_controller = None - manager.output_grammar_compiler = None - manager.req_id_to_out_inf = {} - manager._log_stage_timing = MagicMock() - manager._log_req_header = AsyncMock() - manager._encode = AsyncMock(return_value=[10, 11, 12]) - manager._check_and_repair_length = AsyncMock(side_effect=lambda prompt_ids, _params: prompt_ids) - manager._release_multimodal_resources = AsyncMock() - manager.abort = AsyncMock() - manager._register_running_request = AsyncMock() - manager._unregister_running_request = AsyncMock() - manager.metric_client = MagicMock() - manager._run_reqs_count_lock = asyncio.Lock() - manager.run_reqs_count_mark = _ValueMark() - manager.shm_req_manager = SimpleNamespace( - async_alloc_req_indexes=AsyncMock(side_effect=RuntimeError("alloc failed")) - ) - return manager - - -def _sampling_params(): - sampling_params = SamplingParams() - sampling_params.group_request_id = 123 - sampling_params.n = 1 - sampling_params.max_new_tokens = 1 - sampling_params.pd_node_resource_wait_timeout_seconds = 20 - return sampling_params - - -def _multimodal_params(): - return SimpleNamespace(audios=[], images=[], verify_and_preload=AsyncMock()) - - -@pytest.mark.parametrize("mode", [NodeRole.NORMAL, NodeRole.P, NodeRole.D]) -@pytest.mark.parametrize("warm_cache", [False, True]) -def test_http_publishes_grammar_for_all_choices_before_forwarding(compiler, mode, warm_cache): - import xgrammar as xgr - - async def run(): - manager = _make_manager(mode) - manager.output_grammar_compiler = compiler - manager.tokenizer = None - manager.enable_multimodal = False - manager.args.chunked_prefill_size = 1 - manager._alloc_shm_req_indexes = AsyncMock(return_value=[0, 1]) - reqs = [MagicMock(), MagicMock()] - manager.shm_req_manager.async_get_req_obj_by_index = AsyncMock(side_effect=reqs) - - async def forward(*args): - payloads = [] - for req in reqs: - req.init.assert_called_once() - req.set_compiled_grammar.assert_called_once() - payload = req.set_compiled_grammar.call_args.args[0] - payloads.append(payload) - grammar = compiler.output_constraint_state_factory.get_grammar(payload) - matcher = xgr.GrammarMatcher(grammar) - assert matcher.accept_token(ord("a")) and matcher.accept_token(ord("b")) - assert matcher.accept_token(256) and matcher.is_terminated() - assert payloads[0] is payloads[1] - - manager.transfer_to_next_module_or_node = AsyncMock(side_effect=forward) - - async def results(*args): - for req_id in (123, 124): - yield req_id, "", {"id": 256}, FinishStatus(FinishStatus.FINISHED_STOP) - - manager._wait_to_token_package = results - params = SamplingParams() - params.init(None, regular_constraint="ab", n=2, best_of=2, max_new_tokens=1, group_request_id=123) - with patch.object(compiler, "_compile", wraps=compiler._compile) as compile_grammar: - if warm_cache: - await params.verify_cpu_async(compiler) - params = pickle.loads(pickle.dumps(SamplingParams.from_buffer_copy(params))) - outputs = [result async for result in manager.generate("prompt", params, _multimodal_params(), None)] - assert len(outputs) == 2 - compile_grammar.assert_called_once_with("regex", "ab") - manager.transfer_to_next_module_or_node.assert_awaited_once() - assert not hasattr(params, "compiled_grammar") - - asyncio.run(run()) - - -@pytest.mark.parametrize("mode", [NodeRole.NORMAL, NodeRole.P, NodeRole.D]) -def test_grammar_get_failure_precedes_request_admission_and_pd_handshake(mode): - async def run(): - manager = _make_manager(mode) - manager.output_grammar_compiler = SimpleNamespace(compile=AsyncMock(side_effect=ValueError("compile failed"))) - manager._alloc_shm_req_indexes = AsyncMock() - params = SamplingParams() - params.init(None, regular_constraint="ab", max_new_tokens=1, group_request_id=123) - multimodal_params = _multimodal_params() - websocket = AsyncMock() if mode == NodeRole.P else None - pd_event = asyncio.Event() if mode == NodeRole.P else None - with pytest.raises(ValueError, match="compile failed"): - await _drain_generate(manager, params, multimodal_params, websocket, pd_event) - manager.alloc_req_id.assert_not_called() - manager._encode.assert_not_awaited() - multimodal_params.verify_and_preload.assert_not_awaited() - manager._alloc_shm_req_indexes.assert_not_awaited() - manager._register_running_request.assert_not_awaited() - manager._unregister_running_request.assert_not_awaited() - manager.abort.assert_not_awaited() - if websocket is not None: - websocket.send.assert_not_awaited() - - asyncio.run(run()) - - -@pytest.mark.parametrize("mode", [NodeRole.NORMAL, NodeRole.D, NodeRole.P]) -@pytest.mark.parametrize("include", [False, True]) -@pytest.mark.parametrize("parser", [None, "qwen3"]) -@pytest.mark.parametrize("stop_mode", [None, StopStrMatchMode.ALL_TEXT, StopStrMatchMode.CONTENT_ONLY]) -def test_http_forwards_packaged_tokens_and_preserves_stop_policy(mode, include, parser, stop_mode): - async def run(): - manager = _make_manager(mode) - manager.tokenizer = SimpleNamespace(encode=lambda text, **kwargs: [ord(c) for c in text]) - manager.enable_multimodal = False - manager.args.chunked_prefill_size = 1 - manager.args.reasoning_parser = parser - manager._alloc_shm_req_indexes = AsyncMock(return_value=[0]) - manager.shm_req_manager.async_get_req_obj_by_index = AsyncMock(return_value=MagicMock()) - manager.transfer_to_next_module_or_node = AsyncMock() - - async def results(*args): - yield 123, "hello", {"id": 1}, FinishStatus() - yield 123, "END" if include else "", {"id": 2}, FinishStatus(FinishStatus.FINISHED_STOP) - - manager._wait_to_token_package = results - params = _sampling_params() - params.init( - manager.tokenizer, - stop_sequences="END", - include_stop_str_in_output=include, - max_new_tokens=1, - group_request_id=123, - pd_node_resource_wait_timeout_seconds=20, - ) - # P/D nodes keep the matching policy supplied by Master. - if stop_mode is not None: - params.stop_str_match_mode.set_mode(stop_mode) - websocket = AsyncMock() if mode == NodeRole.P else None - pd_event = None - if mode == NodeRole.P: - pd_event = asyncio.Event() - pd_event.decode_node_info = SimpleNamespace(ready_kv_len=1) - pd_event.set() - - results = [ - result - async for result in manager.generate("prompt", params, _multimodal_params(), None, websocket, pd_event) - ] - expected_mode = StopStrMatchMode.ALL_TEXT if stop_mode is None else stop_mode - assert params.stop_str_match_mode.get_mode() == expected_mode - expected = "helloEND" if include else "hello" - assert "".join(result[1] for result in results) == expected - assert [result[2]["id"] for result in results] == [1, 2] - assert results[0][2]["input_usage"]["input_text_tokens"] == 3 - assert results[-1][3].get_finish_reason() == "stop" - manager.shm_req_manager.async_get_req_obj_by_index.return_value.set_compiled_grammar.assert_not_called() - manager._unregister_running_request.assert_awaited_once() - - asyncio.run(run()) - - -def _req_status(reqs): - return ReqStatus(123, None, reqs, 0) - - -async def _drain_generate(manager, sampling_params, multimodal_params, websocket=None, pd_event=None): - async for _ in manager.generate( - prompt="prompt", - sampling_params=sampling_params, - multimodal_params=multimodal_params, - request=None, - pd_upload_websocket=websocket, - pd_event=pd_event, - ): - pass - - -@pytest.mark.parametrize("mode", [NodeRole.NORMAL, NodeRole.D]) -def test_non_prefill_request_is_counted_before_early_failure_and_always_unregistered( - mode, -): - async def run(): - manager = _make_manager(mode) - manager._encode.side_effect = RuntimeError("encode failed") - - with pytest.raises(RuntimeError, match="encode failed"): - await _drain_generate(manager, _sampling_params(), _multimodal_params()) - - manager._register_running_request.assert_awaited_once() - manager._unregister_running_request.assert_awaited_once() - - asyncio.run(run()) - - -def test_prefill_encode_failure_is_outside_running_request_count(): - async def run(): - manager = _make_manager(NodeRole.P) - manager._encode.side_effect = RuntimeError("encode failed") - - with pytest.raises(RuntimeError, match="encode failed"): - await _drain_generate( - manager, - _sampling_params(), - _multimodal_params(), - AsyncMock(), - asyncio.Event(), - ) - - manager._register_running_request.assert_not_awaited() - manager._unregister_running_request.assert_not_awaited() - - asyncio.run(run()) - - -def test_prefill_full_cache_hit_never_enters_running_request_count(): - async def run(): - manager = _make_manager(NodeRole.P) - websocket = AsyncMock() - pd_event = asyncio.Event() - pd_event.decode_node_info = SimpleNamespace(ready_kv_len=2) - pd_event.set() - - with pytest.raises(PDPrefillNodeStopGenToken): - await _drain_generate(manager, _sampling_params(), _multimodal_params(), websocket, pd_event) - - manager._register_running_request.assert_not_awaited() - manager._unregister_running_request.assert_not_awaited() - websocket.send.assert_awaited_once() - - asyncio.run(run()) - - -def test_prefill_upload_preserves_int64_multimodal_prompt_ids(): - async def run(): - manager = _make_manager(NodeRole.P) - multimodal_token_id = 2 ** 32 - manager._encode.return_value = [10, multimodal_token_id, 12] - websocket = AsyncMock() - pd_event = asyncio.Event() - pd_event.decode_node_info = SimpleNamespace(ready_kv_len=2) - pd_event.set() - - with pytest.raises(PDPrefillNodeStopGenToken): - await _drain_generate(manager, _sampling_params(), _multimodal_params(), websocket, pd_event) - - obj_type, group_request_id, prompt_ids = pickle.loads(websocket.send.await_args.args[0]) - assert obj_type == ObjType.PD_UPLOAD_PREFILL_PROMPT_IDS - assert group_request_id == 123 - assert prompt_ids.itemsize == 8 - assert prompt_ids.tolist() == [10, multimodal_token_id, 12] - - asyncio.run(run()) - - -def test_prefill_waiting_for_decode_assignment_can_be_cancelled_without_touching_count(): - async def run(): - manager = _make_manager(NodeRole.P) - websocket = AsyncMock() - pd_event = asyncio.Event() - task = asyncio.create_task( - _drain_generate(manager, _sampling_params(), _multimodal_params(), websocket, pd_event) - ) - - for _ in range(100): - if websocket.send.await_count: - break - await asyncio.sleep(0) - assert websocket.send.await_count == 1 - assert not task.done() - manager._register_running_request.assert_not_awaited() - - task.cancel() - with pytest.raises(asyncio.CancelledError): - await task - - manager._register_running_request.assert_not_awaited() - manager._unregister_running_request.assert_not_awaited() - - asyncio.run(run()) - - -def test_prefill_is_counted_after_decode_assignment_and_unregistered_on_following_failure(): - async def run(): - manager = _make_manager(NodeRole.P) - websocket = AsyncMock() - pd_event = asyncio.Event() - pd_event.decode_node_info = SimpleNamespace(ready_kv_len=1) - pd_event.set() - - with pytest.raises(RuntimeError, match="alloc failed"): - await _drain_generate(manager, _sampling_params(), _multimodal_params(), websocket, pd_event) - - manager._register_running_request.assert_awaited_once() - manager._unregister_running_request.assert_awaited_once() - - asyncio.run(run()) - - -@pytest.mark.parametrize("mode", [NodeRole.P, NodeRole.D]) -def test_pd_node_returns_busy_when_shm_req_allocation_times_out(mode): - async def run(): - manager = _make_manager(mode) - manager.shm_req_manager = SimpleNamespace( - async_alloc_req_indexes=AsyncMock(return_value=None), - async_release_req_index=AsyncMock(), - ) - - websocket = AsyncMock() if mode == NodeRole.P else None - pd_event = None - if mode == NodeRole.P: - pd_event = asyncio.Event() - pd_event.decode_node_info = SimpleNamespace(ready_kv_len=1) - pd_event.set() - - with patch("lightllm.server.httpserver.manager.time.monotonic", side_effect=[0, 21]): - with pytest.raises(ServerBusyError, match=f"PD {mode.value} node is busy"): - await _drain_generate( - manager, - _sampling_params(), - _multimodal_params(), - websocket, - pd_event, - ) - - manager.shm_req_manager.async_alloc_req_indexes.assert_awaited_once_with(1) - manager.shm_req_manager.async_release_req_index.assert_not_awaited() - manager._register_running_request.assert_awaited_once() - manager._unregister_running_request.assert_awaited_once() - - asyncio.run(run()) - - -@pytest.mark.parametrize("mode", [NodeRole.P, NodeRole.D]) -def test_pd_node_returns_busy_while_first_token_request_waits_in_router(mode): - async def run(): - manager = _make_manager(mode) - req = SimpleNamespace( - request_id=123, - is_aborted=False, - router_arrival_time=0, - infer_start_time=0, - sample_params=SimpleNamespace( - high_priority_request=False, - pd_node_resource_wait_timeout_seconds=20, - ), - ) - req_status = _req_status([req]) - req_status.event.set() - sampling_params = _sampling_params() - - with patch("lightllm.server.httpserver.manager.time.monotonic", return_value=21): - req.router_arrival_time = 1.0 - output_generator = manager._wait_to_token_package( - start_time=0, - prompt_ids=[], - group_request_id=123, - sampling_params=sampling_params, - req_status=req_status, - request=None, - ) - with pytest.raises(ServerBusyError, match="request did not enter inference"): - await output_generator.__anext__() - - asyncio.run(run()) - - -def test_multinode_tp_slave_does_not_apply_router_wait_timeout(): - async def run(): - manager = _make_manager(NodeRole.D) - manager.is_multinode_tp_slave = True - req = SimpleNamespace( - router_arrival_time=1.0, - infer_start_time=0.0, - sample_params=SimpleNamespace(pd_node_resource_wait_timeout_seconds=0), - ) - req_status = _req_status([req]) - req_status.event.set() - request = SimpleNamespace(is_disconnected=AsyncMock(return_value=True)) - - with patch("lightllm.server.httpserver.manager.time.monotonic", return_value=2): - # ReqStatus 只判断时间条件；是否允许 slave 执行超时策略由 manager 在调用处决定。 - assert req_status.has_timed_out_waiting_for_inference() is True - output_generator = manager._wait_to_token_package( - start_time=0, - prompt_ids=[], - group_request_id=123, - sampling_params=_sampling_params(), - req_status=req_status, - request=request, - ) - with pytest.raises(ClientDisconnected): - await output_generator.__anext__() - - asyncio.run(run()) - - -def test_httpserver_keeps_started_requests_and_requests_with_remaining_master_timeout(): - waiting_req = SimpleNamespace( - router_arrival_time=1.0, - infer_start_time=0.0, - sample_params=SimpleNamespace( - high_priority_request=False, - pd_node_resource_wait_timeout_seconds=60, - ), - ) - started_req = SimpleNamespace( - router_arrival_time=1.0, - infer_start_time=2.0, - sample_params=SimpleNamespace(high_priority_request=False), - ) - req_status = _req_status([waiting_req, started_req]) - - with patch("lightllm.server.httpserver.manager.time.monotonic", return_value=30): - assert req_status.has_timed_out_waiting_for_inference() is False - - started_req.infer_start_time = 0.0 - assert req_status.has_timed_out_waiting_for_inference() is False - - -def test_pd_node_allocation_failure_does_not_hold_partial_slots(): - async def run(): - manager = _make_manager(NodeRole.D) - manager.shm_req_manager = SimpleNamespace( - async_alloc_req_indexes=AsyncMock(side_effect=RuntimeError("alloc failed")), - async_release_req_index=AsyncMock(), - ) - sampling_params = _sampling_params() - sampling_params.n = 2 - - with pytest.raises(RuntimeError, match="alloc failed"): - await _drain_generate(manager, sampling_params, _multimodal_params()) - - manager.shm_req_manager.async_alloc_req_indexes.assert_awaited_once_with(2) - manager.shm_req_manager.async_release_req_index.assert_not_awaited() - - asyncio.run(run()) - - -@pytest.mark.parametrize("high_priority", [False, True]) -def test_batch_allocation_retries_with_priority_backoff(high_priority): - async def run(): - manager = _make_manager(NodeRole.D) - manager.shm_req_manager.async_alloc_req_indexes = AsyncMock(side_effect=[None, None, [0, 1]]) - with patch("lightllm.server.httpserver.manager.asyncio.sleep", new_callable=AsyncMock) as sleep: - assert await manager._alloc_shm_req_indexes(2, high_priority_request=high_priority) == [0, 1] - factor = 0.2 if high_priority else 1 - assert [call.args[0] for call in sleep.await_args_list] == pytest.approx([0.1 * factor, 0.11 * factor]) - assert [call.args for call in manager.shm_req_manager.async_alloc_req_indexes.await_args_list] == [(2,)] * 3 - - asyncio.run(run()) - - -def test_batch_allocation_rejects_requests_larger_than_capacity(): - async def run(): - manager = _make_manager(NodeRole.NORMAL) - with pytest.raises(ValueError, match="must not exceed running_max_req_size"): - await manager._alloc_shm_req_indexes(3) - manager.shm_req_manager.async_alloc_req_indexes.assert_not_awaited() - - asyncio.run(run()) - - -def test_multinode_tp_slave_does_not_apply_allocation_wait_timeout(): - async def run(): - manager = _make_manager(NodeRole.D) - manager.is_multinode_tp_slave = True - manager.shm_req_manager.async_alloc_req_indexes = AsyncMock(side_effect=[None, [0, 1]]) - with patch("lightllm.server.httpserver.manager.asyncio.sleep", new_callable=AsyncMock): - with patch("lightllm.server.httpserver.manager.time.monotonic") as monotonic: - assert await manager._alloc_shm_req_indexes(2, pd_node_resource_wait_timeout_seconds=0) == [0, 1] - monotonic.assert_not_called() - - asyncio.run(run()) - - -def test_cancelled_batch_allocation_does_not_hold_slots(): - async def run(): - manager = _make_manager(NodeRole.NORMAL) - manager.shm_req_manager = SimpleNamespace( - async_alloc_req_indexes=AsyncMock(return_value=None), - async_release_req_index=AsyncMock(), - ) - task = asyncio.create_task(manager._alloc_shm_req_indexes(2)) - await asyncio.sleep(0) - task.cancel() - with pytest.raises(asyncio.CancelledError): - await task - manager.shm_req_manager.async_alloc_req_indexes.assert_awaited_once_with(2) - manager.shm_req_manager.async_release_req_index.assert_not_awaited() - - asyncio.run(run()) - - -def test_running_request_helpers_are_atomic_and_refresh_timestamp_only_on_idle_to_running_transition(): - async def run(): - manager = HttpServerManager.__new__(HttpServerManager) - manager._run_reqs_count_lock = asyncio.Lock() - manager.run_reqs_count_mark = _ValueMark() - manager.latest_success_infer_time_mark = _ValueMark(7) - - with patch("lightllm.server.httpserver.manager.time.time", side_effect=[1000, 2000]): - await asyncio.gather(*(manager._register_running_request() for _ in range(100))) - assert manager.run_reqs_count_mark.get_value() == 100 - assert manager.latest_success_infer_time_mark.get_value() == 1000 - - await asyncio.gather(*(manager._unregister_running_request() for _ in range(100))) - assert manager.run_reqs_count_mark.get_value() == 0 - - await manager._register_running_request() - assert manager.run_reqs_count_mark.get_value() == 1 - assert manager.latest_success_infer_time_mark.get_value() == 2000 - await manager._unregister_running_request() - assert manager.run_reqs_count_mark.get_value() == 0 - - asyncio.run(run()) - - -def test_real_supported_max_req_total_len_reserves_mtp_overlap_tokens(): - manager = HttpServerManager.__new__(HttpServerManager) - manager.shm_max_total_token_num = _ValueMark(1000) - manager.max_req_total_len = 2000 - manager.args = SimpleNamespace(mtp_step=2, page_size=4) - - assert manager.get_real_supported_max_req_total_len() == 951 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 deleted file mode 100644 index 1b2675737a..0000000000 --- a/unit_tests/server/multi_level_kv_cache/test_cache_placement_controller.py +++ /dev/null @@ -1,367 +0,0 @@ -from types import SimpleNamespace -from concurrent.futures import ThreadPoolExecutor -import inspect -import pytest - -from lightllm.server.api_cli import make_argument_parser -from lightllm.server.multi_level_kv_cache import ( - AdaptiveCachePlacementController, - CacheCapacityConfig, - CachePlacementController, - CacheTier, - GpuOnlyCachePlacementController, - LegacyCachePlacementController, - create_cache_placement_controller, -) - - -def make_controller( - gpu_tokens: int, - cpu_tokens: int, - disk_tokens: int, - history_size: int = AdaptiveCachePlacementController.MAX_HISTORY_SIZE, - initial_history_size: int = None, - enable_cpu_cache: bool = None, - enable_disk_cache: bool = None, -) -> CachePlacementController: - if enable_cpu_cache is None: - enable_cpu_cache = cpu_tokens > 0 - if enable_disk_cache is None: - enable_disk_cache = disk_tokens > 0 - if initial_history_size is None: - initial_history_size = min(AdaptiveCachePlacementController.INITIAL_HISTORY_SIZE, history_size) - return AdaptiveCachePlacementController( - capacity=CacheCapacityConfig( - gpu_tokens=gpu_tokens, - cpu_tokens=cpu_tokens, - disk_tokens=disk_tokens, - ), - args=SimpleNamespace( - enable_cpu_cache=enable_cpu_cache, - enable_disk_cache=enable_disk_cache, - ), - max_history_size=history_size, - initial_history_size=initial_history_size, - ) - - -def make_req(request_id: int, input_len: int, group_req_id: int = None, cache_tiers=(CacheTier.GPU,)): - if group_req_id is None: - group_req_id = request_id - return SimpleNamespace( - cache_tiers=cache_tiers, - shm_req=SimpleNamespace( - request_id=request_id, - group_req_id=group_req_id, - input_len=input_len, - ), - ) - - -def test_controller_exposes_one_public_operation(): - controller = make_controller(1, 1, 1) - - public_methods = [ - name for name, method in inspect.getmembers(controller, predicate=inspect.ismethod) if not name.startswith("_") - ] - - assert public_methods == ["set_req_cache_way"] - - -def test_controller_is_an_interface(): - assert inspect.isabstract(CachePlacementController) - - -@pytest.mark.parametrize( - "cache_tiers", - [ - (CacheTier.DISK,), - (CacheTier.GPU, CacheTier.DISK), - ], -) -def test_controller_rejects_illegal_existing_cache_tiers(cache_tiers): - req = make_req(0, 100, cache_tiers=cache_tiers) - - with pytest.raises(AssertionError): - make_controller(1, 1, 1).set_req_cache_way([req]) - - -def test_gpu_only_controller_always_selects_gpu(): - reqs = [make_req(0, 100), make_req(1, 1000)] - - GpuOnlyCachePlacementController().set_req_cache_way(reqs) - - assert tuple(req.cache_tiers for req in reqs) == ((CacheTier.GPU,), (CacheTier.GPU,)) - - -def test_legacy_controller_selects_all_enabled_cache_levels(): - req = make_req(0, 100) - - LegacyCachePlacementController(enable_cpu_cache=True, enable_disk_cache=True).set_req_cache_way([req]) - - assert req.cache_tiers == (CacheTier.GPU, CacheTier.CPU, CacheTier.DISK) - - -def test_legacy_controller_does_not_select_disabled_disk_cache(): - req = make_req(0, 100) - - LegacyCachePlacementController(enable_cpu_cache=True, enable_disk_cache=False).set_req_cache_way([req]) - - assert req.cache_tiers == (CacheTier.GPU, CacheTier.CPU) - - -def test_legacy_controller_keeps_diverse_slave_request_on_gpu(): - req = make_req(request_id=2, group_req_id=1, input_len=100) - - LegacyCachePlacementController(enable_cpu_cache=True, enable_disk_cache=True).set_req_cache_way([req]) - - assert req.cache_tiers == (CacheTier.GPU,) - - -def test_cache_placement_strategy_cli_defaults_to_adaptive_and_accepts_legacy(): - parser = make_argument_parser() - - assert parser.parse_args([]).cache_placement_strategy == "adaptive" - assert parser.parse_args(["--cache_placement_strategy", "legacy"]).cache_placement_strategy == "legacy" - - -def test_factory_creates_gpu_only_controller_when_cpu_cache_is_disabled(): - args = SimpleNamespace(enable_cpu_cache=False) - - controller = create_cache_placement_controller(args=args, radix_cache=None) - - assert isinstance(controller, GpuOnlyCachePlacementController) - - -def test_factory_creates_legacy_controller_for_legacy_strategy(): - args = SimpleNamespace( - enable_cpu_cache=True, - enable_disk_cache=True, - cache_placement_strategy="legacy", - ) - - controller = create_cache_placement_controller(args=args, radix_cache=None) - - assert isinstance(controller, LegacyCachePlacementController) - - -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, 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( - enable_cpu_cache=True, - enable_disk_cache=True, - cache_placement_strategy="adaptive", - cpu_cache_token_page_size=16, - disk_cache_storage_size=1, - dp=4, - nnodes=2, - ) - - controller = create_cache_placement_controller( - args=args, - radix_cache=SimpleNamespace(total_token_num=80), - ) - - assert isinstance(controller, AdaptiveCachePlacementController) - assert controller._capacity == CacheCapacityConfig( - gpu_tokens=80, - cpu_tokens=160, - disk_tokens=int(1024 ** 3) // 64 * 16, - ) - - -@pytest.mark.parametrize( - ("configured_ratio", "expected_ratio"), - [ - (None, 0.8), - ("0.5", 0.5), - ("1", 1.0), - ], -) -def test_cache_placement_gpu_capacity_ratio_from_environment(monkeypatch, configured_ratio, expected_ratio): - from lightllm.utils.envs_utils import get_cache_placement_gpu_capacity_ratio - - if configured_ratio is None: - monkeypatch.delenv("LIGHTLLM_CACHE_PLACEMENT_GPU_CAPACITY_RATIO", raising=False) - else: - monkeypatch.setenv("LIGHTLLM_CACHE_PLACEMENT_GPU_CAPACITY_RATIO", configured_ratio) - - assert get_cache_placement_gpu_capacity_ratio() == expected_ratio - - -@pytest.mark.parametrize("configured_ratio", ["0", "-0.1", "1.1", "nan"]) -def test_cache_placement_gpu_capacity_ratio_rejects_out_of_range_values(monkeypatch, configured_ratio): - from lightllm.utils.envs_utils import get_cache_placement_gpu_capacity_ratio - - monkeypatch.setenv("LIGHTLLM_CACHE_PLACEMENT_GPU_CAPACITY_RATIO", configured_ratio) - - with pytest.raises(AssertionError): - get_cache_placement_gpu_capacity_ratio() - - -def test_set_req_cache_way_uses_recent_input_lengths_and_capacity_ratio(): - controller = make_controller(gpu_tokens=1, cpu_tokens=1, disk_tokens=2, history_size=4) - warmup_reqs = [ - make_req(0, 400), - make_req(1, 100), - make_req(2, 300), - make_req(3, 200), - ] - controller.set_req_cache_way(warmup_reqs) - reqs = [ - make_req(4, 400), - make_req(5, 100), - make_req(6, 300), - make_req(7, 200), - ] - - result = controller.set_req_cache_way(reqs) - - assert result is None - assert tuple(req.cache_tiers for req in warmup_reqs) == ((CacheTier.GPU, CacheTier.CPU, CacheTier.DISK),) * 4 - assert tuple(req.cache_tiers for req in reqs) == ( - (CacheTier.CPU, CacheTier.DISK), - (CacheTier.GPU,), - (CacheTier.GPU,), - (CacheTier.GPU,), - ) - - -def test_adaptive_controller_uses_enable_flags_instead_of_capacity_to_select_disk_path(): - controller = make_controller( - gpu_tokens=1, - cpu_tokens=1, - disk_tokens=0, - enable_cpu_cache=True, - enable_disk_cache=True, - ) - warmup_req = make_req(0, 100) - - controller.set_req_cache_way([warmup_req]) - - assert warmup_req.cache_tiers == (CacheTier.GPU, CacheTier.CPU, CacheTier.DISK) - - -def test_length_boundary_uses_max_cpu_and_disk_capacity(): - controller = make_controller(gpu_tokens=1, cpu_tokens=1, disk_tokens=3, history_size=4) - controller.set_req_cache_way([make_req(0, 100), make_req(1, 200), make_req(2, 300), make_req(3, 400)]) - assert controller._gpu_max_input_len == 200 - req = make_req(4, 150) - - controller.set_req_cache_way([req]) - - assert req.cache_tiers == (CacheTier.GPU,) - - -def test_zero_gpu_capacity_does_not_place_requests_in_gpu_after_warmup(): - controller = make_controller(gpu_tokens=0, cpu_tokens=1, disk_tokens=0, history_size=2) - controller.set_req_cache_way([make_req(0, 100), make_req(1, 200)]) - req = make_req(2, 100) - - controller.set_req_cache_way([req]) - - assert controller._gpu_max_input_len == 0 - assert req.cache_tiers == (CacheTier.CPU,) - - -def test_set_req_cache_way_does_not_use_disabled_tiers(): - controller = make_controller(gpu_tokens=1, cpu_tokens=3, disk_tokens=0, history_size=4) - warmup_reqs = [make_req(index, length) for index, length in enumerate((100, 200, 300, 400))] - controller.set_req_cache_way(warmup_reqs) - reqs = [make_req(index + 4, length) for index, length in enumerate((100, 200, 300, 400))] - - controller.set_req_cache_way(reqs) - - assert tuple(req.cache_tiers for req in warmup_reqs) == ((CacheTier.GPU, CacheTier.CPU),) * 4 - assert tuple(req.cache_tiers for req in reqs) == ( - (CacheTier.GPU,), - (CacheTier.GPU,), - (CacheTier.CPU,), - (CacheTier.CPU,), - ) - - -def test_length_boundary_keeps_history_after_initial_update(): - controller = make_controller(1, 1, 1, history_size=3) - - controller.set_req_cache_way([make_req(0, 100), make_req(1, 200)]) - assert controller._gpu_max_input_len is None - assert list(controller._recent_input_lengths) == [100, 200] - - controller.set_req_cache_way([make_req(2, 300)]) - assert controller._gpu_max_input_len == 300 - assert list(controller._recent_input_lengths) == [100, 200, 300] - - latest_req = make_req(3, 1000) - controller.set_req_cache_way([latest_req]) - - assert controller._gpu_max_input_len == 300 - assert list(controller._recent_input_lengths) == [200, 300, 1000] - assert latest_req.cache_tiers == (CacheTier.CPU, CacheTier.DISK) - - -def test_length_boundary_updates_every_36_steps_after_initial_update(): - controller = make_controller(1, 1, 1, history_size=4, initial_history_size=2) - - controller.set_req_cache_way([make_req(0, 100), make_req(1, 200)]) - assert controller._gpu_max_input_len == 200 - assert controller._steps_since_last_update == 0 - assert list(controller._recent_input_lengths) == [100, 200] - - for step in range(35): - controller.set_req_cache_way([make_req(step + 2, 1000)]) - assert controller._gpu_max_input_len == 200 - assert controller._steps_since_last_update == 35 - assert list(controller._recent_input_lengths) == [1000, 1000, 1000, 1000] - - controller.set_req_cache_way([make_req(37, 1000)]) - assert controller._gpu_max_input_len == 1000 - assert controller._steps_since_last_update == 0 - assert list(controller._recent_input_lengths) == [1000, 1000, 1000, 1000] - - -def test_history_window_keeps_latest_requests_when_single_extend_exceeds_limit(): - controller = make_controller(1, 1, 1, history_size=3) - reqs = [make_req(index, length) for index, length in enumerate((100, 200, 300, 400, 500))] - - controller.set_req_cache_way(reqs) - - assert controller._gpu_max_input_len == 400 - assert list(controller._recent_input_lengths) == [300, 400, 500] - - -def test_repeated_assignment_is_rejected(): - controller = make_controller(1, 1, 1) - req = make_req(0, 100) - controller.set_req_cache_way([req]) - - req.shm_req.input_len = 10000 - with pytest.raises(AssertionError): - controller.set_req_cache_way([req]) - - assert list(controller._recent_input_lengths) == [100] - - -def test_concurrent_reassignment_is_rejected(): - controller = make_controller(1, 1, 1) - req = make_req(0, 100) - - with pytest.raises(AssertionError): - with ThreadPoolExecutor(max_workers=2) as executor: - list(executor.map(lambda _: controller.set_req_cache_way([req]), range(2))) - - assert list(controller._recent_input_lengths) == [100] - - -def test_diverse_slave_request_stays_on_gpu_and_is_not_recorded(): - controller = make_controller(0, 1, 1) - slave_req = make_req(request_id=2, group_req_id=1, input_len=1000) - - controller.set_req_cache_way([slave_req]) - - assert slave_req.cache_tiers == (CacheTier.GPU,) - assert list(controller._recent_input_lengths) == [] 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 deleted file mode 100644 index 4238e199f5..0000000000 --- a/unit_tests/server/multi_level_kv_cache/test_multimodal_cache_matching.py +++ /dev/null @@ -1,120 +0,0 @@ -import time -from types import SimpleNamespace - -import pytest -import torch - -from lightllm.server.multi_level_kv_cache.manager import MultiLevelKVCacheManager -from lightllm.utils.kv_cache_utils import CpuKVCacheMeta - - -def test_disk_roundtrip_preserves_raw_pages_before_mm_boundary_trim(cpu_cache_factory, tmp_path, monkeypatch): - light_mem = pytest.importorskip("light_mem") - from lightllm.server.multi_level_kv_cache.disk_cache_worker import DiskCacheWorker - - 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) - client.update_pages_status_to_ready(pages, disk_offload_enable=True, token_num_in_page_list=768) - expected = client.cpu_kv_cache_tensor.clone() - worker = DiskCacheWorker.__new__(DiskCacheWorker) - worker.cpu_cache_client = client - worker.max_concurrent_write_tasks = 8 - worker.service = light_mem.PyLocalCacheService( - worker._prepare_tensor(client.cpu_kv_cache_tensor), - str(tmp_path / "cache"), - storage_size=16 * 1024 * 1024, - num_shard=1, - num_worker=2, - ) - for payloads in worker._gather_offload_payloads(): - worker._persist_pages_to_disk(payloads) - deadline = time.monotonic() + 10 - while not all(worker.service.query([11, 22, 33])): - assert time.monotonic() < deadline, "Disk write did not finish" - time.sleep(0.01) - - # Evict CPU copies and overwrite their bytes before restoring from disk. - replaced, _ = client.allocate_pages([44, 55, 66], False) - client.recycle_pages(replaced) - client.cpu_kv_cache_tensor.fill_(165) - manager = MultiLevelKVCacheManager.__new__(MultiLevelKVCacheManager) - manager.args = SimpleNamespace(enable_disk_cache=True) - manager.cpu_cache_client = client - manager.disk_cache_worker = worker - loaded, disk_pages = manager._disk_cache_match([11, 22, 33], []) - assert len(loaded) == disk_pages == 3 - assert torch.equal(client.cpu_kv_cache_tensor[loaded], expected) - - from lightllm.server.router.model_infer.mode_backend import multi_level_kv_cache as common - - copied_pages = [] - req = SimpleNamespace( - req_idx=0, - cur_kv_len=0, - hold_kv_len=0, - mm_input_spans=[(600, 800)], - sampling_param=SimpleNamespace(shm_param=SimpleNamespace(prompt_logprobs=-1)), - shm_req=SimpleNamespace( - input_len=1024, - disk_prompt_cache_len=768, - cpu_cache_match_page_indexes=SimpleNamespace(get_all=lambda: loaded), - token_hash_page_len_list=SimpleNamespace(get_all=lambda: [256, 512, 768]), - ), - ) - - def allocate(alloc_token_num): - if alloc_token_num == 0: - return None - start = req.cur_kv_len - req.hold_kv_len += alloc_token_num - return torch.arange(start, start + alloc_token_num, dtype=torch.int32) - - req.args = SimpleNamespace(page_size=1) - req.token_index = SimpleNamespace( - get_gpu_token_indexes=lambda start, end: torch.empty(end - start, dtype=torch.int32) - ) - req.kv_cache_alloc_need = lambda target: max(0, target - req.hold_kv_len) - req.alloc_kv_mem = allocate - - from collections import deque - - module = common.MultiLevelKvCacheModule.__new__(common.MultiLevelKvCacheModule) - module.load_tasks = deque() - 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, - dp_world_size=1, - is_hybrid_att_model=False, - model=SimpleNamespace( - req_manager=SimpleNamespace( - req_to_token_indexs=torch.arange(1024).view(1, -1), - alloc_token_indexes=allocate, - wait_token_indexes=lambda: None, - ), - att_state_manager=SimpleNamespace( - 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 - ), - ), - full_att_allocator=None, - ), - ), - ) - monkeypatch.setattr(common.g_infer_context, "get_can_alloc_token_num", lambda: 1024) - monkeypatch.setattr(common.dist, "barrier", lambda **kwargs: None) - monkeypatch.setattr(torch.Tensor, "cuda", lambda self, **kwargs: self) - module.load_cpu_cache_to_reqs([req]) - for task in module.load_tasks: - task.event.synchronize() - module.update_cpu_cache_load_task_states() - assert copied_pages == loaded[:2] - assert req.cur_kv_len == req.hold_kv_len == req.shm_req.disk_prompt_cache_len == 512 - assert req.shm_req.cpu_prompt_cache_len == 0 - assert all(client.page_items.linked_items[page].ref_count == 0 for page in loaded) diff --git a/unit_tests/server/pd_master_fault_launcher.py b/unit_tests/server/pd_master_fault_launcher.py deleted file mode 100644 index 1b804bedf6..0000000000 --- a/unit_tests/server/pd_master_fault_launcher.py +++ /dev/null @@ -1,94 +0,0 @@ -"""Test-only launcher that injects failures at real startup boundaries.""" - -import argparse -import json -from pathlib import Path -import tempfile - -import psutil - -from lightllm.server.api_cli import add_cli_args -from lightllm.server.api_server import launch_server -from lightllm.server.core.objs.start_args_type import StartArgs -from lightllm.server.httpserver_for_pd_master import multiprocess -from lightllm.utils.envs_utils import get_unique_server_name - - -def observe_cleanup(observation_path): - original_stop = multiprocess.process_manager.terminate_all_processes - - def stop(): - prefix = "lightllm-pd-master-" + get_unique_server_name() + "-*" - directories = list(Path(tempfile.gettempdir()).glob(prefix)) - before = [directory.exists() for directory in directories] - processes = list(multiprocess.process_manager.processes) - original_stop() - - def process_stopped(process): - try: - return not process.is_running() - except psutil.NoSuchProcess: - return True - - observation_path.write_text( - json.dumps( - { - "runtime_paths": [str(directory) for directory in directories], - "runtime_before_stop": before, - "runtime_after_stop": [directory.exists() for directory in directories], - "backends_stopped": all(process_stopped(process) for process in processes), - } - ), - encoding="utf-8", - ) - - multiprocess.process_manager.terminate_all_processes = stop - - -def inject_failure(stage, haproxy_binary): - if stage == "backend_start": - original_get_context = multiprocess.multiprocessing.get_context - - def get_context(method): - context = original_get_context(method) - process_count = 0 - - class ProcessContext: - def Process(self, *args, **kwargs): - nonlocal process_count - process_count += 1 - if process_count == 2: - process = context.Process(*args, **kwargs) - - def fail_to_start(): - raise RuntimeError("injected backend startup failure") - - process.start = fail_to_start - return process - return context.Process(*args, **kwargs) - - return ProcessContext() - - multiprocess.multiprocessing.get_context = get_context - elif stage == "proxy_start": - original_popen = multiprocess.subprocess.Popen - - def popen(command, *args, **kwargs): - if command[0] == haproxy_binary and kwargs.get("start_new_session"): - raise OSError("injected HAProxy startup failure") - return original_popen(command, *args, **kwargs) - - multiprocess.subprocess.Popen = popen - - -if __name__ == "__main__": - parser = add_cli_args(argparse.ArgumentParser()) - parser.add_argument("--failure-stage", required=True) - parser.add_argument("--cleanup-observations", type=Path, required=True) - options = vars(parser.parse_args()) - failure_stage = options.pop("failure_stage") - observe_cleanup(options.pop("cleanup_observations")) - args = StartArgs(**options) - haproxy_binary = multiprocess.validate_multiprocess_pd_master_args(args) - inject_failure(failure_stage, haproxy_binary) - launch_server(args) diff --git a/unit_tests/server/pd_master_test_utils.py b/unit_tests/server/pd_master_test_utils.py deleted file mode 100644 index 075cd2765e..0000000000 --- a/unit_tests/server/pd_master_test_utils.py +++ /dev/null @@ -1,188 +0,0 @@ -"""Helpers for real CPU-only PD master integration tests.""" - -from contextlib import contextmanager -from dataclasses import dataclass -import json -import os -from pathlib import Path -import re -import signal -import socket -import subprocess -import sys -import time -from urllib.error import HTTPError, URLError -from urllib.request import ProxyHandler, build_opener - -import psutil -import pytest - -from lightllm.utils.start_utils import kill_recursive - - -def get_test_haproxy_binary(): - try: - from ray_haproxy import get_haproxy_binary - except ImportError: - pytest.skip("Install ray-haproxy to run native proxy integration tests") - return get_haproxy_binary() - - -def create_tiny_tokenizer(model_dir: Path) -> Path: - from tokenizers import Tokenizer - from tokenizers.models import WordLevel - from tokenizers.pre_tokenizers import Whitespace - from transformers import PreTrainedTokenizerFast - - model_dir.mkdir(parents=True, exist_ok=True) - tokenizer = Tokenizer(WordLevel({"[UNK]": 0, "[EOS]": 1, "hello": 2, "world": 3}, unk_token="[UNK]")) - tokenizer.pre_tokenizer = Whitespace() - fast_tokenizer = PreTrainedTokenizerFast(tokenizer_object=tokenizer, unk_token="[UNK]", eos_token="[EOS]") - fast_tokenizer.save_pretrained(model_dir) - (model_dir / "config.json").write_text( - json.dumps( - { - "model_type": "llama", - "architectures": ["LlamaForCausalLM"], - "max_position_embeddings": 128, - "vocab_size": 4, - "eos_token_id": 1, - } - ), - encoding="utf-8", - ) - (model_dir / "generation_config.json").write_text('{"eos_token_id":1}', encoding="utf-8") - return model_dir - - -@dataclass -class RunningPDMaster: - process: subprocess.Popen - public_url: str - workers: list[dict] - opener: object - log_path: Path - - def get_json(self, path, base_url=None): - with self.opener.open((base_url or self.public_url) + path, timeout=3) as response: - return json.load(response) - - def descendants(self): - try: - return psutil.Process(self.process.pid).children(recursive=True) - except psutil.NoSuchProcess: - return [] - - def haproxy_runtime_dir(self): - match = re.search(r"PD master HAProxy configuration file (.+)/haproxy.cfg", self.log_path.read_text()) - assert match is not None, self.log_path.read_text() - return Path(match.group(1)) - - -def assert_processes_exit(processes, timeout=8): - def is_alive(process): - try: - return process.is_running() and process.status() != psutil.STATUS_ZOMBIE - except psutil.NoSuchProcess: - return False - - deadline = time.monotonic() + timeout - while True: - alive = [process for process in processes if is_alive(process)] - if not alive: - return - assert time.monotonic() < deadline, f"Processes survived shutdown: {[process.pid for process in alive]}" - time.sleep(0.1) - - -def discover_test_workers(opener, public_url, timeout=10): - """Build local worker identities from the ports published by the launcher.""" - deadline = time.monotonic() + timeout - while True: - try: - with opener.open(public_url + "/get_server_info", timeout=3) as response: - public_info = json.load(response) - break - except (HTTPError, OSError, URLError): - if time.monotonic() >= deadline: - raise - time.sleep(0.1) - host = public_url.split("://", 1)[1].rsplit(":", 1)[0] - single_worker = public_info["httpserver_workers"] == 1 - workers = [] - for worker_port in public_info["pd_master_multi_ports"]: - host_ip_port = f"{host}:{worker_port}" - workers.append( - { - "node_id": public_info["pd_node_id"] if single_worker else worker_port, - "host_ip_port": host_ip_port, - "port": worker_port, - } - ) - return workers - - -@contextmanager -def running_pd_master(runtime_dir: Path, *, extra_args=(), wait_for_startup=True, launcher_script=None): - """Run real metric/scheduler/API workers and HAProxy without model weights.""" - get_test_haproxy_binary() - runtime_dir.mkdir(parents=True, exist_ok=True) - model_dir = create_tiny_tokenizer(runtime_dir / "tiny-tokenizer") - with socket.socket() as listener: - listener.bind(("", 0)) - public_port = listener.getsockname()[1] - command = [sys.executable, "-m", "lightllm.server.api_server"] - if launcher_script is not None: - command = [sys.executable, str(launcher_script)] - command.extend( - [ - "--run_mode", - "pd_master", - "--host", - "0.0.0.0", - "--port", - str(public_port), - "--httpserver_workers", - "2", - "--model_dir", - str(model_dir), - "--max_req_total_len", - "128", - ] - ) - command.extend(extra_args) - env = os.environ.copy() - env.update(CUDA_VISIBLE_DEVICES="", HF_HUB_OFFLINE="1", TRANSFORMERS_OFFLINE="1", TOKENIZERS_PARALLELISM="false") - env["PYTHONPATH"] = str(Path(__file__).resolve().parents[2]) + os.pathsep + env.get("PYTHONPATH", "") - opener = build_opener(ProxyHandler({})) - public_url = f"http://127.0.0.1:{public_port}" - log_path = runtime_dir / "launcher.log" - process = None - with log_path.open("wb") as log: - try: - process = subprocess.Popen(command, env=env, stdout=log, stderr=subprocess.STDOUT) - if not wait_for_startup: - yield RunningPDMaster(process, public_url, [], opener, log_path) - return - deadline = time.monotonic() + 90 - while True: - assert process.poll() is None, log_path.read_text() - try: - with opener.open(public_url + "/liveness", timeout=0.5) as response: - assert response.status == 200 - break - except (OSError, URLError): - if time.monotonic() >= deadline: - pytest.fail("PD master launcher did not become live:\n" + log_path.read_text()) - time.sleep(0.1) - workers = discover_test_workers(opener, public_url) - yield RunningPDMaster(process, public_url, workers, opener, log_path) - finally: - if process is not None: - if process.poll() is None: - process.send_signal(signal.SIGTERM) - try: - process.wait(timeout=15) - except subprocess.TimeoutExpired: - kill_recursive(process) - process.wait(timeout=5) diff --git a/unit_tests/server/router/dynamic_prompt/test_radix_cache.py b/unit_tests/server/router/dynamic_prompt/test_radix_cache.py deleted file mode 100644 index c2d5a67c19..0000000000 --- a/unit_tests/server/router/dynamic_prompt/test_radix_cache.py +++ /dev/null @@ -1,722 +0,0 @@ -from types import SimpleNamespace - -from sortedcontainers import SortedDict - -import pytest -import torch - -from lightllm.server.router.dynamic_prompt.radix_cache import RadixCache, TreeNode -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(): - monkeypatch = pytest.MonkeyPatch() - monkeypatch.setattr(shm_utils, "get_unique_server_name", lambda: "test_radix_cache_service_0") - yield - monkeypatch.undo() - - -def test_case1(): - tree = RadixCache(100, 0) - ans, _ = tree.insert(torch.tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], dtype=torch.int64, device="cpu")) - assert ans == 0 - tree.print_self() - ans, _ = tree.insert(torch.tensor([0, 1, 2, 3, 4, 7, 8, 9], dtype=torch.int64, device="cpu")) - assert ans == 5 - tree.print_self() - ans, _ = tree.insert(torch.tensor([0, 1, 2, 3, 4, 7, 8, 9], dtype=torch.int64, device="cpu")) - assert ans == 8 - tree.print_self() - - assert tree.get_refed_tokens_num() == 0 - assert tree.get_tree_total_tokens_num() == 13 - - # print("evict") - tree.evict(9, lambda x: x) - tree.print_self() - assert tree.get_refed_tokens_num() == 0 and tree.get_tree_total_tokens_num() == 0 - - -def test_case2(): - tree = RadixCache(100, 1) - 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")) - tree.print_self() - - tree_node, size, values = tree.match_prefix( - torch.tensor([0, 1, 2, 3, 4], dtype=torch.int64, device="cpu"), update_refs=False - ) - assert tree_node.node_prefix_total_len == 5 and size == 5 and len(values) == 5 - tree_node, size, values = tree.match_prefix( - torch.tensor([0, 1, 2, 3, 4, 9], dtype=torch.int64, device="cpu"), update_refs=False - ) - assert tree_node.node_prefix_total_len == 5 and size == 5 and len(values) == 5 - tree_node, size, values = tree.match_prefix( - torch.tensor([0, 1, 2, 3, 4, 7, 8], dtype=torch.int64, device="cpu"), update_refs=False - ) - assert tree_node.node_prefix_total_len == 7 and size == 7 and len(values) == 7 - tree_node, size, values = tree.match_prefix( - torch.tensor([0, 1, 2, 3, 4, 7, 9], dtype=torch.int64, device="cpu"), update_refs=False - ) - assert tree_node.node_prefix_total_len == 6 and size == 6 and len(values) == 6 - print(ans) - return - - -def test_case3(): - 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")) - tree.print_self() - - tree_node, size, values = tree.match_prefix( - torch.tensor([0, 1, 2, 3, 4], dtype=torch.int64, device="cpu"), update_refs=True - ) - assert tree_node.node_prefix_total_len == 5 and size == 5 and len(values) == 5 - assert tree.get_refed_tokens_num() == 5 and tree.get_tree_total_tokens_num() == 13 - - tree_node, size, values = tree.match_prefix( - torch.tensor([0, 1, 2, 3, 4, 7, 9], dtype=torch.int64, device="cpu"), update_refs=True - ) - assert tree_node.node_prefix_total_len == 6 and size == 6 and len(values) == 6 - assert tree.get_refed_tokens_num() == 6 and tree.get_tree_total_tokens_num() == 13 - - tree.print_self() - tree.evict(2, lambda x: x) - assert tree.get_refed_tokens_num() == 6 and tree.get_tree_total_tokens_num() == 8 - tree.print_self() - - tree.dec_node_ref_counter(tree_node) - tree.print_self() - print(ans) - return - - -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")) - tree.print_self() - - tree.clear_tree_nodes() - print(ans) - return - - -def test_case5(): - """ - 测试场景：一个简单的父子节点链 (A -> B)，在 ref_counter 都为 0 时，应该成功合并。 - """ - print("\nTest Case 5: Merging simple parent-child nodes when ref_counter is 0\n") - tree = RadixCache(100, 0) - - _, node_a = tree.insert(torch.tensor([1, 2, 3], dtype=torch.int64)) - _, node_b = tree.insert(torch.tensor([1, 2, 3, 4, 5], dtype=torch.int64)) - tree.print_self() - - # 验证初始状态：A -> B 结构，且 ref_counter 均为 0 - assert node_b.parent == node_a - assert torch.equal(node_a.token_id_key, torch.tensor([1, 2, 3], dtype=torch.int64)) - assert len(node_a.children) == 1 - assert node_a.ref_counter == 0 - assert node_b.ref_counter == 0 - assert tree.get_tree_total_tokens_num() == 5 - - # 执行合并 - tree.merge_unreferenced_nodes() - tree.print_self() - - assert torch.equal(node_b.token_id_key, torch.tensor([1, 2, 3, 4, 5], dtype=torch.int64)) - assert node_b.is_leaf() - assert tree.get_tree_total_tokens_num() == 5 - assert tree.root_node.children[1] is node_b - - -def test_case6(): - """ - 测试场景：一个长的节点链 (A -> B -> C)，在 ref_counter 都为 0 时，应该级联合并成一个节点。 - """ - print("\nTest Case 6: Merging long nodes when ref_counter is 0\n") - tree = RadixCache(100, 0) - _, node_a = tree.insert(torch.tensor([1], dtype=torch.int64)) - _, node_b = tree.insert(torch.tensor([1, 2], dtype=torch.int64)) - _, node_c = tree.insert(torch.tensor([1, 2, 3, 4], dtype=torch.int64)) - tree.print_self() - - assert node_c.parent == node_b - assert node_b.parent == node_a - assert tree.get_tree_total_tokens_num() == 4 - tree.merge_unreferenced_nodes() - tree.print_self() - - assert len(tree.root_node.children) == 1 - # 节点 C 的 key 应该是完整的 [1, 2, 3, 4] - assert torch.equal(node_c.token_id_key, torch.tensor([1, 2, 3, 4], dtype=torch.int64)) - assert node_c.is_leaf() - assert tree.get_tree_total_tokens_num() == 4 - - -def test_case7(): - """ - 测试场景：由于父节点或子节点的 ref_counter > 0，合并不应该发生。 - """ - print("\nTest Case 7: Merging when parent or child ref_counter > 0\n") - tree = RadixCache(100, 0) - - _, node_a = tree.insert(torch.tensor([1, 2, 3], dtype=torch.int64)) - _, node_b = tree.insert(torch.tensor([1, 2, 3, 4, 5], dtype=torch.int64)) - tree.print_self() - - matched_node, _, _ = tree.match_prefix(torch.tensor([1, 2, 3], dtype=torch.int64), update_refs=True) - assert matched_node is node_a - assert node_a.ref_counter == 1 - assert node_b.ref_counter == 0 - - tree.merge_unreferenced_nodes() - tree.print_self() - - assert torch.equal(node_a.token_id_key, torch.tensor([1, 2, 3], dtype=torch.int64)) - assert not node_a.is_leaf() - assert node_b.parent is node_a - - -def test_case8(): - """ - 测试场景：由于父节点有多个子节点，合并不应该发生。 - """ - print("\nTest Case 8: Merging when parent has multiple children\n") - tree = RadixCache(100, 0) - - _, node_a = tree.insert(torch.tensor([1, 2], dtype=torch.int64)) - _, node_b = tree.insert(torch.tensor([1, 2, 3], dtype=torch.int64)) - _, node_c = tree.insert(torch.tensor([1, 2, 4], dtype=torch.int64)) - tree.print_self() - - assert len(node_a.children) == 2 - assert node_a.ref_counter == 0 - assert node_b.ref_counter == 0 - assert node_c.ref_counter == 0 - - tree.merge_unreferenced_nodes() - tree.print_self() - - assert len(node_a.children) == 2 - assert torch.equal(node_a.token_id_key, torch.tensor([1, 2], dtype=torch.int64)) - assert tree.root_node.children[1].children[3] is node_b - assert tree.root_node.children[1].children[4] is node_c - - -def test_case9(): - """ - 测试场景：在一个复杂的树中，只有满足条件的分支被合并。 - """ - print("\nTest Case 9: Merging in a complex tree with mixed conditions\n") - tree = RadixCache(100, 0) - - # 分支1: 可合并的链 A -> B - _, node_a = tree.insert(torch.tensor([1, 2], dtype=torch.int64)) - _, node_b = tree.insert(torch.tensor([1, 2, 3], dtype=torch.int64)) - - # 分支2: 不可合并的链 C -> D (因为 C 被引用) - _, node_c = tree.insert(torch.tensor([4, 5], dtype=torch.int64)) - _, node_d = tree.insert(torch.tensor([4, 5, 6], dtype=torch.int64)) - - # 增加 C 的引用计数 - tree.match_prefix(torch.tensor([4, 5], dtype=torch.int64), update_refs=True) - assert node_c.ref_counter == 1 - tree.print_self() - - tree.merge_unreferenced_nodes() - tree.print_self() - - merged_node_b = tree.root_node.children[1] - assert torch.equal(merged_node_b.token_id_key, torch.tensor([1, 2, 3], dtype=torch.int64)) - assert merged_node_b.is_leaf() - - unmerged_node_c = tree.root_node.children[4] - assert torch.equal(unmerged_node_c.token_id_key, torch.tensor([4, 5], dtype=torch.int64)) - assert not unmerged_node_c.is_leaf() - assert len(unmerged_node_c.children) == 1 - - unmerged_node_d = unmerged_node_c.children[6] - assert torch.equal(unmerged_node_d.token_id_key, torch.tensor([6], dtype=torch.int64)) - - -def test_flush_cache_requires_released_references_and_returns_allocator_pages(): - from lightllm.common.attention_state.allocators import FullAttKVAllocator - - allocator = FullAttKVAllocator(100, pin_memory=False) - 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]) - 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() - 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 - - -@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 - - monkeypatch.setattr( - "lightllm.common.attention_state.bounded_att.utils.get_env_start_args", - lambda: SimpleNamespace(enable_dp_prompt_cache_fetch=False), - ) - resources = composition(device="cuda") - 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_runtime_to_pages( - req_indices=[0], - state_token_lens=[8], - page_kind="big", - page_slot_ids=[big], - ) - resources.save_bounded_att_runtime_to_pages( - req_indices=[1], - state_token_lens=[12], - page_kind="small", - page_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) - - prefix_len, _ = tree.insert(torch.tensor([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]), values) - assert prefix_len == 0 - assert tree.get_tree_total_tokens_num() == 8 - - # A mismatch inside a page cannot produce a partial-page cache hit. - node, matched_len, matched_values = tree.match_prefix(torch.tensor([1, 2, 0, 4, 5, 6, 7, 8])) - assert node is None - assert matched_len == 0 - assert matched_values is None - - node, matched_len, matched_values = tree.match_prefix(torch.tensor([1, 2, 3, 4, 5, 6, 7, 8, 11])) - assert node is not None - assert matched_len == 8 - assert matched_values.tolist() == list(range(100, 108)) - - # The second sequence shares exactly one page and then branches by its - # complete second-page key. - prefix_len, _ = tree.insert( - torch.tensor([1, 2, 3, 4, 5, 6, 0, 8]), - torch.arange(200, 208, dtype=torch.int64), - ) - assert prefix_len == 4 - assert tree.get_tree_total_tokens_num() == 12 - - -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_state_manager=SimpleNamespace(full_att_allocator=mem_manager), page_size=4) - - -@pytest.mark.parametrize("update_refs", [False, True]) -@pytest.mark.parametrize( - "spans,expected", - [ - ([], 24), - ([(10, 21), (22, 30)], 0), - ([(0, 10), (11, 30)], 0), - ([(8, 24)], 24), - ([(24, 30)], 24), - ], -) -def test_mm_gpu_match_uses_existing_ancestors_and_preserves_other_refs(update_refs, spans, expected): - tree = RadixCache(64, 103, page_size=4) - key = torch.arange(32) - values = key + 100 - tree.insert(key, values) - held_node, _, _ = tree.match_prefix(key[:24], update_refs=True) - - node, matched_len, matched_values = tree.match_prefix(key[:27], update_refs=update_refs, mm_input_spans=spans) - assert matched_len == expected - if expected: - assert node.node_prefix_total_len == expected - assert torch.equal(matched_values, values[:expected]) - else: - assert node is None and matched_values is None - assert tree.get_refed_tokens_num() == 24 - - tree.dec_node_ref_counter(held_node) - assert tree.get_refed_tokens_num() == (expected if update_refs else 0) - if update_refs: - tree.dec_node_ref_counter(node) - assert tree.get_refed_tokens_num() == 0 - assert tree.root_node.ref_counter == 1 - freed = [] - tree.evict(32, lambda indexes: freed.extend(indexes.tolist())) - assert sorted(freed) == values.tolist() - - -@pytest.mark.parametrize("update_refs", [False, True]) -@pytest.mark.parametrize("block_start,expected", [(1, 0), (5, 4), (9, 4), (13, 12)]) -def test_mm_gpu_fallback_walks_parents_without_revisiting(monkeypatch, update_refs, block_start, expected): - tree = RadixCache(64, 105, page_size=4) - tree.insert(torch.arange(40, 44), torch.arange(200, 204)) - key = torch.arange(24) - values = key + 100 - cached_nodes = {} - for length in (4, 12, 24): - _, cached_nodes[length] = tree.insert(key[:length], values[:length]) - - # Retreat across existing nodes, even when the image starts inside an ancestor. - visited = [] - update_time = TreeNode.update_time - - def record_visit(node): - visited.append(node) - update_time(node) - - monkeypatch.setattr(TreeNode, "update_time", record_visit) - node, matched_len, matched_values = tree.match_prefix( - key, update_refs=update_refs, mm_input_spans=[(block_start, 30)] - ) - assert matched_len == expected - assert node is cached_nodes.get(expected) - assert [len(n.token_id_key) for n in cached_nodes.values()] == [4, 8, 12] - assert len(visited) == len(set(visited)) == 4 - assert tree.get_refed_tokens_num() == (expected if update_refs else 0) - if expected: - assert node.node_prefix_total_len == expected - assert torch.equal(matched_values, values[:expected]) - else: - assert node is None and matched_values is None - - # The untouched branch is still oldest; changed tail refs must preserve eviction ordering. - freed = [] - tree.evict(4, lambda indexes: freed.extend(indexes.tolist())) - assert freed == list(range(200, 204)) - if update_refs: - tree.dec_node_ref_counter(node) - assert tree.get_refed_tokens_num() == 0 and tree.root_node.ref_counter == 1 - tree.evict(24, lambda indexes: freed.extend(indexes.tolist())) - assert sorted(freed) == sorted(values.tolist() + list(range(200, 204))) - assert not tree.root_node.children - - -def test_page_key_bytes_does_not_share_tensor_memory(): - token_ids = torch.tensor([1, 2, 3, 4], dtype=torch.int64) - expected_key = token_ids.clone().numpy().tobytes() - key = TreeNode(page_size=4).get_child_key(token_ids) - - token_ids[0] = 100 - - assert isinstance(key, bytes) - assert key == expected_key - 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 = [] - tree = HybridAttPagedRadixCache( - total_token_num=100, - rank_in_node=101, - hash_page_size=4, - big_page_num=2, - 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) - tree.insert( - key, - values, - block_hashs=[11, 22, 33], - block_state_idxs=[None, None, 7], - len_to_big_page_id=SortedDict({8: 5}), - ) - - big_page = tree.root_node.children[22] - tail_page = big_page.children[33] - node, matched_len, _ = tree.match_prefix(key, block_hashs=[11, 22, 33], update_refs=True) - assert node is tail_page and matched_len == 12 - assert tree.get_refed_tokens_num() == 12 - tree.dec_node_ref_counter(node) - assert tree.get_refed_tokens_num() == 0 - - tree.free_one_small_page_buffer() - assert freed_state_ids == [7] - node, matched_len, matched_values = tree.match_prefix(key, block_hashs=[11, 22, 33], update_refs=True) - assert node is big_page and matched_len == 8 - assert matched_values.tolist() == list(range(100, 108)) - assert big_page.ref_counter == 1 and tail_page.ref_counter == 0 - assert tree.get_refed_tokens_num() == 8 - tree.dec_node_ref_counter(node) - assert big_page.ref_counter == tail_page.ref_counter == tree.get_refed_tokens_num() == 0 - assert tree.root_node.ref_counter == 1 - - -@pytest.fixture -def hybrid_tree(): - freed_big_states = [] - tree = HybridAttPagedRadixCache( - total_token_num=100, - rank_in_node=102, - hash_page_size=4, - big_page_num=2, - att_state_manager=hybrid_resources(freed_big=freed_big_states), - ) - return tree, freed_big_states - - -def test_hybrid_match_traverses_pages_inside_image(hybrid_tree): - tree, freed = hybrid_tree - key = torch.arange(20) - tree.insert( - key, - key + 100, - block_hashs=[1, 2, 3, 4, 5], - block_state_idxs=[None, None, None, None, 7], - len_to_big_page_id=SortedDict({8: 5, 16: 6}), - ) - first_page = tree.root_node.children[2] - second_page = first_page.children[4] - assert tree.get_big_page_ids_by_node(second_page) == [5, 6] - - for length, expected in [(8, 0), (12, 0), (16, 16), (20, 20), (24, 20)]: - node, matched, values = tree.match_prefix( - torch.arange(length), - block_hashs=list(range(1, length // 4 + 1)), - update_refs=True, - mm_input_spans=[(4, 12)], - ) - assert matched == expected - if expected: - assert not 4 < node.node_prefix_total_len < 12 - assert values.tolist() == list(range(100, 100 + expected)) - assert tree.get_refed_tokens_num() == expected - tree.dec_node_ref_counter(node) - assert tree.get_refed_tokens_num() == 0 - - # Duplicate requests release their slots; no per-page validity metadata is needed. - tree.insert(key[:8], key[:8] + 100, [1, 2], [None, None], SortedDict({8: 8})) - tree.insert(key[:16], key[:16] + 100, [1, 2, 3, 4], [None] * 4, SortedDict({8: 9, 16: 10})) - assert first_page.big_page_buffer_idx == 5 - assert second_page.big_page_buffer_idx == 6 - assert freed == [8, 9, 10] - tree._evict(20, lambda *args: None) - assert sorted(freed) == [5, 6, 8, 9, 10] - assert not tree.root_node.children - - -@pytest.mark.parametrize( - "small_state_valid,mm_spans,match_len,expected", - [ - (True, (), 640, 512), - (True, [(448, 640)], 384, 384), - (True, [(448, 640)], 640, 256), - (False, [(448, 640)], 640, 256), - (True, [(448, 640)], 768, 768), - (True, [(448, 896)], 896, 256), - (False, [(448, 896)], 896, 256), - (True, [(448, 896)], 1024, 1024), - (True, [(700, 800)], 896, 512), - ], -) -def test_hybrid_match_falls_back_along_matched_branch(monkeypatch, small_state_valid, mm_spans, match_len, expected): - freed_big_states = [] - tree = HybridAttPagedRadixCache( - total_token_num=2048, - rank_in_node=104, - hash_page_size=64, - big_page_num=4, - 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) - tree.insert( - torch.arange(384), - old_values, - block_hashs=list(range(1, 7)), - block_state_idxs=[None] * 5 + [7], - len_to_big_page_id=SortedDict({256: 5}), - ) - # A longer request keeps the old small branch and owns a separate big-page tail. - new_values = torch.cat([old_values[:256], torch.arange(2000, 2768)]) - big_slots = {512: 6, 768: 8, 1024: 9} - prefix_len, _ = tree.insert( - torch.arange(1024), - new_values, - block_hashs=list(range(1, 17)), - block_state_idxs=[None] * 16, - len_to_big_page_id=SortedDict(big_slots), - ) - assert prefix_len == 256 - if not small_state_valid: - tree.free_one_small_page_buffer() - - common = tree.root_node.children[4] - small_320 = common.children[5] - small_384 = small_320.children[6] - big_512 = common.children[8] - big_768 = big_512.children[12] - big_1024 = big_768.children[16] - visited = [] - update_time = type(common).update_time - - def record_visit(node): - visited.append(node) - update_time(node) - - monkeypatch.setattr(type(common), "update_time", record_visit) - node, matched_len, values = tree.match_prefix( - torch.arange(match_len), - block_hashs=list(range(1, match_len // 64 + 1)), - update_refs=True, - mm_input_spans=mm_spans, - ) - assert matched_len == expected - # Retain small pages only when already on the matched path; never rematch sibling branches. - assert len(visited) == len(set(visited)) - assert not any(start < matched_len < end for start, end in mm_spans) - expected_values = old_values if expected <= 384 else new_values - assert torch.equal(values, expected_values[:expected]) - assert tree.get_refed_tokens_num() == expected - assert tree.root_node.ref_counter == 2 and common.ref_counter == 1 - assert small_320.ref_counter == small_384.ref_counter == int(expected == 384) - for big in (big_512, big_768, big_1024): - assert big.ref_counter == int(big.node_prefix_total_len <= expected) - - tree.dec_node_ref_counter(node) - assert tree.get_refed_tokens_num() == 0 and tree.root_node.ref_counter == 1 - assert all(n.ref_counter == 0 for n in (common, small_320, small_384, big_512, big_768, big_1024)) - freed_indexes = [] - tree._evict(tree.get_tree_total_tokens_num(), lambda indexes, _: freed_indexes.extend(indexes.tolist())) - assert sorted(freed_indexes) == sorted(old_values.tolist() + new_values[256:].tolist()) - assert sorted(freed_big_states) == [5, 6, 8, 9] - assert not tree.root_node.children - - -@pytest.mark.parametrize("first_outside_image", [False, True]) -@pytest.mark.parametrize("require_state", [False, True]) -def test_hybrid_fallback_distinguishes_kv_ownership_from_state_restore(hybrid_tree, first_outside_image, require_state): - tree, freed = hybrid_tree - tree.insert( - torch.arange(20), - block_hashs=[1, 2, 3, 4, 5], - block_state_idxs=[None] * 4 + [7], - len_to_big_page_id=SortedDict({8: 6, 16: 8}), - ) - spans = [(12, 18)] if first_outside_image else [(4, 18)] - node, _, _ = tree.match_prefix( - torch.arange(20), block_hashs=[1, 2, 3, 4, 5], update_refs=True, mm_input_spans=spans - ) - node = tree.deref_to_first_big_page_node(node, mm_input_spans=spans if require_state else ()) - expected = 16 if not require_state else (8 if first_outside_image else 0) - assert (0 if node is None else node.node_prefix_total_len) == expected - assert tree.get_refed_tokens_num() == expected - tree.dec_node_ref_counter(node) - assert tree.get_refed_tokens_num() == 0 - tree._evict(20, lambda *args: None) - assert sorted(freed) == [6, 8] # Invalid state slots have the same ownership and lifetime. - - -@pytest.mark.parametrize( - "computed,cache_len,retained,expected", - [(12, 20, 0, 0), (18, 20, 16, 16), (20, 20, 20, 20), (26, 20, 20, 20), (26, 24, 24, 16)], -) -def test_hybrid_request_release_frees_or_transfers_image_page_slots( - monkeypatch, hybrid_tree, computed, cache_len, retained, expected -): - from lightllm.server.router.model_infer import infer_batch - from lightllm.server.router.model_infer.infer_batch import InferReq, InferenceContext - - tree, freed = hybrid_tree - context = InferenceContext() - context.is_hybrid_att_model = True - context.args = SimpleNamespace(linear_att_hash_page_size=4, linear_att_page_block_num=2) - 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_state_manager = tree.att_state_manager - table = torch.arange(100, 132).view(1, -1) - context.req_manager = SimpleNamespace() - req = InferReq.__new__(InferReq) - req.req_idx = req.cur_output_len = 0 - req.token_index = SimpleNamespace( - get_cpu_token_indexes=lambda start, end: table[0, start:end].clone(), - get_free_token_indexes=lambda start, end: table[0, start:end].clone(), - reset=lambda: None, - ) - req.cur_kv_len, req.hold_kv_len, req.hybrid_cache_len = computed, max(24, computed), cache_len - req.shared_kv_node = None - req.hybrid_len_to_big_page_id = SortedDict({end: end // 8 + 4 for end in range(8, min(computed, cache_len) + 1, 8)}) - big_state_ids = list(req.hybrid_len_to_big_page_id.values()) - req.mm_input_spans = [(4, 12)] + ([(20, cache_len + 1)] if cache_len >= 24 else []) - req.tail_small_page_buffer_id = 7 if computed >= cache_len and cache_len % 8 else None - block_hashs = list(range(1, cache_len // 4 + 1)) - req.shm_req = SimpleNamespace( - input_len=cache_len + 1, - shm_prompt_ids=SimpleNamespace(arr=list(range(cache_len + 1))), - hybrid_token_hash_list=SimpleNamespace(get_all=lambda: block_hashs), - ) - released = [] - context._hybrid_att_free_req(released, req) - assert [index for block in released for index in block.tolist()] == list( - range(100 + retained, 100 + max(24, computed)) - ) - assert not req.hybrid_len_to_big_page_id - assert req.tail_small_page_buffer_id is None - assert tree.get_tree_total_tokens_num() == retained - assert freed == ([] if retained else big_state_ids) - node, length, _ = tree.match_prefix( - torch.arange(cache_len), block_hashs=block_hashs, update_refs=True, mm_input_spans=req.mm_input_spans - ) - assert length == expected - assert not any(start < length < end for start, end in req.mm_input_spans) - tree.dec_node_ref_counter(node) - tree._evict(retained, lambda *args: None) - assert sorted(freed) == big_state_ids - - -if __name__ == "__main__": - pytest.main() 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 deleted file mode 100644 index 4bf7aea081..0000000000 --- a/unit_tests/server/router/model_infer/mode_backend/test_dp_big_pages.py +++ /dev/null @@ -1,286 +0,0 @@ -"""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.attention_state.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, - _copy_req_token_indexes_to_gpu=lambda reqs: 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), (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) - 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_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: - 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), - full_att_allocator=None, - 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( - is_hybrid_att_model=hybrid, - is_master_in_dp=True, - node_nccl_group="node", - model=SimpleNamespace(att_state_manager=manager), - dp_att_state_managers={}, - dp_req_to_token_indices={}, - _copy_req_token_indexes_to_gpu=lambda reqs: None, - ) - 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, - 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, - req_manager=SimpleNamespace(get_need_alloc_token_num=allocation_need, alloc_token_indexes=allocate_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, - ) - - 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_state_manager=SimpleNamespace( - bounded_att_big_page_allocator=allocator, full_att_allocator=SimpleNamespace(page_size=4) - ) - ), - dp_req_to_token_indices={1: torch.arange(64, dtype=torch.int32).reshape(2, 32)}, - ) - 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]]])) - requests = SimpleNamespace( - get_need_alloc_token_num=lambda req, end, page_size=None: max(0, ((end + 3) // 4) * 4 - req.hold_kv_len), - ) - req.kv_cache_alloc_need = lambda target: requests.get_need_alloc_token_num(req, target) - req.alloc_kv_mem = lambda alloc_token_num: allocate(req, alloc_token_num) - monkeypatch.setattr( - dp_shared_kv_trans, - "g_infer_context", - SimpleNamespace(get_can_alloc_token_num=lambda: 100, req_manager=requests, radix_cache=None), - ) - 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) - requests.get_need_alloc_token_num = lambda req, end, page_size=None: max(0, ((end + 3) // 4) * 4 - req.hold_kv_len) - - def allocate_tokens(req, end, allocator, page_size, radix_cache=None): - count = requests.get_need_alloc_token_num(req, end) - radix_cache.free_radix_cache_to_get_enough_token(count) - indexes = full.alloc(count) - start = req.hold_kv_len - requests.req_to_token_indexs[req.req_idx, start : start + count] = indexes - req.hold_kv_len += count - return requests.req_to_token_indexs[req.req_idx, req.cur_kv_len : end].clone() - - req.kv_cache_alloc_need = lambda target: requests.get_need_alloc_token_num(req, target) - req.alloc_kv_mem = lambda alloc_token_num: allocate_tokens(req, req.cur_kv_len + alloc_token_num, None, 4, radix) - 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 diff --git a/unit_tests/server/router/model_infer/mode_backend/test_dp_overlap_spec_engine.py b/unit_tests/server/router/model_infer/mode_backend/test_dp_overlap_spec_engine.py deleted file mode 100644 index fdc29fb87a..0000000000 --- a/unit_tests/server/router/model_infer/mode_backend/test_dp_overlap_spec_engine.py +++ /dev/null @@ -1,491 +0,0 @@ -from contextlib import nullcontext -from types import SimpleNamespace - -import pytest -import torch - -from lightllm.server.router.model_infer.mode_backend.dp_backend import ( - impl as dp_backend_impl, -) -from lightllm.server.router.model_infer.mode_backend.dp_backend.impl import ( - DPChunkedPrefillBackend, -) -from lightllm.server.router.model_infer.mode_backend.base_backend import ModeBackend -from lightllm.server.router.model_infer.mode_backend.chunked_prefill.impl import ( - ChunkedPrefillBackend, -) -from lightllm.server.router.model_infer.mtp_speculative.dp_overlap_engine import ( - DPOverlapSpecEngine, -) -from lightllm.server.router.model_infer.mtp_speculative.engine import SpecEngine -from lightllm.server.router.model_infer.mtp_speculative.planner import ( - lightspec as lightspec_planner_impl, -) -from lightllm.server.router.model_infer.mtp_speculative.planner import ( - FixedSpecPlanner, - LightSpecPlanner, - SpecDecodePlan, -) -from lightllm.server.router.model_infer.mtp_speculative.proposers.base import ( - SpecProposal, -) - - -def test_dp_backend_reuses_common_engine_outside_overlap(): - args = SimpleNamespace(mtp_mode="eagle3", mtp_dynamic_verify=False, dp=1) - backend = ChunkedPrefillBackend.__new__(ChunkedPrefillBackend) - backend.args = args - backend.max_draft_step = 2 - backend.init_spec_engine() - dp_backend = DPChunkedPrefillBackend.__new__(DPChunkedPrefillBackend) - dp_backend.args = args - dp_backend.max_draft_step = 2 - dp_backend.dp_size = 1 - dp_backend.enable_prefill_microbatch_overlap = False - dp_backend.enable_decode_microbatch_overlap = True - dp_backend.init_spec_engine() - - assert "spec_engine_class" not in ModeBackend.__dict__ - assert type(backend.spec_engine) is SpecEngine - assert type(dp_backend.spec_engine) is SpecEngine - assert not dp_backend.spec_engine.proposer.enable_dynmaic_mtp - assert type(dp_backend.spec_engine.planner) is FixedSpecPlanner - assert type(dp_backend.dp_overlap_spec_engine) is DPOverlapSpecEngine - assert dp_backend.dp_overlap_spec_engine.common_engine is dp_backend.spec_engine - assert dp_backend.prefill_draft_engine is dp_backend.spec_engine - assert dp_backend.decode_draft_engine is dp_backend.dp_overlap_spec_engine - assert not issubclass(DPOverlapSpecEngine, SpecEngine) - - -def test_dp_backend_dspark_uses_common_engine_without_overlap(): - backend = DPChunkedPrefillBackend.__new__(DPChunkedPrefillBackend) - backend.args = SimpleNamespace(mtp_mode="dspark", mtp_dynamic_verify=False, dp=2) - backend.max_draft_step = 5 - backend.dp_size = 2 - backend.draft_models = [SimpleNamespace(block_size=5, graph=None)] - backend.enable_prefill_microbatch_overlap = False - backend.enable_decode_microbatch_overlap = False - - backend.init_spec_engine() - - assert type(backend.spec_engine) is SpecEngine - assert backend.dp_overlap_spec_engine is None - assert backend.prefill_draft_engine is backend.spec_engine - assert backend.decode_draft_engine is backend.spec_engine - - -def test_lightspec_planner_reduces_draft_step_with_max(monkeypatch): - group = object() - planner = LightSpecPlanner.__new__(LightSpecPlanner) - planner._draft_step_group = group - planner._draft_step_tensor = torch.zeros((1,), dtype=torch.int32) - planner._draft_step_stream = object() - planner.draft_steps = (1, 2, 3) - planner.pre_draft_step = 1 - planner.backend = SimpleNamespace(dp_size=2) - - def all_reduce(tensor, op, group, async_op): - assert op == torch.distributed.ReduceOp.MAX - assert group is planner._draft_step_group - assert not async_op - tensor.fill_(3) - - monkeypatch.setattr(lightspec_planner_impl.dist, "all_reduce", all_reduce) - monkeypatch.setattr(lightspec_planner_impl.torch.cuda, "stream", lambda stream: nullcontext()) - - plan = planner.plan(decode_reqs=[], origin_batch_size=2) - - assert plan.dynamic_batch_size == 2 - assert plan.draft_step == 3 - assert planner.pre_draft_step == 3 - - -def test_dp_backend_builds_dedicated_global_nccl_group(monkeypatch): - created_group = object() - created_stream = object() - new_group_args = {} - real_torch_zeros = torch.zeros - - def new_group(*, ranks, backend): - new_group_args.update(ranks=ranks, backend=backend) - return created_group - - def zeros_without_cuda(*args, **kwargs): - kwargs.pop("device", None) - return real_torch_zeros(*args, **kwargs) - - monkeypatch.setattr(lightspec_planner_impl, "get_global_world_size", lambda: 4) - monkeypatch.setattr(lightspec_planner_impl.dist, "new_group", new_group) - monkeypatch.setattr(lightspec_planner_impl.torch, "zeros", zeros_without_cuda) - monkeypatch.setattr(lightspec_planner_impl.torch.cuda, "Stream", lambda: created_stream) - - backend = DPChunkedPrefillBackend.__new__(DPChunkedPrefillBackend) - backend.args = SimpleNamespace(mtp_mode="eagle3", mtp_dynamic_verify=True, dp=2) - backend.max_draft_step = 3 - backend.dp_size = 2 - backend.model = SimpleNamespace(graph=None) - backend.draft_models = [SimpleNamespace(graph=None)] - backend.enable_prefill_microbatch_overlap = False - backend.enable_decode_microbatch_overlap = False - backend.init_spec_engine() - - assert new_group_args == {"ranks": [0, 1, 2, 3], "backend": "nccl"} - planner = backend.spec_engine.planner - assert planner._draft_step_group is created_group - assert planner._draft_step_tensor.shape == (1,) - assert planner._draft_step_stream is created_stream - - -def test_dp_backend_does_not_build_group_for_single_draft_step_mode(monkeypatch): - monkeypatch.setattr( - lightspec_planner_impl.dist, - "new_group", - lambda **kwargs: pytest.fail(f"unexpected NCCL group: {kwargs}"), - ) - - backend = DPChunkedPrefillBackend.__new__(DPChunkedPrefillBackend) - backend.args = SimpleNamespace(mtp_mode="vanilla_with_att", mtp_dynamic_verify=True, dp=2) - backend.max_draft_step = 3 - backend.dp_size = 2 - backend.model = SimpleNamespace(graph=None) - backend.draft_models = [SimpleNamespace(graph=None, block_size=4)] - backend.enable_prefill_microbatch_overlap = False - backend.enable_decode_microbatch_overlap = False - - backend.init_spec_engine() - - assert backend.spec_engine.planner._draft_step_group is None - assert backend.spec_engine.planner._draft_step_tensor is None - assert backend.spec_engine.planner._draft_step_stream is None - - -def test_dp_backend_does_not_build_group_for_fixed_planner(monkeypatch): - monkeypatch.setattr( - lightspec_planner_impl.dist, - "new_group", - lambda **kwargs: pytest.fail(f"unexpected NCCL group: {kwargs}"), - ) - - backend = DPChunkedPrefillBackend.__new__(DPChunkedPrefillBackend) - backend.args = SimpleNamespace(mtp_mode="eagle3", mtp_dynamic_verify=False, dp=2) - backend.max_draft_step = 3 - backend.dp_size = 2 - backend.enable_prefill_microbatch_overlap = False - backend.enable_decode_microbatch_overlap = False - - backend.init_spec_engine() - - assert not hasattr(backend.spec_engine.planner, "_draft_step_group") - - -def test_dp_backend_overlap_decode_reuses_common_lightspec_group(monkeypatch): - created_group = object() - created_stream = object() - real_torch_zeros = torch.zeros - - monkeypatch.setattr(lightspec_planner_impl, "get_global_world_size", lambda: 2) - monkeypatch.setattr(lightspec_planner_impl.dist, "new_group", lambda **kwargs: created_group) - monkeypatch.setattr( - lightspec_planner_impl.torch, - "zeros", - lambda *args, **kwargs: real_torch_zeros( - *args, **{key: value for key, value in kwargs.items() if key != "device"} - ), - ) - monkeypatch.setattr(lightspec_planner_impl.torch.cuda, "Stream", lambda: created_stream) - - backend = DPChunkedPrefillBackend.__new__(DPChunkedPrefillBackend) - backend.args = SimpleNamespace(mtp_mode="eagle3", mtp_dynamic_verify=True, dp=2) - backend.max_draft_step = 3 - backend.dp_size = 2 - backend.model = SimpleNamespace(graph=None) - backend.draft_models = [SimpleNamespace(graph=None)] - backend.enable_prefill_microbatch_overlap = False - backend.enable_decode_microbatch_overlap = True - - backend.init_spec_engine() - - assert backend.spec_engine.planner._draft_step_group is created_group - assert backend.spec_engine.planner._draft_step_stream is created_stream - assert backend.decode_draft_engine is backend.dp_overlap_spec_engine - assert backend.decode_draft_engine.common_engine.planner is backend.spec_engine.planner - - -def test_dp_backend_keeps_dynamic_planning_before_global_reduction(): - backend = DPChunkedPrefillBackend.__new__(DPChunkedPrefillBackend) - backend.args = SimpleNamespace(mtp_mode="eagle3", mtp_dynamic_verify=True, dp=1) - backend.max_draft_step = 3 - backend.dp_size = 1 - backend.model = SimpleNamespace(graph=None) - backend.draft_models = [SimpleNamespace(graph=None)] - backend.enable_prefill_microbatch_overlap = False - backend.enable_decode_microbatch_overlap = False - - backend.init_spec_engine() - - assert isinstance(backend.spec_engine.planner, LightSpecPlanner) - assert backend.spec_engine.proposer.enable_dynmaic_mtp - assert backend.spec_engine.planner._draft_step_group is None - - -def test_dp_prefill_and_decode_select_overlap_engine_independently(): - args = SimpleNamespace(mtp_mode="eagle3", mtp_dynamic_verify=False, dp=1) - - for prefill_overlap in (False, True): - for decode_overlap in (False, True): - backend = DPChunkedPrefillBackend.__new__(DPChunkedPrefillBackend) - backend.args = args - backend.max_draft_step = 2 - backend.dp_size = 1 - backend.enable_prefill_microbatch_overlap = prefill_overlap - backend.enable_decode_microbatch_overlap = decode_overlap - backend.init_spec_engine() - - expected_prefill_engine = backend.dp_overlap_spec_engine if prefill_overlap else backend.spec_engine - expected_decode_engine = backend.dp_overlap_spec_engine if decode_overlap else backend.spec_engine - assert backend.prefill_draft_engine is expected_prefill_engine - assert backend.decode_draft_engine is expected_decode_engine - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") -def test_dp_decode_mtp_runs_common_engine_for_empty_batch(monkeypatch): - device = "cuda" - empty_i32 = torch.empty((0,), dtype=torch.int32, device=device) - model_input = SimpleNamespace( - batch_size=0, - b_req_idx=empty_i32, - b_mtp_index=empty_i32, - ) - model_output = SimpleNamespace(logits=torch.empty((0, 8), device=device)) - calls = [] - - class _CommonEngine: - def plan_decode(self, **kwargs): - calls.append("plan") - return SpecDecodePlan(0, 0, 2, 2) - - def prepare_decode_model_input(self, **kwargs): - calls.append("prepare") - return kwargs["model_input"], None - - def propose_next(self, **kwargs): - calls.append("propose") - assert kwargs["target_next_token_ids"].shape == (0,) - assert kwargs["b_req_mtp_start_loc"].shape == (0,) - assert kwargs["accept_len"].shape == (0,) - return SpecProposal(token_ids=torch.empty((0, 2), dtype=torch.int64, device=device)) - - backend = DPChunkedPrefillBackend.__new__(DPChunkedPrefillBackend) - backend.spec_engine = _CommonEngine() - backend.args = SimpleNamespace(penalty_counter_mode="gpu_counter") - backend.model = SimpleNamespace(forward=lambda _: model_output) - event_pack = SimpleNamespace( - notify_post_handle_and_wait_pre_post_handle=lambda: calls.append("post_wait"), - notify_forward_and_wait_post_handle=lambda: calls.append("forward_wait"), - notify_pre_post_handle=lambda: calls.append("pre_post"), - ) - monkeypatch.setattr( - dp_backend_impl, - "prepare_decode_inputs", - lambda req_objs: (model_input, []), - ) - monkeypatch.setattr( - dp_backend_impl, - "g_infer_context", - SimpleNamespace(get_overlap_stream=lambda: torch.cuda.current_stream()), - ) - backend.decode_mtp(event_pack=event_pack, decode_reqs=[]) - - assert calls == [ - "plan", - "prepare", - "post_wait", - "propose", - "forward_wait", - "pre_post", - ] - - -def test_dp_overlap_engine_delegates_raw_verify_layout_to_proposer(): - calls = {} - - class _Proposer: - def propose_next_overlap(self, **kwargs): - calls.update(kwargs) - return SpecProposal(token_ids=kwargs["target_next_token_ids0"].new_empty((3, 7))) - - engine = DPOverlapSpecEngine.__new__(DPOverlapSpecEngine) - engine.proposer = _Proposer() - model_input0 = SimpleNamespace(batch_size=8) - model_input1 = SimpleNamespace(batch_size=16) - model_output0 = SimpleNamespace() - model_output1 = SimpleNamespace() - target_next_token_ids0 = torch.arange(8, dtype=torch.int64) - target_next_token_ids1 = torch.arange(8, 24, dtype=torch.int64) - accept_len0 = torch.tensor([2], dtype=torch.int32) - accept_len1 = torch.tensor([3, 4], dtype=torch.int32) - - proposal = engine.propose_next_overlap( - target_model_input0=model_input0, - target_model_output0=model_output0, - target_next_token_ids0=target_next_token_ids0, - accept_len0=accept_len0, - target_model_input1=model_input1, - target_model_output1=model_output1, - target_next_token_ids1=target_next_token_ids1, - accept_len1=accept_len1, - draft_step=7, - ) - - assert calls["target_model_input0"] is model_input0 - assert calls["target_model_input1"] is model_input1 - assert calls["target_model_output0"] is model_output0 - assert calls["target_model_output1"] is model_output1 - assert calls["target_next_token_ids0"] is target_next_token_ids0 - assert calls["target_next_token_ids1"] is target_next_token_ids1 - assert calls["accept_len0"] is accept_len0 - assert calls["accept_len1"] is accept_len1 - assert calls["draft_step"] == 7 - assert proposal.token_ids.shape == (3, 7) - - -def test_dp_overlap_engine_distributes_dynamic_verify_budget_by_request_count(): - prepare_calls = [] - - class _CommonEngine: - def prepare_decode_model_input(self, **kwargs): - prepare_calls.append(kwargs) - model_input = kwargs["model_input"] - model_input.batch_size = kwargs["plan"].dynamic_batch_size - return model_input, f"mask{len(prepare_calls)}" - - engine = DPOverlapSpecEngine.__new__(DPOverlapSpecEngine) - engine.common_engine = _CommonEngine() - model_input0 = SimpleNamespace(batch_size=12) - model_input1 = SimpleNamespace(batch_size=8) - plan = SpecDecodePlan( - origin_batch_size=20, - dynamic_batch_size=11, - draft_step=2, - pre_draft_step=3, - ) - - compacted_input0, mask0, compacted_input1, mask1 = engine.prepare_decode_model_inputs( - model_input0=model_input0, - req_num0=3, - model_input1=model_input1, - req_num1=2, - plan=plan, - ) - - assert compacted_input0.batch_size == 7 - assert compacted_input1.batch_size == 4 - assert mask0 == "mask1" - assert mask1 == "mask2" - assert prepare_calls[0]["req_num"] == 3 - assert prepare_calls[0]["plan"] == SpecDecodePlan(12, 7, 2, 3) - assert prepare_calls[1]["req_num"] == 2 - assert prepare_calls[1]["plan"] == SpecDecodePlan(8, 4, 2, 3) - - -@pytest.mark.parametrize( - ("batch_size0", "batch_size1", "dynamic_batch_size", "expected_batch_sizes"), - ( - (4, 8, 9, (4, 5)), - (8, 4, 10, (6, 4)), - ), -) -def test_dp_overlap_engine_moves_verify_rows_to_the_side_with_capacity( - batch_size0, - batch_size1, - dynamic_batch_size, - expected_batch_sizes, -): - dynamic_batch_sizes = [] - - class _CommonEngine: - def prepare_decode_model_input(self, **kwargs): - dynamic_batch_sizes.append(kwargs["plan"].dynamic_batch_size) - return kwargs["model_input"], None - - engine = DPOverlapSpecEngine.__new__(DPOverlapSpecEngine) - engine.common_engine = _CommonEngine() - engine.prepare_decode_model_inputs( - model_input0=SimpleNamespace(batch_size=batch_size0), - req_num0=2, - model_input1=SimpleNamespace(batch_size=batch_size1), - req_num1=2, - plan=SpecDecodePlan( - origin_batch_size=batch_size0 + batch_size1, - dynamic_batch_size=dynamic_batch_size, - draft_step=2, - pre_draft_step=3, - ), - ) - - assert tuple(dynamic_batch_sizes) == expected_batch_sizes - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") -def test_dp_overlap_decode_delegates_empty_layout(monkeypatch): - device = "cuda" - empty_i32 = torch.empty((0,), dtype=torch.int32, device=device) - model_input0 = SimpleNamespace(batch_size=0, b_req_idx=empty_i32, b_mtp_index=empty_i32) - model_input1 = SimpleNamespace(batch_size=0, b_req_idx=empty_i32, b_mtp_index=empty_i32) - model_output0 = SimpleNamespace(logits=torch.empty((0, 8), device=device)) - model_output1 = SimpleNamespace(logits=torch.empty((0, 8), device=device)) - calls = [] - - class _OverlapEngine: - def plan_decode(self, **kwargs): - calls.append("plan") - return SpecDecodePlan(0, 0, 2, 2) - - def prepare_decode_model_inputs(self, **kwargs): - calls.append("prepare") - return kwargs["model_input0"], None, kwargs["model_input1"], None - - def propose_next_overlap(self, **kwargs): - calls.append("propose") - assert kwargs["target_next_token_ids0"].shape == (0,) - assert kwargs["target_next_token_ids0"].dtype == torch.int64 - assert kwargs["target_next_token_ids1"].shape == (0,) - assert kwargs["target_next_token_ids1"].dtype == torch.int64 - assert kwargs["accept_len0"].shape == (0,) - assert kwargs["accept_len0"].dtype == torch.int32 - assert kwargs["accept_len1"].shape == (0,) - assert kwargs["accept_len1"].dtype == torch.int32 - assert kwargs["draft_step"] == 2 - return SpecProposal(token_ids=torch.empty((0, 2), dtype=torch.int64, device=device)) - - backend = DPChunkedPrefillBackend.__new__(DPChunkedPrefillBackend) - backend.decode_draft_engine = _OverlapEngine() - backend.args = SimpleNamespace(penalty_counter_mode="gpu_counter") - backend.model = SimpleNamespace( - microbatch_overlap_decode=lambda input0, input1: (model_output0, model_output1), - ) - event_pack = SimpleNamespace( - notify_post_handle_and_wait_pre_post_handle=lambda: calls.append("post_wait"), - notify_forward_and_wait_post_handle=lambda: calls.append("forward_wait"), - notify_pre_post_handle=lambda: calls.append("pre_post"), - ) - monkeypatch.setattr( - dp_backend_impl, - "overlap_prepare_decode_inputs", - lambda req_objs: (model_input0, [], [], model_input1, [], []), - ) - monkeypatch.setattr( - dp_backend_impl, - "g_infer_context", - SimpleNamespace(get_overlap_stream=lambda: torch.cuda.current_stream()), - ) - backend.decode_overlap_mtp(event_pack=event_pack, decode_reqs=[]) - - assert calls == [ - "plan", - "prepare", - "post_wait", - "propose", - "forward_wait", - "pre_post", - ] diff --git a/unit_tests/server/router/model_infer/mode_backend/test_generic_pre_process.py b/unit_tests/server/router/model_infer/mode_backend/test_generic_pre_process.py deleted file mode 100644 index 42519a7415..0000000000 --- a/unit_tests/server/router/model_infer/mode_backend/test_generic_pre_process.py +++ /dev/null @@ -1,176 +0,0 @@ -from types import SimpleNamespace - -import torch -from lightllm.server.router.model_infer.mode_backend import generic_pre_process - - -def _patch_empty_input_context(monkeypatch): - mem_manager = SimpleNamespace( - HOLD_TOKEN_MEMINDEXES=(-1,), - alloc=lambda size: torch.empty((size,), dtype=torch.int32), - ) - infer_context = SimpleNamespace( - args=SimpleNamespace(page_size=1), - req_manager=SimpleNamespace( - HOLD_REQUEST_ID=-1, - mem_manager=mem_manager, - req_to_token_indexs=torch.zeros((16, 32), dtype=torch.int32), - ), - radix_cache=None, - ) - monkeypatch.setattr(generic_pre_process, "g_infer_context", infer_context) - return infer_context - - -def _patch_overlap_input_context(monkeypatch): - return _patch_empty_input_context(monkeypatch) - - -def _make_prefill_req(req_idx: int, token_num: int): - input_token_ids = [req_idx] * token_num - return SimpleNamespace( - req_idx=req_idx, - cur_kv_len=0, - hold_kv_len=token_num, - multimodal_params={"images": [], "audios": []}, - get_chuncked_input_token_ids=lambda: input_token_ids, - get_input_token_ids=lambda: input_token_ids, - get_cur_total_len=lambda: token_num, - ) - - -def _make_decode_req(req_idx: int): - return SimpleNamespace( - req_idx=req_idx, - cur_kv_len=3, - hold_kv_len=4, - mtp_step=0, - multimodal_params={"images": [], "audios": []}, - shared_kv_node=None, - get_cur_total_len=lambda: 4, - get_radix_cache_shared_len=lambda: 0, - ) - - -def test_prepare_prefill_inputs_allows_empty_batch(monkeypatch): - _patch_empty_input_context(monkeypatch) - - model_input, run_reqs = generic_pre_process.prepare_prefill_inputs([], is_chuncked_mode=True) - - assert run_reqs == [] - assert model_input.batch_size == 0 - assert model_input.input_ids.shape == (0,) - assert model_input.b_req_idx.shape == (0,) - assert model_input.b_prefill_start_loc.shape == (0,) - assert model_input.b_prefill_has_output_cpu == [] - assert model_input.max_q_seq_len == 0 - assert model_input.max_kv_seq_len == 0 - - -def test_prepare_decode_inputs_allows_empty_batch(monkeypatch): - _patch_empty_input_context(monkeypatch) - - model_input, run_reqs = generic_pre_process.prepare_decode_inputs([]) - - assert run_reqs == [] - assert model_input.batch_size == 0 - assert model_input.input_ids is None - assert model_input.b_req_idx.shape == (0,) - assert model_input.b_position_delta.shape == (0,) - assert model_input.b_shared_seq_len.shape == (0,) - assert model_input.b_shared_radix_node_id.shape == (0,) - assert model_input.max_q_seq_len == 1 - assert model_input.max_kv_seq_len == 0 - - -def test_overlap_prefill_balances_request_token_load_without_padding(monkeypatch): - _patch_overlap_input_context(monkeypatch) - reqs = [ - _make_prefill_req(req_idx=0, token_num=8), - _make_prefill_req(req_idx=1, token_num=7), - _make_prefill_req(req_idx=2, token_num=6), - _make_prefill_req(req_idx=3, token_num=5), - _make_prefill_req(req_idx=4, token_num=1), - ] - - ( - model_input0, - run_reqs0, - model_input1, - run_reqs1, - ) = generic_pre_process.overlap_prepare_prefill_inputs(reqs) - - assert [req.req_idx for req in run_reqs0] == [0, 3, 4] - assert [req.req_idx for req in run_reqs1] == [1, 2] - assert model_input0.b_req_idx.tolist() == [0, 3, 4] - assert model_input1.b_req_idx.tolist() == [1, 2] - assert model_input0.input_ids.shape == (14,) - assert model_input1.input_ids.shape == (13,) - assert model_input0.batch_size == 3 - assert model_input1.batch_size == 2 - - -def test_overlap_prefill_balances_single_token_request_normally(monkeypatch): - _patch_overlap_input_context(monkeypatch) - req = _make_prefill_req(req_idx=7, token_num=1) - - ( - model_input0, - run_reqs0, - model_input1, - run_reqs1, - ) = generic_pre_process.overlap_prepare_prefill_inputs([req]) - - assert run_reqs0 == [req] - assert model_input0.batch_size == 1 - assert model_input0.input_ids.tolist() == [7] - assert model_input0.b_req_idx.tolist() == [7] - assert run_reqs1 == [] - assert model_input1.batch_size == 0 - assert model_input1.input_ids.shape == (0,) - assert model_input1.b_req_idx.shape == (0,) - - -def test_overlap_decode_builds_two_unpadded_inputs(monkeypatch): - _patch_overlap_input_context(monkeypatch) - reqs = [_make_decode_req(req_idx=index) for index in range(3)] - - ( - model_input0, - run_reqs0, - decode_reqs0, - model_input1, - run_reqs1, - decode_reqs1, - ) = generic_pre_process.overlap_prepare_decode_inputs(reqs) - - assert decode_reqs0 == reqs[:2] - assert decode_reqs1 == reqs[2:] - assert run_reqs0 == reqs[:2] - assert run_reqs1 == reqs[2:] - assert model_input0.batch_size == 2 - assert model_input1.batch_size == 1 - assert model_input0.b_req_idx.tolist() == [0, 1] - assert model_input1.b_req_idx.tolist() == [2] - - -def test_overlap_decode_preserves_empty_microbatch(monkeypatch): - _patch_overlap_input_context(monkeypatch) - req = _make_decode_req(req_idx=7) - - ( - model_input0, - run_reqs0, - decode_reqs0, - model_input1, - run_reqs1, - decode_reqs1, - ) = generic_pre_process.overlap_prepare_decode_inputs([req]) - - assert decode_reqs0 == [req] - assert decode_reqs1 == [] - assert run_reqs0 == [req] - assert model_input0.batch_size == 1 - assert run_reqs1 == [] - assert model_input1.batch_size == 0 - assert model_input1.b_req_idx.shape == (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 deleted file mode 100644 index ece87927e0..0000000000 --- a/unit_tests/server/router/model_infer/mode_backend/test_multi_level_kv_cache.py +++ /dev/null @@ -1,507 +0,0 @@ -from types import SimpleNamespace -from collections import deque - -import pytest -import torch - -from lightllm.server.multi_level_kv_cache import ( - AdaptiveCachePlacementController, - CacheCapacityConfig, - CacheTier, -) -from lightllm.server.router.model_infer.mode_backend.multi_level_kv_cache import MultiLevelKvCacheModule -from lightllm.server.router.model_infer.mode_backend import multi_level_kv_cache as multi_level_kv_cache_impl -from lightllm.server.router.model_infer.infer_batch import InferenceContext - - -def test_running_cpu_cache_tasks_are_not_ready(monkeypatch): - from lightllm.server.router.model_infer.infer_batch import InferReq - from lightllm.server.router.model_infer.mode_backend.base_backend import ModeBackend - from lightllm.server.router.model_infer.mode_backend import base_backend as base_backend_impl - - status = InferReq._CpuCacheTaskStatus - - def req(load_status, offload_status): - return SimpleNamespace( - cpu_cache_load_task_status=load_status, - cpu_cache_offload_task_status=offload_status, - ) - - running_load = req(status.RUNNING, status.NOT_STARTED) - running_offload = req(status.FINISHED, status.RUNNING) - ready = req(status.NOT_STARTED, status.NOT_STARTED) - backend = ModeBackend.__new__(ModeBackend) - backend.args = SimpleNamespace(enable_cpu_cache=True) - monkeypatch.setattr( - base_backend_impl, - "g_infer_context", - SimpleNamespace(requests_mapping={1: running_load, 2: running_offload, 3: ready}), - ) - - assert backend._filter_not_ready_reqs([1, 2, 3]) == [ready] - backend.args.enable_cpu_cache = False - assert backend._filter_not_ready_reqs([1, 2, 3]) == [running_load, running_offload, ready] - - -def test_cache_tiers_reassignment_is_rejected(): - context = InferenceContext() - context.cache_placement_controller = AdaptiveCachePlacementController( - capacity=CacheCapacityConfig(1, 1, 1), - args=SimpleNamespace(enable_cpu_cache=True, enable_disk_cache=True), - ) - reqs = [ - SimpleNamespace( - cache_tiers=(CacheTier.GPU,), - shm_req=SimpleNamespace(group_req_id=index, request_id=index, input_len=length), - ) - for index, length in enumerate((100, 200, 300)) - ] - - context.cache_placement_controller.set_req_cache_way(reqs) - context.cache_placement_controller = AdaptiveCachePlacementController( - capacity=CacheCapacityConfig(0, 1, 0), - args=SimpleNamespace(enable_cpu_cache=True, enable_disk_cache=False), - ) - - with pytest.raises(AssertionError): - context.cache_placement_controller.set_req_cache_way(reqs) - - -def test_non_gpu_cache_tiers_release_owned_tokens_without_radix_insert(): - released_refs = [] - context = InferenceContext() - context.is_hybrid_att_model = False - context.req_manager = SimpleNamespace() - context.radix_cache = SimpleNamespace(dec_node_ref_counter=released_refs.append) - shared_node = SimpleNamespace(node_prefix_total_len=2) - req = SimpleNamespace( - req_idx=0, - cur_kv_len=5, - hold_kv_len=7, - shared_kv_node=shared_node, - token_index=SimpleNamespace( - get_free_token_indexes=lambda start, end: torch.arange(10 + start, 10 + end), - get_cpu_token_indexes=lambda start, end: torch.arange(10 + start, 10 + end), - reset=lambda: None, - ), - ) - free_token_index_blocks = [] - - context._free_req_mem_without_radix_insert(free_token_index_blocks, req) - - assert free_token_index_blocks[0].tolist() == [12, 13, 14, 15, 16] - assert released_refs == [shared_node] - assert req.shared_kv_node is None - - -def test_legacy_cache_tiers_still_insert_gpu_radix_cache(): - context = InferenceContext() - context.radix_cache = object() - context.is_hybrid_att_model = False - inserted_reqs = [] - context._full_att_free_req = lambda free_token_index, req: inserted_reqs.append(req) - req = SimpleNamespace( - cache_tiers=(CacheTier.GPU, CacheTier.CPU, CacheTier.DISK), - cur_kv_len=3, - shm_req=SimpleNamespace(shm_cur_kv_len=3), - token_index=SimpleNamespace(reset=lambda: None), - ) - - context.free_a_req_mem([], req) - - assert inserted_reqs == [req] - - -def test_finished_batch_routes_cpu_and_disk_offloads_separately(monkeypatch): - class NotStartedStatus: - @staticmethod - def is_finished(): - return False - - @staticmethod - def is_running(): - return False - - @staticmethod - def is_not_started(): - return True - - module = MultiLevelKvCacheModule.__new__(MultiLevelKvCacheModule) - module.args = SimpleNamespace(cpu_cache_token_page_size=64, linear_att_hash_page_size=64) - module.backend = SimpleNamespace(radix_cache=object()) - module.offload_tasks = deque() - module.need_sync_compute_stream = lambda: False - offload_calls = [] - - def start_offload(req): - offload_calls.append((req.shm_req.request_id, CacheTier.DISK in req.cache_tiers)) - return SimpleNamespace(req=req) - - module.write_req_to_cpu = start_offload - current = torch.cuda.current_stream() - offload_stream = torch.cuda.Stream() - waited = [] - offload_stream.wait_stream = lambda stream: waited.append(stream) - module.get_cpu_offload_stream = lambda: offload_stream - monkeypatch.setattr(multi_level_kv_cache_impl.g_infer_context, "is_hybrid_att_model", False) - cache_tiers = ( - (CacheTier.GPU,), - (CacheTier.CPU,), - (CacheTier.CPU, CacheTier.DISK), - ) - reqs = [ - SimpleNamespace( - cache_tiers=req_cache_tiers, - cur_kv_len=length, - cpu_cache_offload_task_status=NotStartedStatus(), - shm_req=SimpleNamespace( - group_req_id=index, - request_id=index, - input_len=length, - ), - ) - for index, (length, req_cache_tiers) in enumerate(zip((100, 200, 300), cache_tiers)) - ] - offload_reqs = [req for req in reqs if CacheTier.CPU in req.cache_tiers or CacheTier.DISK in req.cache_tiers] - - true_finished_reqs = module.offload_finished_reqs_to_cpu_cache(offload_reqs) - - assert waited == [current] - assert true_finished_reqs == [] - assert offload_calls == [(1, False), (2, True)] - assert len(module.offload_tasks) == 2 - - overlap = torch.cuda.Stream() - waited_by_overlap = [] - overlap.wait_stream = lambda stream: waited_by_overlap.append(stream) - monkeypatch.setattr(multi_level_kv_cache_impl.g_infer_context, "get_overlap_stream", lambda: overlap) - module.need_sync_compute_stream = lambda: True - waited.clear() - assert module.offload_finished_reqs_to_cpu_cache([]) == [] - assert waited == [] and waited_by_overlap == [] - short_req = SimpleNamespace( - cache_tiers=(CacheTier.CPU,), - cur_kv_len=1, - cpu_cache_offload_task_status=NotStartedStatus(), - shm_req=SimpleNamespace(group_req_id=0, request_id=0, input_len=1), - ) - assert module.offload_finished_reqs_to_cpu_cache([short_req]) == [short_req] - assert waited == [current, overlap] - assert waited_by_overlap == [offload_stream] - - -def test_non_gpu_linear_cache_tiers_release_pending_state_pages(): - freed_small_pages = [] - freed_big_pages = [] - context = InferenceContext() - context.is_hybrid_att_model = True - context.req_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), - ) - req = SimpleNamespace( - req_idx=0, - cur_kv_len=3, - hold_kv_len=3, - token_index=SimpleNamespace( - get_free_token_indexes=lambda start, end: torch.arange(10 + start, 10 + end), - reset=lambda: None, - ), - shared_kv_node=None, - tail_small_page_buffer_id=7, - hybrid_len_to_big_page_id={128: 8, 256: 9}, - ) - free_token_index_blocks = [] - - context._free_req_mem_without_radix_insert(free_token_index_blocks, req) - - assert free_token_index_blocks[0].tolist() == [10, 11, 12] - assert freed_small_pages == [7] - assert freed_big_pages == [8, 9] - assert req.tail_small_page_buffer_id is None - assert req.hybrid_len_to_big_page_id == {} - - -def test_cpu_cache_load_uses_exact_aligned_size(monkeypatch): - table = torch.full((1, 136), -1, dtype=torch.int32) - table[0, :4] = torch.arange(4, dtype=torch.int32) - next_index = 4 - alloc_sizes = [] - loaded_indexes = [] - - def alloc(need_size): - nonlocal next_index - alloc_sizes.append(need_size) - mem_indexes = torch.arange(next_index, next_index + need_size, dtype=torch.int32) - next_index += need_size - return mem_indexes - - def alloc_token_indexes(req, end, allocator, page_size, radix_cache=None): - alloc_token_num = end - req.hold_kv_len - mem_indexes = alloc(alloc_token_num) - new_hold_kv_len = req.hold_kv_len + alloc_token_num - table[req.req_idx, req.hold_kv_len : new_hold_kv_len] = mem_indexes - req.hold_kv_len = new_hold_kv_len - return mem_indexes - - resources = SimpleNamespace( - 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() - ), - ), - full_att_allocator=None, - ) - backend = SimpleNamespace( - is_master_in_dp=False, - dp_world_size=1, - is_hybrid_att_model=False, - radix_cache=None, - model=SimpleNamespace( - att_state_manager=resources, - req_manager=SimpleNamespace( - req_to_token_indexs=table, alloc_token_indexes=alloc_token_indexes, wait_token_indexes=lambda: None - ), - ), - ) - module = MultiLevelKvCacheModule.__new__(MultiLevelKvCacheModule) - module.load_tasks = deque() - module.args = SimpleNamespace(page_size=4, cpu_cache_token_page_size=132) - module.backend = backend - module.init_sync_group = object() - module.need_sync_compute_stream = lambda: False - module.cpu_cache_client = SimpleNamespace(cpu_kv_cache_tensor=None) - context = SimpleNamespace( - req_manager=SimpleNamespace(req_to_token_indexs=table, alloc_token_indexes=alloc_token_indexes), - radix_cache=None, - get_can_alloc_token_num=lambda: 200, - ) - req = SimpleNamespace( - req_idx=0, - cur_kv_len=4, - hold_kv_len=4, - mm_input_spans=[], - shm_req=SimpleNamespace( - input_len=256, - disk_prompt_cache_len=0, - cpu_cache_match_page_indexes=SimpleNamespace(get_all=lambda: [7]), - token_hash_page_len_list=SimpleNamespace(get_all=lambda: [132]), - ), - sampling_param=SimpleNamespace(shm_param=SimpleNamespace(prompt_logprobs=-1)), - args=SimpleNamespace(page_size=4), - ) - req.token_index = SimpleNamespace(get_gpu_token_indexes=lambda start, end: table[req.req_idx, start:end]) - req.kv_cache_alloc_need = lambda target: max(0, (target + 3) // 4 * 4 - req.hold_kv_len) - req.alloc_kv_mem = lambda alloc_token_num: alloc_token_indexes(req, req.hold_kv_len + alloc_token_num, None, 4) - monkeypatch.setattr(multi_level_kv_cache_impl, "g_infer_context", context) - monkeypatch.setattr(multi_level_kv_cache_impl.dist, "barrier", lambda group: None) - monkeypatch.setattr(torch.Tensor, "cuda", lambda self, **kwargs: self) - - module.load_cpu_cache_to_reqs([req]) - - assert alloc_sizes == [128] - assert req.cur_kv_len == 132 - assert req.hold_kv_len == 132 - assert loaded_indexes == list(range(132)) - assert table[0, :132].tolist() == list(range(132)) - - -def test_standard_load_releases_unloaded_pages_immediately(monkeypatch): - module = MultiLevelKvCacheModule.__new__(MultiLevelKvCacheModule) - module.backend = SimpleNamespace(is_master_in_dp=True, dp_world_size=1) - module.init_sync_group = object() - module.need_sync_compute_stream = lambda: False - events = [] - module.cpu_cache_client = SimpleNamespace( - lock=SimpleNamespace( - acquire_sleep1ms=lambda: events.append("lock"), - release=lambda: events.append("unlock"), - ), - deref_pages=lambda page_list: events.append(("deref", list(page_list))), - ) - reqs = [ - SimpleNamespace( - cur_kv_len=0, - mm_input_spans=[], - sampling_param=SimpleNamespace(shm_param=SimpleNamespace(prompt_logprobs=logprobs)), - shm_req=SimpleNamespace( - input_len=64, - disk_prompt_cache_len=0, - cpu_cache_match_page_indexes=SimpleNamespace(get_all=lambda page=page: [page]), - token_hash_page_len_list=SimpleNamespace(get_all=lambda: [64]), - ), - ) - for page, logprobs in ((4, 0), (7, -1)) - ] - monkeypatch.setattr(multi_level_kv_cache_impl.g_infer_context, "get_can_alloc_token_num", lambda: 128) - - module.load_cpu_cache_to_reqs(reqs) - - assert events == ["lock", ("deref", [4]), ("deref", [7]), "unlock"] - assert not getattr(module, "load_tasks", None) - assert all(req.cpu_cache_load_task_status.is_finished() for req in reqs) - - -def _load_module(monkeypatch, *, sync, master=True, world=1): - events = [] - - def alloc_token_indexes(req, end, allocator, page_size, radix_cache=None): - req.hold_kv_len = end - return torch.arange(req.cur_kv_len, end, dtype=torch.int32) - - module = MultiLevelKvCacheModule.__new__(MultiLevelKvCacheModule) - module.args = SimpleNamespace(page_size=4, cpu_cache_token_page_size=128) - module.backend = SimpleNamespace( - is_master_in_dp=master, - dp_world_size=world, - is_hybrid_att_model=False, - model=SimpleNamespace( - req_manager=SimpleNamespace( - req_to_token_indexs=torch.arange(512, dtype=torch.int32).view(1, -1), - alloc_token_indexes=alloc_token_indexes, - wait_token_indexes=lambda: None, - ), - att_state_manager=SimpleNamespace( - get_cpu_cache_operator=lambda: SimpleNamespace( - get_page_layouts=lambda: {}, - read_req=lambda **kwargs: events.append("copy"), - ), - full_att_allocator=None, - ), - ), - ) - module.need_sync_compute_stream = lambda: sync - module.init_sync_group = module.filter_group = object() - module.load_sync_tensor = object() - module.load_tasks = deque() - module.cpu_cache_client = SimpleNamespace( - cpu_kv_cache_tensor=None, - lock=SimpleNamespace(acquire_sleep1ms=lambda: events.append("lock"), release=lambda: events.append("unlock")), - deref_pages=lambda page_list: events.append(("deref", list(page_list))), - ) - req = SimpleNamespace( - req_idx=0, - cur_kv_len=0, - hold_kv_len=0, - mm_input_spans=[], - sampling_param=SimpleNamespace(shm_param=SimpleNamespace(prompt_logprobs=-1)), - shm_req=SimpleNamespace( - input_len=256, - disk_prompt_cache_len=0, - cpu_prompt_cache_len=0, - shm_cur_kv_len=0, - cpu_cache_match_page_indexes=SimpleNamespace(get_all=lambda: [3, 4]), - token_hash_page_len_list=SimpleNamespace(get_all=lambda: [128, 256]), - ), - args=SimpleNamespace(page_size=4), - ) - token_table = module.backend.model.req_manager.req_to_token_indexs - req.token_index = SimpleNamespace(get_gpu_token_indexes=lambda start, end: token_table[req.req_idx, start:end]) - req.kv_cache_alloc_need = lambda target: max(0, (target + 3) // 4 * 4 - req.hold_kv_len) - req.alloc_kv_mem = lambda alloc_token_num: alloc_token_indexes(req, req.hold_kv_len + alloc_token_num, None, 4) - monkeypatch.setattr(multi_level_kv_cache_impl.g_infer_context, "get_can_alloc_token_num", lambda: 512) - monkeypatch.setattr(multi_level_kv_cache_impl.g_infer_context, "get_overlap_stream", torch.cuda.Stream) - monkeypatch.setattr(multi_level_kv_cache_impl.dist, "barrier", lambda group: events.append("barrier")) - monkeypatch.setattr(torch.Tensor, "cuda", lambda self, **kwargs: self) - return module, req, events - - -def test_async_load_defers_deref_until_the_event_is_notified(monkeypatch): - module, req, events = _load_module(monkeypatch, sync=False) - current = torch.cuda.current_stream() - overlap = torch.cuda.Stream() - waited_by_load = [] - overlap.wait_stream = lambda stream: events.append("wait") - load_stream = module.get_cpu_load_stream() - load_stream.wait_stream = lambda stream: waited_by_load.append(stream) - monkeypatch.setattr(multi_level_kv_cache_impl.g_infer_context, "get_overlap_stream", lambda: overlap) - recorded = [] - - class Event: - def __init__(self): - self.ready = False - - def record(self): - recorded.append(self) - - def query(self): - return self.ready - - monkeypatch.setattr(multi_level_kv_cache_impl.torch.cuda, "Event", Event) - module.load_cpu_cache_to_reqs([req]) - assert waited_by_load == [current] - assert events == ["copy"] and len(recorded) == 1 and len(module.load_tasks) == 1 - module.update_cpu_cache_load_task_states() - assert "deref" not in [item[0] for item in events if isinstance(item, tuple)] - recorded[0].ready = True - module.update_cpu_cache_load_task_states() - assert events[-3:] == ["lock", ("deref", [3, 4]), "unlock"] and not module.load_tasks - - -def test_synchronous_backend_waits_overlap_for_the_load_stream(monkeypatch): - module, req, events = _load_module(monkeypatch, sync=True) - current = torch.cuda.current_stream() - overlap = torch.cuda.Stream() - waited_by_overlap = [] - waited_by_load = [] - overlap.wait_stream = lambda stream: waited_by_overlap.append(stream) - load_stream = module.get_cpu_load_stream() - load_stream.wait_stream = lambda stream: waited_by_load.append(stream) - monkeypatch.setattr(multi_level_kv_cache_impl.g_infer_context, "get_overlap_stream", lambda: overlap) - module.load_cpu_cache_to_reqs([req]) - assert waited_by_load == [current, overlap] - assert waited_by_overlap == [load_stream] - assert events == ["copy"] and len(module.load_tasks) == 1 - module.load_tasks[0].event.synchronize() - module.update_cpu_cache_load_task_states() - assert events[-3:] == ["lock", ("deref", [3, 4]), "unlock"] and not module.load_tasks - - -def test_async_load_notifies_every_rank_before_deref(monkeypatch): - module, req, events = _load_module(monkeypatch, sync=False, world=2) - collectives = [] - - class Event: - def __init__(self): - self.ready = True - - def record(self): - return None - - def query(self): - return self.ready - - monkeypatch.setattr(multi_level_kv_cache_impl.torch.cuda, "Event", Event) - monkeypatch.setattr( - multi_level_kv_cache_impl.dist, - "all_reduce", - lambda tensor, group: collectives.append(("all_reduce", group)), - ) - monkeypatch.setattr( - multi_level_kv_cache_impl.dist, - "broadcast_object_list", - lambda data, group, group_src: collectives.append(("broadcast", data[0], group)), - ) - module.load_cpu_cache_to_reqs([req]) - assert collectives == [("all_reduce", module.init_sync_group)] and not any( - isinstance(item, tuple) and item[0] == "deref" for item in events - ) - module.update_cpu_cache_load_task_states() - assert collectives[1][0] == "broadcast" and collectives[1][1] == 1 - assert ("deref", [3, 4]) in events - - -def test_fa3_cpu_cache_sync_defaults_off(): - from lightllm.common.basemodel.attention.base_att import BaseAttBackend - from lightllm.common.basemodel.attention.fa3.fp import Fa3AttBackend - from lightllm.common.basemodel.attention.fa3.fp8 import Fp8Fa3AttBackend - from lightllm.common.basemodel.attention.fa3.mla import MlaFa3AttBackend - from lightllm.common.basemodel.attention.fa3.neo import NeoFa3AttBackend - - assert BaseAttBackend.requires_cpu_cache_sync is False - assert all( - backend.requires_cpu_cache_sync is False - for backend in (Fa3AttBackend, Fp8Fa3AttBackend, MlaFa3AttBackend, NeoFa3AttBackend) - ) 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 deleted file mode 100644 index 17d564987d..0000000000 --- a/unit_tests/server/router/model_infer/mode_backend/test_nccl_control_channel.py +++ /dev/null @@ -1,151 +0,0 @@ -from __future__ import annotations - -import concurrent.futures -import socket -import threading -from types import SimpleNamespace -from unittest.mock import Mock - -import rpyc -import pytest - -from lightllm.server.router.model_infer.mode_backend.pd import nccl_kv_transporter as module - - -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.mode == "send" - assert handle.finish_cuda_event is event - event.query.return_value = True - assert handle.check_status() == "DONE" - - 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 - 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): - 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 deleted file mode 100644 index 623f9564b0..0000000000 --- a/unit_tests/server/router/model_infer/mode_backend/test_paged_side_allocations.py +++ /dev/null @@ -1,253 +0,0 @@ -from types import SimpleNamespace - -import numpy as np -import pytest -import torch - -# 先加载通用预处理模块，沿用生产初始化顺序，避免 basemodel 与 req_manager -# 在独立运行本测试文件时形成模块级循环导入。 -from lightllm.server.router.model_infer.mode_backend import generic_pre_process # noqa: F401 -from lightllm.server.router.model_infer.infer_batch import InferReq -from lightllm.server.router.model_infer.mode_backend.dp_backend import dp_shared_kv_trans -from lightllm.server.router.model_infer.mode_backend.pd.decode_node_impl import decode_impl -from lightllm.server.router.model_infer.mode_backend.pd.prefill_node_impl import prefill_impl - - -@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) - table[0, :4] = torch.arange(4, dtype=torch.int32) - alloc_sizes = [] - next_index = [4] - tasks = [] - backend = decode_impl.PDDecodeNode.__new__(decode_impl.PDDecodeNode) - backend.args = SimpleNamespace(pd_kv_page_size=4, page_size=4) - req_manager = _make_req_manager(table, alloc_sizes, next_index) - backend.model = SimpleNamespace( - req_manager=req_manager, - att_state_manager=SimpleNamespace(full_att_allocator=None), - ) - backend.is_master_in_dp = False - backend.radix_cache = None - - def create_task(**kwargs): - tasks.append( - ( - kwargs["kv_start_index"], - kwargs["kv_end_index"], - kwargs["mem_indexes"], - kwargs.get("page_kind", "kv"), - ) - ) - kwargs["group"].task_list.append(SimpleNamespace()) - - backend._create_pd_trans_task = create_task - monkeypatch.setattr(decode_impl.g_infer_context, "is_hybrid_att_model", is_hybrid_att_model) - req = _bind_alloc_need( - SimpleNamespace( - req_idx=0, - cur_kv_len=4, - hold_kv_len=4, - pd_trans_kv_start_index=0, - shm_req=SimpleNamespace(input_len=10), - ), - page_size=4, - ) - req.kv_cache_alloc_need = lambda target: req_manager.get_need_alloc_token_num(req, target) - req.alloc_kv_mem = lambda alloc_token_num: req_manager.alloc_token_indexes( - req, req.cur_kv_len + alloc_token_num, None, req.args.page_size - ) - - backend._decode_node_gen_trans_tasks(req) - - assert alloc_sizes == [8] - assert req.cur_kv_len == 10 - assert req.hold_kv_len == 12 - expected_tasks = [ - (4, 8, [4, 5, 6, 7], "kv"), - (8, 10, [8, 9], "kv"), - ] - if is_hybrid_att_model: - expected_tasks.append((10, 10, [], "att_state")) - assert tasks == expected_tasks - assert table[0, :12].tolist() == list(range(12)) - - -def test_pd_decode_rejects_unaligned_transfer_page_size(): - backend = decode_impl.PDDecodeNode.__new__(decode_impl.PDDecodeNode) - backend.args = SimpleNamespace(pd_kv_page_size=3, page_size=4) - 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._decode_node_gen_trans_tasks(req) - - -def test_pd_prefill_rejects_unaligned_transfer_page_size(): - backend = prefill_impl.PDChunkedPrefillForPrefillNode.__new__(prefill_impl.PDChunkedPrefillForPrefillNode) - backend.args = SimpleNamespace(pd_kv_page_size=3, page_size=4) - 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_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)]) -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 = [] - next_index = [4] - req = _bind_alloc_need( - SimpleNamespace( - req_idx=0, - cur_kv_len=4, - hold_kv_len=4, - get_cur_total_len=lambda: 20, - ), - page_size=4, - ) - source_table = torch.arange(32, dtype=torch.int32).reshape(2, 16) - backend = SimpleNamespace( - model=SimpleNamespace( - mem_manager=SimpleNamespace(), - att_state_manager=SimpleNamespace(full_att_allocator=SimpleNamespace(page_size=4)), - ), - dp_world_size=tp_size, - rank_in_dp=tp_rank, - node_nccl_group=None, - dp_req_to_token_indices={0: source_table, 1: source_table}, - is_hybrid_att_model=False, - ) - - module = dp_shared_kv_trans.DPKVSharedMoudle.__new__(dp_shared_kv_trans.DPKVSharedMoudle) - module.backend = backend - module.dp_rank_in_node = 0 - module.shared_req_infos = SimpleNamespace(arr=np.array([[[4, 0], [10, 0]]], dtype=np.int64)) - req_manager = _make_req_manager(table, alloc_sizes, next_index) - req.kv_cache_alloc_need = lambda target: req_manager.get_need_alloc_token_num(req, target) - req.alloc_kv_mem = lambda alloc_token_num: req_manager.alloc_token_indexes( - req, req.cur_kv_len + alloc_token_num, None, req.args.page_size - ) - context = SimpleNamespace( - req_manager=req_manager, - radix_cache=None, - get_can_alloc_token_num=lambda: 100, - ) - monkeypatch.setattr(dp_shared_kv_trans, "g_infer_context", context) - monkeypatch.setattr(dp_shared_kv_trans.dist, "barrier", lambda group: None) - - tasks = module.build_shared_kv_trans_tasks([req], req_dp_ranks=[0]) - - assert alloc_sizes == [8] - assert req.cur_kv_len == 4 - assert req.hold_kv_len == 12 - assert len(tasks) == 1 - assert tasks[0].src_dp_rank == 1 - assert tasks[0].dst_token_indices.tolist() == list(range(4, 10)) - assert len(tasks[0].src_token_indices) == 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( - get_dp_shared_operator=lambda rank=rank: SimpleNamespace( - register_page_pools_pin_memory=lambda: 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_state_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_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_state_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) - 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 - - -def _make_req_manager(table, alloc_sizes, next_index): - def need(req, end, page_size=None): - page_size = req.args.page_size if page_size is None else page_size - return max(0, (end + page_size - 1) // page_size * page_size - req.hold_kv_len) - - def allocate(req, end, allocator, page_size, radix_cache=None): - _reserve_into_table(table, req, alloc_sizes, next_index, need(req, end)) - return table[req.req_idx, req.cur_kv_len : end].clone() - - return SimpleNamespace(req_to_token_indexs=table, alloc_token_indexes=allocate, get_need_alloc_token_num=need) diff --git a/unit_tests/server/router/model_infer/mode_backend/test_pd_dynamic_split.py b/unit_tests/server/router/model_infer/mode_backend/test_pd_dynamic_split.py deleted file mode 100644 index 38e60dd255..0000000000 --- a/unit_tests/server/router/model_infer/mode_backend/test_pd_dynamic_split.py +++ /dev/null @@ -1,244 +0,0 @@ -from types import MethodType, SimpleNamespace -from unittest.mock import MagicMock - -import pytest - -from lightllm.common.req_manager import ReqManager -from lightllm.server.core.objs import FinishStatus -from lightllm.server.core.objs import req as req_module -from lightllm.server.router.model_infer.mode_backend import base_backend -from lightllm.server.router.model_infer.mode_backend.pd.decode_node_impl import ( - decode_impl as pd_decode_impl, -) -from lightllm.server.router.req_queue import _get_req_queue_class -from lightllm.server.router.req_queue.chunked_prefill.impl_for_pd_decode import PDDecodeQueue -from lightllm.server.router.req_queue.chunked_prefill.impl_for_pd_prefill import PDPrefillQueue - - -def _make_infer_req(cur_output_len: int, shm_output_len: int): - return SimpleNamespace( - req_id=1, - filter_mark=False, - wait_pause=False, - paused=False, - infer_aborted=False, - finish_status=FinishStatus(), - cpu_cache_offload_task_status=SimpleNamespace(is_not_started=MagicMock(return_value=True)), - cur_kv_len=10, - hold_kv_len=8, - cur_output_len=cur_output_len, - shm_req=SimpleNamespace(shm_cur_output_len=shm_output_len), - sampling_param=SimpleNamespace( - shm_param=SimpleNamespace(max_new_tokens=65535), - ), - get_cur_total_len=MagicMock(return_value=11), - args=SimpleNamespace(page_size=4), - decode_need_token_num=MagicMock(return_value=(1, 4)), - kv_cache_alloc_need=lambda target: max(0, (target + 3) // 4 * 4 - 8), - alloc_kv_mem=lambda alloc_token_num: None, - ) - - -def _classify_without_token_capacity(monkeypatch, req, support_overlap=True): - reqs = req if isinstance(req, list) else [req] - backend = pd_decode_impl.PDDecodeNode.__new__(pd_decode_impl.PDDecodeNode) - backend.args = SimpleNamespace( - enable_cpu_cache=False, - enable_prefill_decode_mixed=False, - run_mode="decode", - page_size=4, - ) - backend.support_overlap = support_overlap - backend.is_master_in_dp = True - logger = MagicMock() - backend.logger = logger - backend._timer_merge_radix_tree = MagicMock() - backend._filter_not_ready_reqs = MagicMock(return_value=reqs) - backend.prefill_queue_strategy = SimpleNamespace(reorder=MagicMock(side_effect=lambda reqs, **kwargs: reqs)) - - backend.batch_max_tokens = 16 - infer_context = base_backend.g_infer_context - manager = ReqManager.__new__(ReqManager) - manager.page_size = 4 - monkeypatch.setattr(infer_context, "req_manager", manager) - monkeypatch.setattr( - infer_context, - "att_state_manager", - SimpleNamespace(full_att_allocator=SimpleNamespace(page_size=4, can_use_mem_size=0)), - ) - monkeypatch.setattr(infer_context, "get_can_alloc_token_num", MagicMock(return_value=0)) - monkeypatch.setattr( - infer_context, - "cache_placement_controller", - SimpleNamespace(set_req_cache_way=MagicMock()), - ) - filter_reqs = MagicMock() - monkeypatch.setattr(infer_context, "filter_reqs", filter_reqs) - monkeypatch.setattr(infer_context, "pause_reqs", MagicMock()) - - backend._get_classed_reqs(req_ids=[req.req_id for req in reqs]) - return filter_reqs, logger - - -def test_pd_decode_capacity_shortage_is_delayed_for_overlap(monkeypatch): - req = _make_infer_req(cur_output_len=5, shm_output_len=4) - - filter_reqs, logger = _classify_without_token_capacity(monkeypatch, req, support_overlap=True) - - assert req.filter_mark - assert req.finished_by_pd_decode_capacity - assert not req.finish_status.is_finished() - assert req.sampling_param.shm_param.max_new_tokens == 65535 - assert not req.wait_pause - filter_reqs.assert_called_once_with(finished_reqs=[]) - assert logger.info.call_args.args[0] == ( - "force early finish for PD decode req_id=1 because token capacity is insufficient" - ) - - filter_reqs, _ = _classify_without_token_capacity(monkeypatch, req, support_overlap=True) - filter_reqs.assert_called_once_with(finished_reqs=[req]) - - -def test_pd_decode_capacity_shortage_is_filtered_without_overlap(monkeypatch): - req = _make_infer_req(cur_output_len=5, shm_output_len=4) - - filter_reqs, logger = _classify_without_token_capacity(monkeypatch, req, support_overlap=False) - - assert not req.filter_mark - assert req.finished_by_pd_decode_capacity - assert not req.finish_status.is_finished() - assert req.sampling_param.shm_param.max_new_tokens == 65535 - assert not req.wait_pause - filter_reqs.assert_called_once_with(finished_reqs=[req]) - assert logger.info.call_args.args[0] == ( - "force early finish for PD decode req_id=1 because token capacity is insufficient" - ) - - -def test_pd_decode_capacity_shortage_only_handles_two_requests_per_iteration(monkeypatch): - reqs = [_make_infer_req(cur_output_len=5, shm_output_len=4) for _ in range(3)] - for req_id, req in enumerate(reqs, start=1): - req.req_id = req_id - - filter_reqs, _ = _classify_without_token_capacity(monkeypatch, reqs, support_overlap=True) - - assert [req.filter_mark for req in reqs] == [True, True, False] - filter_reqs.assert_called_once_with(finished_reqs=[]) - - -def test_pd_decode_capacity_finish_status_is_written_to_shm(): - shm_req = SimpleNamespace(mark_simulated_finished=MagicMock()) - req = SimpleNamespace( - finish_status=FinishStatus(), - finished_by_pd_decode_capacity=True, - shm_req=shm_req, - cur_output_len=3, - ) - - pd_decode_impl.InferReq.mark_shm_aborted_finished(req) - - shm_req.mark_simulated_finished.assert_called_once_with( - FinishStatus.FINISHED_PD_DECODE_CAPACITY, - output_len=3, - ) - - -def test_request_finish_status_takes_priority_over_pd_decode_capacity_marker(): - shm_req = SimpleNamespace(mark_simulated_finished=MagicMock()) - req = SimpleNamespace( - finish_status=FinishStatus(FinishStatus.FINISHED_STOP), - finished_by_pd_decode_capacity=True, - shm_req=shm_req, - cur_output_len=3, - ) - - pd_decode_impl.InferReq.mark_shm_aborted_finished(req) - - shm_req.mark_simulated_finished.assert_not_called() - - -def test_pd_decode_capacity_limit_never_extends_original_length(monkeypatch): - req = _make_infer_req(cur_output_len=5, shm_output_len=1) - req.sampling_param.shm_param.max_new_tokens = 3 - - _classify_without_token_capacity(monkeypatch, req) - assert req.sampling_param.shm_param.max_new_tokens == 3 - - -@pytest.mark.parametrize( - "overrides", - [{}, {"diverse_mode": True}, {"output_constraint_mode": "xgrammar"}], -) -def test_pd_nodes_use_pd_queue(overrides): - base_args = { - "diverse_mode": False, - "token_healing_mode": False, - "output_constraint_mode": "none", - "disable_chunked_prefill": False, - } - - base_args.update(overrides) - prefill_args = SimpleNamespace(**base_args, run_mode="prefill") - decode_args = SimpleNamespace(**base_args, run_mode="decode") - - assert _get_req_queue_class(prefill_args, router=None, dp_size_in_node=1) is PDPrefillQueue - assert _get_req_queue_class(decode_args, router=None, dp_size_in_node=1) is PDDecodeQueue - - -@pytest.mark.parametrize("page_size", [1, 16]) -def test_pd_decode_queue_uses_ema_for_prefill_stage_output_length(page_size, monkeypatch): - monkeypatch.setattr(req_module, "get_env_start_args", lambda: SimpleNamespace(page_size=page_size, mtp_step=0)) - queue = PDDecodeQueue.__new__(PDDecodeQueue) - queue.args = SimpleNamespace(run_mode="decode", page_size=page_size) - queue.dp_index = 0 - queue.max_total_tokens = 4096 - queue.running_max_req_size = 8 - queue.router = SimpleNamespace( - router_statics=SimpleNamespace(ema_req_out_len=128), - shared_token_load=SimpleNamespace( - set_estimated_peak_token_count=MagicMock(), - set_dynamic_max_load=MagicMock(), - ), - ) - queue.is_busy = MagicMock(return_value=False) - - req = SimpleNamespace( - input_len=10, - shm_cur_output_len=0, - shm_cur_kv_len=0, - sample_params=SimpleNamespace(suggested_dp_index=0, max_new_tokens=1024, ignore_eos=False), - is_infer_decode=MagicMock(return_value=False), - ) - req.get_tuple_tokens = MagicMock(wraps=MethodType(req_module.ChunkedPrefillReq.get_tuple_tokens, req)) - batch = SimpleNamespace(reqs=[req]) - - assert queue._caclu_batch_estimated_peak_token_num(batch) == 156 + page_size - req.get_tuple_tokens.assert_called_once_with(False, 128) - req.get_tuple_tokens.reset_mock() - assert queue._can_add_new_req(req, estimated_peak_token_num=0, batch_req_num=0) == (True, 156 + page_size, 1) - req.get_tuple_tokens.assert_called_once_with(False, 128) - - # 接近上下文上限的请求可能只剩很少输出额度，预计输出长度仍受 max_new_tokens 约束， - # 并在此基础上计入分页、MTP 和异步退出余量，避免因 EMA 偏大而永久滞留在 Decode 等待队列。 - req.sample_params.max_new_tokens = 20 - assert queue._caclu_batch_estimated_peak_token_num(batch) == 48 + page_size - assert queue._can_add_new_req(req, estimated_peak_token_num=0, batch_req_num=0) == (True, 48 + page_size, 1) - - req.sample_params.max_new_tokens = 1024 - req.get_tuple_tokens.reset_mock() - prefill_queue = PDPrefillQueue.__new__(PDPrefillQueue) - prefill_queue.args = SimpleNamespace(run_mode="prefill", page_size=page_size) - prefill_queue.dp_index = queue.dp_index - prefill_queue.max_total_tokens = queue.max_total_tokens - prefill_queue.running_max_req_size = queue.running_max_req_size - prefill_queue.router = queue.router - prefill_queue.is_busy = MagicMock(return_value=False) - assert prefill_queue._caclu_batch_estimated_peak_token_num(batch) == 156 + page_size - req.get_tuple_tokens.assert_called_once_with(False, 128) - req.get_tuple_tokens.reset_mock() - assert prefill_queue._can_add_new_req(req, estimated_peak_token_num=0, batch_req_num=0) == ( - True, - 156 + page_size, - 1, - ) - req.get_tuple_tokens.assert_called_once_with(False, 128) diff --git a/unit_tests/server/router/model_infer/mode_backend/test_pd_manager_startup.py b/unit_tests/server/router/model_infer/mode_backend/test_pd_manager_startup.py deleted file mode 100644 index 9dbf0c8025..0000000000 --- a/unit_tests/server/router/model_infer/mode_backend/test_pd_manager_startup.py +++ /dev/null @@ -1,174 +0,0 @@ -import ctypes -import importlib -import multiprocessing as mp -import os -import signal -import sys -import time -import traceback -from types import ModuleType, SimpleNamespace - -import psutil -import pytest - -from lightllm.utils import process_check - - -pytestmark = pytest.mark.skipif(sys.platform != "linux", reason="requires Linux subreaper and fork") - - -def _kill_and_reap_children(): - for child in psutil.Process().children(recursive=True): - try: - os.kill(child.pid, signal.SIGKILL) - except ProcessLookupError: - pass - while True: - try: - os.waitpid(-1, 0) - except ChildProcessError: - break - - -def _run_startup_failure(mode, late_monitor, result_connection): - """Own orphaned workers in an isolated subreaper so the test can always reap them.""" - result = None - try: - libc = ctypes.CDLL(None, use_errno=True) - if libc.prctl(36, 1, 0, 0, 0) != 0: # PR_SET_CHILD_SUBREAPER - raise OSError(ctypes.get_errno(), "could not become a child subreaper") - - ctx = mp.get_context("fork") - events, event_writer = ctx.Pipe(duplex=False) - start_monitor = ctx.Event() - if not late_monitor: - start_monitor.set() - - # Speed up the real monitor without changing the shared time module's sleep. - process_check.time = SimpleNamespace(sleep=lambda _: time.sleep(0.01)) - - def worker_main(started): - if late_monitor: - started.set() - assert start_monitor.wait(timeout=10) - event_writer.send(("monitor_parent", os.getpid(), os.getppid())) - process_check.start_parent_check_thread() - started.set() - while True: - time.sleep(100) - - def start_worker(*args, **kwargs): - started = ctx.Event() - worker = ctx.Process(target=worker_main, args=(started,)) - worker.start() - event_writer.send(("worker", worker.pid)) - assert not worker.daemon - assert started.wait(timeout=5) - - module_name = "lightllm.server.router.model_infer.mode_backend.pd." f"{mode}_node_impl.{mode}_kv_move_manager" - manager_module = importlib.import_module(module_name) - manager_module.get_unique_server_name = lambda: "startup-regression" - original_logger = manager_module.logger - - def log_failure(message): - original_logger.exception(message) - event_writer.send(("init_error",)) - - manager_module.logger = SimpleNamespace(exception=log_failure) - - def fail_initialization(*args, **kwargs): - start_worker() - raise RuntimeError("injected transport initialization failure") - - manager_class = "PrefillKVMoveManager" if mode == "prefill" else "DecodeKVMoveManager" - setattr(manager_module, manager_class, fail_initialization) - - # Exercise the real entry point while avoiding CUDA and network initialization. - package_name = module_name.rsplit(".", 1)[0] - transport = ModuleType(f"{package_name}.{mode}_trans_process") - setattr(transport, f"start_{mode}_trans_process", lambda *args, **kwargs: None) - sys.modules[transport.__name__] = transport - sys.modules["lightllm.utils.rpyc_fix_utils"] = ModuleType("lightllm.utils.rpyc_fix_utils") - if mode == "decode": - up_status = ModuleType(f"{package_name}.up_status") - up_status.start_up_kv_status_process = start_worker - sys.modules[up_status.__name__] = up_status - - ready = ctx.Event() - manager = ctx.Process(target=manager_module._init_env, args=(SimpleNamespace(), None, ready)) - manager.start() - manager.join(timeout=5) - assert not manager.is_alive(), "manager hung while joining its non-daemon children" - assert manager.exitcode == 1 - assert not ready.is_set() - - # In the late case workers are already reparented when they start the monitor. - start_monitor.set() - worker_count = 2 if mode == "decode" else 1 - workers = set() - monitor_parents = {} - logged_failure = False - deadline = time.monotonic() + 5 - while len(workers) < worker_count or len(monitor_parents) < worker_count or not logged_failure: - assert time.monotonic() < deadline, "missing worker or initialization failure report" - if not events.poll(0.05): - continue - event = events.recv() - if event[0] == "worker": - workers.add(event[1]) - elif event[0] == "monitor_parent": - monitor_parents[event[1]] = event[2] - elif event[0] == "init_error": - logged_failure = True - - expected_parent = os.getpid() if late_monitor else manager.pid - assert set(monitor_parents) == workers - assert set(monitor_parents.values()) == {expected_parent} - - for pid in workers: - while True: - reaped_pid, status = os.waitpid(pid, os.WNOHANG) - if reaped_pid: - assert os.WIFSIGNALED(status) - assert os.WTERMSIG(status) == signal.SIGKILL - break - assert time.monotonic() < deadline, "worker did not detect its dead parent" - time.sleep(0.01) - result = "ok" - except BaseException: - result = traceback.format_exc() - finally: - _kill_and_reap_children() - result_connection.send(result) - - -@pytest.mark.parametrize("mode", ["prefill", "decode"]) -@pytest.mark.parametrize("late_monitor", [False, True], ids=["monitor-running", "monitor-after-parent-exit"]) -def test_failed_manager_exits_and_workers_detect_parent_death(mode, late_monitor): - # Import before forking so costly framework imports happen only once per test run. - importlib.import_module( - f"lightllm.server.router.model_infer.mode_backend.pd.{mode}_node_impl.{mode}_kv_move_manager" - ) - ctx = mp.get_context("fork") - result_reader, result_writer = ctx.Pipe(duplex=False) - controller = ctx.Process(target=_run_startup_failure, args=(mode, late_monitor, result_writer)) - controller.start() - try: - controller.join(timeout=15) - assert not controller.is_alive(), "startup regression controller timed out" - assert controller.exitcode == 0 - assert result_reader.poll(timeout=1), "startup regression controller returned no result" - assert result_reader.recv() == "ok" - finally: - if controller.is_alive(): - # The controller normally owns cleanup; retain a bounded test-runner fallback. - for child in psutil.Process(controller.pid).children(recursive=True): - try: - os.kill(child.pid, signal.SIGKILL) - except ProcessLookupError: - pass - controller.kill() - controller.join(timeout=5) - controller.close() - result_reader.close() - result_writer.close() 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 deleted file mode 100644 index 8f99c6add3..0000000000 --- a/unit_tests/server/router/model_infer/mode_backend/test_pd_notify.py +++ /dev/null @@ -1,223 +0,0 @@ -import pickle -import subprocess -import sys -from types import SimpleNamespace -from unittest.mock import Mock - -import pytest - -from lightllm.server.pd_io_struct import PDChunckedTransTask -from lightllm.server.router.model_infer.mode_backend.pd.decode_node_impl.decode_trans_process import ( - _SharedTimeOutMark, -) -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, -) - - -def make_task(page_kind): - task = PDChunckedTransTask( - request_id=123, - start_kv_index=4, - end_kv_index=5 if page_kind == "kv" else 4, - time_out_secs=10, - pd_master_node_id=2, - prefill_dp_index=0, - decode_dp_index=1, - src_device_id=0, - dst_device_id=1, - mem_indexes=[7] if page_kind == "kv" else [], - prefill_agent_name="prefill", - prefill_agent_metadata=b"prefill metadata", - prefill_num_pages=8, - prefill_page_reg_desc=b"prefill pages", - decode_agent_name="decode", - decode_agent_metadata=b"decode metadata", - decode_num_pages=9, - decode_page_reg_desc=b"decode pages", - first_gen_token_id=42, - first_gen_token_logprob=-0.5, - page_kind=page_kind, - req_idx=17, - src_page_index=3, - dst_page_index=6, - xfer_handle=object(), - ) - return task - - -@pytest.mark.parametrize("page_kind", ["kv", "att_state"]) -@pytest.mark.parametrize("with_mark", [False, True]) -def test_copy_for_notify_preserves_local_task(page_kind, with_mark): - task = make_task(page_kind) - if with_mark: - task.shared_time_out_mark = _SharedTimeOutMark() - time_mark = task.shared_time_out_mark.time_mark - original = vars(task).copy() - mem_indexes = list(task.mem_indexes) - - notify = task.copy_for_notify() - - assert notify is not task - assert not hasattr(notify, "shared_time_out_mark") - assert notify.mem_indexes is None - assert notify.xfer_handle is None - assert vars(notify) == { - **{key: value for key, value in original.items() if key != "shared_time_out_mark"}, - "mem_indexes": None, - "xfer_handle": None, - } - assert vars(task) == original - assert all(vars(task)[key] is value for key, value in original.items()) - assert task.mem_indexes == mem_indexes - if with_mark: - assert task.shared_time_out_mark.time_mark == time_mark - - -_UNPICKLE_WITHOUT_DECODE = """ -import importlib.abc -import pickle -import sys -from lightllm.server.pd_io_struct import PDChunckedTransTask - -class BlockDecode(importlib.abc.MetaPathFinder): - def find_spec(self, fullname, path=None, target=None): - if 'decode_node_impl' in fullname: - raise ImportError('decode modules are unavailable to the receiver') - -assert not any('decode_node_impl' in name for name in sys.modules) -sys.meta_path.insert(0, BlockDecode()) -task = pickle.loads(sys.stdin.buffer.read()) -assert isinstance(task, PDChunckedTransTask) -assert not hasattr(task, 'shared_time_out_mark') -assert task.mem_indexes is None -assert task.xfer_handle is None -assert not any('decode_node_impl' in name for name in sys.modules) -""" - - -@pytest.mark.parametrize("backend", [NcclKVTransporter, NixlKVTransporter]) -@pytest.mark.parametrize("page_kind", ["kv", "att_state"]) -@pytest.mark.parametrize( - "method,stage,destination", - [ - ("send_write_done_task_to_decode_node", "done", "decode"), - ("send_write_request_task_to_decode_node", "request", "decode"), - ("send_write_ready_task_to_prefill_node", "ready", "prefill"), - ("send_error_info_to_prefill_node", "error", "prefill"), - ("send_error_info_to_decode_node", "error", "decode"), - ], -) -def test_notification_payload(backend, page_kind, method, stage, destination): - task = make_task(page_kind) - task.error_info = "transfer failed" if stage == "error" else None - task.shared_time_out_mark = _SharedTimeOutMark() - original = vars(task).copy() - time_mark = task.shared_time_out_mark.time_mark - transporter = backend.__new__(backend) - transporter.node_id, transporter.tp_idx = 10, 0 - transporter.num_pages = 12 - transporter.remote_agents = {"prefill": object(), "decode": object()} - 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="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: - transporter.page_reg_desc = object() - transporter.nixl_agent = SimpleNamespace( - get_agent_metadata=lambda: b"local metadata", - get_serialized_descs=lambda _: b"local pages", - send_notif=Mock(), - ) - - getattr(transporter, method)(task) - - if backend is NcclKVTransporter: - transporter.control_channel.send_notif.assert_called_once() - ( - peer, - host, - port, - payload, - ) = transporter.control_channel.send_notif.call_args.args - assert (peer, host, port) == (destination, "127.0.0.2", 4321) - else: - transporter.nixl_agent.send_notif.assert_called_once() - kwargs = transporter.nixl_agent.send_notif.call_args.kwargs - assert kwargs["remote_agent_name"] == destination - payload = kwargs["notif_msg"] - result = subprocess.run( - [sys.executable, "-c", _UNPICKLE_WITHOUT_DECODE], - input=payload, - capture_output=True, - timeout=30, - ) - assert result.returncode == 0, result.stderr.decode() - notify = pickle.loads(payload) - sender = "prefill" if destination == "decode" else "decode" - expected = {key: value for key, value in original.items() if key != "shared_time_out_mark"} - expected.update(mem_indexes=None, xfer_handle=None, write_stage=stage) - expected.update( - { - f"{sender}_agent_name": transporter.agent_name, - f"{sender}_agent_metadata": transporter.agent_metadata, - f"{sender}_num_pages": transporter.num_pages, - f"{sender}_page_reg_desc": transporter.local_page_mem_desc, - } - ) - if backend is NixlKVTransporter and stage == "done": - expected.update(decode_agent_metadata=None, decode_page_reg_desc=None) - assert vars(notify) == expected - assert vars(task) == original - assert all(vars(task)[key] is value for key, value in original.items()) - assert task.shared_time_out_mark.time_mark == time_mark - - -def test_receiver_rejects_payload_with_decode_local_mark(): - task = make_task("kv") - task.shared_time_out_mark = _SharedTimeOutMark() - result = subprocess.run( - [sys.executable, "-c", _UNPICKLE_WITHOUT_DECODE], - input=pickle.dumps(task), - capture_output=True, - timeout=30, - ) - 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_bytes = 128 - transporter.nixl_agent = SimpleNamespace(get_xfer_descs=Mock(), prep_xfer_dlist=Mock()) - 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_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)) - 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_bytes": 256}), - ) - with pytest.raises(AssertionError, match="same page byte capacity"): - transporter.connect_add_remote_agent(peer) - assert not transporter.remote_agents - peer.page_reg_desc = transporter.local_page_mem_desc - 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_pd_page_io.py b/unit_tests/server/router/model_infer/mode_backend/test_pd_page_io.py deleted file mode 100644 index d2adf388f4..0000000000 --- a/unit_tests/server/router/model_infer/mode_backend/test_pd_page_io.py +++ /dev/null @@ -1,53 +0,0 @@ -from types import SimpleNamespace - -import torch -import pytest - - -@pytest.mark.parametrize("role", ["prefill", "decode"]) -def test_hybrid_pd_uses_the_shared_att_state_task_without_model_flags(monkeypatch, role): - from lightllm.server.router.model_infer.infer_batch import g_infer_context - from lightllm.server.router.model_infer.mode_backend.pd.prefill_node_impl.prefill_impl import ( - PDChunkedPrefillForPrefillNode, - ) - from lightllm.server.router.model_infer.mode_backend.pd.decode_node_impl import decode_impl - from lightllm.server.router.model_infer.mode_backend.pd.decode_node_impl.decode_impl import PDDecodeNode - - monkeypatch.setattr(g_infer_context, "is_hybrid_att_model", True) - cls = PDChunkedPrefillForPrefillNode if role == "prefill" else PDDecodeNode - backend = cls.__new__(cls) - backend.args = SimpleNamespace(pd_kv_page_size=256, page_size=256) - backend.is_master_in_dp = role == "decode" - groups, tasks = [], [] - backend.info_queue = SimpleNamespace(put=groups.append) - req = SimpleNamespace( - shm_req=SimpleNamespace(input_len=512), - cur_kv_len=512 if role == "prefill" else 0, - hold_kv_len=0, - pd_trans_kv_start_index=0, - ) - - def create(**kwargs): - task = SimpleNamespace( - page_kind=kwargs.get("page_kind", "kv"), first_gen_token_id=None, first_gen_token_logprob=None - ) - tasks.append(task) - if "group" in kwargs: - kwargs["group"].task_list.append(task) - return task - - backend._create_pd_trans_task = create - if role == "prefill": - backend._prefill_chuncked_handle_func(req, next_token_id=17, next_token_prob=-0.2, output_len=1) - assert tasks[-1].first_gen_token_id == 17 - else: - backend.radix_cache = None - backend.model = SimpleNamespace( - att_state_manager=SimpleNamespace(full_att_allocator=SimpleNamespace(page_size=backend.args.page_size)), - ) - req.args = SimpleNamespace(page_size=backend.args.page_size) - req.kv_cache_alloc_need = lambda target: target - req.alloc_kv_mem = lambda alloc_token_num: torch.arange(alloc_token_num) - backend._decode_node_gen_trans_tasks(req) - assert len(groups) == 1 - assert [task.page_kind for task in tasks] == ["kv", "kv", "att_state"] diff --git a/unit_tests/server/router/model_infer/mode_backend/test_pd_prefill_abort.py b/unit_tests/server/router/model_infer/mode_backend/test_pd_prefill_abort.py deleted file mode 100644 index 5fff4ddc2e..0000000000 --- a/unit_tests/server/router/model_infer/mode_backend/test_pd_prefill_abort.py +++ /dev/null @@ -1,199 +0,0 @@ -import queue -import threading -from types import SimpleNamespace -from unittest.mock import MagicMock - -import pytest - -from lightllm.server.pd_io_struct import PDAbortReq -from lightllm.server.router.model_infer.mode_backend.pd.decode_node_impl import decode_impl -from lightllm.server.router.model_infer.mode_backend.pd.decode_node_impl.decode_impl import PDDecodeNode -from lightllm.server.router.model_infer.mode_backend.pd.prefill_node_impl import prefill_impl, prefill_trans_process -from lightllm.server.router.model_infer.mode_backend.pd.prefill_node_impl.prefill_impl import ( - PDChunkedPrefillForPrefillNode, -) -from lightllm.server.router.model_infer.mode_backend.pd.prefill_node_impl.prefill_kv_move_manager import ( - PrefillKVMoveManager, -) -from lightllm.server.router.model_infer.mode_backend.pd.prefill_node_impl.prefill_trans_process import ( - _PrefillTransModule, -) - - -class _StopLoop(Exception): - pass - - -class _SequenceQueue: - def __init__(self, *values): - self.values = list(values) - - def get(self): - if self.values: - return self.values.pop(0) - raise _StopLoop() - - -def test_prefill_infer_forwards_abort_to_kv_move_manager(monkeypatch): - request_id = 123 - req = SimpleNamespace( - req_id=request_id, - infer_aborted=True, - pd_task_num=2, - pd_task_failed_num=0, - pd_task_success_num=0, - pd_trans_device_id=1, - cur_kv_len=4, - shm_req=SimpleNamespace(input_len=8), - ) - monkeypatch.setattr(prefill_impl.g_infer_context, "requests_mapping", {request_id: req}) - - backend = PDChunkedPrefillForPrefillNode.__new__(PDChunkedPrefillForPrefillNode) - backend.is_master_in_dp = True - backend.args = SimpleNamespace(enable_cpu_cache=False) - backend.info_queue = queue.Queue() - - ready_reqs = backend._filter_not_ready_reqs([request_id]) - - assert ready_reqs == [] - abort_req = backend.info_queue.get_nowait() - assert abort_req == PDAbortReq(request_id=request_id, device_id=1) - assert req.pd_abort_req_send_count == 1 - - -def test_prefill_infer_sends_abort_at_most_six_times(monkeypatch): - request_id = 123 - req = SimpleNamespace( - req_id=request_id, - infer_aborted=True, - pd_task_num=2, - pd_task_failed_num=0, - pd_task_success_num=0, - pd_trans_device_id=1, - cur_kv_len=4, - shm_req=SimpleNamespace(input_len=8), - ) - monkeypatch.setattr(prefill_impl.g_infer_context, "requests_mapping", {request_id: req}) - - backend = PDChunkedPrefillForPrefillNode.__new__(PDChunkedPrefillForPrefillNode) - backend.is_master_in_dp = True - backend.args = SimpleNamespace(enable_cpu_cache=False) - backend.info_queue = queue.Queue() - - for _ in range(8): - assert backend._filter_not_ready_reqs([request_id]) == [] - - abort_reqs = [backend.info_queue.get_nowait() for _ in range(6)] - assert all(abort_req == PDAbortReq(request_id=request_id, device_id=1) for abort_req in abort_reqs) - assert backend.info_queue.empty() - assert req.pd_abort_req_send_count == 6 - - -def test_decode_infer_sends_abort_at_most_six_times(monkeypatch): - request_id = 123 - req = SimpleNamespace( - req_id=request_id, - infer_aborted=True, - pd_task_num=2, - pd_task_failed_num=0, - pd_task_success_num=0, - pd_trans_device_id=1, - ) - monkeypatch.setattr(decode_impl.g_infer_context, "requests_mapping", {request_id: req}) - - backend = PDDecodeNode.__new__(PDDecodeNode) - backend.is_master_in_dp = True - backend.info_queue = queue.Queue() - - for _ in range(8): - assert backend._filter_not_ready_reqs([request_id]) == [] - - abort_reqs = [backend.info_queue.get_nowait() for _ in range(6)] - assert all(abort_req == PDAbortReq(request_id=request_id, device_id=1) for abort_req in abort_reqs) - assert backend.info_queue.empty() - assert req.pd_abort_req_send_count == 6 - - -def test_prefill_filter_skips_request_while_cpu_cache_load_is_running(monkeypatch): - from lightllm.server.router.model_infer.infer_batch import InferReq - - request_id = 123 - status = InferReq._CpuCacheTaskStatus - req = SimpleNamespace( - req_id=request_id, - infer_aborted=False, - pd_task_num=0, - pd_task_failed_num=0, - pd_task_success_num=0, - cur_kv_len=8192, - shm_req=SimpleNamespace(input_len=8192), - cpu_cache_load_task_status=status.RUNNING, - cpu_cache_offload_task_status=status.NOT_STARTED, - ) - monkeypatch.setattr(prefill_impl.g_infer_context, "requests_mapping", {request_id: req}) - - backend = PDChunkedPrefillForPrefillNode.__new__(PDChunkedPrefillForPrefillNode) - backend.is_master_in_dp = True - backend.args = SimpleNamespace(enable_cpu_cache=True) - backend.info_queue = queue.Queue() - - assert backend._filter_not_ready_reqs([request_id]) == [] - - req.cpu_cache_load_task_status = status.FINISHED - assert backend._filter_not_ready_reqs([request_id]) == [req] - - -def test_prefill_kv_move_manager_routes_abort_to_target_device(): - abort_req = PDAbortReq(request_id=123, device_id=1) - target_queue = queue.Queue() - manager = PrefillKVMoveManager.__new__(PrefillKVMoveManager) - manager.info_queue = _SequenceQueue(abort_req) - manager.kv_trans_processes = [ - SimpleNamespace(task_in_queue=queue.Queue()), - SimpleNamespace(task_in_queue=target_queue), - ] - - with pytest.raises(_StopLoop): - manager.task_dispatcher_loop() - - assert target_queue.get_nowait() is abort_req - assert manager.kv_trans_processes[0].task_in_queue.empty() - - -def test_prefill_trans_process_handles_abort_without_allocating_page(monkeypatch): - abort_req = PDAbortReq(request_id=123, device_id=0) - module = _PrefillTransModule.__new__(_PrefillTransModule) - module.device_id = 0 - module.task_in_queue = _SequenceQueue(abort_req) - module.page_index_queue = MagicMock() - module._abort = MagicMock() - monkeypatch.setattr(prefill_trans_process.torch.cuda, "set_device", lambda _device_id: None) - - with pytest.raises(_StopLoop): - module.recv_task_loop() - - module._abort.assert_called_once_with(request_id=123) - module.page_index_queue.get.assert_not_called() - - -def test_prefill_abort_only_fails_tasks_not_submitted_to_transporter(): - pending_task = SimpleNamespace(request_id=123, xfer_handle=None, error_info=None) - active_task = SimpleNamespace(request_id=123, xfer_handle=object(), error_info=None) - unrelated_task = SimpleNamespace(request_id=456, xfer_handle=None, error_info=None) - module = _PrefillTransModule.__new__(_PrefillTransModule) - module.waiting_dict_lock = threading.Lock() - module.waiting_dict = { - "pending": pending_task, - "active": active_task, - "unrelated": unrelated_task, - } - module.failed_queue = queue.Queue() - - module._abort(request_id=123) - - assert module.failed_queue.get_nowait() is pending_task - assert pending_task.error_info == "aborted req" - assert module.waiting_dict == { - "active": active_task, - "unrelated": unrelated_task, - } 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 deleted file mode 100644 index 726df8dbee..0000000000 --- a/unit_tests/server/router/model_infer/mode_backend/test_pd_token_indices.py +++ /dev/null @@ -1,126 +0,0 @@ -"""PD 索引由普通 CPU Tensor 上传，使用实际拷贝流并跳过运行态任务。""" - -from types import SimpleNamespace -from unittest.mock import Mock -from queue import Queue - -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_operator=lambda: SimpleNamespace(get_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, - xfer_handle=None, - ) - 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_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) - - -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() diff --git a/unit_tests/server/router/model_infer/mode_backend/test_pd_transfer_progress.py b/unit_tests/server/router/model_infer/mode_backend/test_pd_transfer_progress.py deleted file mode 100644 index b3ca2f9b64..0000000000 --- a/unit_tests/server/router/model_infer/mode_backend/test_pd_transfer_progress.py +++ /dev/null @@ -1,176 +0,0 @@ -import pickle -import queue -import threading -from types import SimpleNamespace -from unittest.mock import Mock - -import pytest - -from lightllm.server.pd_io_struct import PDChunckedTransTask, PDChunckedTransTaskGroup -from lightllm.server.router.model_infer.mode_backend.pd.decode_node_impl import decode_trans_process - - -class _StopLoop(Exception): - pass - - -class _SequenceQueue: - def __init__(self, *values): - self.values = list(values) - - def get(self): - if self.values: - return self.values.pop(0) - raise _StopLoop() - - -def make_task(request_id, page, empty=False): - return PDChunckedTransTask( - request_id=request_id, - start_kv_index=page, - end_kv_index=page if empty else page + 1, - time_out_secs=10, - pd_master_node_id=1, - prefill_dp_index=0, - decode_dp_index=0, - src_device_id=0, - dst_device_id=0, - mem_indexes=[] if empty else [page], - prefill_agent_name="prefill", - prefill_agent_metadata=b"prefill metadata", - prefill_num_pages=1, - prefill_page_reg_desc=b"prefill page", - decode_agent_name=None, - decode_agent_metadata=None, - decode_num_pages=None, - decode_page_reg_desc=None, - first_gen_token_id=42, - first_gen_token_logprob=-0.5, - ) - - -def make_module(*tasks): - module = decode_trans_process._DecodeTransModule.__new__(decode_trans_process._DecodeTransModule) - module.device_id = 0 - module.args = SimpleNamespace(pd_node_id=2, pd_detail_log=False) - module.waiting_dict_lock = threading.Lock() - module.waiting_dict = {task.get_key(): task for task in tasks} - module.transporter = SimpleNamespace( - agent_name="decode", - agent_metadata=b"decode metadata", - num_pages=1, - local_page_mem_desc=b"decode page", - send_error_info_to_prefill_node=Mock(), - ) - for name in ( - "request_page_task_queue", - "ready_page_task_queue", - "success_queue", - "failed_queue", - "up_status_in_queue", - ): - setattr(module, name, queue.Queue()) - return module - - -def test_dispatch_shares_timeout_mark_with_all_pages_in_each_group(monkeypatch): - monkeypatch.setattr(decode_trans_process.time, "time", lambda: 100.0) - first, second, empty, unrelated = make_task(1, 0), make_task(1, 1), make_task(1, 2, empty=True), make_task(2, 0) - module = make_module() - module.recv_task_group_queue = _SequenceQueue( - PDChunckedTransTaskGroup([first, second, empty]), PDChunckedTransTaskGroup([unrelated]) - ) - - with pytest.raises(_StopLoop): - module.dispatch_task_loop() - - assert first.shared_time_out_mark is second.shared_time_out_mark is empty.shared_time_out_mark - assert unrelated.shared_time_out_mark is not first.shared_time_out_mark - assert first.shared_time_out_mark.time_mark == 100.0 - assert module.waiting_dict == {task.get_key(): task for task in (first, second, unrelated)} - assert module.success_queue.get_nowait() == (None, None, empty) - assert empty.start_trans_time == 100.0 - restored = pickle.loads(pickle.dumps(PDChunckedTransTaskGroup([first, second, empty]))).task_list - assert restored[0].shared_time_out_mark is restored[1].shared_time_out_mark is restored[2].shared_time_out_mark - - -def test_waiting_pages_timeout_from_shared_progress_and_initial_deadline(monkeypatch): - now = [100.0] - monkeypatch.setattr(decode_trans_process.time, "time", lambda: now[0]) - first, second, unrelated = make_task(1, 0), make_task(1, 1), make_task(2, 0) - first.shared_time_out_mark = second.shared_time_out_mark = decode_trans_process._SharedTimeOutMark() - unrelated.shared_time_out_mark = decode_trans_process._SharedTimeOutMark() - for task in (first, second, unrelated): - task.create_time = 0.0 - module = make_module(first, second, unrelated) - - now[0] = 109.0 - module._check_tasks_time_out() - assert len(module.waiting_dict) == 3 - first.shared_time_out_mark.update() - now[0] = 111.0 - module._check_tasks_time_out() - assert module.failed_queue.get_nowait() is unrelated - assert len(module.waiting_dict) == 2 - now[0] = 120.0 - module._check_tasks_time_out() - assert [module.failed_queue.get_nowait(), module.failed_queue.get_nowait()] == [first, second] - assert not module.waiting_dict - - -@pytest.mark.parametrize("stage", ["request", "done"]) -def test_matched_notifications_refresh_local_siblings_only(monkeypatch, stage): - now = [100.0] - monkeypatch.setattr(decode_trans_process.time, "time", lambda: now[0]) - monkeypatch.setattr(decode_trans_process.torch.cuda, "set_device", lambda _device: None) - first, sibling, unrelated = make_task(1, 0), make_task(1, 1), make_task(2, 0) - first.shared_time_out_mark = sibling.shared_time_out_mark = decode_trans_process._SharedTimeOutMark() - unrelated.shared_time_out_mark = decode_trans_process._SharedTimeOutMark() - module = make_module(first, sibling, unrelated) - remote = make_task(1, 0) - remote.write_stage = stage - remote.shared_time_out_mark = decode_trans_process._SharedTimeOutMark() - remote.shared_time_out_mark.time_mark = -100.0 - module.transporter.get_new_notifs = lambda: {"prefill": [pickle.dumps(remote)]} - - now[0] = 109.0 - with monkeypatch.context() as patch: - patch.setattr(module, "_check_tasks_time_out", Mock(side_effect=_StopLoop)) - with pytest.raises(_StopLoop): - module.accept_peer_task_loop() - - assert sibling.shared_time_out_mark.time_mark == 109.0 - assert unrelated.shared_time_out_mark.time_mark == 100.0 - target_queue = module.request_page_task_queue if stage == "request" else module.ready_page_task_queue - assert target_queue.get_nowait() is first - - # The repeated notify cannot keep siblings alive once its waiting task was removed. - now[0] = 110.0 - with monkeypatch.context() as patch: - patch.setattr(module, "_check_tasks_time_out", Mock(side_effect=_StopLoop)) - with pytest.raises(_StopLoop): - module.accept_peer_task_loop() - assert sibling.shared_time_out_mark.time_mark == 109.0 - module.transporter.send_error_info_to_prefill_node.assert_called_once() - - -def test_started_transfer_keeps_real_per_page_timeout(monkeypatch): - now = [100.0] - monkeypatch.setattr(decode_trans_process.time, "time", lambda: now[0]) - task = make_task(1, 0) - task.start_trans_time = 50.0 - task.shared_time_out_mark = decode_trans_process._SharedTimeOutMark() - module = make_module(task) - - now[0] = 111.0 - module._check_tasks_time_out() - assert module.waiting_dict - now[0] = 116.0 - task.shared_time_out_mark.update() - module._check_tasks_time_out() - assert module.waiting_dict - now[0] = 117.0 - module._check_tasks_time_out() - assert module.failed_queue.get_nowait() is task - assert task.error_info == "time out in accept_peer_task_loop" - assert not module.waiting_dict diff --git a/unit_tests/server/router/model_infer/mode_backend/test_pd_transfer_readiness.py b/unit_tests/server/router/model_infer/mode_backend/test_pd_transfer_readiness.py deleted file mode 100644 index 1ad7a981dc..0000000000 --- a/unit_tests/server/router/model_infer/mode_backend/test_pd_transfer_readiness.py +++ /dev/null @@ -1,38 +0,0 @@ -import queue -from unittest.mock import Mock - -import pytest - -from lightllm.server.router.model_infer.mode_backend.pd.trans_process_obj import KVTransProcess - - -@pytest.mark.parametrize("status, expected", [("module_ready", True), ("init_failed", False)]) -def test_transfer_process_requires_module_ready(status, expected): - worker = KVTransProcess(process=Mock(), task_out_queue=queue.Queue(), device_id=0) - worker.task_out_queue.put(status) - assert worker.wait_until_ready() is expected - - -def test_transfer_process_waits_through_slow_initialization(): - output = Mock() - output.get.side_effect = [queue.Empty(), queue.Empty(), "module_ready"] - worker = KVTransProcess(process=Mock(), task_out_queue=output, device_id=0) - worker.process.is_alive.return_value = True - assert worker.wait_until_ready() - assert output.get.call_count == 3 - - -def test_transfer_process_detects_exit_before_ready(): - worker = KVTransProcess(process=Mock(), task_out_queue=Mock(), device_id=0) - worker.task_out_queue.get.side_effect = queue.Empty() - worker.process.is_alive.return_value = False - assert not worker.wait_until_ready() - assert worker.task_out_queue.get.call_count == 1 - - -def test_transfer_process_initialization_has_bounded_wait(): - worker = KVTransProcess(process=Mock(), task_out_queue=Mock(), device_id=0) - worker.task_out_queue.get.side_effect = queue.Empty() - worker.process.is_alive.return_value = True - assert not worker.wait_until_ready() - assert worker.task_out_queue.get.call_count == 600 diff --git a/unit_tests/server/router/model_infer/mode_backend/test_prefill_queue_strategy.py b/unit_tests/server/router/model_infer/mode_backend/test_prefill_queue_strategy.py deleted file mode 100644 index 606e616aa3..0000000000 --- a/unit_tests/server/router/model_infer/mode_backend/test_prefill_queue_strategy.py +++ /dev/null @@ -1,401 +0,0 @@ -from types import SimpleNamespace - -import pytest - -from lightllm.server.api_cli import make_argument_parser -from lightllm.server.core.objs import FinishStatus -from lightllm.server.router.model_infer.mode_backend.base_backend import ModeBackend -from lightllm.server.router.model_infer.infer_batch import InferReq -from lightllm.server.router.model_infer.mode_backend.prefill_queue_strategy import ( - PREFILL_QUEUE_STRATEGIES, - PrefillQueueStrategy, - create_prefill_queue_strategy, -) - - -def req( - group_id, - input_len, - kv_len=0, - priority=0, - router_priority=False, - total_len=None, - finish_status=FinishStatus.NO_FINISH, - infer_aborted=False, - filter_mark=False, - wait_pause=False, - paused=False, -): - request = InferReq.__new__(InferReq) - request.req_id = group_id - request.cur_kv_len = kv_len - request.finish_status = FinishStatus(finish_status) - request.infer_aborted = infer_aborted - request.filter_mark = filter_mark - request.wait_pause = wait_pause - request.paused = paused - request.shm_req = SimpleNamespace( - input_len=input_len, - group_req_id=group_id, - sample_params=SimpleNamespace( - high_priority_request=router_priority, - infer_high_priority=priority, - ), - ) - request.get_cur_total_len = lambda: input_len if total_len is None else total_len - return request - - -def strategy(name="default", processed_prefill_tokens=0): - backend = ModeBackend.__new__(ModeBackend) - backend.args = SimpleNamespace(prefill_queue_strategy=name) - backend.processed_prefill_tokens = processed_prefill_tokens - return create_prefill_queue_strategy(backend) - - -@pytest.mark.parametrize("name", PREFILL_QUEUE_STRATEGIES) -def test_empty_and_singleton(name): - policy = strategy(name) - assert isinstance(policy, PrefillQueueStrategy) - assert policy.reorder([]) == [] - request = req(1, 10) - assert policy.reorder([request]) == [request] - - -@pytest.mark.parametrize( - "name, expected", - [ - ("default", [2, 4, 1, 3]), - ("promote_shortest", [2, 4, 3, 1]), - ("hrrn", [2, 4, 1, 3]), - ], -) -def test_ordering_and_input_preservation(name, expected): - requests = [req(1, 20), req(2, 50, priority=-2), req(3, 102, 100), req(4, 5, priority=-1)] - original = list(requests) - result = strategy(name).reorder(requests) - assert [r.shm_req.group_req_id for r in result] == expected - assert all(a is b for a, b in zip(requests, original)) - assert sorted(map(id, result)) == sorted(map(id, original)) - - -def test_inference_priority_is_independent_from_router_priority(): - router_high_priority = req(1, 10, priority=0, router_priority=True) - infer_high_priority = req(2, 10, priority=-1, router_priority=False) - - assert strategy().reorder([router_high_priority, infer_high_priority]) == [ - infer_high_priority, - router_high_priority, - ] - - -def test_promote_shortest_clamps_remaining_tokens_and_preserves_ties(): - requests = [req(9, 5, 6), req(8, 5, 10), req(7, 5, 5), req(6, 10, 8), req(1, 2)] - assert strategy("promote_shortest").reorder(requests) == requests - - -def test_promote_shortest_ignores_finished_request_when_selecting_shortest(): - finished = req(1, 5, kv_len=5, finish_status=FinishStatus.FINISHED_LENGTH) - long_request = req(2, 20) - short_request = req(3, 2) - - assert strategy("promote_shortest").reorder([finished, long_request, short_request], no_decode=True) == [ - short_request, - long_request, - finished, - ] - - -def test_promote_shortest_preserves_arrival_order_for_equal_positive_remaining_tokens(): - long_request = req(1, 20) - first_short = req(3, 10, kv_len=8) - second_short = req(2, 2) - - assert strategy("promote_shortest").reorder([long_request, first_short, second_short]) == [ - first_short, - long_request, - second_short, - ] - - -@pytest.mark.parametrize("name", PREFILL_QUEUE_STRATEGIES) -@pytest.mark.parametrize("mode", [{}, {"no_decode": True}, {"strict_prefill": True}]) -def test_finished_aborted_and_cleanup_requests_stay_last(name, mode): - finished_normal = req(9, 5, kv_len=5, finish_status=FinishStatus.FINISHED_LENGTH) - finished_high = req(8, 2, priority=-10, finish_status=FinishStatus.FINISHED_STOP) - finished_decode = req(7, 5, kv_len=4, finish_status=FinishStatus.FINISHED_ERROR) - aborted = req(6, 2, priority=-10, infer_aborted=True) - cleanup = req(5, 5, kv_len=4, filter_mark=True) - long_request = req(1, 20) - short_request = req(2, 2) - high_priority = req(3, 100, priority=-1) - decode = req(4, 5, kv_len=4) - active = [long_request, short_request, high_priority, decode] - finished = [finished_normal, finished_high, finished_decode, aborted, cleanup] - requests = [ - finished_normal, - long_request, - finished_high, - short_request, - finished_decode, - high_priority, - aborted, - decode, - cleanup, - ] - original = list(requests) - policy = strategy(name) - - expected_active = policy.reorder(active, **mode) - result = policy.reorder(requests, **mode) - - assert result == expected_active + finished - assert all(a is b for a, b in zip(requests, original)) - assert sorted(map(id, result)) == sorted(map(id, original)) - for request in finished: - assert not hasattr(request, "arrival_processed_prefill_tokens") - assert not hasattr(request, "uncached_prefill_tokens") - - -@pytest.mark.parametrize("name", PREFILL_QUEUE_STRATEGIES) -def test_all_finished_requests_preserve_arrival_order(name): - requests = [ - req(6, 2, finish_status=FinishStatus.FINISHED_ABORTED), - req(5, 5, kv_len=5, finish_status=FinishStatus.FINISHED_LENGTH), - req(4, 10, priority=-10, finish_status=FinishStatus.FINISHED_STOP), - req(3, 5, kv_len=4, infer_aborted=True), - req(2, 2, filter_mark=True), - req(1, 5, kv_len=4, finish_status=FinishStatus.FINISHED_PD_DECODE_CAPACITY), - ] - - assert strategy(name).reorder(requests) == requests - - -@pytest.mark.parametrize("pause_state", ["wait_pause", "paused"]) -def test_promote_shortest_ignores_paused_request_when_selecting_shortest(pause_state): - suspended = req(1, 2, **{pause_state: True}) - long_request = req(2, 20) - short_request = req(3, 5) - - assert strategy("promote_shortest").reorder([suspended, long_request, short_request]) == [ - short_request, - long_request, - suspended, - ] - assert getattr(suspended, pause_state) is True - - -@pytest.mark.parametrize("name", PREFILL_QUEUE_STRATEGIES) -@pytest.mark.parametrize("mode", [{}, {"no_decode": True}, {"strict_prefill": True}]) -def test_paused_requests_keep_arrival_order_at_tail(name, mode): - paused_high = req(5, 2, priority=-10, paused=True) - waiting = req(6, 5, kv_len=4, wait_pause=True) - paused_normal = req(4, 2, paused=True) - finished = req(7, 5, finish_status=FinishStatus.FINISHED_STOP) - active = [req(1, 20), req(2, 5), req(3, 100, priority=-1)] - requests = [paused_high, active[0], waiting, active[1], finished, paused_normal, active[2]] - original = list(requests) - policy = strategy(name) - - assert policy.reorder(requests, **mode) == policy.reorder(active, **mode) + [ - paused_high, - waiting, - finished, - paused_normal, - ] - assert all(a is b for a, b in zip(requests, original)) - - -@pytest.mark.parametrize("pause_state", ["wait_pause", "paused"]) -def test_hrrn_initializes_paused_request_only_after_recovery(pause_state): - suspended = req(1, 100, **{pause_state: True}) - active = req(2, 20) - policy = strategy("hrrn", processed_prefill_tokens=1000) - - assert policy.reorder([suspended, active]) == [active, suspended] - assert not hasattr(suspended, "arrival_processed_prefill_tokens") - assert not hasattr(suspended, "uncached_prefill_tokens") - - setattr(suspended, pause_state, False) - suspended.cur_kv_len = 90 - policy.backend.processed_prefill_tokens = 2000 - assert policy.reorder([suspended, active]) == [active, suspended] - assert suspended.arrival_processed_prefill_tokens == 2000 - assert suspended.uncached_prefill_tokens == 10 - - -def test_hrrn_preserves_existing_aging_snapshot_across_pause_and_recovery(): - request = req(1, 100) - policy = strategy("hrrn", processed_prefill_tokens=1000) - policy.reorder([request]) - - request.paused = True - request.cur_kv_len = 0 - policy.backend.processed_prefill_tokens = 2000 - assert policy.reorder([request]) == [request] - assert request.paused is True - - request.paused = False - request.cur_kv_len = 90 - policy.backend.processed_prefill_tokens = 3000 - assert policy.reorder([request]) == [request] - assert request.arrival_processed_prefill_tokens == 1000 - assert request.uncached_prefill_tokens == 100 - - -def test_promote_shortest_keeps_high_priority_requests_first(): - long_high_priority = req(1, 1000, priority=-1) - short_normal_priority = req(2, 2, priority=0) - - assert strategy("promote_shortest").reorder([short_normal_priority, long_high_priority]) == [ - long_high_priority, - short_normal_priority, - ] - - -def test_promote_shortest_preserves_high_priority_arrival_order(): - first_high_priority = req(1, 100, priority=-1) - second_high_priority = req(2, 100, priority=-2) - - assert strategy("promote_shortest").reorder([first_high_priority, second_high_priority]) == [ - first_high_priority, - second_high_priority, - ] - - -def test_promote_shortest_moves_only_one_normal_request_and_preserves_other_order(): - first_normal = req(1, 20) - shortest_normal = req(2, 2) - third_normal = req(3, 10) - fourth_normal = req(4, 5) - - assert strategy("promote_shortest").reorder([first_normal, shortest_normal, third_normal, fourth_normal]) == [ - shortest_normal, - first_normal, - third_normal, - fourth_normal, - ] - - -def test_hrrn_prioritizes_short_work_for_equally_aged_requests(): - requests = [req(1, 100), req(2, 10), req(3, 2)] - policy = strategy("hrrn") - policy.reorder(requests) - policy.backend.processed_prefill_tokens = 1000 - - assert policy.reorder(requests) == [ - requests[2], - requests[1], - requests[0], - ] - - -def test_hrrn_aging_allows_long_request_to_overtake_new_short_request(): - long_old = req(1, 100) - short_new = req(2, 2) - policy = strategy("hrrn") - policy.reorder([long_old]) - policy.backend.processed_prefill_tokens = 100_000 - - assert policy.reorder([short_new, long_old]) == [ - long_old, - short_new, - ] - - -def test_hrrn_aging_starts_at_first_sort_and_survives_later_sorts(): - long_request = req(1, 100) - short_request = req(2, 2) - requests = [short_request, long_request] - policy = strategy("hrrn", processed_prefill_tokens=1000) - - # 首轮两个请求的等待量均为零，使用请求 ID 打破平局。 - assert policy.reorder(requests) == [long_request, short_request] - policy.backend.processed_prefill_tokens = 2000 - # 后续排序保留首次快照，等龄请求中短请求的响应比更高。 - assert policy.reorder(requests) == [short_request, long_request] - - -def test_hrrn_binds_uncached_tokens_once_and_uses_request_id_tiebreaker(): - less_cached = req(2, 100, kv_len=0) - equally_cached_later_id = req(3, 100, kv_len=90) - more_cached = req(1, 100, kv_len=90) - policy = strategy("hrrn") - - requests = [less_cached, equally_cached_later_id, more_cached] - policy.reorder(requests) - policy.backend.processed_prefill_tokens = 1000 - assert policy.reorder(requests) == [more_cached, equally_cached_later_id, less_cached] - assert [request.uncached_prefill_tokens for request in requests] == [100, 10, 10] - - less_cached.cur_kv_len = 95 - more_cached.cur_kv_len = 0 - policy.backend.processed_prefill_tokens = 2000 - assert policy.reorder(requests) == [more_cached, equally_cached_later_id, less_cached] - - -def test_hrrn_keeps_internal_high_priority_requests_first(): - shortest_normal = req(1, 2, priority=0) - long_high_priority = req(2, 1000, priority=-1) - - assert strategy("hrrn", processed_prefill_tokens=1000).reorder([shortest_normal, long_high_priority]) == [ - long_high_priority, - shortest_normal, - ] - - -@pytest.mark.parametrize("name", PREFILL_QUEUE_STRATEGIES) -def test_decode_requests_stay_left_and_are_not_reordered(name): - decode_normal = req(1, 5, kv_len=4, priority=0) - prefill_high = req(2, 100, priority=-2) - decode_high = req(3, 5, kv_len=4, priority=-10) - prefill_normal = req(4, 2, priority=0) - - assert strategy(name).reorder([decode_normal, prefill_normal, decode_high, prefill_high]) == [ - decode_normal, - decode_high, - prefill_high, - prefill_normal, - ] - - -def test_no_decode_and_strict_prefill_match_backend_classification_rules(): - boundary_req = req(1, 5, kv_len=4, priority=0) - high_priority_prefill = req(2, 10, priority=-1) - - policy = strategy() - assert policy.reorder([boundary_req, high_priority_prefill]) == [boundary_req, high_priority_prefill] - assert policy.reorder([boundary_req, high_priority_prefill], no_decode=True) == [ - high_priority_prefill, - boundary_req, - ] - assert policy.reorder([boundary_req, high_priority_prefill], strict_prefill=True) == [ - high_priority_prefill, - boundary_req, - ] - - -@pytest.mark.parametrize("name", PREFILL_QUEUE_STRATEGIES) -def test_cli_strategy_selection(name): - args = make_argument_parser().parse_args(["--prefill_queue_strategy", name]) - backend = SimpleNamespace(args=args) - policy = create_prefill_queue_strategy(backend) - assert isinstance(policy, PREFILL_QUEUE_STRATEGIES[name]) - assert policy.backend is backend - - -def test_cli_default_and_invalid_strategy(): - parser = make_argument_parser() - assert parser.parse_args([]).prefill_queue_strategy == "auto" - assert parser.parse_args(["--prefill_queue_strategy", "auto"]).prefill_queue_strategy == "auto" - assert set(PREFILL_QUEUE_STRATEGIES) == {"default", "promote_shortest", "hrrn"} - for name in ("unknown", "none"): - with pytest.raises(SystemExit): - parser.parse_args(["--prefill_queue_strategy", name]) - with pytest.raises(ValueError, match="Unknown prefill queue strategy"): - strategy(name) - - -def test_strategy_base_is_abstract(): - with pytest.raises(TypeError): - PrefillQueueStrategy(SimpleNamespace()) 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 deleted file mode 100644 index a760f44a35..0000000000 --- a/unit_tests/server/router/model_infer/mode_backend/test_rl_backend_ops.py +++ /dev/null @@ -1,57 +0,0 @@ -from types import SimpleNamespace -from unittest.mock import Mock - -import numpy as np -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 _ReqIndexPool -from lightllm.server.router.model_infer.mode_backend.rl_backend_ops import RlBackendOps -from lightllm.utils.torch_memory_saver_utils import MemoryTag - - -@pytest.mark.parametrize( - ("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, kv_allocator_environment -): - allocator = FullAttKVAllocator(16, page_size=4, pin_memory=False) - req_manager = ReqManager.__new__(ReqManager) - req_manager.req_list = _ReqIndexPool(2) - req_manager.HOLD_REQUEST_ID = 2 - req_manager.page_size = allocator.page_size - req_manager._cpu_token_index_blocks = np.empty((3, 2), dtype=np.int32) - 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=resume), - req_manager=req_manager, - att_state_manager=SimpleNamespace(full_att_allocator=allocator), - ) - backend = SimpleNamespace(model=model, logger=Mock()) - ops = RlBackendOps(backend) - ops._clear_cuda_cache = lambda: None - - success, message = ops.resume_memory_occupation(tags) - - 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() 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 deleted file mode 100644 index 784848cb13..0000000000 --- a/unit_tests/server/router/model_infer/mtp_speculative/test_dflash.py +++ /dev/null @@ -1,180 +0,0 @@ -from types import SimpleNamespace - -import pytest -import torch - -from lightllm.server.router.model_infer.mtp_speculative.proposers.dflash import DFlashProposer -from lightllm.server.router.model_infer.mtp_speculative.proposers.proposal_type import DFlashSpecProposal -from lightllm.server.router.model_infer.pin_mem_manager import g_pin_mem_manager - - -def test_dflash_prefill_uses_a_shallow_copy_for_target_hidden(): - forwarded_inputs = [] - draft_model = SimpleNamespace(forward=forwarded_inputs.append) - proposer = DFlashProposer( - backend=SimpleNamespace(draft_models=[draft_model]), - enable_dynmaic_mtp=False, - ) - model_input = SimpleNamespace( - is_prefill=True, - b_position_delta=None, - b_req_idx=torch.tensor([3, 5], dtype=torch.int32), - input_ids=torch.arange(5, dtype=torch.int64), - mtp_draft_input_hiddens=None, - ) - target_hidden = torch.empty((5, 8)) - - proposer.fill_draft_model_kv_state( - target_model_input=model_input, - target_model_output=SimpleNamespace( - mtp_collector=SimpleNamespace(spec_hidden=target_hidden), - ), - target_next_token_ids=torch.tensor([11, 13], dtype=torch.int64), - ) - - assert len(forwarded_inputs) == 1 - assert forwarded_inputs[0] is not model_input - assert forwarded_inputs[0].mtp_draft_input_hiddens is target_hidden - assert model_input.mtp_draft_input_hiddens is None - - -def test_dflash_commits_verify_kv_and_builds_parallel_block(monkeypatch): - block_size = 3 - flat_draft_token_ids = torch.tensor([30, 31, 32, 40, 41, 42], dtype=torch.int64) - flat_draft_token_probs = torch.tensor([0.9, 0.8, 0.7, 0.6, 0.5, 0.4], dtype=torch.float16) - forwarded_inputs = [] - - def forward(model_input): - forwarded_inputs.append(model_input) - return SimpleNamespace() - - draft_model = SimpleNamespace( - block_size=block_size, - mask_token_id=99, - forward=forward, - ) - proposer = DFlashProposer( - backend=SimpleNamespace( - draft_models=[draft_model], - _gen_argmax_token_ids_and_prob=lambda _: (flat_draft_token_ids, flat_draft_token_probs), - ), - enable_dynmaic_mtp=True, - ) - monkeypatch.setattr(torch.Tensor, "cuda", lambda self, non_blocking=False: self) - monkeypatch.setattr( - g_pin_mem_manager, - "get_const_gpu_tensor", - lambda *, shape, fill_value, dtype, **_: torch.full(shape, fill_value, dtype=dtype), - ) - - target_hidden = torch.empty((5, 8)) - model_input = SimpleNamespace( - is_prefill=False, - batch_size=5, - total_token_num=40, - max_q_seq_len=1, - max_kv_seq_len=9, - input_ids=torch.tensor([10, 11, 12, 20, 21], dtype=torch.int64), - b_req_idx=torch.tensor([7, 7, 7, 9, 9], dtype=torch.int32), - b_mtp_index=torch.arange(5, dtype=torch.int32), - b_seq_len=torch.tensor([4, 5, 6, 8, 9], dtype=torch.int32), - b_position_delta=torch.tensor([1, 1, 1, 2, 2], dtype=torch.int32), - b_shared_seq_len=torch.tensor([3, 3, 3, 6, 6], dtype=torch.int32), - b_shared_radix_node_id=torch.tensor([17, 17, 17, 19, 19], dtype=torch.int64), - multimodal_params=[{"images": [], "audios": []} for _ in range(5)], - mtp_draft_input_hiddens=None, - ) - - proposal = proposer.propose_next( - target_model_input=model_input, - target_model_output=SimpleNamespace( - mtp_collector=SimpleNamespace(spec_hidden=target_hidden), - ), - target_next_token_ids=model_input.input_ids, - b_req_mtp_start_loc=torch.tensor([0, 3], dtype=torch.int32), - draft_step=2, - accept_len=torch.tensor([2, 2], dtype=torch.int32), - ) - - assert isinstance(proposal, DFlashSpecProposal) - assert len(forwarded_inputs) == 2 - verify_draft_input, block_draft_input = forwarded_inputs - assert verify_draft_input is not model_input - assert verify_draft_input.mtp_draft_input_hiddens is target_hidden - assert model_input.mtp_draft_input_hiddens is None - - assert torch.equal(block_draft_input.input_ids, torch.tensor([11, 99, 99, 21, 99, 99])) - assert torch.equal(block_draft_input.b_req_idx, torch.tensor([7, 7, 7, 9, 9, 9], dtype=torch.int32)) - assert torch.equal(block_draft_input.b_mtp_index, torch.zeros(6, dtype=torch.int32)) - assert torch.equal(block_draft_input.b_seq_len, torch.tensor([6, 7, 8, 10, 11, 12], dtype=torch.int32)) - assert torch.equal(block_draft_input.b_position_delta, torch.tensor([1, 1, 1, 2, 2, 2], dtype=torch.int32)) - assert block_draft_input.mtp_draft_input_hiddens is None - - assert torch.equal(proposal.token_ids, torch.tensor([[30, 31], [40, 41]], dtype=torch.int64)) - torch.testing.assert_close( - 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() diff --git a/unit_tests/server/router/model_infer/mtp_speculative/test_dflash2.py b/unit_tests/server/router/model_infer/mtp_speculative/test_dflash2.py deleted file mode 100644 index 021b6d47fa..0000000000 --- a/unit_tests/server/router/model_infer/mtp_speculative/test_dflash2.py +++ /dev/null @@ -1,117 +0,0 @@ -from types import SimpleNamespace - -import pytest -import torch - -from lightllm.common.basemodel.batch_objs import ModelMtpOutputCollector, ModelOutput, PostLayerOutput -from lightllm.server.router.model_infer.mtp_speculative.proposers.proposal_type import DFlashSpecProposal -from lightllm.server.router.model_infer.mtp_speculative.proposers.dflash2 import DFlash2Proposer - - -@pytest.mark.parametrize("dynamic_verify", [False, True]) -@pytest.mark.parametrize("draft_step", [1, 3]) -def test_dflash2_proposal_returns_greedy_tokens(dynamic_verify, draft_step): - tokens = torch.arange(6).view(2, 3) - physical_tokens = torch.tensor([100, 0, 1, 2, 200, 3, 4, 5]) - proposer = DFlash2Proposer(backend=SimpleNamespace(), enable_dynmaic_mtp=dynamic_verify) - confidence = torch.tensor([[0.0, 1.0, -1.0], [float("inf"), -100.0, 2.0]]) - output = SimpleNamespace( - mtp_collector=ModelMtpOutputCollector(draft_token_ids=physical_tokens, confidence_logits=confidence) - ) - proposal = proposer._build_proposal(output, 2, 4, draft_step) - assert type(proposal) is DFlashSpecProposal - torch.testing.assert_close(proposal.token_ids, tokens[:, :draft_step]) - if dynamic_verify: - torch.testing.assert_close(proposal.schedule_scores, confidence[:, :draft_step].sigmoid().clamp(0.01, 0.99)) - else: - assert proposal.schedule_scores is None - - -@pytest.mark.parametrize("dynamic_verify", [False, True]) -@pytest.mark.parametrize("req_num", [0, 1, 2]) -def test_dflash2_collector_unpads_physical_token_rows(dynamic_verify, req_num): - tokens = torch.arange(6).view(2, 3) - physical_tokens = torch.tensor([100, 0, 1, 2, 200, 3, 4, 5]) - confidence = tokens.float() if dynamic_verify else None - collector = ModelMtpOutputCollector(draft_token_ids=physical_tokens, confidence_logits=confidence) - unpadded = collector.unpad_decode(padded_batch_size=8, origin_batch_size=req_num * 4) - torch.testing.assert_close(unpadded.draft_token_ids, physical_tokens[: req_num * 4]) - if dynamic_verify: - torch.testing.assert_close(unpadded.confidence_logits, confidence[:req_num]) - else: - assert unpadded.confidence_logits is None - assert collector.draft_token_ids.shape == (8,) - - proposer = DFlash2Proposer(backend=SimpleNamespace(), enable_dynmaic_mtp=dynamic_verify) - proposal = proposer._build_proposal(SimpleNamespace(mtp_collector=unpadded), req_num, 4, 3) - torch.testing.assert_close(proposal.token_ids, tokens[:req_num]) - if dynamic_verify: - torch.testing.assert_close(proposal.schedule_scores, confidence[:req_num].sigmoid().clamp(0.01, 0.99)) - - -def test_dflash2_dynamic_verification_requires_confidence(): - proposer = DFlash2Proposer(backend=SimpleNamespace(), enable_dynmaic_mtp=True) - output = ModelOutput( - logits=torch.empty(4, 1), - mtp_collector=ModelMtpOutputCollector(draft_token_ids=torch.ones(4, dtype=torch.int64)), - ) - with pytest.raises(RuntimeError, match="selector confidence"): - proposer._build_proposal(output, 1, 4, 3) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("width", [1, 3, 7]) -@pytest.mark.parametrize("dynamic_verify", [False, True]) -def test_greedy_selector_conditions_on_selected_predecessor(width, dynamic_verify): - from lightllm.models.qwen3_dflash2.triton_kernel.selector_walk import greedy_selector_walk - - scores = torch.randn(3, width, 16, 16, generator=torch.Generator().manual_seed(42)) - scores[0].zero_() # Ties select the first candidate. - candidates = torch.arange(3 * width * 16).view(3, width, 16) - expected = torch.empty(3, width, dtype=torch.int64) - expected_probs = torch.empty(3, width) - for req in range(3): - previous = 0 - for slot in range(width): - row = scores[req, slot, previous] - previous = int(row.argmax()) - expected_probs[req, slot] = row.softmax(-1)[previous] - expected[req, slot] = candidates[req, slot, previous] - scores, candidates = scores.cuda(), candidates.cuda() - rng_before = torch.cuda.get_rng_state() - confidence = torch.empty(3, width, device="cuda") if dynamic_verify else None - actual = greedy_selector_walk(scores, candidates, confidence_logits=confidence) - if dynamic_verify: - torch.testing.assert_close(confidence.sigmoid().cpu(), expected_probs) - torch.testing.assert_close(actual.cpu(), expected) - assert torch.equal(rng_before, torch.cuda.get_rng_state()) - - -@pytest.mark.parametrize("is_prefill", [True, False]) -def test_dflash2_post_layer_satisfies_base_model_output_contract(is_prefill): - from lightllm.models.qwen3_dflash2.layer_infer.post_layer_infer import Qwen3DFlash2PostLayerInfer - - layer = Qwen3DFlash2PostLayerInfer.__new__(Qwen3DFlash2PostLayerInfer) - layer.block_size_ = 4 - layer.selector_top_k_ = 2 - hidden = torch.randn(4, 8) - ids = torch.arange(6).view(3, 2) - unary = torch.randn(3, 2) - selected = ids[:, 0].view(1, 3) - layer._slice_get_last_input = lambda *_: (hidden, 4) - layer._norm = lambda x, *_: x - layer._compute_candidates = lambda **_: (ids, unary) - confidence = torch.randn(1, 3) - layer._select_path = lambda **_: (selected, confidence) - collected = {} - state = SimpleNamespace( - is_prefill=is_prefill, - input_ids=torch.tensor([99, 0, 0, 0]), - hidden_collector=SimpleNamespace(add_mtp_outputs=lambda **kwargs: collected.update(kwargs)), - ) - output = layer.token_forward(hidden, state, None) - assert isinstance(output, PostLayerOutput) - assert output.logits.shape == ((0,) if is_prefill else (4, 1)) - if not is_prefill: - torch.testing.assert_close(collected["draft_token_ids"], torch.tensor([99, 0, 2, 4])) - torch.testing.assert_close(collected["confidence_logits"], confidence) diff --git a/unit_tests/server/router/model_infer/mtp_speculative/test_dspark.py b/unit_tests/server/router/model_infer/mtp_speculative/test_dspark.py deleted file mode 100644 index 63e78e4e5d..0000000000 --- a/unit_tests/server/router/model_infer/mtp_speculative/test_dspark.py +++ /dev/null @@ -1,129 +0,0 @@ -from types import SimpleNamespace - -import torch - -from lightllm.server.router.model_infer.mtp_speculative.proposers.dspark import DSparkProposer -from lightllm.server.router.model_infer.pin_mem_manager import g_pin_mem_manager - - -def test_dspark_prefill_uses_a_shallow_copy_for_target_hidden(): - forwarded_inputs = [] - draft_model = SimpleNamespace(forward=forwarded_inputs.append) - proposer = DSparkProposer( - backend=SimpleNamespace(draft_models=[draft_model]), - enable_dynmaic_mtp=False, - ) - model_input = SimpleNamespace( - is_prefill=True, - b_position_delta=None, - b_req_idx=torch.tensor([3, 5], dtype=torch.int32), - input_ids=torch.arange(5, dtype=torch.int64), - mtp_draft_input_hiddens=None, - ) - target_hidden = torch.empty((5, 8)) - - proposer.fill_draft_model_kv_state( - target_model_input=model_input, - target_model_output=SimpleNamespace( - mtp_collector=SimpleNamespace(spec_hidden=target_hidden), - ), - target_next_token_ids=torch.tensor([11, 13], dtype=torch.int64), - ) - - assert len(forwarded_inputs) == 1 - assert forwarded_inputs[0] is not model_input - assert forwarded_inputs[0].mtp_draft_input_hiddens is target_hidden - assert model_input.mtp_draft_input_hiddens is None - - -def test_dspark_commits_verify_kv_and_builds_parallel_block(monkeypatch): - block_size = 3 - draft_token_ids = torch.tensor([30, 31, 32, 40, 41, 42], dtype=torch.int64) - confidence_logits = torch.tensor( - [ - [-10.0, 0.0, 10.0], - [10.0, 0.0, -10.0], - ] - ) - block_output = SimpleNamespace( - mtp_collector=SimpleNamespace( - draft_token_ids=draft_token_ids, - confidence_logits=confidence_logits, - ) - ) - forwarded_inputs = [] - - def forward(model_input): - forwarded_inputs.append(model_input) - return block_output if len(forwarded_inputs) == 2 else SimpleNamespace() - - draft_model = SimpleNamespace( - block_size=block_size, - mask_token_id=99, - forward=forward, - ) - proposer = DSparkProposer( - backend=SimpleNamespace(draft_models=[draft_model]), - enable_dynmaic_mtp=True, - ) - monkeypatch.setattr(torch.Tensor, "cuda", lambda self, non_blocking=False: self) - monkeypatch.setattr( - g_pin_mem_manager, - "get_const_gpu_tensor", - lambda *, shape, fill_value, dtype, **_: torch.full(shape, fill_value, dtype=dtype), - ) - monkeypatch.setattr( - g_pin_mem_manager, - "async_copy_from_gpu_tensor", - lambda *, gpu_tensor, **_: gpu_tensor.clone(), - ) - - target_hidden = torch.empty((5, 8)) - model_input = SimpleNamespace( - is_prefill=False, - batch_size=5, - total_token_num=40, - max_q_seq_len=1, - max_kv_seq_len=9, - input_ids=torch.tensor([10, 11, 12, 20, 21], dtype=torch.int64), - b_req_idx=torch.tensor([7, 7, 7, 9, 9], dtype=torch.int32), - b_mtp_index=torch.arange(5, dtype=torch.int32), - b_seq_len=torch.tensor([4, 5, 6, 8, 9], dtype=torch.int32), - b_position_delta=torch.tensor([1, 1, 1, 2, 2], dtype=torch.int32), - b_shared_seq_len=torch.tensor([3, 3, 3, 6, 6], dtype=torch.int32), - b_shared_radix_node_id=torch.tensor([17, 17, 17, 19, 19], dtype=torch.int64), - multimodal_params=[{"images": [], "audios": []} for _ in range(5)], - mtp_draft_input_hiddens=None, - ) - - proposal = proposer.propose_next( - target_model_input=model_input, - target_model_output=SimpleNamespace( - mtp_collector=SimpleNamespace(spec_hidden=target_hidden), - ), - target_next_token_ids=model_input.input_ids, - b_req_mtp_start_loc=torch.tensor([0, 3], dtype=torch.int32), - draft_step=2, - accept_len=torch.tensor([2, 2], dtype=torch.int32), - ) - - assert len(forwarded_inputs) == 2 - verify_draft_input, block_draft_input = forwarded_inputs - assert verify_draft_input is not model_input - assert verify_draft_input.mtp_draft_input_hiddens is target_hidden - assert model_input.mtp_draft_input_hiddens is None - - assert torch.equal(block_draft_input.input_ids, torch.tensor([11, 99, 99, 21, 99, 99])) - assert torch.equal(block_draft_input.b_req_idx, torch.tensor([7, 7, 7, 9, 9, 9], dtype=torch.int32)) - assert torch.equal(block_draft_input.b_mtp_index, torch.zeros(6, dtype=torch.int32)) - assert torch.equal(block_draft_input.b_seq_len, torch.tensor([6, 7, 8, 10, 11, 12], dtype=torch.int32)) - assert torch.equal(block_draft_input.b_position_delta, torch.tensor([1, 1, 1, 2, 2, 2], dtype=torch.int32)) - assert block_draft_input.mtp_draft_input_hiddens is None - - assert torch.equal(proposal.token_ids, torch.tensor([[30, 31], [40, 41]], dtype=torch.int64)) - torch.testing.assert_close( - proposal.schedule_scores, - torch.tensor([[0.01, 0.5], [0.99, 0.5]], dtype=torch.float32), - ) - assert torch.equal(proposal.schedule_scores_cpu, proposal.schedule_scores) - assert proposal.schedule_scores_cpu is not proposal.schedule_scores diff --git a/unit_tests/server/router/model_infer/mtp_speculative/test_eagle_no_att.py b/unit_tests/server/router/model_infer/mtp_speculative/test_eagle_no_att.py deleted file mode 100644 index 71de619b9a..0000000000 --- a/unit_tests/server/router/model_infer/mtp_speculative/test_eagle_no_att.py +++ /dev/null @@ -1,93 +0,0 @@ -from types import SimpleNamespace - -import pytest -import torch - -from lightllm.server.router.model_infer.mtp_speculative.proposers.eagle_no_att import ( - EagleNoAttProposer, -) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") -def test_eagle_no_att_recurrently_proposes_from_accepted_tails(): - device = "cuda" - draft_calls = [] - draft_outputs = [ - SimpleNamespace( - token_ids=torch.tensor([21, 24], device=device), - token_probs=torch.tensor([0.8, 0.7], device=device), - mtp_collector=SimpleNamespace(spec_hidden=torch.tensor([[102.0, 103.0], [108.0, 109.0]], device=device)), - ), - SimpleNamespace( - token_ids=torch.tensor([31, 34], device=device), - token_probs=torch.tensor([0.6, 0.5], device=device), - mtp_collector=SimpleNamespace(spec_hidden=torch.tensor([[202.0, 203.0], [208.0, 209.0]], device=device)), - ), - ] - - def forward(model_input): - draft_calls.append( - { - "batch_size": model_input.batch_size, - "input_ids": model_input.input_ids.cpu(), - "draft_hidden": model_input.mtp_draft_input_hiddens.cpu(), - "b_req_idx": model_input.b_req_idx.cpu(), - "b_seq_len": model_input.b_seq_len.cpu(), - } - ) - return draft_outputs[len(draft_calls) - 1] - - backend = SimpleNamespace( - draft_models=[SimpleNamespace(forward=forward)], - _gen_argmax_token_ids=lambda output: output.token_ids, - _gen_argmax_token_ids_and_prob=lambda output: ( - output.token_ids, - output.token_probs, - ), - ) - proposer = EagleNoAttProposer(backend=backend, enable_dynmaic_mtp=True) - target_model_input = SimpleNamespace( - batch_size=5, - input_ids=torch.tensor([10, 11, 12, 13, 14], device=device), - b_req_idx=torch.tensor([7, 7, 7, 9, 9], dtype=torch.int32, device=device), - b_mtp_index=torch.tensor([0, 1, 2, 0, 1], dtype=torch.int32, device=device), - b_seq_len=torch.tensor([10, 11, 12, 20, 21], dtype=torch.int32, device=device), - b_position_delta=torch.tensor([0, 1, 2, 3, 4], dtype=torch.int32, device=device), - b_shared_seq_len=torch.tensor([8, 8, 8, 6, 6], dtype=torch.int32, device=device), - b_shared_radix_node_id=torch.tensor([70, 70, 70, 90, 90], dtype=torch.int64, device=device), - multimodal_params=[{"images": [], "audios": []} for _ in range(5)], - ) - target_model_output = SimpleNamespace( - mtp_collector=SimpleNamespace(spec_hidden=torch.arange(10, dtype=torch.float32, device=device).reshape(5, 2)) - ) - - proposal = proposer.propose_next( - target_model_input=target_model_input, - target_model_output=target_model_output, - target_next_token_ids=target_model_input.input_ids, - b_req_mtp_start_loc=torch.tensor([0, 3], dtype=torch.int32, device=device), - draft_step=2, - accept_len=torch.tensor([2, 2], dtype=torch.int32, device=device), - ) - - torch.testing.assert_close(proposal.token_ids, torch.tensor([[21, 31], [24, 34]], device=device)) - torch.testing.assert_close(proposal.schedule_scores, torch.tensor([[0.8, 0.6], [0.7, 0.5]], device=device)) - assert len(draft_calls) == 2 - assert draft_calls[0]["batch_size"] == 2 - torch.testing.assert_close(draft_calls[0]["input_ids"], torch.tensor([11, 14])) - torch.testing.assert_close(draft_calls[0]["draft_hidden"], torch.tensor([[2.0, 3.0], [8.0, 9.0]])) - torch.testing.assert_close(draft_calls[0]["b_req_idx"], torch.tensor([7, 9], dtype=torch.int32)) - torch.testing.assert_close(draft_calls[0]["b_seq_len"], torch.tensor([11, 21], dtype=torch.int32)) - torch.testing.assert_close(draft_calls[1]["input_ids"], torch.tensor([21, 24])) - torch.testing.assert_close( - draft_calls[1]["draft_hidden"], - torch.tensor([[102.0, 103.0], [108.0, 109.0]]), - ) - assert target_model_input.batch_size == 5 - torch.testing.assert_close(target_model_input.input_ids.cpu(), torch.tensor([10, 11, 12, 13, 14])) - - -def test_eagle_no_att_fill_hook_is_noop(): - proposer = EagleNoAttProposer(backend=SimpleNamespace(draft_models=[]), enable_dynmaic_mtp=False) - - proposer.fill_draft_model_kv_state(None, None, None) diff --git a/unit_tests/server/router/model_infer/mtp_speculative/test_eagle_overlap.py b/unit_tests/server/router/model_infer/mtp_speculative/test_eagle_overlap.py deleted file mode 100644 index d5dae90b4b..0000000000 --- a/unit_tests/server/router/model_infer/mtp_speculative/test_eagle_overlap.py +++ /dev/null @@ -1,363 +0,0 @@ -from types import SimpleNamespace - -import pytest -import torch - -from lightllm.common.basemodel.batch_objs import ModelMtpOutputCollector, ModelOutput -from lightllm.server.router.model_infer.mtp_speculative.dp_overlap_proposers import eagle_with_att -from lightllm.server.router.model_infer.mtp_speculative.dp_overlap_proposers.utils import ( - get_dp_overlap_req_start_rows, -) -from lightllm.server.router.model_infer.mtp_speculative.dp_overlap_proposers.eagle3 import ( - DpOverlapEagle3Proposer, -) -from lightllm.server.router.model_infer.mtp_speculative.dp_overlap_proposers.eagle_no_att import ( - DpOverlapEagleNoAttProposer, -) -from lightllm.server.router.model_infer.mtp_speculative.dp_overlap_proposers.eagle_with_att import ( - DpOverlapEagleWithAttProposer, -) -from lightllm.server.router.model_infer.mtp_speculative.proposers.eagle3 import ( - Eagle3Proposer, -) - - -class _DraftModel: - def __init__(self): - self.extend_batch_sizes = None - self.extend_inputs = None - self.decode_batch_sizes = [] - self.decode_inputs = [] - - def _microbatch_overlap_prefill_cuda(self, input0, input1): - self.extend_batch_sizes = (input0.batch_size, input1.batch_size) - self.extend_inputs = (input0, input1) - return tuple( - ModelOutput( - logits=torch.arange(model_input.batch_size, dtype=torch.float32).view(-1, 1), - mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.ones((model_input.batch_size, 2))), - ) - for model_input in (input0, input1) - ) - - def _microbatch_overlap_decode_cuda(self, input0, input1): - self.decode_batch_sizes.append((input0.batch_size, input1.batch_size)) - self.decode_inputs.append((input0, input1)) - return tuple( - ModelOutput( - logits=torch.arange( - model_input.batch_size, - dtype=torch.float32, - device=model_input.input_ids.device, - ).view(-1, 1), - mtp_collector=ModelMtpOutputCollector( - spec_hidden=torch.ones( - (model_input.batch_size, 2), - device=model_input.input_ids.device, - ) - ), - ) - for model_input in (input0, input1) - ) - - def map_draft_vocab_to_main_vocab(self, token_ids): - return token_ids - - -def _target_input(batch_size, b_mtp_index=None, device="cpu"): - if b_mtp_index is None: - b_mtp_index = torch.arange(batch_size, dtype=torch.int32, device=device) % 3 - return SimpleNamespace( - batch_size=batch_size, - total_token_num=batch_size, - input_ids=torch.arange(batch_size, dtype=torch.int64, device=device), - b_seq_len=torch.arange(batch_size, dtype=torch.int32, device=device) + 4, - b_req_idx=torch.arange(batch_size, dtype=torch.int32, device=device), - b_mtp_index=b_mtp_index, - b_position_delta=torch.zeros(batch_size, dtype=torch.int32, device=device), - b_shared_seq_len=torch.zeros(batch_size, dtype=torch.int32, device=device), - b_shared_radix_node_id=torch.arange(batch_size, dtype=torch.int64, device=device), - max_kv_seq_len=16, - max_cache_len=16, - is_prefill=False, - multimodal_params=[{"images": [], "audios": []}] * batch_size, - ) - - -def _patch_cpu_req_start_rows(monkeypatch): - def get_cpu_req_start_rows(b_mtp_index, req_num): - req_start_rows = torch.nonzero(b_mtp_index == 0, as_tuple=False).flatten().to(dtype=torch.int32) - assert req_start_rows.shape == (req_num,) - return req_start_rows - - monkeypatch.setattr(eagle_with_att, "get_dp_overlap_req_start_rows", get_cpu_req_start_rows) - - -def test_dp_overlap_req_start_rows_rejects_nonempty_cpu_input(): - with pytest.raises(AssertionError, match="must be a CUDA tensor"): - get_dp_overlap_req_start_rows( - b_mtp_index=torch.tensor([0, 1], dtype=torch.int32), - req_num=1, - ) - - -def test_overlap_eagle_supports_variable_verify_layout(monkeypatch): - _patch_cpu_req_start_rows(monkeypatch) - draft_model = _DraftModel() - backend = SimpleNamespace( - max_draft_step=2, - draft_models=[draft_model], - model=SimpleNamespace( - req_manager=SimpleNamespace( - mem_manager=SimpleNamespace(HOLD_TOKEN_MEMINDEXES=(99,), page_size=1), - ) - ), - _gen_argmax_token_ids=lambda output: output.logits[:, 0].to(torch.int64), - ) - proposer = DpOverlapEagleWithAttProposer(backend=backend, enable_dynmaic_mtp=False) - model_input0 = _target_input(batch_size=3) - model_input1 = _target_input(batch_size=6) - - proposal = proposer.propose_next_overlap( - target_model_input0=model_input0, - target_model_output0=ModelOutput( - logits=torch.empty((3, 1)), - mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.ones((3, 2))), - ), - target_next_token_ids0=torch.arange(3, dtype=torch.int64), - accept_len0=torch.tensor([2], dtype=torch.int32), - target_model_input1=model_input1, - target_model_output1=ModelOutput( - logits=torch.empty((6, 1)), - mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.ones((6, 2))), - ), - target_next_token_ids1=torch.arange(10, 16, dtype=torch.int64), - accept_len1=torch.tensor([1, 3], dtype=torch.int32), - draft_step=2, - ) - - assert draft_model.extend_batch_sizes is None - assert draft_model.decode_batch_sizes == [(3, 6), (1, 2)] - assert proposal.token_ids.shape == (3, 2) - assert torch.equal(proposal.token_ids, torch.tensor([[1, 0], [0, 0], [5, 1]])) - - -def test_overlap_eagle_supports_empty_verify_rows(): - draft_model = _DraftModel() - backend = SimpleNamespace( - max_draft_step=2, - draft_models=[draft_model], - model=SimpleNamespace( - req_manager=SimpleNamespace( - mem_manager=SimpleNamespace(HOLD_TOKEN_MEMINDEXES=(99,), page_size=1), - ) - ), - _gen_argmax_token_ids=lambda output: output.logits[:, 0].to(torch.int64), - ) - proposer = DpOverlapEagleWithAttProposer(backend=backend, enable_dynmaic_mtp=False) - - proposal = proposer.propose_next_overlap( - target_model_input0=_target_input(batch_size=0), - target_model_output0=ModelOutput( - logits=torch.empty((0, 1)), - mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.ones((0, 2))), - ), - target_next_token_ids0=torch.empty((0,), dtype=torch.int64), - accept_len0=torch.empty((0,), dtype=torch.int32), - target_model_input1=_target_input(batch_size=0), - target_model_output1=ModelOutput( - logits=torch.empty((0, 1)), - mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.ones((0, 2))), - ), - target_next_token_ids1=torch.empty((0,), dtype=torch.int64), - accept_len1=torch.empty((0,), dtype=torch.int32), - draft_step=2, - ) - - assert proposal.token_ids.shape == (0, 2) - assert draft_model.decode_batch_sizes == [(0, 0), (0, 0)] - - -def test_overlap_eagle_returns_dynamic_schedule_scores(monkeypatch): - _patch_cpu_req_start_rows(monkeypatch) - draft_model = _DraftModel() - backend = SimpleNamespace( - max_draft_step=2, - draft_models=[draft_model], - model=SimpleNamespace( - req_manager=SimpleNamespace( - mem_manager=SimpleNamespace(HOLD_TOKEN_MEMINDEXES=(99,), page_size=1), - ) - ), - _gen_argmax_token_ids_and_prob=lambda output: ( - output.logits[:, 0].to(torch.int64), - output.logits[:, 0] / 10 + 0.5, - ), - ) - proposer = DpOverlapEagleWithAttProposer(backend=backend, enable_dynmaic_mtp=True) - - proposal = proposer.propose_next_overlap( - target_model_input0=_target_input( - batch_size=2, - b_mtp_index=torch.tensor([0, 1], dtype=torch.int32), - ), - target_model_output0=ModelOutput( - logits=torch.empty((2, 1)), - mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.ones((2, 2))), - ), - target_next_token_ids0=torch.arange(2, dtype=torch.int64), - accept_len0=torch.tensor([2], dtype=torch.int32), - target_model_input1=_target_input( - batch_size=3, - b_mtp_index=torch.tensor([0, 1, 0], dtype=torch.int32), - ), - target_model_output1=ModelOutput( - logits=torch.empty((3, 1)), - mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.ones((3, 2))), - ), - target_next_token_ids1=torch.arange(2, 5, dtype=torch.int64), - accept_len1=torch.tensor([1, 1], dtype=torch.int32), - draft_step=2, - ) - - assert torch.equal(proposal.token_ids, torch.tensor([[1, 0], [0, 0], [2, 1]])) - assert torch.allclose( - proposal.schedule_scores, - torch.tensor([[0.6, 0.5], [0.5, 0.5], [0.7, 0.6]]), - ) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") -def test_overlap_eagle_no_att_supports_dynamic_draft_step(): - device = "cuda" - draft_model = _DraftModel() - backend = SimpleNamespace( - max_draft_step=3, - draft_models=[draft_model], - _gen_argmax_token_ids_and_prob=lambda output: ( - output.logits[:, 0].to(torch.int64), - output.logits[:, 0] / 10 + 0.5, - ), - ) - proposer = DpOverlapEagleNoAttProposer(backend=backend, enable_dynmaic_mtp=True) - model_input0 = _target_input( - batch_size=2, - b_mtp_index=torch.tensor([0, 1], dtype=torch.int32, device=device), - device=device, - ) - model_input1 = _target_input( - batch_size=3, - b_mtp_index=torch.tensor([0, 1, 0], dtype=torch.int32, device=device), - device=device, - ) - - proposal = proposer.propose_next_overlap( - target_model_input0=model_input0, - target_model_output0=ModelOutput( - logits=torch.empty((2, 1), device=device), - mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.ones((2, 2), device=device)), - ), - target_next_token_ids0=torch.arange(2, dtype=torch.int64, device=device), - accept_len0=torch.tensor([2], dtype=torch.int32, device=device), - target_model_input1=model_input1, - target_model_output1=ModelOutput( - logits=torch.empty((3, 1), device=device), - mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.ones((3, 2), device=device)), - ), - target_next_token_ids1=torch.arange(2, 5, dtype=torch.int64, device=device), - accept_len1=torch.tensor([1, 1], dtype=torch.int32, device=device), - draft_step=2, - ) - - assert proposal.token_ids.tolist() == [[0, 0], [0, 0], [1, 1]] - assert torch.allclose( - proposal.schedule_scores, - torch.tensor( - [[0.5, 0.5], [0.5, 0.5], [0.6, 0.6]], - device=device, - ), - ) - assert draft_model.decode_batch_sizes == [(1, 2), (1, 2)] - - -def test_autoregressive_eagle_reuses_overlap_inputs(monkeypatch): - _patch_cpu_req_start_rows(monkeypatch) - draft_model = _DraftModel() - backend = SimpleNamespace( - max_draft_step=2, - draft_models=[draft_model], - model=SimpleNamespace( - req_manager=SimpleNamespace( - mem_manager=SimpleNamespace(HOLD_TOKEN_MEMINDEXES=(99,), page_size=1), - ) - ), - _gen_argmax_token_ids=lambda output: output.logits[:, 0].to(torch.int64), - ) - proposer = DpOverlapEagle3Proposer(backend=backend, enable_dynmaic_mtp=False) - model_input0 = _target_input(batch_size=3) - model_input1 = _target_input(batch_size=6) - - proposal = proposer.propose_next_overlap( - target_model_input0=model_input0, - target_model_output0=ModelOutput( - logits=torch.empty((3, 1)), - mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.ones((3, 2))), - ), - target_next_token_ids0=torch.arange(3, dtype=torch.int64), - accept_len0=torch.tensor([2], dtype=torch.int32), - target_model_input1=model_input1, - target_model_output1=ModelOutput( - logits=torch.empty((6, 1)), - mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.ones((6, 2))), - ), - target_next_token_ids1=torch.arange(10, 16, dtype=torch.int64), - accept_len1=torch.tensor([1, 3], dtype=torch.int32), - draft_step=2, - ) - - assert draft_model.extend_inputs is None - assert len(draft_model.decode_inputs) == 2 - assert draft_model.decode_inputs[0][0] is not model_input0 - assert draft_model.decode_inputs[0][1] is not model_input1 - assert draft_model.extend_batch_sizes is None - assert draft_model.decode_batch_sizes == [(3, 6), (1, 2)] - assert proposal.token_ids.shape == (3, 2) - assert torch.equal(proposal.token_ids, torch.tensor([[1, 0], [0, 0], [5, 1]])) - - -def test_eagle3_maps_draft_token_ids_in_proposer(): - proposer = Eagle3Proposer.__new__(Eagle3Proposer) - proposer.backend = SimpleNamespace( - draft_models=[SimpleNamespace(map_draft_vocab_to_main_vocab=lambda token_ids: token_ids + 100)], - _gen_argmax_token_ids=lambda _: torch.tensor([1, 2]), - _gen_argmax_token_ids_and_prob=lambda _: ( - torch.tensor([3, 4]), - torch.tensor([0.8, 0.7]), - ), - ) - - token_ids = proposer._gen_argmax_token_ids(ModelOutput(logits=torch.empty(0))) - token_ids_with_prob, probs = proposer._gen_argmax_token_ids_and_prob(ModelOutput(logits=torch.empty(0))) - - assert torch.equal(token_ids, torch.tensor([101, 102])) - assert torch.equal(token_ids_with_prob, torch.tensor([103, 104])) - assert torch.equal(probs, torch.tensor([0.8, 0.7])) - - -def test_dp_overlap_eagle3_maps_draft_token_ids_in_proposer(): - proposer = DpOverlapEagle3Proposer.__new__(DpOverlapEagle3Proposer) - proposer.backend = SimpleNamespace( - draft_models=[SimpleNamespace(map_draft_vocab_to_main_vocab=lambda token_ids: token_ids + 100)], - _gen_argmax_token_ids=lambda _: torch.tensor([1, 2]), - _gen_argmax_token_ids_and_prob=lambda _: ( - torch.tensor([3, 4]), - torch.tensor([0.8, 0.7]), - ), - ) - - token_ids = proposer._gen_argmax_token_ids(ModelOutput(logits=torch.empty(0))) - token_ids_with_prob, probs = proposer._gen_argmax_token_ids_and_prob(ModelOutput(logits=torch.empty(0))) - - assert torch.equal(token_ids, torch.tensor([101, 102])) - assert torch.equal(token_ids_with_prob, torch.tensor([103, 104])) - assert torch.equal(probs, torch.tensor([0.8, 0.7])) diff --git a/unit_tests/server/router/model_infer/mtp_speculative/test_eagle_with_att.py b/unit_tests/server/router/model_infer/mtp_speculative/test_eagle_with_att.py deleted file mode 100644 index 365acb30b2..0000000000 --- a/unit_tests/server/router/model_infer/mtp_speculative/test_eagle_with_att.py +++ /dev/null @@ -1,213 +0,0 @@ -from types import SimpleNamespace - -import pytest -import torch - -from lightllm.common.basemodel.batch_objs import ModelMtpOutputCollector, ModelOutput -from lightllm.server.router.model_infer.mtp_speculative.proposers.eagle3 import Eagle3Proposer -from lightllm.server.router.model_infer.mtp_speculative.proposers.eagle_with_att import EagleWithAttProposer -from lightllm.server.router.model_infer.mtp_speculative.proposers.proposal_type import EagleSpecProposal - - -def test_eagle3_reuses_attention_flow_and_maps_proposal_tokens(): - target_input_ids = torch.tensor([10, 20], dtype=torch.int64) - target_hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) - - def forward(model_input): - assert model_input.input_ids is target_input_ids - assert model_input.mtp_draft_input_hiddens is target_hidden - return ModelOutput( - logits=torch.tensor([[1.0], [2.0]]), - mtp_collector=ModelMtpOutputCollector(spec_hidden=target_hidden + 10), - ) - - draft_model = SimpleNamespace( - forward=forward, - map_draft_vocab_to_main_vocab=lambda token_ids: token_ids + 100, - ) - proposer = Eagle3Proposer( - backend=SimpleNamespace( - draft_models=[draft_model], - _gen_argmax_token_ids=lambda output: output.logits[:, 0].long(), - ), - enable_dynmaic_mtp=False, - ) - target_input = SimpleNamespace( - is_prefill=False, - batch_size=2, - input_ids=None, - b_position_delta=torch.zeros(2, dtype=torch.int32), - mtp_draft_input_hiddens=None, - ) - - proposal = proposer.propose_next( - target_model_input=target_input, - target_model_output=SimpleNamespace(mtp_collector=SimpleNamespace(spec_hidden=target_hidden)), - target_next_token_ids=target_input_ids, - b_req_mtp_start_loc=torch.tensor([0, 1], dtype=torch.int32), - draft_step=1, - accept_len=torch.ones(2, dtype=torch.int32), - ) - - assert isinstance(proposal, EagleSpecProposal) - torch.testing.assert_close(proposal.token_ids, torch.tensor([[101], [102]])) - assert proposal.schedule_scores is None - assert target_input.input_ids is None - assert target_input.mtp_draft_input_hiddens is None - - -def test_eagle_with_att_rejects_zero_draft_steps(): - proposer = EagleWithAttProposer( - backend=SimpleNamespace(draft_models=[]), - enable_dynmaic_mtp=False, - ) - - with pytest.raises(AssertionError, match="requires draft_step to be greater than 0"): - proposer.propose_next( - target_model_input=None, - target_model_output=None, - target_next_token_ids=torch.tensor([10]), - b_req_mtp_start_loc=torch.tensor([0], dtype=torch.int32), - draft_step=0, - ) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") -def test_eagle_with_att_prefill_builds_draft_kv_without_mutating_target(): - device = "cuda" - original_input_ids = torch.tensor([10, 11, 12, 20, 21, 22], dtype=torch.int64, device=device) - target_hidden = torch.arange(12, dtype=torch.float32, device=device).reshape(6, 2) - target_input = SimpleNamespace( - is_prefill=True, - batch_size=2, - input_ids=original_input_ids, - b_req_idx=torch.tensor([7, 9], dtype=torch.int32, device=device), - b_seq_len=torch.tensor([3, 3], dtype=torch.int32, device=device), - b_ready_cache_len=torch.zeros(2, dtype=torch.int32, device=device), - b_position_delta=None, - b_is_decode_req=None, - mtp_draft_input_hiddens=None, - ) - forwarded = [] - - def forward(model_input): - forwarded.append( - ( - model_input, - model_input.input_ids.clone(), - model_input.mtp_draft_input_hiddens, - model_input.b_is_decode_req, - ) - ) - return ModelOutput(logits=torch.empty((0,), device=device)) - - proposer = EagleWithAttProposer( - backend=SimpleNamespace(draft_models=[SimpleNamespace(forward=forward)]), - enable_dynmaic_mtp=False, - ) - proposer.fill_draft_model_kv_state( - target_model_input=target_input, - target_model_output=SimpleNamespace(mtp_collector=SimpleNamespace(spec_hidden=target_hidden)), - target_next_token_ids=torch.tensor([13, 23], dtype=torch.int64, device=device), - ) - - assert len(forwarded) == 1 - assert forwarded[0][0] is not target_input - torch.testing.assert_close( - forwarded[0][1], - torch.tensor([11, 12, 13, 21, 22, 23], dtype=torch.int64, device=device), - ) - assert forwarded[0][2] is target_hidden - assert not forwarded[0][3].any() - assert target_input.input_ids is original_input_ids - assert target_input.b_is_decode_req is None - assert target_input.mtp_draft_input_hiddens is None - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") -def test_eagle_with_att_commits_verify_kv_then_recurrently_decodes(monkeypatch): - device = "cuda" - original_input_ids = torch.tensor([10, 11, 12, 20, 21, 22], dtype=torch.int64, device=device) - target_hidden = torch.arange(12, dtype=torch.float32, device=device).reshape(6, 2) - extend_hidden = target_hidden + 100 - decode_hidden = torch.tensor([[201.0, 202.0], [203.0, 204.0]], device=device) - draft_calls = [] - - def forward(model_input): - draft_calls.append( - { - "model_input": model_input, - "batch_size": model_input.batch_size, - "input_ids": model_input.input_ids.clone(), - "draft_hidden": model_input.mtp_draft_input_hiddens.clone(), - "b_req_idx": model_input.b_req_idx.clone(), - "b_mtp_index": model_input.b_mtp_index.clone(), - "b_seq_len": model_input.b_seq_len.clone(), - "max_kv_seq_len": model_input.max_kv_seq_len, - "total_token_num": model_input.total_token_num, - } - ) - if len(draft_calls) == 1: - return ModelOutput( - logits=torch.arange(30, 36, dtype=torch.float32, device=device).unsqueeze(1), - mtp_collector=ModelMtpOutputCollector(spec_hidden=extend_hidden), - ) - return ModelOutput( - logits=torch.tensor([[40.0], [41.0]], device=device), - mtp_collector=ModelMtpOutputCollector(spec_hidden=decode_hidden), - ) - - backend = SimpleNamespace( - draft_models=[SimpleNamespace(forward=forward)], - _gen_argmax_token_ids=lambda output: output.logits[:, 0].long(), - _gen_argmax_token_ids_and_prob=lambda output: ( - output.logits[:, 0].long(), - output.logits[:, 0] / 100, - ), - ) - proposer = EagleWithAttProposer(backend=backend, enable_dynmaic_mtp=True) - target_input = SimpleNamespace( - is_prefill=False, - batch_size=6, - total_token_num=66, - max_kv_seq_len=22, - input_ids=original_input_ids, - b_req_idx=torch.tensor([7, 7, 7, 9, 9, 9], dtype=torch.int32, device=device), - b_mtp_index=torch.tensor([0, 1, 2, 0, 1, 2], dtype=torch.int32, device=device), - b_seq_len=torch.tensor([10, 11, 12, 20, 21, 22], dtype=torch.int32, device=device), - b_position_delta=torch.tensor([0, 1, 2, 3, 4, 5], dtype=torch.int32, device=device), - b_shared_seq_len=torch.tensor([8, 8, 8, 6, 6, 6], dtype=torch.int32, device=device), - b_shared_radix_node_id=torch.tensor([70, 70, 70, 90, 90, 90], dtype=torch.int64, device=device), - multimodal_params=[{"images": [], "audios": []} for _ in range(6)], - mtp_draft_input_hiddens=None, - ) - proposal = proposer.propose_next( - target_model_input=target_input, - target_model_output=SimpleNamespace(mtp_collector=SimpleNamespace(spec_hidden=target_hidden)), - target_next_token_ids=original_input_ids, - b_req_mtp_start_loc=torch.tensor([0, 3], dtype=torch.int32, device=device), - draft_step=2, - accept_len=torch.tensor([3, 2], dtype=torch.int32, device=device), - ) - - assert isinstance(proposal, EagleSpecProposal) - torch.testing.assert_close(proposal.token_ids, torch.tensor([[32, 40], [34, 41]], device=device)) - torch.testing.assert_close(proposal.schedule_scores, torch.tensor([[0.32, 0.40], [0.34, 0.41]], device=device)) - assert len(draft_calls) == 2 - assert draft_calls[0]["model_input"] is not target_input - assert draft_calls[0]["batch_size"] == 6 - torch.testing.assert_close(draft_calls[0]["input_ids"], original_input_ids) - torch.testing.assert_close(draft_calls[0]["draft_hidden"], target_hidden) - assert draft_calls[1]["batch_size"] == 2 - torch.testing.assert_close(draft_calls[1]["input_ids"], torch.tensor([32, 34], device=device)) - torch.testing.assert_close( - draft_calls[1]["draft_hidden"], - extend_hidden.index_select(0, torch.tensor([2, 4], device=device)), - ) - torch.testing.assert_close(draft_calls[1]["b_req_idx"], torch.tensor([7, 9], dtype=torch.int32, device=device)) - torch.testing.assert_close(draft_calls[1]["b_mtp_index"], torch.zeros(2, dtype=torch.int32, device=device)) - torch.testing.assert_close(draft_calls[1]["b_seq_len"], torch.tensor([13, 22], dtype=torch.int32, device=device)) - assert draft_calls[1]["max_kv_seq_len"] == 23 - assert draft_calls[1]["total_token_num"] == 46 - assert target_input.input_ids is original_input_ids - assert target_input.mtp_draft_input_hiddens is None diff --git a/unit_tests/server/router/model_infer/mtp_speculative/test_planner.py b/unit_tests/server/router/model_infer/mtp_speculative/test_planner.py deleted file mode 100644 index 2e488ebb30..0000000000 --- a/unit_tests/server/router/model_infer/mtp_speculative/test_planner.py +++ /dev/null @@ -1,964 +0,0 @@ -from types import SimpleNamespace - -import numpy as np -import pytest -import torch - -from lightllm.server.router.model_infer.mtp_speculative.engine import SpecEngine -from lightllm.server.router.model_infer.mtp_speculative.dp_overlap_proposers import ( - build_dp_overlap_spec_proposer, -) -from lightllm.server.router.model_infer.mtp_speculative.dp_overlap_proposers.base import ( - BaseDpOverlapProposer, -) -from lightllm.server.router.model_infer.mtp_speculative.dp_overlap_proposers.eagle3 import ( - DpOverlapEagle3Proposer, -) -from lightllm.server.router.model_infer.mtp_speculative.dp_overlap_proposers.eagle_no_att import ( - DpOverlapEagleNoAttProposer, -) -from lightllm.server.router.model_infer.mtp_speculative.dp_overlap_proposers.eagle_with_att import ( - DpOverlapEagleWithAttProposer, -) -from lightllm.server.router.model_infer.mtp_speculative.dp_overlap_proposers.vanilla_no_att import ( - DpOverlapVanillaNoAttProposer, -) -from lightllm.server.router.model_infer.mtp_speculative.dp_overlap_proposers.vanilla_with_att import ( - DpOverlapVanillaWithAttProposer, -) -from lightllm.server.router.model_infer.mtp_speculative.planner import ( - BaseMtpPlanner, - DSparkPlanner, - FixedSpecPlanner, - LightSpecPlanner, - SpecDecodePlan, -) -from lightllm.server.router.model_infer.mtp_speculative.planner.base import ( - _InferCostMsTable, -) -from lightllm.server.router.model_infer.mtp_speculative.proposers import ( - build_spec_proposer, -) -from lightllm.server.router.model_infer.mtp_speculative.proposers.base import ( - BaseSpecProposer, - SpecProposal, -) -from lightllm.server.router.model_infer.mtp_speculative.proposers.dflash import ( - DFlashProposer, -) -from lightllm.server.router.model_infer.mtp_speculative.proposers.dflash2 import ( - DFlash2Proposer, -) -from lightllm.server.router.model_infer.mtp_speculative.proposers.dspark import ( - DSparkProposer, -) -from lightllm.server.router.model_infer.mtp_speculative.proposers.eagle3 import ( - Eagle3Proposer, -) -from lightllm.server.router.model_infer.mtp_speculative.proposers.eagle_no_att import ( - EagleNoAttProposer, -) -from lightllm.server.router.model_infer.mtp_speculative.proposers.eagle_with_att import ( - EagleWithAttProposer, -) -from lightllm.server.router.model_infer.mtp_speculative.proposers.proposal_type import ( - DFlashSpecProposal, - DSparkSpecProposal, - EagleSpecProposal, - VanillaSpecProposal, -) -from lightllm.server.router.model_infer.mtp_speculative.proposers.vanilla_no_att import ( - VanillaNoAttProposer, -) -from lightllm.server.router.model_infer.mtp_speculative.proposers.vanilla_with_att import ( - VanillaWithAttProposer, -) -from lightllm.server.router.model_infer.mtp_speculative import utils as mtp_utils - - -def build_lightspec_planner( - max_draft_step: int = 3, - spec_mode: str = "vanilla_with_att", - block_size: int = 3, - enable_decode_microbatch_overlap: bool = False, -): - backend = SimpleNamespace( - args=SimpleNamespace(dp=1), - enable_decode_microbatch_overlap=enable_decode_microbatch_overlap, - max_draft_step=max_draft_step, - model=SimpleNamespace(graph=None), - draft_models=[SimpleNamespace(block_size=block_size, graph=None)], - ) - return LightSpecPlanner( - spec_mode=spec_mode, - backend=backend, - ) - - -def build_dspark_planner(max_draft_step: int = 3, block_size: int = 3): - backend = SimpleNamespace( - max_draft_step=max_draft_step, - model=SimpleNamespace(graph=None), - draft_models=[SimpleNamespace(block_size=block_size, graph=None)], - ) - return DSparkPlanner(backend=backend) - - -def build_decode_reqs(req_num: int, req_num_with_proposals: int | None = None): - if req_num_with_proposals is None: - req_num_with_proposals = req_num - return [SimpleNamespace(cur_output_len=2)] * req_num_with_proposals + [SimpleNamespace(cur_output_len=1)] * ( - req_num - req_num_with_proposals - ) - - -def build_planner(spec_mode: str, enable_dynmaic_mtp: bool = True): - engine = SpecEngine.__new__(SpecEngine) - engine.backend = SimpleNamespace( - args=SimpleNamespace(dp=1), - enable_decode_microbatch_overlap=False, - max_draft_step=3, - model=SimpleNamespace(graph=None), - draft_models=[SimpleNamespace(block_size=3, graph=None)], - ) - return engine._build_mtp_planner( - spec_mode=spec_mode, - enable_dynmaic_mtp=enable_dynmaic_mtp, - ) - - -def test_fixed_planner_returns_static_plan(): - planner = FixedSpecPlanner(max_draft_step=3) - plan = planner.plan(decode_reqs=build_decode_reqs(4), origin_batch_size=16) - - assert isinstance(planner, BaseMtpPlanner) - assert plan.origin_batch_size == 16 - assert plan.dynamic_batch_size == plan.origin_batch_size - assert plan.draft_step == plan.pre_draft_step == 3 - assert not plan.skip_verify_sync - - -def test_common_engine_accepts_empty_fixed_decode_batch(): - engine = SpecEngine.__new__(SpecEngine) - engine.planner = FixedSpecPlanner(max_draft_step=3) - - plan = engine.plan_decode(model_input=SimpleNamespace(batch_size=0), decode_reqs=[]) - - assert plan == SpecDecodePlan( - origin_batch_size=0, - dynamic_batch_size=0, - draft_step=3, - pre_draft_step=3, - ) - - -def test_common_engine_delegates_empty_dp_batch_to_lightspec_planner(): - engine = SpecEngine.__new__(SpecEngine) - engine.backend = SimpleNamespace(max_draft_step=3, dp_size=2) - engine.planner = build_lightspec_planner(spec_mode="eagle3") - engine.planner.backend.args.dp = 2 - engine.planner.pre_draft_step = 2 - - plan = engine.plan_decode(model_input=SimpleNamespace(batch_size=0), decode_reqs=[]) - - assert plan == SpecDecodePlan( - origin_batch_size=0, - dynamic_batch_size=0, - draft_step=1, - pre_draft_step=2, - ) - - -def test_spec_engine_only_exposes_planning_and_proposal_interfaces(): - public_methods = { - name for name, value in SpecEngine.__dict__.items() if callable(value) and not name.startswith("_") - } - - assert public_methods == { - "fill_draft_model_kv_state", - "plan_decode", - "prepare_decode_model_input", - "propose_next", - "update_planner_statics", - } - - -def test_mode_proposals_own_their_schedule_metadata(): - assert "schedule_scores" not in SpecProposal.__dataclass_fields__ - assert "schedule_scores_cpu" not in SpecProposal.__dataclass_fields__ - - for proposal_type in ( - VanillaSpecProposal, - EagleSpecProposal, - DFlashSpecProposal, - DSparkSpecProposal, - ): - assert "schedule_scores" in proposal_type.__dataclass_fields__ - assert "schedule_scores_cpu" not in VanillaSpecProposal.__dataclass_fields__ - assert "schedule_scores_cpu" not in EagleSpecProposal.__dataclass_fields__ - assert "schedule_scores_cpu" not in DFlashSpecProposal.__dataclass_fields__ - assert "schedule_scores_cpu" in DSparkSpecProposal.__dataclass_fields__ - - -def test_scatter_mtp_next_tokens_consumes_mode_proposal(monkeypatch): - scatter_args = {} - monkeypatch.setattr( - mtp_utils, - "mtp_scatter_next_token_ids", - lambda **kwargs: scatter_args.update(kwargs), - ) - req_to_next_token_ids = torch.empty((4, 2), dtype=torch.int64) - req_to_next_token_scores = torch.empty((4, 1), dtype=torch.float32) - backend = SimpleNamespace( - model=SimpleNamespace( - req_manager=SimpleNamespace( - req_sampling_params_manager=SimpleNamespace( - req_to_bitmask_enabled=None, - req_to_next_token_ids=req_to_next_token_ids, - req_to_next_token_scores=req_to_next_token_scores, - ) - ) - ) - ) - proposal = DFlashSpecProposal( - token_ids=torch.arange(2, dtype=torch.int64).view(2, 1), - schedule_scores=torch.arange(2, dtype=torch.float32).view(2, 1), - ) - next_token_ids = torch.tensor([10, 11], dtype=torch.int64) - - mtp_utils.scatter_mtp_next_tokens( - backend=backend, - proposal=proposal, - target_next_token_ids=next_token_ids, - b_req_mtp_start_loc=torch.tensor([0, 1], dtype=torch.int32), - b_req_idx=torch.tensor([0, 1], dtype=torch.int32), - mtp_accept_len=torch.ones(2, dtype=torch.int32), - ) - - assert scatter_args["req_to_next_token_ids"] is req_to_next_token_ids - assert scatter_args["req_to_next_token_scores"] is req_to_next_token_scores - assert torch.equal(scatter_args["target_next_token_ids"], next_token_ids) - assert torch.equal(scatter_args["draft_token_ids"], proposal.token_ids) - assert torch.equal(scatter_args["schedule_scores"], proposal.schedule_scores) - - -def test_scatter_mtp_next_tokens_ignores_empty_schedule_scores(monkeypatch): - scatter_args = {} - monkeypatch.setattr( - mtp_utils, - "mtp_scatter_next_token_ids", - lambda **kwargs: scatter_args.update(kwargs), - ) - backend = SimpleNamespace( - model=SimpleNamespace( - req_manager=SimpleNamespace( - req_sampling_params_manager=SimpleNamespace( - req_to_bitmask_enabled=None, - req_to_next_token_ids=torch.empty((4, 2), dtype=torch.int64), - req_to_next_token_scores=torch.empty((4, 2), dtype=torch.float32), - ) - ) - ) - ) - proposal = VanillaSpecProposal( - token_ids=torch.empty((2, 0), dtype=torch.int64), - schedule_scores=torch.empty((2, 0), dtype=torch.float32), - ) - - mtp_utils.scatter_mtp_next_tokens( - backend=backend, - proposal=proposal, - target_next_token_ids=torch.tensor([10, 11], dtype=torch.int64), - b_req_mtp_start_loc=torch.tensor([0, 1], dtype=torch.int32), - b_req_idx=torch.tensor([0, 1], dtype=torch.int32), - mtp_accept_len=torch.ones(2, dtype=torch.int32), - ) - - assert scatter_args["req_to_next_token_scores"] is None - assert scatter_args["schedule_scores"] is None - - -def test_infer_cost_candidates_include_feasible_boundaries(): - costs = _InferCostMsTable() - costs.update(batch_size=4, infer_cost_ms=1.0) - costs.update(batch_size=8, infer_cost_ms=2.0) - - assert costs.get_batch_size_keys_between(5, 10) == [5, 8, 10] - - -def test_engine_routes_only_dspark_to_the_confidence_planner(): - fixed_planner = build_planner("eagle3", enable_dynmaic_mtp=False) - dspark_planner = build_planner("dspark") - assert isinstance(fixed_planner, FixedSpecPlanner) - assert isinstance(dspark_planner, DSparkPlanner) - assert isinstance(fixed_planner, BaseMtpPlanner) - assert isinstance(dspark_planner, BaseMtpPlanner) - - vanilla_planner = build_planner("vanilla_with_att") - assert isinstance(vanilla_planner, LightSpecPlanner) - assert isinstance(vanilla_planner, BaseMtpPlanner) - assert vanilla_planner.draft_steps == (3,) - - vanilla_no_att_planner = build_planner("vanilla_no_att") - assert vanilla_no_att_planner.draft_steps == (0, 1, 2, 3) - - eagle_no_att_planner = build_planner("eagle_no_att") - assert eagle_no_att_planner.draft_steps == (0, 1, 2, 3) - - dflash_planner = build_planner("dflash") - assert isinstance(dflash_planner, LightSpecPlanner) - assert dflash_planner.draft_steps == (3,) - - dflash2_planner = build_planner("dflash2", enable_dynmaic_mtp=False) - assert isinstance(dflash2_planner, FixedSpecPlanner) - dflash2_dynamic_planner = build_planner("dflash2") - assert isinstance(dflash2_dynamic_planner, LightSpecPlanner) - assert dflash2_dynamic_planner.draft_steps == (3,) - - eagle_planner = build_planner("eagle3") - assert isinstance(eagle_planner, LightSpecPlanner) - assert eagle_planner.draft_steps == (1, 2, 3) - - eagle_with_att_planner = build_planner("eagle_with_att") - assert eagle_with_att_planner.draft_steps == (1, 2, 3) - - -def test_dynamic_planner_registers_cuda_graph_costs_from_backend(): - backend = SimpleNamespace( - args=SimpleNamespace(dp=1), - enable_decode_microbatch_overlap=False, - max_draft_step=3, - model=SimpleNamespace(graph=SimpleNamespace(infer_cost_ms_by_batch_size={4: 1.2})), - draft_models=[ - SimpleNamespace( - block_size=3, - graph=SimpleNamespace(infer_cost_ms_by_batch_size={4: 0.3}), - ) - ], - ) - - planner = LightSpecPlanner(spec_mode="vanilla_with_att", backend=backend) - - assert planner.target_infer_costs.estimate(4) == 1.2 - assert planner.draft_infer_costs.estimate(4) == 0.3 - - -def test_overlap_lightspec_consumes_combined_cuda_graph_batch_size(): - backend = SimpleNamespace( - args=SimpleNamespace(dp=1), - enable_decode_microbatch_overlap=True, - max_draft_step=3, - model=SimpleNamespace(graph=SimpleNamespace(infer_cost_ms_by_batch_size={8: 1.2})), - draft_models=[ - SimpleNamespace( - block_size=3, - graph=SimpleNamespace(infer_cost_ms_by_batch_size={8: 0.3}), - ) - ], - ) - - planner = LightSpecPlanner(spec_mode="vanilla_with_att", backend=backend) - - assert planner.target_infer_costs.estimate(7) == 1.2 - assert planner.target_infer_costs.estimate(8) == 1.2 - assert planner.draft_infer_costs.estimate(7) == 0.3 - assert planner.draft_infer_costs.estimate(8) == 0.3 - - -def test_each_mode_proposer_inherits_its_expected_implementation_base(): - proposer_types = ( - VanillaWithAttProposer, - VanillaNoAttProposer, - EagleWithAttProposer, - EagleNoAttProposer, - DFlashProposer, - DSparkProposer, - ) - dp_overlap_proposer_types = ( - DpOverlapVanillaWithAttProposer, - DpOverlapVanillaNoAttProposer, - DpOverlapEagleWithAttProposer, - DpOverlapEagleNoAttProposer, - ) - - for proposer_type in proposer_types: - assert proposer_type.__bases__ == (BaseSpecProposer,) - assert Eagle3Proposer.__bases__ == (EagleWithAttProposer,) - assert DFlash2Proposer.__bases__ == (DFlashProposer,) - for proposer_type in dp_overlap_proposer_types: - assert proposer_type.__bases__ == (BaseDpOverlapProposer,) - assert DpOverlapEagle3Proposer.__bases__ == (DpOverlapEagleWithAttProposer,) - - -def test_each_mtp_mode_builds_its_own_proposer(): - backend = SimpleNamespace() - proposer_types = { - "vanilla_with_att": VanillaWithAttProposer, - "vanilla_no_att": VanillaNoAttProposer, - "eagle_with_att": EagleWithAttProposer, - "eagle_no_att": EagleNoAttProposer, - "eagle3": Eagle3Proposer, - "dflash": DFlashProposer, - "dflash2": DFlash2Proposer, - "dspark": DSparkProposer, - } - - for spec_mode, proposer_type in proposer_types.items(): - proposer = build_spec_proposer(spec_mode=spec_mode, backend=backend, enable_dynmaic_mtp=False) - assert type(proposer) is proposer_type - - -def test_each_supported_dp_overlap_mtp_mode_builds_its_own_proposer(): - backend = SimpleNamespace() - proposer_types = { - "vanilla_with_att": DpOverlapVanillaWithAttProposer, - "vanilla_no_att": DpOverlapVanillaNoAttProposer, - "eagle_with_att": DpOverlapEagleWithAttProposer, - "eagle_no_att": DpOverlapEagleNoAttProposer, - "eagle3": DpOverlapEagle3Proposer, - } - - for spec_mode, proposer_type in proposer_types.items(): - proposer = build_dp_overlap_spec_proposer(spec_mode=spec_mode, backend=backend, enable_dynmaic_mtp=False) - assert type(proposer) is proposer_type - assert isinstance(proposer, BaseDpOverlapProposer) - - -def test_eagle_proposer_skips_draft_forward_for_zero_steps(): - proposer = EagleNoAttProposer( - backend=SimpleNamespace(draft_models=[]), - enable_dynmaic_mtp=True, - ) - next_token_ids = torch.tensor([10, 11], dtype=torch.int64) - - proposal = proposer.propose_next( - target_model_input=None, - target_model_output=None, - target_next_token_ids=next_token_ids, - b_req_mtp_start_loc=torch.tensor([0, 1], dtype=torch.int32), - draft_step=0, - ) - - assert isinstance(proposal, EagleSpecProposal) - assert proposal.token_ids.shape == (2, 0) - assert proposal.schedule_scores.shape == (2, 0) - - -def test_records_request_mtp_metrics_in_one_pass(): - class MetricReq: - def __init__(self, req_idx: int, mtp_step: int): - self.req_idx = req_idx - self.mtp_step = mtp_step - self.accepted = 0 - self.verified = 0 - self.verify_steps = 0 - - def update_mtp_accepted_token_num(self, accept_token_num: int): - self.accepted += accept_token_num - - def update_mtp_verify_token_num(self, verify_token_num: int): - self.verified += verify_token_num - - def update_mtp_verify_step_num(self, verify_step_num: int): - self.verify_steps += verify_step_num - - backend = SimpleNamespace(is_master_in_dp=True) - req0 = MetricReq(req_idx=10, mtp_step=3) - req1 = MetricReq(req_idx=11, mtp_step=3) - - mtp_utils.record_request_mtp_metrics( - backend=backend, - decode_reqs=[req0, req1], - accept_lengths_cpu=torch.tensor([2, 1], dtype=torch.int32), - verify_run_reqs=[req0, req0, req1], - ) - - assert (req0.accepted, req0.verified, req0.verify_steps) == (1, 2, 1) - assert (req1.accepted, req1.verified, req1.verify_steps) == (0, 1, 1) - - fixed_req = MetricReq(req_idx=12, mtp_step=3) - mtp_utils.record_request_mtp_metrics( - backend=backend, - decode_reqs=[fixed_req], - accept_lengths_cpu=torch.tensor([3], dtype=torch.int32), - verify_run_reqs=[fixed_req] * 4, - ) - assert (fixed_req.accepted, fixed_req.verified, fixed_req.verify_steps) == (2, 4, 1) - - -def test_lightspec_stays_full_width_until_costs_are_profiled(): - plan = build_lightspec_planner().plan( - decode_reqs=build_decode_reqs(2), - origin_batch_size=8, - ) - - assert plan.origin_batch_size == 8 - assert plan.dynamic_batch_size == 8 - assert plan.draft_step == plan.pre_draft_step == 3 - - -def test_lightspec_collects_full_width_progress_before_adapting(): - planner = build_lightspec_planner(spec_mode="eagle_with_att") - for batch_size in (2, 4, 8): - planner.target_infer_costs.update(batch_size=batch_size, infer_cost_ms=float(batch_size)) - planner.draft_infer_costs.update(batch_size=batch_size, infer_cost_ms=float(batch_size)) - - plan = planner.plan(decode_reqs=build_decode_reqs(2), origin_batch_size=8) - - assert plan.dynamic_batch_size == 8 - assert plan.draft_step == plan.pre_draft_step == 3 - - -def test_lightspec_selects_eagle_draft_depth_and_verify_capacity(): - planner = build_lightspec_planner(spec_mode="eagle_with_att") - for batch_size, target_cost in ((2, 1.0), (4, 1.1), (8, 10.0)): - planner.target_infer_costs.update(batch_size=batch_size, infer_cost_ms=target_cost) - for batch_size, draft_cost in ((2, 0.1), (4, 0.1), (8, 0.2)): - planner.draft_infer_costs.update(batch_size=batch_size, infer_cost_ms=draft_cost) - for _ in range(8): - planner._update_verified_batch( - accept_lengths=[4, 4], - req_num=2, - dynamic_batch_size=8, - verified_draft_step=3, - ) - - plan = planner.plan(decode_reqs=build_decode_reqs(2), origin_batch_size=8) - - assert plan.dynamic_batch_size == 4 - assert plan.pre_draft_step == 3 - assert plan.draft_step == 1 - - -def test_lightspec_overlap_prices_combined_batch_size(): - planner = build_lightspec_planner( - spec_mode="eagle_with_att", - enable_decode_microbatch_overlap=True, - ) - planner.target_infer_costs.update(batch_size=4, infer_cost_ms=1.0) - planner.target_infer_costs.update(batch_size=8, infer_cost_ms=3.0) - planner.draft_infer_costs.update(batch_size=4, infer_cost_ms=0.25) - planner.draft_infer_costs.update(batch_size=8, infer_cost_ms=1.0) - - assert planner.target_infer_costs.estimate(7) == 3.0 - assert planner.target_infer_costs.get_batch_size_keys_between(4, 8) == [4, 8] - assert planner._get_cost_ms(req_num=4, dynamic_batch_size=8, draft_step=1) == 0.5 - - -def test_lightspec_overlap_selects_dynamic_draft_step(): - planner = build_lightspec_planner( - spec_mode="eagle_with_att", - enable_decode_microbatch_overlap=True, - ) - for batch_size, target_cost in ((2, 1.0), (4, 1.1), (8, 10.0)): - planner.target_infer_costs.update(batch_size=batch_size, infer_cost_ms=target_cost) - for batch_size, draft_cost in ((2, 0.1), (4, 0.1), (8, 0.2)): - planner.draft_infer_costs.update(batch_size=batch_size, infer_cost_ms=draft_cost) - planner._update_verified_batch( - accept_lengths=[4, 4], - req_num=2, - dynamic_batch_size=8, - verified_draft_step=3, - ) - - plan = planner.plan(decode_reqs=build_decode_reqs(2), origin_batch_size=8) - - assert plan.dynamic_batch_size == 4 - assert plan.pre_draft_step == 3 - assert plan.draft_step == 1 - - -def test_lightspec_keeps_vanilla_attention_chained_depth_fixed(): - planner = build_lightspec_planner(spec_mode="vanilla_with_att") - for batch_size, target_cost in ((2, 1.0), (4, 1.1), (8, 10.0)): - planner.target_infer_costs.update(batch_size=batch_size, infer_cost_ms=target_cost) - for batch_size, draft_cost in ((2, 0.1), (4, 0.1), (8, 0.2)): - planner.draft_infer_costs.update(batch_size=batch_size, infer_cost_ms=draft_cost) - planner._update_verified_batch( - accept_lengths=[4, 4], - req_num=2, - dynamic_batch_size=8, - verified_draft_step=3, - ) - - plan = planner.plan(decode_reqs=build_decode_reqs(2), origin_batch_size=8) - - assert planner.draft_steps == (3,) - assert plan.draft_step == plan.pre_draft_step == 3 - - -@pytest.mark.parametrize("spec_mode,block_size", [("dflash", 7), ("dflash2", 8)]) -def test_lightspec_compacts_block_verify_without_changing_draft_shape(spec_mode, block_size): - planner = build_lightspec_planner( - max_draft_step=7, - spec_mode=spec_mode, - block_size=block_size, - ) - for batch_size, target_cost in ((2, 1.0), (4, 1.1), (8, 3.0), (16, 8.0)): - planner.target_infer_costs.update(batch_size=batch_size, infer_cost_ms=target_cost) - planner.draft_infer_costs.update(batch_size=14, infer_cost_ms=0.5) - for _ in range(8): - planner._update_verified_batch( - accept_lengths=[3, 3], - req_num=2, - dynamic_batch_size=16, - verified_draft_step=7, - ) - - plan = planner.plan(decode_reqs=build_decode_reqs(2), origin_batch_size=16) - - assert plan.dynamic_batch_size < 16 - assert plan.draft_step == plan.pre_draft_step == 7 - - -def test_lightspec_bounds_verify_to_existing_proposals(): - planner = build_lightspec_planner() - - cold_start = planner.plan(decode_reqs=build_decode_reqs(2, 0), origin_batch_size=8) - mixed_batch = planner.plan(decode_reqs=build_decode_reqs(2, 1), origin_batch_size=8) - ready_batch = planner.plan(decode_reqs=build_decode_reqs(2, 2), origin_batch_size=8) - - assert cold_start.dynamic_batch_size == 2 - assert mixed_batch.dynamic_batch_size == 5 - assert ready_batch.dynamic_batch_size == 8 - assert not cold_start.all_reqs_have_proposals - assert not mixed_batch.all_reqs_have_proposals - assert ready_batch.all_reqs_have_proposals - - -def test_engine_lets_planner_count_requests_with_a_previous_proposal(): - engine = SpecEngine.__new__(SpecEngine) - engine.planner = build_lightspec_planner() - model_input = SimpleNamespace(batch_size=8) - - mixed_plan = engine.plan_decode( - model_input=model_input, - decode_reqs=[ - SimpleNamespace(cur_output_len=1), - SimpleNamespace(cur_output_len=2), - ], - ) - ready_plan = engine.plan_decode( - model_input=model_input, - decode_reqs=[ - SimpleNamespace(cur_output_len=2), - SimpleNamespace(cur_output_len=2), - ], - ) - - assert mixed_plan.dynamic_batch_size == 5 - assert not mixed_plan.all_reqs_have_proposals - assert ready_plan.dynamic_batch_size == 8 - assert ready_plan.all_reqs_have_proposals - - -def test_dynamic_prepare_delegates_row_compaction(monkeypatch): - from lightllm.common.basemodel.triton_kernel import dynamic_mtp_utils - from lightllm.server.router.model_infer.mtp_speculative import ( - engine as engine_module, - ) - - selected_row_mask = torch.tensor([1, 0, 1, 0], dtype=torch.int32) - monkeypatch.setattr( - dynamic_mtp_utils, - "prepare_dynamic_mtp_model_input", - lambda model_input, **kwargs: (model_input, selected_row_mask), - ) - async_mask = object() - monkeypatch.setattr( - engine_module.g_pin_mem_manager, - "async_copy_from_gpu_tensor_with_event", - lambda **kwargs: async_mask, - ) - - engine = SpecEngine.__new__(SpecEngine) - engine.backend = SimpleNamespace( - model=SimpleNamespace( - req_manager=SimpleNamespace( - req_sampling_params_manager=SimpleNamespace(req_to_next_token_scores=torch.empty(0)) - ) - ) - ) - model_input = SimpleNamespace( - batch_size=4, - ) - plan = SpecDecodePlan(origin_batch_size=4, dynamic_batch_size=2, draft_step=1, pre_draft_step=1) - - compacted_input, selected_mask_cpu = engine.prepare_decode_model_input( - model_input=model_input, - req_num=1, - plan=plan, - ) - - assert compacted_input is model_input - assert selected_mask_cpu is async_mask - - -def test_lightspec_eagle_draft_always_keeps_the_extend_candidate(): - planner = build_lightspec_planner(spec_mode="eagle_with_att") - for batch_size in (2, 4, 8): - planner.target_infer_costs.update(batch_size=batch_size, infer_cost_ms=float(batch_size)) - planner.draft_infer_costs.update(batch_size=batch_size, infer_cost_ms=float(batch_size)) - planner._update_verified_batch( - accept_lengths=[2, 2], - req_num=2, - dynamic_batch_size=8, - verified_draft_step=3, - ) - - plan = planner.plan(decode_reqs=build_decode_reqs(2), origin_batch_size=8) - - assert planner.draft_steps == (1, 2, 3) - assert plan.draft_step >= 1 - - -def test_vanilla_with_attention_planner_prices_extend_then_normal_batches(): - planner = build_lightspec_planner(spec_mode="vanilla_with_att") - planner.draft_infer_costs.update(batch_size=2, infer_cost_ms=0.1) - - draft_cost_ms = planner._get_draft_cost_ms( - req_num=4, - verify_batch_size=8, - draft_step=3, - ) - - assert np.isclose(draft_cost_ms, 0.8) - with pytest.raises(AssertionError, match="requires draft_step to be greater than 0"): - planner._get_draft_cost_ms(req_num=4, verify_batch_size=8, draft_step=0) - - -def test_no_attention_planner_prices_only_normal_request_batches(): - for spec_mode in ("vanilla_no_att", "eagle_no_att"): - planner = build_lightspec_planner(spec_mode=spec_mode) - planner.draft_infer_costs.update(batch_size=2, infer_cost_ms=0.1) - - draft_cost_ms = planner._get_draft_cost_ms( - req_num=4, - verify_batch_size=8, - draft_step=3, - ) - - assert np.isclose(draft_cost_ms, 0.6) - - -def test_eagle_planner_prices_extend_and_decode_separately(): - planner = build_lightspec_planner(spec_mode="eagle_with_att") - planner.draft_infer_costs.update(batch_size=2, infer_cost_ms=0.25) - - draft_cost_ms = planner._get_draft_cost_ms( - req_num=2, - verify_batch_size=8, - draft_step=3, - ) - - assert draft_cost_ms == 1.5 - - -def test_autoregressive_eagle_planner_prices_extend_and_decode_rows(): - for spec_mode in ("eagle_with_att", "eagle3"): - planner = build_lightspec_planner( - max_draft_step=7, - spec_mode=spec_mode, - ) - for batch_size in (2, 4, 8, 16): - planner.draft_infer_costs.update(batch_size=batch_size, infer_cost_ms=float(batch_size)) - - draft_cost_ms = planner._get_draft_cost_ms( - req_num=8, - verify_batch_size=16, - draft_step=7, - ) - - assert draft_cost_ms == 64.0 - - -@pytest.mark.parametrize("spec_mode,block_size,expected_cost", [("dflash", 7, 1.5), ("dflash2", 8, 1.6)]) -def test_block_planner_prices_commit_and_complete_block(spec_mode, block_size, expected_cost): - planner = build_lightspec_planner( - max_draft_step=7, - spec_mode=spec_mode, - block_size=block_size, - ) - planner.draft_infer_costs.update(batch_size=8, infer_cost_ms=0.4) - - draft_cost_ms = planner._get_draft_cost_ms( - req_num=2, - verify_batch_size=16, - draft_step=7, - ) - - assert np.isclose(draft_cost_ms, expected_cost) - - -def test_dspark_planner_prices_commit_and_complete_block(): - planner = build_dspark_planner(max_draft_step=3, block_size=3) - planner.draft_infer_costs.update(batch_size=8, infer_cost_ms=0.8) - - draft_cost_ms = planner._get_draft_cost_ms( - req_num=4, - verify_batch_size=8, - draft_step=3, - ) - - assert np.isclose(draft_cost_ms, 2.0) - - -def test_lightspec_records_one_batch_observation_per_configuration(): - planner = build_lightspec_planner() - - planner._update_verified_batch( - accept_lengths=[2, 2], - req_num=2, - dynamic_batch_size=4, - verified_draft_step=1, - ) - planner._update_verified_batch( - accept_lengths=[1, 1], - req_num=2, - dynamic_batch_size=4, - verified_draft_step=3, - ) - - assert planner.progress_ema_by_config[(2, 4, 1)].get() == 1.0 - assert planner.progress_ema_by_config[(2, 4, 3)].get() == 0.5 - assert planner.progress_ema_by_config[(2, 4, 1)].get_count() == 1 - assert planner.progress_ema_by_config[(2, 4, 3)].get_count() == 1 - - -def test_lightspec_high_concurrency_does_not_multiply_ema_updates(): - planner = build_lightspec_planner() - - planner._update_verified_batch( - accept_lengths=[1] * 128, - req_num=128, - dynamic_batch_size=128, - verified_draft_step=1, - ) - - assert planner.progress_ema_by_config[(128, 128, 1)].get_count() == 1 - - -def test_lightspec_estimates_unseen_shapes_from_prefix_survival(): - planner = build_lightspec_planner() - planner._update_verified_batch( - accept_lengths=[2, 1], - req_num=2, - dynamic_batch_size=8, - verified_draft_step=3, - ) - - assert planner._estimate_progress(req_num=2, dynamic_batch_size=4, draft_step=3) == 0.75 - assert planner._estimate_progress(req_num=2, dynamic_batch_size=2, draft_step=3) == 1.0 - assert planner._estimate_progress(req_num=2, dynamic_batch_size=4, draft_step=2) == 0.75 - - -def test_lightspec_does_not_transfer_deep_progress_to_short_drafts(): - planner = build_lightspec_planner() - planner._update_verified_batch( - accept_lengths=[4, 2], - req_num=2, - dynamic_batch_size=8, - verified_draft_step=3, - ) - - assert np.isclose(planner._estimate_progress(req_num=2, dynamic_batch_size=6, draft_step=2), 5 / 6) - - -def test_lightspec_short_current_proposal_can_recover_to_a_deeper_draft(): - planner = build_lightspec_planner(spec_mode="eagle_with_att") - for batch_size, target_cost in ((2, 1.0), (4, 1.1), (6, 1.2), (8, 1.3)): - planner.target_infer_costs.update(batch_size=batch_size, infer_cost_ms=target_cost) - planner.draft_infer_costs.update(batch_size=batch_size, infer_cost_ms=0.01) - planner._update_verified_batch( - accept_lengths=[4, 4], - req_num=2, - dynamic_batch_size=8, - verified_draft_step=3, - ) - planner.pre_draft_step = 1 - - plan = planner.plan(decode_reqs=build_decode_reqs(2), origin_batch_size=8) - - assert plan.dynamic_batch_size <= 4 - assert plan.pre_draft_step == 1 - assert plan.draft_step == 3 - - -def test_engine_skips_feedback_for_a_mixed_proposal_batch(): - engine = SpecEngine.__new__(SpecEngine) - engine.planner = build_lightspec_planner() - plan = SpecDecodePlan( - origin_batch_size=8, - dynamic_batch_size=5, - draft_step=3, - pre_draft_step=3, - all_reqs_have_proposals=False, - ) - - engine.update_planner_statics( - plan=plan, - proposal=SpecProposal( - token_ids=torch.empty((0,), dtype=torch.int64), - ), - req_num=2, - accept_lengths_cpu=torch.tensor([1, 2], dtype=torch.int32), - ) - - assert not engine.planner.progress_ema_by_config - - -def test_dspark_applies_confidence_capacity_after_two_step_delay(): - planner = build_dspark_planner() - for batch_size, target_cost in ((2, 1.0), (4, 1.1), (8, 10.0)): - planner.target_infer_costs.update(batch_size=batch_size, infer_cost_ms=target_cost) - planner.draft_infer_costs.update(batch_size=6, infer_cost_ms=0.5) - confidence_probs = np.asarray([[0.9, 0.9, 0.9]] * 2, dtype=np.float64) - plan = SpecDecodePlan(origin_batch_size=8, dynamic_batch_size=8, draft_step=3, pre_draft_step=3) - proposal = DSparkSpecProposal( - token_ids=torch.empty((0,), dtype=torch.int64), - schedule_scores_cpu=torch.from_numpy(confidence_probs), - ) - engine = SpecEngine.__new__(SpecEngine) - engine.planner = planner - - engine.update_planner_statics( - plan=plan, - proposal=proposal, - req_num=2, - accept_lengths_cpu=torch.tensor([1, 1], dtype=torch.int32), - ) - first_plan = planner.plan(decode_reqs=build_decode_reqs(2), origin_batch_size=8) - engine.update_planner_statics( - plan=plan, - proposal=proposal, - req_num=2, - accept_lengths_cpu=torch.tensor([1, 1], dtype=torch.int32), - ) - second_plan = planner.plan(decode_reqs=build_decode_reqs(2), origin_batch_size=8) - - assert first_plan.dynamic_batch_size == 8 - assert second_plan.dynamic_batch_size == 4 - assert second_plan.draft_step == second_plan.pre_draft_step == 3 - - -def test_dspark_uses_confidence_scores_without_acceptance_ema(): - planner = build_dspark_planner() - for batch_size, target_cost in ((2, 1.0), (4, 1.2), (8, 10.0)): - planner.target_infer_costs.update(batch_size=batch_size, infer_cost_ms=target_cost) - planner.draft_infer_costs.update(batch_size=6, infer_cost_ms=0.5) - - low_confidence = np.cumprod(np.full((2, 3), 0.1), axis=1) - high_confidence = np.cumprod(np.full((2, 3), 0.9), axis=1) - - assert planner._select_dynamic_batch_size_from_survival_scores(2, low_confidence) == 2 - assert planner._select_dynamic_batch_size_from_survival_scores(2, high_confidence) == 4 - - -def test_topk_prefix_sums_only_computes_requested_counts(): - result = DSparkPlanner._topk_prefix_sums( - values=np.asarray([0.1, 0.9, 0.4, 0.7]), - counts=[0, 2, 4], - ) - - assert set(result) == {0, 2, 4} - np.testing.assert_allclose([result[0], result[2], result[4]], [0.0, 1.6, 2.1]) diff --git a/unit_tests/server/router/model_infer/mtp_speculative/test_vanilla_no_att.py b/unit_tests/server/router/model_infer/mtp_speculative/test_vanilla_no_att.py deleted file mode 100644 index bedf17a5b0..0000000000 --- a/unit_tests/server/router/model_infer/mtp_speculative/test_vanilla_no_att.py +++ /dev/null @@ -1,211 +0,0 @@ -from types import SimpleNamespace - -import pytest -import torch - -from lightllm.common.basemodel.triton_kernel.select_mtp_rows import select_accepted_tail_rows -from lightllm.server.router.model_infer.mtp_speculative.dp_overlap_proposers.vanilla_no_att import ( - DpOverlapVanillaNoAttProposer, -) -from lightllm.server.router.model_infer.mtp_speculative.proposers.vanilla_no_att import VanillaNoAttProposer - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") -def test_select_accepted_tail_rows_triton_matches_index_select(): - device = "cuda" - b_req_mtp_start_loc = torch.tensor([0, 3, 6], dtype=torch.int32, device=device) - accept_len = torch.tensor([2, 3, 2], dtype=torch.int32, device=device) - expected_rows = torch.tensor([1, 5, 7], dtype=torch.int64, device=device) - - # Exercise non-contiguous source strides as well as hidden widths spanning - # multiple Triton column blocks. - input_ids = torch.arange(16, dtype=torch.int64, device=device)[::2] - hidden = torch.arange(8 * 514, dtype=torch.float32, device=device).reshape(8, 514)[:, ::2] - b_req_idx = (torch.arange(16, dtype=torch.int32, device=device) + 10)[::2] - b_mtp_index = torch.arange(8, dtype=torch.int32, device=device) - b_seq_len = (torch.arange(16, dtype=torch.int32, device=device) + 20)[::2] - b_shared_seq_len = (torch.arange(16, dtype=torch.int32, device=device) + 30)[::2] - b_shared_radix_node_id = (torch.arange(16, dtype=torch.int64, device=device) + 1000)[::2] - b_position_delta = (torch.arange(16, dtype=torch.int32, device=device) - 8)[::2] - - selected = select_accepted_tail_rows( - b_req_mtp_start_loc=b_req_mtp_start_loc, - accept_len=accept_len, - input_ids=input_ids, - hidden=hidden, - b_req_idx=b_req_idx, - b_mtp_index=b_mtp_index, - b_seq_len=b_seq_len, - b_shared_seq_len=b_shared_seq_len, - b_shared_radix_node_id=b_shared_radix_node_id, - b_position_delta=b_position_delta, - ) - - torch.testing.assert_close(selected.input_ids, input_ids.index_select(0, expected_rows)) - torch.testing.assert_close(selected.hidden, hidden.index_select(0, expected_rows)) - torch.testing.assert_close(selected.b_req_idx, b_req_idx.index_select(0, expected_rows)) - torch.testing.assert_close(selected.b_mtp_index, b_mtp_index.index_select(0, expected_rows)) - torch.testing.assert_close(selected.b_seq_len, b_seq_len.index_select(0, expected_rows)) - torch.testing.assert_close(selected.b_shared_seq_len, b_shared_seq_len.index_select(0, expected_rows)) - torch.testing.assert_close( - selected.b_shared_radix_node_id, - b_shared_radix_node_id.index_select(0, expected_rows), - ) - torch.testing.assert_close(selected.b_position_delta, b_position_delta.index_select(0, expected_rows)) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") -def test_vanilla_no_att_runs_all_draft_steps_for_empty_dp_batch(): - device = "cuda" - draft_batch_sizes = [] - - def draft_forward(model_input): - draft_batch_sizes.append(model_input.batch_size) - return SimpleNamespace( - token_ids=torch.empty((0,), dtype=torch.int64, device=device), - mtp_collector=SimpleNamespace( - spec_hidden=torch.empty((0, 2), dtype=torch.float32, device=device), - ), - ) - - backend = SimpleNamespace( - draft_models=[SimpleNamespace(forward=draft_forward) for _ in range(2)], - _gen_argmax_token_ids=lambda output: output.token_ids, - ) - proposer = VanillaNoAttProposer(backend=backend, enable_dynmaic_mtp=False) - empty_i32 = torch.empty((0,), dtype=torch.int32, device=device) - target_model_input = SimpleNamespace( - batch_size=0, - b_req_idx=empty_i32, - b_mtp_index=empty_i32, - b_seq_len=empty_i32, - b_position_delta=empty_i32, - b_shared_seq_len=empty_i32, - b_shared_radix_node_id=torch.empty((0,), dtype=torch.int64, device=device), - multimodal_params=[], - ) - target_model_output = SimpleNamespace( - mtp_collector=SimpleNamespace( - spec_hidden=torch.empty((0, 2), dtype=torch.float32, device=device), - ) - ) - - proposal = proposer.propose_next( - target_model_input=target_model_input, - target_model_output=target_model_output, - target_next_token_ids=torch.empty((0,), dtype=torch.int64, device=device), - b_req_mtp_start_loc=empty_i32, - draft_step=2, - accept_len=empty_i32, - ) - - assert draft_batch_sizes == [0, 0] - assert proposal.token_ids.shape == (0, 2) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") -def test_vanilla_no_att_proposes_from_one_accepted_tail_per_request(): - device = "cuda" - draft_calls = [] - draft_outputs = [ - SimpleNamespace( - token_ids=torch.tensor([21, 24]), - token_probs=torch.tensor([0.8, 0.7]), - mtp_collector=SimpleNamespace(spec_hidden=torch.tensor([[102.0, 103.0], [108.0, 109.0]])), - ), - SimpleNamespace( - token_ids=torch.tensor([31, 34]), - token_probs=torch.tensor([0.6, 0.5]), - mtp_collector=SimpleNamespace(spec_hidden=torch.tensor([[202.0, 203.0], [208.0, 209.0]])), - ), - ] - - def build_draft_model(step): - def forward(model_input): - draft_calls.append( - { - "batch_size": model_input.batch_size, - "input_ids": model_input.input_ids.cpu(), - "draft_hidden": model_input.mtp_draft_input_hiddens.cpu(), - "b_req_idx": model_input.b_req_idx.cpu(), - "b_mtp_index": model_input.b_mtp_index.cpu(), - "b_seq_len": model_input.b_seq_len.cpu(), - } - ) - return draft_outputs[step] - - return SimpleNamespace(forward=forward) - - backend = SimpleNamespace( - draft_models=[build_draft_model(0), build_draft_model(1)], - _gen_argmax_token_ids=lambda output: output.token_ids, - _gen_argmax_token_ids_and_prob=lambda output: (output.token_ids, output.token_probs), - ) - proposer = VanillaNoAttProposer(backend=backend, enable_dynmaic_mtp=True) - target_model_input = SimpleNamespace( - batch_size=5, - input_ids=torch.tensor([10, 11, 12, 13, 14], device=device), - b_req_idx=torch.tensor([7, 7, 7, 9, 9], dtype=torch.int32, device=device), - b_mtp_index=torch.tensor([0, 1, 2, 0, 1], dtype=torch.int32, device=device), - b_seq_len=torch.tensor([10, 11, 12, 20, 21], dtype=torch.int32, device=device), - b_position_delta=torch.tensor([0, 1, 2, 3, 4], dtype=torch.int32, device=device), - b_shared_seq_len=torch.tensor([8, 8, 8, 6, 6], dtype=torch.int32, device=device), - b_shared_radix_node_id=torch.tensor([70, 70, 70, 90, 90], dtype=torch.int64, device=device), - multimodal_params=[{"images": [], "audios": []} for _ in range(5)], - ) - target_model_output = SimpleNamespace( - mtp_collector=SimpleNamespace(spec_hidden=torch.arange(10, dtype=torch.float32, device=device).reshape(5, 2)) - ) - - proposal = proposer.propose_next( - target_model_input=target_model_input, - target_model_output=target_model_output, - target_next_token_ids=target_model_input.input_ids, - b_req_mtp_start_loc=torch.tensor([0, 3], dtype=torch.int32, device=device), - draft_step=2, - accept_len=torch.tensor([2, 2], dtype=torch.int32, device=device), - ) - - torch.testing.assert_close(proposal.token_ids, torch.tensor([[21, 31], [24, 34]])) - torch.testing.assert_close(proposal.schedule_scores, torch.tensor([[0.8, 0.6], [0.7, 0.5]])) - assert len(draft_calls) == 2 - assert draft_calls[0]["batch_size"] == 2 - torch.testing.assert_close(draft_calls[0]["input_ids"], torch.tensor([11, 14])) - torch.testing.assert_close(draft_calls[0]["draft_hidden"], torch.tensor([[2.0, 3.0], [8.0, 9.0]])) - torch.testing.assert_close(draft_calls[0]["b_req_idx"], torch.tensor([7, 9], dtype=torch.int32)) - torch.testing.assert_close(draft_calls[0]["b_mtp_index"], torch.tensor([1, 1], dtype=torch.int32)) - torch.testing.assert_close(draft_calls[0]["b_seq_len"], torch.tensor([11, 21], dtype=torch.int32)) - torch.testing.assert_close(draft_calls[1]["input_ids"], torch.tensor([21, 24])) - torch.testing.assert_close( - draft_calls[1]["draft_hidden"], - torch.tensor([[102.0, 103.0], [108.0, 109.0]]), - ) - assert target_model_input.batch_size == 5 - torch.testing.assert_close(target_model_input.input_ids.cpu(), torch.tensor([10, 11, 12, 13, 14])) - - -def test_vanilla_no_att_skips_draft_forward_for_zero_steps(): - proposer = VanillaNoAttProposer( - backend=SimpleNamespace(draft_models=[]), - enable_dynmaic_mtp=True, - ) - - proposal = proposer.propose_next( - target_model_input=None, - target_model_output=None, - target_next_token_ids=torch.tensor([10, 11]), - b_req_mtp_start_loc=torch.tensor([0, 1], dtype=torch.int32), - draft_step=0, - ) - - assert proposal.token_ids.shape == (2, 0) - assert proposal.schedule_scores.shape == (2, 0) - - -def test_vanilla_no_att_fill_hooks_are_noops(): - backend = SimpleNamespace(draft_models=[]) - proposer = VanillaNoAttProposer(backend=backend, enable_dynmaic_mtp=False) - overlap_proposer = DpOverlapVanillaNoAttProposer(backend=backend, enable_dynmaic_mtp=False) - - proposer.fill_draft_model_kv_state(None, None, None) - overlap_proposer.fill_draft_model_kv_state_overlap(None, None, None, None, None, None) diff --git a/unit_tests/server/router/model_infer/mtp_speculative/test_vanilla_overlap.py b/unit_tests/server/router/model_infer/mtp_speculative/test_vanilla_overlap.py deleted file mode 100644 index 23a6672fd2..0000000000 --- a/unit_tests/server/router/model_infer/mtp_speculative/test_vanilla_overlap.py +++ /dev/null @@ -1,147 +0,0 @@ -from types import SimpleNamespace - -import pytest -import torch - -from lightllm.common.basemodel.batch_objs import ModelMtpOutputCollector, ModelOutput -from lightllm.server.router.model_infer.mtp_speculative.dp_overlap_proposers.vanilla_no_att import ( - DpOverlapVanillaNoAttProposer, -) -from lightllm.server.router.model_infer.mtp_speculative.dp_overlap_proposers.vanilla_with_att import ( - DpOverlapVanillaWithAttProposer, -) - - -class _DraftModel: - def __init__(self): - self.decode_batch_sizes = [] - - def _microbatch_overlap_decode_cuda(self, input0, input1): - self.decode_batch_sizes.append((input0.batch_size, input1.batch_size)) - return tuple( - ModelOutput( - logits=torch.arange( - model_input.batch_size, - dtype=torch.float32, - device=model_input.input_ids.device, - ).view(-1, 1), - mtp_collector=ModelMtpOutputCollector( - spec_hidden=torch.ones( - (model_input.batch_size, 2), - device=model_input.input_ids.device, - ) - ), - ) - for model_input in (input0, input1) - ) - - -def test_dp_vanilla_no_att_supports_zero_dynamic_draft_step(): - proposer = DpOverlapVanillaNoAttProposer( - backend=SimpleNamespace(draft_models=[]), - enable_dynmaic_mtp=True, - ) - target_next_token_ids0 = torch.tensor([10], dtype=torch.int64) - target_next_token_ids1 = torch.tensor([11], dtype=torch.int64) - - proposal = proposer.propose_next_overlap( - target_model_input0=None, - target_model_output0=None, - target_next_token_ids0=target_next_token_ids0, - accept_len0=torch.ones((1,), dtype=torch.int32), - target_model_input1=None, - target_model_output1=None, - target_next_token_ids1=target_next_token_ids1, - accept_len1=torch.ones((1,), dtype=torch.int32), - draft_step=0, - ) - - assert proposal.token_ids.shape == (2, 0) - assert proposal.schedule_scores.shape == (2, 0) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") -def test_dp_vanilla_proposer_owns_overlap_decode(): - device = "cuda" - draft_models = [_DraftModel(), _DraftModel()] - backend = SimpleNamespace( - max_draft_step=2, - draft_models=draft_models, - _gen_argmax_token_ids=lambda output: output.logits[:, 0].to(torch.int64), - ) - proposer = DpOverlapVanillaWithAttProposer(backend=backend, enable_dynmaic_mtp=False) - b_mtp_index = torch.arange(3, dtype=torch.int32, device=device) - model_input0 = SimpleNamespace( - batch_size=3, - b_req_idx=torch.arange(3, dtype=torch.int32, device=device), - b_mtp_index=b_mtp_index, - ) - model_input1 = SimpleNamespace( - batch_size=3, - b_req_idx=torch.arange(3, dtype=torch.int32, device=device), - b_mtp_index=b_mtp_index, - ) - - proposal = proposer.propose_next_overlap( - target_model_input0=model_input0, - target_model_output0=ModelOutput( - logits=torch.empty((6, 1), device=device), - mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.ones((6, 2), device=device)), - ), - target_next_token_ids0=torch.tensor([10, 11, 0], dtype=torch.int64, device=device), - accept_len0=torch.tensor([2], dtype=torch.int32, device=device), - target_model_input1=model_input1, - target_model_output1=ModelOutput( - logits=torch.empty((6, 1), device=device), - mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.ones((6, 2), device=device)), - ), - target_next_token_ids1=torch.tensor([20, 21, 22], dtype=torch.int64, device=device), - accept_len1=torch.tensor([1], dtype=torch.int32, device=device), - draft_step=2, - ) - - assert proposal.token_ids.tolist() == [ - [1, 1], - [0, 0], - ] - assert draft_models[0].decode_batch_sizes == [(3, 3)] - assert draft_models[1].decode_batch_sizes == [(3, 3)] - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") -def test_dp_vanilla_proposer_builds_padded_inputs_for_empty_verify_rows(): - device = "cuda" - draft_models = [_DraftModel(), _DraftModel()] - backend = SimpleNamespace( - max_draft_step=2, - draft_models=draft_models, - _gen_argmax_token_ids=lambda output: output.logits[:, 0].to(torch.int64), - ) - proposer = DpOverlapVanillaWithAttProposer(backend=backend, enable_dynmaic_mtp=False) - empty_i32 = torch.empty((0,), dtype=torch.int32, device=device) - model_input0 = SimpleNamespace(batch_size=0, b_req_idx=empty_i32, b_mtp_index=empty_i32) - model_input1 = SimpleNamespace(batch_size=0, b_req_idx=empty_i32, b_mtp_index=empty_i32) - model_output0 = ModelOutput( - logits=torch.empty((0, 1), device=device), - mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.ones((0, 2), device=device)), - ) - model_output1 = ModelOutput( - logits=torch.empty((0, 1), device=device), - mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.ones((0, 2), device=device)), - ) - - proposal = proposer.propose_next_overlap( - target_model_input0=model_input0, - target_model_output0=model_output0, - target_next_token_ids0=torch.empty((0,), dtype=torch.int64, device=device), - accept_len0=torch.empty((0,), dtype=torch.int32, device=device), - target_model_input1=model_input1, - target_model_output1=model_output1, - target_next_token_ids1=torch.empty((0,), dtype=torch.int64, device=device), - accept_len1=torch.empty((0,), dtype=torch.int32, device=device), - draft_step=2, - ) - - assert proposal.token_ids.shape == (0, 2) - assert draft_models[0].decode_batch_sizes == [(0, 0)] - assert draft_models[1].decode_batch_sizes == [(0, 0)] diff --git a/unit_tests/server/router/model_infer/mtp_speculative/test_vanilla_prefill.py b/unit_tests/server/router/model_infer/mtp_speculative/test_vanilla_prefill.py deleted file mode 100644 index d22db894f7..0000000000 --- a/unit_tests/server/router/model_infer/mtp_speculative/test_vanilla_prefill.py +++ /dev/null @@ -1,237 +0,0 @@ -from types import SimpleNamespace - -import pytest -import torch - -from lightllm.server.router.model_infer.mtp_speculative.dp_overlap_proposers.vanilla_with_att import ( - DpOverlapVanillaWithAttProposer, -) -from lightllm.server.router.model_infer.mtp_speculative.proposers.vanilla_with_att import VanillaWithAttProposer - - -def _prefill_input(input_ids, batch_size): - return SimpleNamespace( - is_prefill=True, - b_position_delta=None, - batch_size=batch_size, - input_ids=input_ids, - b_req_idx=torch.arange(batch_size, dtype=torch.int32, device=input_ids.device), - b_seq_len=torch.full( - (batch_size,), input_ids.shape[0] // batch_size, dtype=torch.int32, device=input_ids.device - ), - b_ready_cache_len=torch.zeros(batch_size, dtype=torch.int32, device=input_ids.device), - mtp_draft_input_hiddens=None, - ) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") -def test_chained_prefill_advances_local_input_without_mutating_target(): - device = "cuda" - original_input_ids = torch.tensor([10, 11, 12, 20, 21, 22], dtype=torch.int64, device=device) - target_input = _prefill_input(original_input_ids, batch_size=2) - target_hidden = torch.arange(12, dtype=torch.float32, device=device).reshape(6, 2) - forwarded = [] - - def draft_model(output_tokens, output_hidden): - def forward(model_input): - forwarded.append( - ( - model_input, - model_input.input_ids.clone(), - model_input.mtp_draft_input_hiddens, - model_input.b_is_decode_req, - ) - ) - return SimpleNamespace( - token_ids=output_tokens, - mtp_collector=SimpleNamespace(spec_hidden=output_hidden), - ) - - return SimpleNamespace(forward=forward) - - stage0_hidden = target_hidden + 100 - stage1_hidden = target_hidden + 200 - backend = SimpleNamespace( - draft_models=[ - draft_model(torch.tensor([30, 40], dtype=torch.int64, device=device), stage0_hidden), - draft_model(torch.tensor([31, 41], dtype=torch.int64, device=device), stage1_hidden), - ], - _gen_argmax_token_ids=lambda output: output.token_ids, - ) - proposer = VanillaWithAttProposer(backend=backend, enable_dynmaic_mtp=False) - - proposer.fill_draft_model_kv_state( - target_model_input=target_input, - target_model_output=SimpleNamespace(mtp_collector=SimpleNamespace(spec_hidden=target_hidden)), - target_next_token_ids=torch.tensor([13, 23], dtype=torch.int64, device=device), - ) - - assert forwarded[0][0] is forwarded[1][0] - assert forwarded[0][0] is not target_input - torch.testing.assert_close( - forwarded[0][1], - torch.tensor([11, 12, 13, 21, 22, 23], dtype=torch.int64, device=device), - ) - torch.testing.assert_close( - forwarded[1][1], - torch.tensor([12, 13, 30, 22, 23, 40], dtype=torch.int64, device=device), - ) - assert forwarded[0][2] is target_hidden - assert forwarded[1][2] is stage0_hidden - assert not forwarded[0][3].any() - assert forwarded[0][3].data_ptr() == forwarded[1][3].data_ptr() - assert target_input.input_ids is original_input_ids - assert target_input.mtp_draft_input_hiddens is None - - -def test_overlap_chained_prefill_uses_local_microbatch_inputs(monkeypatch): - def prepare(model_input, b_next_token_ids, mtp_draft_input_hiddens): - model_input.input_ids = model_input.input_ids + b_next_token_ids - model_input.mtp_draft_input_hiddens = mtp_draft_input_hiddens - return model_input - - target_input0 = _prefill_input(torch.tensor([1, 2], dtype=torch.int64), batch_size=2) - target_input1 = _prefill_input(torch.tensor([3, 4], dtype=torch.int64), batch_size=2) - target_hidden0 = torch.tensor([[1.0], [2.0]]) - target_hidden1 = torch.tensor([[3.0], [4.0]]) - forwarded = [] - - class DraftModel: - def __init__(self, token_offset): - self.token_offset = token_offset - - def _microbatch_overlap_prefill_cuda(self, input0, input1): - forwarded.append((input0, input1, input0.input_ids.clone(), input1.input_ids.clone())) - return tuple( - SimpleNamespace( - token_ids=torch.full((2,), self.token_offset + index, dtype=torch.int64), - mtp_collector=SimpleNamespace(spec_hidden=hidden + self.token_offset), - ) - for index, hidden in enumerate((target_hidden0, target_hidden1)) - ) - - backend = SimpleNamespace( - draft_models=[DraftModel(10), DraftModel(20)], - _gen_argmax_token_ids=lambda output: output.token_ids, - ) - proposer = DpOverlapVanillaWithAttProposer(backend=backend, enable_dynmaic_mtp=False) - monkeypatch.setattr(proposer, "_prepare_mtp_prefill_inputs", prepare) - - proposer.fill_draft_model_kv_state_overlap( - target_model_input0=target_input0, - target_model_output0=SimpleNamespace(mtp_collector=SimpleNamespace(spec_hidden=target_hidden0)), - target_next_token_ids0=torch.tensor([5, 6], dtype=torch.int64), - target_model_input1=target_input1, - target_model_output1=SimpleNamespace(mtp_collector=SimpleNamespace(spec_hidden=target_hidden1)), - target_next_token_ids1=torch.tensor([7, 8], dtype=torch.int64), - ) - - assert forwarded[0][0] is forwarded[1][0] - assert forwarded[0][1] is forwarded[1][1] - assert forwarded[0][0] is not target_input0 - assert forwarded[0][1] is not target_input1 - torch.testing.assert_close(forwarded[0][2], torch.tensor([6, 8], dtype=torch.int64)) - torch.testing.assert_close(forwarded[0][3], torch.tensor([10, 12], dtype=torch.int64)) - torch.testing.assert_close(forwarded[1][2], torch.tensor([16, 18], dtype=torch.int64)) - torch.testing.assert_close(forwarded[1][3], torch.tensor([21, 23], dtype=torch.int64)) - torch.testing.assert_close(target_input0.input_ids, torch.tensor([1, 2], dtype=torch.int64)) - torch.testing.assert_close(target_input1.input_ids, torch.tensor([3, 4], dtype=torch.int64)) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") -def test_chained_decode_overlays_verified_tokens_for_the_next_level(): - device = "cuda" - original_input_ids = torch.tensor([10, 11, 12, 20, 21, 22], dtype=torch.int64, device=device) - target_input = SimpleNamespace( - is_prefill=False, - batch_size=6, - input_ids=original_input_ids, - mtp_draft_input_hiddens=None, - ) - target_hidden = torch.arange(12, dtype=torch.float32, device=device).reshape(6, 2) - forwarded = [] - - def draft_model(output_tokens, output_probs, output_hidden): - def forward(model_input): - forwarded.append( - ( - model_input, - model_input.input_ids.clone(), - model_input.mtp_draft_input_hiddens, - ) - ) - return SimpleNamespace( - token_ids=output_tokens, - token_probs=output_probs, - mtp_collector=SimpleNamespace(spec_hidden=output_hidden), - ) - - return SimpleNamespace(forward=forward) - - stage0_tokens = torch.tensor([30, 31, 32, 33, 34, 35], dtype=torch.int64, device=device) - stage1_tokens = torch.tensor([40, 41, 42, 43, 44, 45], dtype=torch.int64, device=device) - stage2_tokens = torch.tensor([50, 51, 52, 53, 54, 55], dtype=torch.int64, device=device) - stage0_probs = torch.tensor([0.30, 0.31, 0.32, 0.33, 0.34, 0.35], device=device) - stage1_probs = torch.tensor([0.40, 0.41, 0.42, 0.43, 0.44, 0.45], device=device) - stage2_probs = torch.tensor([0.50, 0.51, 0.52, 0.53, 0.54, 0.55], device=device) - stage0_hidden = target_hidden + 100 - stage1_hidden = target_hidden + 200 - stage2_hidden = target_hidden + 300 - backend = SimpleNamespace( - max_draft_step=3, - draft_models=[ - draft_model(stage0_tokens, stage0_probs, stage0_hidden), - draft_model(stage1_tokens, stage1_probs, stage1_hidden), - draft_model(stage2_tokens, stage2_probs, stage2_hidden), - ], - _gen_argmax_token_ids_and_prob=lambda output: (output.token_ids, output.token_probs), - ) - proposer = VanillaWithAttProposer(backend=backend, enable_dynmaic_mtp=True) - - proposal = proposer.propose_next( - target_model_input=target_input, - target_model_output=SimpleNamespace(mtp_collector=SimpleNamespace(spec_hidden=target_hidden)), - target_next_token_ids=original_input_ids, - b_req_mtp_start_loc=torch.tensor([0, 3], dtype=torch.int32, device=device), - draft_step=3, - accept_len=torch.tensor([3, 2], dtype=torch.int32, device=device), - ) - - assert forwarded[0][0] is forwarded[1][0] is forwarded[2][0] - assert forwarded[0][0] is not target_input - torch.testing.assert_close(forwarded[0][1], original_input_ids) - # 第一次覆盖：[10, 11, 12] -> [11, 12, 32]， - # [20, 21] -> [21, 34]。 - torch.testing.assert_close( - forwarded[1][1], - torch.tensor([11, 12, 32, 21, 34, 35], dtype=torch.int64, device=device), - ) - # 第二次覆盖：[11, 12, 32] -> [12, 32, 42]， - # [21, 34] -> [34, 44]。 - torch.testing.assert_close( - forwarded[2][1], - torch.tensor([12, 32, 42, 34, 44, 45], dtype=torch.int64, device=device), - ) - assert forwarded[0][2] is target_hidden - assert forwarded[1][2] is stage0_hidden - assert forwarded[2][2] is stage1_hidden - torch.testing.assert_close( - proposal.token_ids, - torch.tensor([[32, 42, 52], [34, 44, 54]], dtype=torch.int64, device=device), - ) - torch.testing.assert_close( - proposal.schedule_scores, - torch.tensor([[0.32, 0.42, 0.52], [0.34, 0.44, 0.54]], device=device), - ) - assert target_input.input_ids is original_input_ids - assert target_input.mtp_draft_input_hiddens is None - - with pytest.raises(AssertionError, match="requires the full chained draft depth"): - proposer.propose_next( - target_model_input=target_input, - target_model_output=SimpleNamespace(mtp_collector=SimpleNamespace(spec_hidden=target_hidden)), - target_next_token_ids=original_input_ids, - b_req_mtp_start_loc=torch.tensor([0, 3], dtype=torch.int32, device=device), - draft_step=2, - accept_len=torch.tensor([3, 2], dtype=torch.int32, device=device), - ) diff --git a/unit_tests/server/router/model_infer/structured_output/test_error_output.py b/unit_tests/server/router/model_infer/structured_output/test_error_output.py deleted file mode 100644 index e8285f555a..0000000000 --- a/unit_tests/server/router/model_infer/structured_output/test_error_output.py +++ /dev/null @@ -1,66 +0,0 @@ -from types import MethodType, SimpleNamespace -from unittest.mock import MagicMock - -import numpy as np -import pytest -import torch -import xgrammar as xgr - -from unit_tests.server.grammar_helpers import init_request, make_req -from lightllm.server.router.model_infer.infer_batch import InferReq, InferReqUpdatePack - - -@pytest.mark.parametrize("failure", ["mask", "commit", "rollback"]) -@pytest.mark.parametrize("exception_type", [RuntimeError, BaseException]) -def test_grammar_failure_finishes_after_publishing_inflight_token(compiler, monkeypatch, failure, exception_type): - req = make_req(regular_constraint="abc") - init_request(compiler, req) - shm_req = req.shm_req - shm_req.shm_prompt_ids.arr = np.array([ord("P")] + [0] * 8, dtype=np.int64) - shm_req.shm_logprobs = SimpleNamespace(arr=np.zeros(9, dtype=[("logprob", np.float32), ("rank", np.int32)])) - req.set_next_gen_token_id = MethodType(InferReq.set_next_gen_token_id, req) - req.update_finish_status = lambda **kwargs: None - InferReqUpdatePack(req, 1).handle(ord("a"), -0.1, -1, [256], is_master_in_dp=True) - - if failure == "mask": - monkeypatch.setattr( - xgr.GrammarMatcher, - "fill_next_token_bitmask", - MagicMock(side_effect=exception_type("injected mask failure")), - ) - assert not req.output_constraint.fill_masks(torch.empty((1, 9), dtype=torch.int32), []) - elif failure == "commit": - monkeypatch.setattr( - xgr.GrammarMatcher, "accept_token", MagicMock(side_effect=exception_type("injected commit failure")) - ) - else: - monkeypatch.setattr( - xgr.GrammarMatcher, "rollback", MagicMock(side_effect=exception_type("injected rollback failure")) - ) - assert req.output_constraint.fill_masks(torch.empty((2, 9), dtype=torch.int32), [ord("b")]) - - next_token = ord("b") if failure == "rollback" else ord("!") - InferReqUpdatePack(req, 2).handle(next_token, -0.2, -1, [256], is_master_in_dp=True) - InferReqUpdatePack(req, 3).handle(ord("z"), -0.3, -1, [256], is_master_in_dp=True) - - # In-flight tokens can update KV history, but output publication stops at the error. - assert req.finish_status.is_finished_error() and shm_req.finish_status.is_finished_error() - assert req.output_constraint.is_terminated() - assert shm_req.shm_prompt_ids.arr[1:4].tolist() == [ord("a"), next_token, ord("z")] - assert req.output_constraint.error == f"injected {failure} failure" - assert shm_req.candetoken_out_len == shm_req.shm_cur_output_len == 2 - assert shm_req.finish_token_index == 2 - - -def test_draft_length_failure_disables_mask(compiler): - req = make_req(regular_constraint="ab") - init_request(compiler, req) - state = req.output_constraint - bitmask = torch.empty((1, 9), dtype=torch.int32) - draft_tokens = MagicMock() - draft_tokens.__len__.side_effect = BaseException("injected draft_length failure") - - assert not state.fill_masks(bitmask, draft_tokens) - assert state.error == "injected draft_length failure" - assert state.is_terminated() - assert state.matcher.accept_token(ord("a")) diff --git a/unit_tests/server/router/model_infer/structured_output/test_mtp.py b/unit_tests/server/router/model_infer/structured_output/test_mtp.py deleted file mode 100644 index 823f3d6049..0000000000 --- a/unit_tests/server/router/model_infer/structured_output/test_mtp.py +++ /dev/null @@ -1,444 +0,0 @@ -from contextlib import nullcontext -from types import SimpleNamespace - -import pytest -import torch - -from unit_tests.server.grammar_helpers import ( - allowed, - commit_token, - build_mask, - make_mask_buffers, - apply_masks, - make_req, - init_request, - prepare_sampling_tensors, -) - - -def verify_mask(compiler, reqs, inputs): - mask = build_mask(reqs, draft_input_ids=inputs) - logits = torch.zeros(len(reqs), compiler.vocab_size) - apply_masks(reqs, logits, mask) - return torch.isfinite(logits) - - -def test_verify_masks_follow_each_draft_prefix_and_roll_back(compiler): - req = make_req(regular_constraint="abc") - init_request(compiler, req) - mask = verify_mask(compiler, [req] * 4, [0, ord("a"), ord("b"), ord("c")]) - assert [row.nonzero().flatten().tolist() for row in mask] == [[ord("a")], [ord("b")], [ord("c")], [256]] - assert allowed([req])[0].nonzero().flatten().tolist() == [ord("a")] - - # Only tokens accepted by target verification become permanent history. - commit_token(req, ord("a")) - commit_token(req, ord("b")) - mask = verify_mask(compiler, [req] * 2, [ord("b"), ord("c")]) - assert mask[0, ord("c")] and mask[1, 256] - assert allowed([req])[0, ord("c")] - - -@pytest.mark.parametrize("draft", [b"!bc", b"a!c", b"ab!"]) -def test_invalid_draft_is_rejected_without_failing_the_request(compiler, draft): - req = make_req(regular_constraint="abc") - init_request(compiler, req) - mask = verify_mask(compiler, [req] * 4, [0, *draft]) - mismatch = draft.index(ord("!")) - assert mask[mismatch, b"abc"[mismatch]] - assert not mask[mismatch, ord("!")] - assert mask[mismatch + 1 :].all() # Unreachable suffix rows remain harmless. - assert req.output_constraint.error is None - assert not req.output_constraint.is_terminated() - assert allowed([req])[0, ord("a")] - - -def test_mixed_compacted_requests_have_independent_verify_states(compiler): - first, second, plain = make_req(regular_constraint="ab"), make_req(regular_constraint="cd"), make_req() - for req in (first, second, plain): - init_request(compiler, req) - reqs = [plain, plain, second, second, second, first, first] - mask = verify_mask(compiler, reqs, [0, 0, 0, ord("c"), ord("d"), 0, ord("a")]) - assert mask[:2].all() - assert [row.nonzero().flatten().tolist() for row in mask[2:]] == [ - [ord("c")], - [ord("d")], - [256], - [ord("a")], - [ord("b")], - ] - assert allowed([first])[0, ord("a")] - assert allowed([second])[0, ord("c")] - - -@pytest.mark.parametrize("reasoning_end", [b"", b"]"]) -def test_stop_token_in_draft_does_not_terminate_the_real_matcher(compiler, reasoning_end): - req = make_req(regular_constraint="a", guided_reasoning_end=tuple(reasoning_end)) - init_request(compiler, req) - draft = [*reasoning_end, ord("a"), 256, ord("!")] - mask = verify_mask(compiler, [req] * (len(draft) + 1), [0, *draft]) - assert mask[: len(reasoning_end)].all() - mask = mask[len(reasoning_end) :] - assert mask[0, ord("a")] and mask[1, 256] and mask[2:].all() - assert not req.output_constraint.matcher.is_terminated() - assert req.output_constraint.error is None - for token in reasoning_end: - assert not commit_token(req, token) - assert allowed([req])[0, ord("a")] - - -def test_completed_grammar_does_not_enable_mask(compiler): - req = make_req(regular_constraint="a") - init_request(compiler, req) - state = req.output_constraint - state.commit(ord("a")) - state.commit(256) - assert state.is_terminated() - - bitmask = torch.zeros((1, 9), dtype=torch.int32) - assert not state.fill_masks(bitmask, []) - assert (bitmask == -1).all() - assert state.is_terminated() and state.error is None - - -def test_long_reasoning_tail_restored_after_repeated_speculation(compiler): - req = make_req(regular_constraint="ab", guided_reasoning_end=tuple(b"]>")) - init_request(compiler, req) - state = req.output_constraint - # 长思考内容只保留结束标记长度的后缀，标记仍能跨已提交和 draft token 匹配。 - assert verify_mask(compiler, [req] * 4, [0, *b"xxx"]).all() - for token in b"x" * 100 + b"]": - commit_token(req, token) - assert state.reasoning_tail == tuple(b"x]") - saved_tail = state.reasoning_tail - - for _ in range(3): - mask = verify_mask(compiler, [req] * 4, [0, *b">ab"]) - assert mask[0].all() - assert [row.nonzero().flatten().tolist() for row in mask[1:]] == [[ord("a")], [ord("b")], [256]] - assert state.in_reasoning and state.reasoning_tail[-1] == ord("]") - # 快照恢复被临时推进替换的后缀，多轮验证始终观察同一已提交状态。 - assert state.reasoning_tail == saved_tail - assert state.error is None - - for token in b">ab": - commit_token(req, token) - # 思考结束标记保留在 tail 中，grammar token 的提交进度由 matcher 持有。 - assert state.reasoning_tail == tuple(b"]>") - assert not state.in_reasoning - assert allowed([req])[0].nonzero().flatten().tolist() == [256] - - -def test_reasoning_boundary_inside_verify_is_temporary_until_commit(compiler): - req = make_req(regular_constraint="ab", guided_reasoning_end=tuple(b"]>")) - init_request(compiler, req) - commit_token(req, ord("]")) - mask = verify_mask(compiler, [req] * 4, [ord("]"), ord(">"), ord("a"), ord("b")]) - assert mask[0].all() - assert [row.nonzero().flatten().tolist() for row in mask[1:]] == [[ord("a")], [ord("b")], [256]] - state = req.output_constraint - assert state.in_reasoning and state.reasoning_tail == (ord("]"),) - matcher = state.matcher - assert matcher is not None - # Repeating an uncommitted draft reuses the matcher at the same prefix. - repeated_mask = verify_mask(compiler, [req] * 4, [ord("]"), ord(">"), ord("a"), ord("b")]) - assert torch.equal(repeated_mask, mask) - assert state.matcher is matcher - - # A single MTP post_handle can commit the delimiter and the whole answer. - for token in b">ab": - assert not commit_token(req, token) - assert not state.in_reasoning - assert allowed([req])[0, 256] - assert state.matcher is matcher - assert commit_token(req, 256) - - -@pytest.mark.parametrize( - "marker,committed,draft,expected", - [ - (b"]>", b"", b"]>ab", [None, None, ord("a"), ord("b"), 256]), - (b"]>", b"]", b"x]>a", [None, None, None, ord("a"), ord("b")]), - (b"]>", b"]", b">!b", [None, ord("a"), None, None]), - (b"]>", b"]", b"x]x", [None, None, None, None]), - (b"aba", b"ab", b"baba", [None, None, None, None, ord("a")]), - (b"aba", b"ab", b"aab", [None, ord("a"), ord("b"), 256]), - ], -) -def test_verify_locates_answer_without_mutating_reasoning_progress(compiler, marker, committed, draft, expected): - req = make_req(regular_constraint="ab", guided_reasoning_end=tuple(marker)) - init_request(compiler, req) - for token in committed: - commit_token(req, token) - state = req.output_constraint - original_tail = state.reasoning_tail - - mask = verify_mask(compiler, [req] * (len(draft) + 1), [0, *draft]) - - for prediction, token in zip(mask, expected): - if token is None: - assert prediction.all() - else: - assert prediction.nonzero().flatten().tolist() == [token] - assert state.reasoning_end_token_ids == tuple(marker) - assert state.reasoning_tail == original_tail - assert state.in_reasoning and state.error is None - # The target may choose to continue thinking instead of accepting the draft. - assert not commit_token(req, ord("?")) - assert allowed([req])[0].all() - # A later real transition must start from the unmodified grammar root. - for token in marker: - assert not commit_token(req, token) - assert allowed([req])[0].nonzero().flatten().tolist() == [ord("a")] - - -@pytest.mark.parametrize("failure", ["fill_next_token_bitmask", "rollback"]) -@pytest.mark.parametrize("exception_type", [RuntimeError, BaseException]) -def test_failure_after_speculative_reasoning_end_records_terminal_error(compiler, monkeypatch, failure, exception_type): - import xgrammar as xgr - - req = make_req(regular_constraint="ab", guided_reasoning_end=(ord("]"),)) - init_request(compiler, req) - original = getattr(xgr.GrammarMatcher, failure) - - def fail_speculative_mask(matcher, *args): - if failure == "rollback" or args[1] == 2: - raise exception_type("injected speculative mask failure") - return original(matcher, *args) - - with monkeypatch.context() as patcher: - patcher.setattr(xgr.GrammarMatcher, failure, fail_speculative_mask) - masks = build_mask([req] * 5, draft_input_ids=[0, *b"]ab", 256]) - if failure == "rollback": - # 回滚失败不影响此前已完整生成的掩码，本轮仍可使用。 - assert masks is not None - else: - assert masks is None - - state = req.output_constraint - assert state.in_reasoning and state.reasoning_tail == () - assert state.error == "injected speculative mask failure" - assert state.is_terminated() - # 后续批次不再使用错误请求的 matcher。 - assert build_mask([req] * 4, draft_input_ids=[0, *b"]ab"]) is None - - -@pytest.mark.parametrize("exception_type", [RuntimeError, BaseException]) -def test_temporary_state_marks_rollback_failure_as_terminal(compiler, monkeypatch, exception_type): - import xgrammar as xgr - - req = make_req(regular_constraint="ab", guided_reasoning_end=tuple(b"]>")) - init_request(compiler, req) - commit_token(req, ord("]")) - state = req.output_constraint - - def fail_rollback(matcher, token_count): - raise exception_type("injected rollback failure") - - monkeypatch.setattr(xgr.GrammarMatcher, "rollback", fail_rollback) - # 直接使用临时状态接口，验证错误标记不依赖 fill_masks 的外层捕获。 - with state._temporary_state(): - assert state._accept_token(ord(">")) - assert state._accept_token(ord("a")) - - assert state.error == "injected rollback failure" - assert state.is_terminated() - assert state.in_reasoning and state.reasoning_tail == (ord("]"),) - - -def test_thinking_verify_does_not_enable_mask_or_advance_matcher(compiler): - req = make_req(regular_constraint="ab", guided_reasoning_end=(ord("]"),)) - init_request(compiler, req) - assert build_mask([req] * 3, draft_input_ids=[0, ord("x"), ord("y")]) is None - assert req.output_constraint.matcher.accept_token(ord("a")) - - -@pytest.mark.parametrize("failure", ["fill_next_token_bitmask", "rollback"]) -def test_speculative_mask_failure_isolated_and_draft_tokens_rolled_back(compiler, monkeypatch, failure): - import xgrammar as xgr - - bad = make_req(regular_constraint="ab") - good = make_req(regular_constraint="cd") - for req in [bad, good]: - init_request(compiler, req) - allowed([bad, good]) - bad_matcher = bad.output_constraint.matcher - original = getattr(xgr.GrammarMatcher, failure) - - def fail_for_bad_request(matcher, *args): - if matcher is bad_matcher and (failure == "rollback" or args[1] == 1): - raise RuntimeError("injected grammar failure") - return original(matcher, *args) - - with monkeypatch.context() as patcher: - patcher.setattr(xgr.GrammarMatcher, failure, fail_for_bad_request) - mask = verify_mask(compiler, [bad, bad, good], [0, ord("a"), 0]) - - if failure == "rollback": - assert [row.nonzero().flatten().tolist() for row in mask[:2]] == [[ord("a")], [ord("b")]] - else: - assert mask[:2].all() - assert mask[2].nonzero().flatten().tolist() == [ord("c")] - assert bad.output_constraint.error == "injected grammar failure" - assert good.output_constraint.error is None - if failure == "fill_next_token_bitmask": - # The failure happened after accepting draft 'a'; cleanup must still - # leave the matcher at its original prefix, which accepts 'a' again. - assert bad_matcher.accept_token(ord("a")) - - -@pytest.mark.parametrize("dp,microbatch_side", [(False, None), (True, None), (True, 0), (True, 1)]) -@pytest.mark.parametrize("dynamic", [False, True]) -def test_mtp_decode_handoff_compaction_and_accepted_commits(compiler, monkeypatch, dp, microbatch_side, dynamic): - from lightllm.server.router.model_infer.mode_backend.chunked_prefill import impl as chunked_impl - from lightllm.server.router.model_infer.mode_backend.dp_backend import impl as dp_impl - - impl = dp_impl if dp else chunked_impl - backend_class = dp_impl.DPChunkedPrefillBackend if dp else chunked_impl.ChunkedPrefillBackend - backend = backend_class.__new__(backend_class) - backend.eos_id = [256] - backend.is_master_in_dp = False - backend.extra_post_req_handle_func = None - req = make_req(regular_constraint="abc") - req.req_idx = 5 - req.out_token_id_count = {} - init_request(compiler, req) - order = [] - - def make_input(tokens): - return SimpleNamespace( - batch_size=len(tokens), - input_ids=None, - b_req_idx=torch.full((len(tokens),), req.req_idx, dtype=torch.int32), - b_mtp_index=torch.arange(len(tokens), dtype=torch.int32), - ) - - model_input = make_input(list(b"abcd") if dynamic else list(b"abc")) - empty = make_input([]) - original_reqs = [req] * model_input.batch_size - - def compact(**kwargs): - value = kwargs["model_input"] - selection = None - if dynamic and value.batch_size: - value.batch_size = 3 - value.b_req_idx = value.b_req_idx[:3] - value.b_mtp_index = value.b_mtp_index[:3] - selection = SimpleNamespace(wait=lambda: None, tensor=torch.tensor([True, True, True, False])) - return value, selection - - def compact_pair(**kwargs): - first, selection0 = compact(model_input=kwargs["model_input0"]) - second, selection1 = compact(model_input=kwargs["model_input1"]) - return first, selection0, second, selection1 - - def forward(value): - assert value.input_ids is None, "Constraint preparation must not materialize model inputs" - order.append("forward") - logits = torch.zeros(value.batch_size, 257) - logits[:, ord("!")] = 10 # Every constrained row must override this. - return SimpleNamespace(logits=logits) - - def forward_pair(first, second): - value = model_input - output = forward(value) - padding = SimpleNamespace(logits=torch.zeros(0, 257)) - return (output, padding) if microbatch_side == 0 else (padding, output) - - def propose(**kwargs): - order.append("propose") - return SimpleNamespace(token_ids=torch.empty(1, 2, dtype=torch.int64)) - - engine = SimpleNamespace( - plan_decode=lambda **kwargs: SimpleNamespace(draft_step=2, skip_verify_sync=False), - prepare_decode_model_input=compact, - prepare_decode_model_inputs=compact_pair, - propose_next=propose, - propose_next_overlap=propose, - update_planner_statics=lambda **kwargs: None, - ) - backend.spec_engine = backend.decode_draft_engine = engine - backend.model = SimpleNamespace(forward=forward, microbatch_overlap_decode=forward_pair) - monkeypatch.setattr(impl, "prepare_decode_inputs", lambda req_objs: (model_input, original_reqs)) - if microbatch_side is not None: - layout = ( - (model_input, original_reqs, [req], empty, [], []) - if microbatch_side == 0 - else (empty, [], [], model_input, original_reqs, [req]) - ) - monkeypatch.setattr(impl, "overlap_prepare_decode_inputs", lambda req_objs: layout) - monkeypatch.setattr(torch.cuda, "stream", lambda stream: nullcontext()) - monkeypatch.setattr(torch.cuda, "Event", lambda: SimpleNamespace(record=lambda: None, synchronize=lambda: None)) - monkeypatch.setattr(impl.g_infer_context, "get_overlap_stream", lambda: None) - monkeypatch.setattr( - impl.g_infer_context, - "req_sampling_manager", - SimpleNamespace(update_reqs_out_token_counter_gpu=lambda **kwargs: None), - raising=False, - ) - monkeypatch.setattr( - impl.g_pin_mem_manager, "async_copy_from_gpu_tensor", lambda key, gpu_tensor: gpu_tensor.clone() - ) - monkeypatch.setattr(impl, "gen_b_req_mtp_start_loc", lambda *args, **kwargs: torch.tensor([0])) - - token_table = torch.full((8, 4), ord("!"), dtype=torch.int64) - - class ReadyTokenTable: - def __getitem__(self, index): - assert "handoff" in order, "Token table read before the existing post_handle handoff" - assert "propose" not in order, "Token table read after the next proposal could overwrite it" - return token_table[index] - - backend._get_next_token_ranks = lambda logits, token_ids: torch.zeros_like(token_ids) - backend._async_copy_next_token_infos_to_pin_mem = lambda next_token_ids, next_token_logprobs, next_token_ranks: ( - next_token_ids, - next_token_logprobs, - next_token_ranks, - ) - - def sample(logits, reqs, eos_ids, b_mtp_index): - assert len(reqs) == 3 - assert req.out_token_id_count == {ord("a"): 1}, "MTP read stale CPU counters" - manager = make_mask_buffers(reqs, 4) - manager.req_to_next_token_ids_cpu = ReadyTokenTable() - assert prepare_sampling_tensors(reqs, manager)[-1] - assert b_mtp_index.tolist() == [0, 1, 2] - apply_masks(reqs, logits, manager, b_mtp_index) - order.append("sample") - result = logits.argmax(-1) - assert result.tolist() == [ord("b"), ord("c"), 256] - return result, torch.zeros(3) - - monkeypatch.setattr(impl, "sample", sample) - - def verify(**kwargs): - order.append("verify") - assert allowed([req])[0, ord("b")], "Draft advanced permanent matcher state" - return torch.tensor([2], dtype=torch.int32), torch.tensor([1, 1, 0], dtype=torch.int32) - - monkeypatch.setattr(impl.mtp_utils, "verify_mtp_tokens", verify) - monkeypatch.setattr(impl.mtp_utils, "scatter_mtp_next_tokens", lambda **kwargs: None) - backend._pre_post_handle = lambda *args, **kwargs: [] - - def post(**kwargs): - assert kwargs["next_token_ids"].tolist() == list(b"bc") - for token in kwargs["next_token_ids"].tolist(): - commit_token(req, token) - order.append("commit") - - backend._post_handle = post - - def handoff(): - order.append("handoff") - token_table[req.req_idx] = torch.tensor(list(b"abcd")) - commit_token(req, ord("a")) - req.out_token_id_count[ord("a")] = 1 - - events = SimpleNamespace( - notify_post_handle_and_wait_pre_post_handle=handoff, - notify_forward_and_wait_post_handle=lambda: None, - notify_pre_post_handle=lambda: None, - ) - method = backend.decode_overlap_mtp if microbatch_side is not None else backend.decode_mtp - method(events, [req]) - assert order == ["forward", "handoff", "sample", "verify", "propose", "commit"] - assert allowed([req])[0, 256] diff --git a/unit_tests/server/router/model_infer/structured_output/test_overlap.py b/unit_tests/server/router/model_infer/structured_output/test_overlap.py deleted file mode 100644 index 930cc5c1e8..0000000000 --- a/unit_tests/server/router/model_infer/structured_output/test_overlap.py +++ /dev/null @@ -1,152 +0,0 @@ -from contextlib import nullcontext -from threading import Event, Thread -from types import SimpleNamespace - -import pytest -import torch - -from unit_tests.server.grammar_helpers import ( - init_request, - make_req, - make_mask_buffers, - apply_masks, - prepare_sampling_tensors, -) -from lightllm.server.router.model_infer.mode_backend.chunked_prefill import impl as chunked_impl -from lightllm.server.router.model_infer.mode_backend.dp_backend import impl as dp_impl -from lightllm.server.router.model_infer.mode_backend.overlap_events import OverlapEventManager - - -@pytest.mark.parametrize("prefill_first", [False, True]) -@pytest.mark.parametrize("constrained", [False, True, "reasoning"]) -@pytest.mark.parametrize( - "dp,microbatch,empty_second", - [(False, False, False), (True, False, False), (True, False, True), (True, True, False), (True, True, True)], -) -def test_overlap_handoff_completes_postprocessing_before_sampling( - compiler, monkeypatch, prefill_first, constrained, dp, microbatch, empty_second -): - """Exercise the real two-thread event protocol and chunked-prefill methods. - - Sampling must see the previous matcher and CPU token counters. - All paths still submit the next forward before releasing CPU postprocessing. - """ - impl = dp_impl if dp else chunked_impl - backend_class = dp_impl.DPChunkedPrefillBackend if dp else chunked_impl.ChunkedPrefillBackend - backend = object.__new__(backend_class) - backend.args = SimpleNamespace(penalty_counter_mode="cpu_counter") - backend.disable_chunked_prefill = False - backend.extra_post_req_handle_func = backend.pd_prefill_chunked_handle_func = None - backend._capture_prompt_logprobs_if_needed = lambda *args: None - req = make_req( - regular_constraint="b" if constrained == "reasoning" else "ab" if constrained else None, - guided_reasoning_end=(ord("a"),) if constrained == "reasoning" else (), - ) - req.out_token_id_count = {} - init_request(compiler, req) - mask_buffers = make_mask_buffers([req]) - model_input = SimpleNamespace( - b_req_idx=torch.tensor([0]), b_mtp_index=torch.tensor([0]), b_prefill_has_output_cpu=[True] - ) - empty_input = SimpleNamespace(b_req_idx=torch.empty(0), b_mtp_index=torch.empty(0), b_prefill_has_output_cpu=[]) - monkeypatch.setattr(impl, "prepare_decode_inputs", lambda req_objs: (model_input, req_objs)) - monkeypatch.setattr(impl, "prepare_prefill_inputs", lambda req_objs, **kwargs: (model_input, req_objs)) - if microbatch: - monkeypatch.setattr( - impl, - "overlap_prepare_prefill_inputs", - lambda reqs: (model_input if reqs else empty_input, reqs, empty_input, []), - ) - monkeypatch.setattr( - impl, - "overlap_prepare_decode_inputs", - lambda req_objs: (model_input if req_objs else empty_input, req_objs, req_objs, empty_input, [], []), - ) - monkeypatch.setattr(impl.torch.cuda, "stream", lambda stream: nullcontext()) - monkeypatch.setattr( - impl.torch.cuda, "Event", lambda: SimpleNamespace(record=lambda: None, synchronize=lambda: None) - ) - monkeypatch.setattr(impl.g_infer_context, "get_overlap_stream", lambda: None) - monkeypatch.setattr(impl.g_infer_context, "save_hybrid_state_to_cache", lambda **kwargs: None) - monkeypatch.setattr(impl.g_infer_context, "is_hybrid_att_model", False) - second_forward = Event() - order, errors = [], [] - - def forward(model_input): - step = sum(event.startswith("forward") for event in order) - order.append(f"forward{step}") - if step == 1: - second_forward.set() - logits = torch.zeros(model_input.b_req_idx.numel(), 257) - logits[:, ord("a") + step] = 1 - return SimpleNamespace(logits=logits, prompt_logics=None) - - def forward_microbatch(first, second): - return forward(first), SimpleNamespace(logits=torch.zeros(0, 257), prompt_logics=None) - - backend.model = SimpleNamespace( - forward=forward, microbatch_overlap_prefill=forward_microbatch, microbatch_overlap_decode=forward_microbatch - ) - - def sample(**kwargs): - if second_forward.is_set(): - assert req.out_token_id_count.get(ord("a"), 0) == 1, "Sampling read stale CPU token counters" - logits = kwargs["logits"] - reqs = kwargs["run_reqs"] - if prepare_sampling_tensors(reqs, mask_buffers)[-1]: - apply_masks(reqs, logits, mask_buffers) - token = logits.argmax(-1) - order.append(f"sample{chr(token.item())}") - return token, token, torch.zeros(1), torch.zeros(1, dtype=torch.int64) - - backend._sample_and_scatter_token = sample - - def pre_post(*args, **kwargs): - req.cur_output_len += 1 - return [] - - backend._pre_post_handle = pre_post - - def post(**kwargs): - token = kwargs["next_token_ids"].item() - if token == ord("a"): - assert second_forward.wait(3), "CPU postprocessing must overlap the next forward" - if req.output_constraint is not None: - req.output_constraint.commit(token) - req.out_token_id_count[token] = req.out_token_id_count.get(token, 0) + 1 - order.append(f"post{chr(token)}") - - backend._post_handle = post - events = OverlapEventManager() - packs = [events.get_overlap_event_pack() for _ in range(3)] - - def run(step): - try: - packs[step].wait_to_forward() - run_reqs = [] if step == 1 and empty_second else [req] - if prefill_first and (step == 0 or empty_second): - method = backend.prefill_overlap if microbatch else backend.prefill_normal - else: - method = backend.decode_overlap if microbatch else backend.decode_normal - method(packs[step], run_reqs) - except BaseException as exc: - errors.append(exc) - - def drain(): - packs[2].wait_to_forward() - packs[2].notify_post_handle_and_wait_pre_post_handle() - - threads = [Thread(target=run, args=(step,), daemon=True) for step in range(2)] - threads.append(Thread(target=drain, daemon=True)) - for thread in threads: - thread.start() - for thread in threads: - thread.join(5) - assert not errors - assert not any(thread.is_alive() for thread in threads), f"Overlap deadlocked: {order}" - expected = ["forward0", "samplea", "forward1"] - if empty_second: - expected += ["posta"] - else: - expected += ["posta", "sampleb", "postb"] - assert order == expected diff --git a/unit_tests/server/router/model_infer/structured_output/test_reasoning.py b/unit_tests/server/router/model_infer/structured_output/test_reasoning.py deleted file mode 100644 index 8c265847db..0000000000 --- a/unit_tests/server/router/model_infer/structured_output/test_reasoning.py +++ /dev/null @@ -1,194 +0,0 @@ -import asyncio -import ctypes -import json -from types import SimpleNamespace -from unittest.mock import AsyncMock - -import pytest -import torch - -from lightllm.server import api_http, api_openai, build_prompt -from lightllm.server.api_models import ChatCompletionRequest -from lightllm.server.core.objs import FinishStatus -from lightllm.server.core.objs.sampling_params import SamplingParams -from lightllm.server.core.objs.start_args_type import StartArgs -from lightllm.server.router.model_infer.infer_batch import InferSamplingParams -from lightllm.server.router.model_infer.structured_output.state import ConstraintState - - -@pytest.mark.parametrize( - "parser,template_kwargs,effort,marker", - [ - ("qwen3", None, None, ""), - ("qwen3", {"enable_thinking": True}, None, ""), - ("qwen3", {"enable_thinking": False}, None, None), - ("qwen3", {"thinking": False}, None, None), - ("qwen3", None, "none", None), - ("deepseek-v3", None, None, None), - ("deepseek-v3", {"thinking": True}, None, ""), - ("qwen3-thinking", None, None, ""), - ("kimi", None, None, "◁/think▷"), - ("gpt-oss", None, None, "<|channel|>final<|message|>"), - ("gemma4", None, None, "answer"), - (None, None, None, None), - ], -) -def test_chat_grammar_activation_matches_template_thinking(monkeypatch, parser, template_kwargs, effort, marker): - monkeypatch.setattr(api_openai, "get_env_start_args", lambda: StartArgs(reasoning_parser=parser)) - monkeypatch.setattr(build_prompt, "tokenizer_supports_force_thinking", lambda: True) - encoded = [] - - def encode(text, add_special_tokens): - assert not add_special_tokens - encoded.append(text) - return [7, 8] - - request = ChatCompletionRequest( - model="test", - messages=[], - chat_template_kwargs=template_kwargs, - reasoning_effort=effort, - ) - result = api_openai._get_guided_reasoning_end(request, SimpleNamespace(encode=encode)) - assert encoded == ([marker] if marker else []) - assert result == ([7, 8] if marker else []) - - -@pytest.mark.parametrize("response_type", ["json_schema", "json_object", "text"]) -@pytest.mark.parametrize("thinking", [False, True]) -def test_chat_passes_reasoning_marker_through_shared_sampling_params(monkeypatch, compiler, response_type, thinking): - monkeypatch.setattr(api_openai, "get_env_start_args", lambda: StartArgs(reasoning_parser="qwen3")) - monkeypatch.setattr(build_prompt, "tokenizer_supports_force_thinking", lambda: True) - - async def prompt(*args): - return "prompt" - - monkeypatch.setattr(api_openai, "build_prompt", prompt) - captured = [] - - class RequestCaptured(Exception): - pass - - async def generate(prompt, params, *args, **kwargs): - # Exercise the real API-to-shared-memory and grammar-state boundary. - copy = SamplingParams.from_buffer_copy(bytes(params)) - payload = await copy.get_compiled_grammar(compiler) - shared_req = SimpleNamespace( - sample_params=copy, - get_compiled_grammar=lambda: payload, - input_len=1, - shm_prompt_ids=SimpleNamespace(arr=[ord("P")]), - ) - infer_params = InferSamplingParams(shared_req, vocab_size=257) - captured.append(compiler.output_constraint_state_factory.create_state(shared_req, infer_params)) - raise RequestCaptured - yield - - tokenizer = SimpleNamespace(encode=lambda *args, **kwargs: [7, 8]) - manager = SimpleNamespace(tokenizer=tokenizer, generate=generate, output_grammar_compiler=compiler) - monkeypatch.setattr(api_http, "g_objs", SimpleNamespace(httpserver_manager=manager)) - response_format = {"type": response_type} - if response_type == "json_schema": - response_format["json_schema"] = {"name": "result", "schema": {"type": "object"}} - request = ChatCompletionRequest( - model="test", - messages=[], - response_format=response_format, - chat_template_kwargs={"enable_thinking": thinking}, - ) - with pytest.raises(RequestCaptured): - asyncio.run(api_openai.chat_completions_impl(request, None)) - expected = (7, 8) if thinking and response_type != "text" else () - state = captured[0] - if response_type == "text": - assert state is None - else: - assert state.error is None - assert state.reasoning_end_token_ids == expected - assert state.in_reasoning == bool(expected) - - -def test_reasoning_marker_survives_serialization_and_parameter_reuse(): - params = SamplingParams() - params.init(None, guided_reasoning_end=[11, 12]) - copy = SamplingParams.from_buffer_copy(ctypes.string_at(ctypes.addressof(params), ctypes.sizeof(params))) - round_trip = SamplingParams() - round_trip.init(None, **copy.to_origin_dict()) - assert round_trip.constraint_out_params.guided_reasoning_end.to_list() == [11, 12] - round_trip.init(None) - assert round_trip.constraint_out_params.guided_reasoning_end.to_list() == [] - - -@pytest.mark.parametrize( - "parser,marker", - [ - ("deepseek-r1", ""), - ("qwen3-thinking", ""), - ("gpt-oss", "<|channel|>final<|message|>"), - ("minimax", ""), - ("qwen3", None), - ], -) -def test_fixed_reasoning_does_not_require_a_template_thinking_switch(monkeypatch, parser, marker): - monkeypatch.setattr(api_openai, "get_env_start_args", lambda: StartArgs(reasoning_parser=parser)) - monkeypatch.setattr(build_prompt, "tokenizer_supports_force_thinking", lambda: False) - tokenizer = SimpleNamespace(encode=lambda text, **kwargs: list(text.encode())) - request = ChatCompletionRequest(model="test", messages=[]) - assert api_openai._is_force_thinking_mode(request) == (marker is not None) - assert api_openai._get_guided_reasoning_end(request, tokenizer) == (list(marker.encode()) if marker else []) - - -@pytest.mark.parametrize("stream", [False, True]) -@pytest.mark.parametrize("response_type", ["json_object", "json_schema"]) -def test_deepseek_r1_json_answer_stays_in_content(monkeypatch, compiler, stream, response_type): - monkeypatch.setattr(api_openai, "get_env_start_args", lambda: StartArgs(reasoning_parser="deepseek-r1")) - monkeypatch.setattr(build_prompt, "tokenizer_supports_force_thinking", lambda: False) - monkeypatch.setattr(api_openai, "build_prompt", AsyncMock(return_value="")) - monkeypatch.setattr("lightllm.server.reasoning_parser.get_token_id", lambda token: ord(">")) - tokenizer = SimpleNamespace(encode=lambda text, **kwargs: list(text.encode())) - reasoning = "Let me think." - prefix = reasoning + "" - answer = '{"ok":true}' - - async def generate(prompt, params, *args, **kwargs): - payload = await params.get_compiled_grammar(compiler) - state = ConstraintState.from_grammar( - compiler.output_constraint_state_factory.get_grammar(payload), params.constraint_out_params - ) - bitmask = torch.empty((1, 9), dtype=torch.int32) - for index, token_id in enumerate([*tokenizer.encode(prefix + answer), 256]): - # The real matcher must allow thinking and its delimiter before constraining JSON. - masked = state.fill_masks(bitmask, []) - assert masked == (index >= len(prefix)) - if masked: - assert (int(bitmask[0, token_id // 32]) >> (token_id % 32)) & 1 - state.commit(token_id) - assert state.error is None - finish = FinishStatus(FinishStatus.FINISHED_STOP) if token_id == 256 else FinishStatus() - yield 0, "" if token_id == 256 else chr(token_id), {"id": token_id, "prompt_tokens": 1}, finish - - manager = SimpleNamespace(tokenizer=tokenizer, generate=generate, output_grammar_compiler=compiler) - monkeypatch.setattr(api_http, "g_objs", SimpleNamespace(httpserver_manager=manager)) - response_format = {"type": response_type} - if response_type == "json_schema": - response_format["json_schema"] = {"name": "result", "schema": {"const": {"ok": True}}} - request = ChatCompletionRequest(model="test", messages=[], response_format=response_format, stream=stream) - - async def run(): - response = await api_openai.chat_completions_impl(request, None) - if not stream: - message = response.choices[0].message - return message.reasoning, message.content - deltas = [] - async for chunk in response.body_iterator: - chunk = chunk.decode() if isinstance(chunk, bytes) else chunk - if chunk.strip() != "data: [DONE]": - for choice in json.loads(chunk.removeprefix("data: "))["choices"]: - deltas.append(choice["delta"]) - return "".join(delta.get("reasoning", "") for delta in deltas), "".join( - delta.get("content", "") for delta in deltas - ) - - actual_reasoning, content = asyncio.run(run()) - assert actual_reasoning == reasoning - assert content == answer diff --git a/unit_tests/server/router/model_infer/structured_output/test_sampling.py b/unit_tests/server/router/model_infer/structured_output/test_sampling.py deleted file mode 100644 index 8043b70ec9..0000000000 --- a/unit_tests/server/router/model_infer/structured_output/test_sampling.py +++ /dev/null @@ -1,320 +0,0 @@ -from types import SimpleNamespace -from unittest.mock import MagicMock - -import pytest -import torch - -from unit_tests.server.grammar_helpers import ( - allowed, - commit_token, - make_req, - init_request, -) -from lightllm.common.req_manager import req_sampling_params -from lightllm.common.basemodel.triton_kernel.mtp_utils import mtp_verify -from lightllm.server.router.model_infer.mode_backend import generic_post_process as sampling -from lightllm.server.router.model_infer.mtp_speculative.utils import scatter_mtp_next_tokens - - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="Sampling integration requires CUDA") - - -@pytest.fixture -def sampling_manager(request, monkeypatch): - args = SimpleNamespace( - penalty_counter_mode=getattr(request, "param", "cpu_counter"), - model_dir="unused", - mtp_step=3, - mtp_dynamic_verify=True, - output_constraint_mode="xgrammar", - ) - monkeypatch.setattr(req_sampling_params, "get_env_start_args", lambda: args) - monkeypatch.setattr(req_sampling_params, "get_vocab_size", lambda model_dir: 257) - result = req_sampling_params.ReqSamplingParamsManager(8) - assert result.req_to_next_token_ids.is_cuda - assert result.req_to_next_token_ids_cpu.device.type == "cpu" and result.req_to_next_token_ids_cpu.is_pinned() - assert result.req_to_next_token_scores.is_cuda - monkeypatch.setattr(sampling.g_infer_context, "req_manager", SimpleNamespace(req_sampling_params_manager=result)) - yield result - torch.cuda.synchronize() - - -def sampling_req(compiler, sampling_manager, req_idx=5, **constraints): - req = make_req(**({"regular_constraint": "abc"} | constraints)) - req.req_idx = req_idx - req.sampling_param.shm_param = SimpleNamespace( - constraint_out_params=req.sampling_param.shm_param.constraint_out_params, - presence_penalty=0.0, - frequency_penalty=0.0, - repetition_penalty=1.0, - exponential_decay_length_penalty=SimpleNamespace(to_tuple=lambda: (0, 1.1)), - input_penalty=False, - min_new_tokens=0, - temperature=1.0, - top_p=1.0, - top_k=1, - ) - req.get_last_gen_token = lambda: req.shm_req.shm_prompt_ids.arr[-1] - init_request(compiler, req) - sampling_manager.init_req_sampling_params(req) - return req - - -def test_sample_disables_stale_constraint_mask_for_ordinary_requests(compiler, sampling_manager, monkeypatch): - req = sampling_req(compiler, sampling_manager, regular_constraint=None) - sampling_manager.req_to_bitmask_enabled[req.req_idx] = True - sampling_manager.req_to_bitmask[req.req_idx].zero_() - - def unexpected_mask(*args, **kwargs): - pytest.fail("Ordinary requests must not launch the constraint mask kernel") - - monkeypatch.setattr(sampling, "apply_constraint_mask", unexpected_mask) - logits = torch.zeros(1, 257, device="cuda") - logits[:, ord("!")] = 10 - next_tokens, logprobs = sampling.sample(logits, [req], [256]) - - assert next_tokens.tolist() == [ord("!")] - assert torch.isfinite(logprobs).all() - assert not sampling_manager.req_to_bitmask_enabled[req.req_idx] - assert not compiler._cache - - -@pytest.mark.parametrize("sampling_manager", ["cpu_counter", "pin_mem_counter", "gpu_counter"], indirect=True) -def test_sample_reads_scattered_prefix_after_existing_handoff(compiler, sampling_manager, monkeypatch): - req = sampling_req(compiler, sampling_manager) - table = sampling_manager.req_to_next_token_ids - stream = torch.cuda.Stream() - stream.wait_stream(torch.cuda.current_stream()) - # This is the existing previous-step completion event, before post_handle. - previous_done = torch.cuda.Event() - with torch.cuda.stream(stream): - scatter_mtp_next_tokens( - backend=SimpleNamespace(model=SimpleNamespace(req_manager=sampling.g_infer_context.req_manager)), - proposal=SimpleNamespace(token_ids=torch.tensor([[ord("b"), ord("!")]], device="cuda")), - b_req_mtp_start_loc=torch.tensor([0], dtype=torch.int32, device="cuda"), - target_next_token_ids=torch.tensor([ord("a")], device="cuda"), - b_req_idx=torch.tensor([req.req_idx], dtype=torch.int32, device="cuda"), - mtp_accept_len=torch.tensor([1], dtype=torch.int32, device="cuda"), - ) - previous_done.record() - previous_done.synchronize() - commit_token(req, ord("a")) - - def unexpected_copy(*args, **kwargs): - pytest.fail("Sampling must read the ready CPU snapshot without starting another D2H") - - monkeypatch.setattr(sampling.g_pin_mem_manager, "async_copy_from_gpu_tensor", unexpected_copy) - monkeypatch.setattr(sampling.g_pin_mem_manager, "async_copy_from_gpu_tensor_with_event", unexpected_copy) - tensor_tolist = torch.Tensor.tolist - - def cpu_tolist(tensor): - assert not tensor.is_cuda, "Sampling synchronized a GPU tensor to read the prefix" - return tensor_tolist(tensor) - - logits = torch.zeros(3, 257, device="cuda") - logits[:, ord("!")] = 10 - # Three selected verify rows use a prefix of the four-column request table. - with monkeypatch.context() as patcher: - patcher.setattr(torch.Tensor, "tolist", cpu_tolist) - next_tokens, logprobs = sampling.sample( - logits, [req] * 3, [256], b_mtp_index=torch.arange(3, dtype=torch.int32, device="cuda") - ) - assert next_tokens.tolist() == [ord("b"), ord("c"), ord("!")] - assert torch.isfinite(logprobs).all() - assert allowed([req])[0].nonzero().flatten().tolist() == [ord("b")] - - accept_len, accepted = mtp_verify( - table, - torch.tensor([0], dtype=torch.int32, device="cuda"), - next_tokens, - torch.full((3,), req.req_idx, dtype=torch.int32, device="cuda"), - ) - assert accept_len.is_cuda and accepted.is_cuda - assert accept_len.tolist() == [2] and accepted.tolist() == [1, 1, 0] - assert table[req.req_idx].tolist() == [ord("a"), ord("b"), ord("!"), 1] - - -def test_sample_does_not_activate_constraints_for_partial_prefill(compiler, sampling_manager): - req = sampling_req(compiler, sampling_manager) - req.output_constraint.fill_masks = MagicMock(side_effect=AssertionError("Partial prefill must not fill masks")) - logits = torch.zeros(1, 257, device="cuda") - logits[:, ord("!")] = 10 - - next_tokens, logprobs = sampling.sample(logits, [req], [256], has_output=[False]) - - assert next_tokens.tolist() == [ord("!")] - assert torch.isfinite(logprobs).all() - req.output_constraint.fill_masks.assert_not_called() - assert req.output_constraint.error is None - - -@pytest.mark.parametrize("failure", ["fill_next_token_bitmask", "rollback"]) -def test_sample_preserves_valid_masks_and_isolates_failed_requests(compiler, sampling_manager, monkeypatch, failure): - import xgrammar as xgr - - bad = sampling_req(compiler, sampling_manager, req_idx=0) - good = sampling_req(compiler, sampling_manager, req_idx=1, regular_constraint="cd") - bad_matcher = bad.output_constraint.matcher - original = getattr(xgr.GrammarMatcher, failure) - - def fail_for_bad_request(matcher, *args): - if matcher is bad_matcher and (failure == "rollback" or args[1] == 1): - if failure == "fill_next_token_bitmask": - args[0].zero_() - raise RuntimeError("injected grammar failure") - return original(matcher, *args) - - monkeypatch.setattr(xgr.GrammarMatcher, failure, fail_for_bad_request) - - sampling_manager.req_to_next_token_ids[bad.req_idx] = torch.tensor([0, *b"ab!"]) - copy_snapshot_and_wait(sampling_manager) - sampling_manager.req_to_bitmask_enabled[bad.req_idx] = True - sampling_manager.req_to_bitmask[bad.req_idx].zero_() - logits = torch.zeros(3, 257, device="cuda") - logits[:, ord("!")] = 10 - tokens, logprobs = sampling.sample( - logits, - [bad, bad, good], - [256], - b_mtp_index=torch.tensor([0, 1, 0], dtype=torch.int32, device="cuda"), - ) - - expected_tokens = b"abc" if failure == "rollback" else b"!!c" - assert tokens.tolist() == list(expected_tokens) - assert torch.isfinite(logprobs).all() - assert bool(sampling_manager.req_to_bitmask_enabled[bad.req_idx]) == (failure == "rollback") - assert sampling_manager.req_to_bitmask_enabled[good.req_idx] - assert bad.output_constraint.error == "injected grammar failure" - assert good.output_constraint.error is None - - -def test_sample_groups_mixed_verify_rows_and_preserves_speculative_reasoning(compiler, sampling_manager): - active = sampling_req(compiler, sampling_manager, req_idx=0) - thinking = sampling_req(compiler, sampling_manager, req_idx=1, guided_reasoning_end=tuple(b"]>")) - plain = sampling_req(compiler, sampling_manager, req_idx=2, regular_constraint=None) - partial = sampling_req(compiler, sampling_manager, req_idx=3) - commit_token(active, ord("a")) - commit_token(thinking, ord("]")) - table = sampling_manager.req_to_next_token_ids - table[active.req_idx] = torch.tensor(list(b"ab!!")) - table[thinking.req_idx] = torch.tensor(list(b"]>a!")) - copy_snapshot_and_wait(sampling_manager) - - partial.output_constraint.fill_masks = MagicMock(side_effect=AssertionError("Partial prefill must not fill masks")) - # Dynamic verification can keep a different prefix length per request. - reqs = [plain, active, active, thinking, thinking, thinking, plain, partial] - logits = torch.zeros(len(reqs), 257, device="cuda") - logits[:, ord("!")] = 10 - next_tokens, logprobs = sampling.sample( - logits, - reqs, - [256], - has_output=[True] * 7 + [False], - b_mtp_index=torch.tensor([0, 0, 1, 0, 1, 2, 0, 0], dtype=torch.int32, device="cuda"), - ) - - assert next_tokens.tolist() == list(b"!bc!ab!!") - assert torch.isfinite(logprobs).all() - assert allowed([active])[0].nonzero().flatten().tolist() == [ord("b")] - assert thinking.output_constraint.in_reasoning - assert thinking.output_constraint.reasoning_tail == (ord("]"),) - partial.output_constraint.fill_masks.assert_not_called() - assert partial.output_constraint.error is None - assert not commit_token(thinking, ord(">")) - assert allowed([thinking])[0].nonzero().flatten().tolist() == [ord("a")] - - -def test_reused_mask_slots_clear_an_unreachable_draft_suffix(compiler, sampling_manager): - req = sampling_req(compiler, sampling_manager) - positions = torch.arange(4, dtype=torch.int32, device="cuda") - sampling_manager.req_to_next_token_ids[req.req_idx] = torch.tensor([0, *b"abc"]) - copy_snapshot_and_wait(sampling_manager) - logits = torch.zeros(4, 257, device="cuda") - logits[:, ord("!")] = 10 - tokens, _ = sampling.sample(logits, [req] * 4, [256], b_mtp_index=positions) - assert tokens.tolist() == [*b"abc", 256] - completed = torch.cuda.Event() - completed.record() - completed.synchronize() - - # The second prefix rejects its draft at '!', so masks from the previous - # iteration must not constrain that unreachable suffix. - sampling_manager.req_to_next_token_ids[req.req_idx] = torch.tensor([0, *b"a!c"]) - copy_snapshot_and_wait(sampling_manager) - logits.zero_() - logits[:, ord("!")] = 10 - tokens, _ = sampling.sample(logits, [req] * 4, [256], b_mtp_index=positions) - assert tokens.tolist() == list(b"ab!!") - assert allowed([req])[0].nonzero().flatten().tolist() == [ord("a")] - - -@pytest.mark.parametrize("replacement", ["plain", "partial", "thinking", "finished"]) -def test_reused_request_slot_disables_stale_mask(compiler, sampling_manager, replacement): - from lightllm.server.core.objs import FinishStatus - - old = sampling_req(compiler, sampling_manager, req_idx=0) - tokens, _ = sampling.sample(torch.zeros(1, 257, device="cuda"), [old], [256]) - assert tokens.tolist() == [ord("a")] - completed = torch.cuda.Event() - completed.record() - completed.synchronize() - assert sampling_manager.req_to_bitmask_enabled[0] - - constraints = {"regular_constraint": None} if replacement == "plain" else {} - if replacement == "thinking": - constraints["guided_reasoning_end"] = (ord("]"),) - new = sampling_req(compiler, sampling_manager, req_idx=0, **constraints) - if replacement == "finished": - new.finish_status.set_status(FinishStatus.FINISHED_STOP) - good = sampling_req(compiler, sampling_manager, req_idx=1, regular_constraint="xy") - logits = torch.zeros(2, 257, device="cuda") - logits[:, ord("!")] = 10 - tokens, _ = sampling.sample(logits, [new, good], [256], has_output=[replacement != "partial", True]) - - assert tokens.tolist() == list(b"!x") - assert not sampling_manager.req_to_bitmask_enabled[0] - assert sampling_manager.req_to_bitmask_enabled[1] - - -@pytest.mark.parametrize("constraints_enabled,draft_step", [(True, 2), (False, 2), (True, 0)]) -def test_mtp_scatter_reuses_snapshot_and_only_copies_when_needed(sampling_manager, constraints_enabled, draft_step): - table = sampling_manager.req_to_next_token_ids - snapshot = sampling_manager.req_to_next_token_ids_cpu - snapshot.fill_(-42) - snapshot_ptr = snapshot.data_ptr() - if not constraints_enabled: - sampling_manager.req_to_bitmask_enabled = None - stream = torch.cuda.Stream() - stream.wait_stream(torch.cuda.current_stream()) - backend = SimpleNamespace(model=SimpleNamespace(req_manager=sampling.g_infer_context.req_manager)) - for first_draft in [10, 20]: - with torch.cuda.stream(stream): - drafts = torch.arange(first_draft, first_draft + draft_step, device="cuda").reshape(1, draft_step) - scatter_mtp_next_tokens( - backend=backend, - proposal=SimpleNamespace(token_ids=drafts), - target_next_token_ids=torch.tensor([7], device="cuda"), - b_req_mtp_start_loc=torch.tensor([0], dtype=torch.int32, device="cuda"), - b_req_idx=torch.tensor([0], dtype=torch.int32, device="cuda"), - mtp_accept_len=torch.tensor([1], dtype=torch.int32, device="cuda"), - ) - completed = torch.cuda.Event() - completed.record() - completed.synchronize() - assert snapshot.data_ptr() == snapshot_ptr - if constraints_enabled: - expected = [7, first_draft, first_draft + 1, 1] if draft_step else [7, 1, 1, 1] - assert snapshot[0].tolist() == expected - # CPU snapshot is independent of subsequent writes to the GPU table. - table[0].fill_(99) - torch.cuda.current_stream().synchronize() - assert snapshot[0].tolist() == expected - else: - assert snapshot.eq(-42).all() - - -def copy_snapshot_and_wait(manager): - manager.req_to_next_token_ids_cpu.copy_(manager.req_to_next_token_ids, non_blocking=True) - completed = torch.cuda.Event() - completed.record() - completed.synchronize() diff --git a/unit_tests/server/router/model_infer/structured_output/test_state.py b/unit_tests/server/router/model_infer/structured_output/test_state.py deleted file mode 100644 index aa09c7df2e..0000000000 --- a/unit_tests/server/router/model_infer/structured_output/test_state.py +++ /dev/null @@ -1,407 +0,0 @@ -import asyncio -from itertools import count -from types import MethodType, SimpleNamespace -from weakref import ref -from unittest.mock import patch - -import pytest -import torch -import xgrammar as xgr - -from lightllm.server.core.objs import FinishStatus, SamplingParams -from lightllm.server.httpserver.grammar import OutputGrammarCompiler -from unit_tests.server.grammar_helpers import ( - make_req, - init_request, - make_mask_buffers, - prepare_sampling_tensors, - build_mask, - apply_masks, - allowed, - commit_token, -) -from lightllm.server.router.model_infer import infer_batch - - -def test_sampling_preparation_without_constraint_buffers(): - req = make_req() - manager = SimpleNamespace(req_to_bitmask=None, req_to_bitmask_enabled=None) - tensors = prepare_sampling_tensors([req, req], manager) - assert tensors[0].tolist() == [req.req_idx, req.req_idx] - assert tensors[1].tolist() == [1.0, 1.0] - assert not tensors[-1] - - -@pytest.mark.parametrize("mode", ["none", "xgrammar"]) -def test_context_initializes_state_factory_with_resolved_eos_at_startup(monkeypatch, mode): - args = SimpleNamespace( - output_constraint_mode=mode, model_dir="unused", tokenizer_mode="auto", trust_remote_code=False, eos_id=[255] - ) - monkeypatch.setattr(infer_batch, "get_env_start_args", lambda: args) - tokenizer = SimpleNamespace(eos_token_id=256) - tokenizer_info = xgr.TokenizerInfo([bytes([i]) for i in range(256)] + [b""], stop_token_ids=[255]) - payload = xgr.GrammarCompiler(tokenizer_info).compile_regex("a").serialize_json().encode() - context = infer_batch.InferenceContext() - with ( - patch("lightllm.server.tokenizer.get_tokenizer", return_value=tokenizer) as load_tokenizer, - patch.object(xgr.TokenizerInfo, "from_huggingface", return_value=tokenizer_info) as build_tokenizer_info, - patch.object(xgr, "GrammarCompiler", side_effect=AssertionError("Inference must not compile")), - ): - context.register( - backend=SimpleNamespace( - dp_rank_in_node=0, model=SimpleNamespace(att_state_manager=None), is_hybrid_att_model=False - ), - req_manager=SimpleNamespace(req_sampling_params_manager=None), - radix_cache=None, - shm_req_manager=None, - vocab_size=257, - ) - if mode == "none": - assert context.output_constraint_state_factory is None - load_tokenizer.assert_not_called() - build_tokenizer_info.assert_not_called() - return - load_tokenizer.assert_called_once_with("unused", "auto", trust_remote_code=False) - build_tokenizer_info.assert_called_once_with(tokenizer, vocab_size=257, stop_token_ids=[255]) - req = make_req(regular_constraint="a") - req.shm_req.get_compiled_grammar = lambda: payload - req.output_constraint = context.output_constraint_state_factory.create_state(req.shm_req, req.sampling_param) - commit_token(req, ord("a")) - assert req.output_constraint.matcher.accept_token(255) - assert req.output_constraint.is_terminated() - - -@pytest.mark.parametrize("previous_output_len", [-1, 2]) -def test_pd_history_outside_prompt_returns_constraint_error(compiler, previous_output_len): - req = make_req(regular_constraint="ab") - req.sampling_param.shm_param.constraint_out_params.pd_previous_output_len = previous_output_len - init_request(compiler, req) - assert "pd_previous_output_len" in req.output_constraint.error - assert req.output_constraint.matcher is None - assert req.output_constraint.is_terminated() - - -def test_thinking_does_not_advance_matcher_or_enable_mask(compiler): - req = make_req(regular_constraint="ab", guided_reasoning_end=(ord("]"),)) - init_request(compiler, req) - # The state owns its marker snapshot even if the shared parameters change. - req.sampling_param.shm_param.constraint_out_params.guided_reasoning_end.initialize([ord("?")]) - manager = make_mask_buffers([req]) - manager.req_to_bitmask.fill_(42) - assert not prepare_sampling_tensors([req], manager)[-1] - assert (manager.req_to_bitmask[req.req_idx] == -1).all() - assert not manager.req_to_bitmask_enabled[req.req_idx] - assert not commit_token(req, ord("?")) - assert req.output_constraint.error is None - assert not commit_token(req, ord("]")) - mask = allowed([req]) - assert mask[0, ord("a")] and not mask[0, ord("?")] - - -@pytest.mark.parametrize("marker", [b"]", b"aba"]) -def test_reasoning_boundary_mixed_batch_and_shared_schema(compiler, marker): - thinking = make_req(guided_json='{"const":{"ok":true}}', guided_reasoning_end=tuple(marker)) - direct = make_req(guided_json='{"const":{"ok":true}}') - for req in (thinking, direct): - init_request(compiler, req) - assert compiler.output_constraint_state_factory.get_grammar.cache_info().currsize == 1 - mask = allowed([thinking, direct]) - assert mask[0].all() and not mask[1, ord("?")] - assert thinking.output_constraint.matcher is not direct.output_constraint.matcher - - # A failed partial marker must not activate constraints. The prompt is ignored. - for token in b"thought aab?" + marker[:-1]: - assert not commit_token(thinking, token) - assert allowed([thinking, direct])[0].all() - # Sampling preparation observes the delimiter committed by post_handle. - assert not commit_token(thinking, marker[-1]) - mask = build_mask([thinking, direct]) - logits = torch.zeros(2, compiler.vocab_size) - apply_masks([thinking, direct], logits, mask) - assert torch.isfinite(logits[:, ord("{")]).all() - assert not torch.isfinite(logits[:, ord("?")]).any() - - for token in b'{"ok":true}': - assert allowed([thinking])[0, token] - assert not commit_token(thinking, token) - assert allowed([thinking])[0, 256] - assert commit_token(thinking, 256) - assert allowed([direct])[0, ord("{")] - - -def test_unconstrained_batch_does_not_compile_or_allocate_mask(compiler): - req = make_req() - init_request(compiler, req) - assert allowed([req]).all() - assert req.output_constraint is None - assert not compiler._cache - - -@pytest.mark.parametrize("current_dp_rank", [0, 1]) -@pytest.mark.parametrize("constraints_enabled", [False, True]) -def test_request_admission_loads_shared_artifacts_without_compilation( - compiler, monkeypatch, current_dp_rank, constraints_enabled -): - payload = asyncio.run(compiler.compile("regex", "ab")) - shared_reqs = {} - requests = [] - for index, req_id in enumerate((1, 2, 3, 4)): - owner_dp_rank = 1 - current_dp_rank if req_id == 4 else current_dp_rank - params = SamplingParams() - params.init(None, regular_constraint="" if req_id == 3 else "ab") - - def read_artifact(req_id=req_id, owner_dp_rank=owner_dp_rank): - # Disabled constraints and KV-only donor requests must not load grammar artifacts. - assert constraints_enabled and owner_dp_rank == current_dp_rank - return None if req_id == 3 else payload - - shared_reqs[index] = SimpleNamespace( - sample_params=params, - input_len=1, - shm_prompt_ids=SimpleNamespace(arr=[ord("P")]), - link_prompt_ids_shm_array=lambda: None, - link_logprobs_shm_array=lambda: None, - get_compiled_grammar=read_artifact, - ) - requests.append((req_id, index, SimpleNamespace(to_dict=lambda: {}), owner_dp_rank)) - monkeypatch.setattr(compiler, "_compile", lambda *args: pytest.fail("Inference must not compile")) - monkeypatch.setattr( - infer_batch, - "get_env_start_args", - lambda: SimpleNamespace( - diverse_mode=False, mtp_step=0, disable_chunked_mm_input=False, req_token_index_mode="pin_mem_cache" - ), - ) - context = infer_batch.InferenceContext( - req_manager=SimpleNamespace( - alloc=count().__next__, - req_to_token_indexs=torch.zeros((len(requests), 1), dtype=torch.int32), - req_sampling_params_manager=SimpleNamespace(init_req_sampling_params=lambda req: None), - ), - shm_req_manager=SimpleNamespace(get_req_obj_by_index=lambda index: shared_reqs[index]), - requests_mapping={}, - output_constraint_state_factory=compiler.output_constraint_state_factory if constraints_enabled else None, - ) - context.backend = SimpleNamespace(dp_rank_in_node=current_dp_rank) - context.infer_req_ids = [] - context.vocab_size = compiler.vocab_size - monkeypatch.setattr(infer_batch, "g_infer_context", context) - # Exercise the real InferReq constructor, including its internal grammar initialization. - context.add_reqs(requests) - reqs = context.requests_mapping - assert [reqs[req_id].dp_rank_in_node for req_id in (1, 2, 3, 4)] == [ - current_dp_rank, - current_dp_rank, - current_dp_rank, - 1 - current_dp_rank, - ] - if not constraints_enabled: - assert all(req.output_constraint is None for req in reqs.values()) - return - assert reqs[3].output_constraint is None and reqs[4].output_constraint is None - assert compiler.output_constraint_state_factory.get_grammar.cache_info().misses == 1 - assert reqs[1].output_constraint.matcher is not reqs[2].output_constraint.matcher - mask = allowed([reqs[1], reqs[2], reqs[3]]) - assert mask[:2, ord("a")].all() and mask[2].all() - commit_token(reqs[1], ord("a")) - mask = allowed([reqs[2], reqs[1]]) - assert mask[0, ord("a")] and not mask[0, ord("b")] - assert mask[1, ord("b")] and not mask[1, ord("a")] - - -@pytest.mark.parametrize("failure", ["read", "deserialize", "empty", "matcher"]) -def test_artifact_load_failure_finishes_only_its_request(compiler, monkeypatch, failure): - bad = make_req(regular_constraint="ab") - good = make_req(regular_constraint="cd") - init_request(compiler, good) - payload = asyncio.run(compiler.compile("regex", "ab")) - - def read_artifact(): - if failure == "read": - raise FileNotFoundError("missing compiled artifact") - if failure == "empty": - return b"" - return b"invalid artifact" if failure == "deserialize" else payload - - def fail_to_create_matcher(*args): - raise RuntimeError("matcher failed") - - bad.shm_req.get_compiled_grammar = read_artifact - with monkeypatch.context() as patcher: - if failure == "matcher": - patcher.setattr(xgr, "GrammarMatcher", fail_to_create_matcher) - bad.output_constraint = compiler.output_constraint_state_factory.create_state(bad.shm_req, bad.sampling_param) - assert bad.output_constraint.error is not None - assert build_mask([bad]) is None - bad.set_next_gen_token_id = lambda *args, **kwargs: None - bad.update_finish_status = lambda **kwargs: None - infer_batch.InferReqUpdatePack(bad, 1).handle(ord("!"), 0.0, -1, [256], is_master_in_dp=False) - assert bad.finish_status.is_finished_error() - assert allowed([good])[0].nonzero().flatten().tolist() == [ord("c")] - - -def test_constraint_state_lifetime_follows_request_cleanup(compiler): - req = make_req(regular_constraint="ab") - req.req_idx = 7 - init_request(compiler, req) - allowed([req]) - state_ref = ref(req.output_constraint) - payload = req.shm_req.get_compiled_grammar() - grammar = compiler.output_constraint_state_factory.get_grammar(payload) - context = infer_batch.InferenceContext( - req_manager=SimpleNamespace(free=lambda *args: None), - requests_mapping={1: req}, - shm_req_manager=SimpleNamespace(put_back_req_obj=lambda req: None), - ) - context.infer_req_ids = [1] - context.args = SimpleNamespace(diverse_mode=False) - context.backend = SimpleNamespace(is_master_in_dp=False) - context.free_a_req_mem = lambda *args: None - - context._filter([1], modify_shm_finish_state=False) - assert context.requests_mapping == {} and context.infer_req_ids == [] - del req - assert state_ref() is None - assert compiler.output_constraint_state_factory.get_grammar(payload) is grammar - - -def test_mixed_batch_partial_prefill_and_finished_requests(compiler): - partial = make_req(regular_constraint="ab") - active = make_req(regular_constraint="cd") - plain = make_req() - for req in [partial, active, plain]: - init_request(compiler, req) - mask = allowed([partial, active, plain], [False, True, True]) - assert mask[0].all() and mask[2].all() - assert mask[1, ord("c")] and not mask[1, ord("a")] - assert allowed([partial])[0].nonzero().flatten().tolist() == [ord("a")] - - active.finish_status.set_status(FinishStatus.FINISHED_STOP) - assert build_mask([active]) is None - - -@pytest.mark.parametrize("is_master_in_dp", [False, True]) -@pytest.mark.parametrize( - "reasoning_end,history,first_token,expected_next", - [ - (b"", b"", ord("a"), ord("b")), - (b"]", b"", ord("]"), ord("a")), - (b"]>", b"", ord("]"), None), - (b"", b"", ord("!"), None), - (b"", b"a", ord("b"), 256), - (b"]>", b"thought]", ord(">"), ord("a")), - ], -) -def test_pd_first_token_updates_constraint_before_sampling( - compiler, monkeypatch, is_master_in_dp, reasoning_end, history, first_token, expected_next -): - from lightllm.server.router.model_infer.mode_backend.base_backend import ModeBackend - - req = make_req(regular_constraint="ab", guided_reasoning_end=tuple(reasoning_end)) - req.shm_req.shm_prompt_ids.arr.extend(history) - req.shm_req.input_len += len(history) - req.sampling_param.shm_param.constraint_out_params.pd_previous_output_len = len(history) - good = make_req(regular_constraint="cd") - for request in (req, good): - init_request(compiler, request) - req.pd_task_success_num = 0 - req.shm_req.shm_prompt_ids.arr.append(0) - req.shm_req.shm_logprobs = SimpleNamespace(arr=[None] * (req.shm_req.input_len + 1)) - req.set_next_gen_token_id = MethodType(infer_batch.InferReq.set_next_gen_token_id, req) - req.update_finish_status = lambda **kwargs: None - sampling_params = make_mask_buffers([req, good]) - transfer = SimpleNamespace( - request_id=req.req_idx, - has_error=False, - first_gen_token_id=first_token, - first_gen_token_logprob=0.0, - ) - backend = SimpleNamespace( - shm_pd_trans_io_buffer=SimpleNamespace(read_obj=lambda: [transfer], sub_state=lambda: None), - is_pd_decode_mode=True, - model=SimpleNamespace(req_manager=SimpleNamespace(req_sampling_params_manager=sampling_params)), - eos_id=[256], - is_master_in_dp=is_master_in_dp, - ) - monkeypatch.setattr(infer_batch.g_infer_context, "requests_mapping", {req.req_idx: req}) - monkeypatch.setattr(torch.cuda, "current_stream", lambda: SimpleNamespace(synchronize=lambda: None)) - - # Exercise the real PD receive handler and InferReqUpdatePack.handle() - # before any sampling preparation can update the constraint. - ModeBackend._read_pd_trans_io_buffer_and_update_req_status(backend) - - assert req.cur_output_len == 1 and req.pd_task_success_num == 1 - assert req.shm_req.shm_prompt_ids.arr == [ord("P")] + list(history) + [first_token] - assert sampling_params.req_to_next_token_ids[req.req_idx, 0].item() == first_token - if is_master_in_dp: - assert req.shm_req.shm_cur_output_len == 1 - state = req.output_constraint - if first_token == ord("!"): - assert req.finish_status.status == FinishStatus.FINISHED_ERROR - assert "Grammar rejected sampled token" in state.error - assert allowed([good])[0].nonzero().flatten().tolist() == [ord("c")] - assert good.output_constraint.error is None - return - - assert state.error is None and not req.finish_status.is_finished() - if expected_next is None: - assert state.in_reasoning and state.reasoning_tail == (first_token,) - assert allowed([req])[0].all() - assert not commit_token(req, ord(">")) - expected_next = ord("a") - - # Repeated mask preparation must observe the same committed prefix. - for _ in range(2): - mask = allowed([req, good]) - assert [row.nonzero().flatten().tolist() for row in mask] == [[expected_next], [ord("c")]] - - -@pytest.mark.parametrize( - "constraints,text", - [ - ({"guided_grammar": 'root ::= "hello"'}, "hello"), - ( - { - "guided_json": '{"type":"object","properties":{"x":{"const":1}},' - '"required":["x"],"additionalProperties":false}' - }, - '{"x":1}', - ), - ({"guided_grammar": "json"}, '{"x":1}'), - ], -) -def test_grammar_and_schema_enforce_each_token(compiler, constraints, text): - req = make_req(**constraints) - init_request(compiler, req) - for token in text.encode(): - mask = allowed([req]) - assert mask[0, token] - req.output_constraint.commit(token) - assert not req.output_constraint.is_terminated() - assert allowed([req])[0, 256] - req.output_constraint.commit(256) - assert req.output_constraint.is_terminated() - - -def test_json_object_does_not_accept_array_root(compiler): - req = make_req(guided_grammar="json") - init_request(compiler, req) - mask = allowed([req]) - assert mask[0, ord("{")] - assert not mask[0, ord("[")] - - -def test_wrapped_tokenizer_uses_underlying_hf_tokenizer(monkeypatch): - tokenizer = object() - wrapper = SimpleNamespace(tokenizer=SimpleNamespace(tokenizer=tokenizer)) - received = [] - - def from_huggingface(value, **kwargs): - received.append(value) - return xgr.TokenizerInfo([bytes([i]) for i in range(256)] + [b""], stop_token_ids=[256]) - - monkeypatch.setattr(xgr.TokenizerInfo, "from_huggingface", from_huggingface) - compiler = OutputGrammarCompiler(wrapper, 257, [256]) - asyncio.run(compiler.compile("regex", "ab")) - compiler.shutdown() - assert received == [tokenizer] 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 deleted file mode 100644 index 21bbb06339..0000000000 --- a/unit_tests/server/router/model_infer/test_cpu_resource_layout.py +++ /dev/null @@ -1,513 +0,0 @@ -"""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.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 只有一个自有层，不能沿用 main 的多层类型列表。 - layer_types=["deepseek_sparse_attention"] if kind == "glm" else ["full_attention"], - ) - 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, - 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 - prepare_model_states(draft, logical_layer_ids=(config["num_hidden_layers"] + step,)) - drafts.append(draft) - 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_state_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=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") - assert main is not draft - 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.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_cache_operator().get_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=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: - 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_cache_operator().get_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_runtime_to_pages( - req_indices=[0], - state_token_lens=[8], - page_kind="big", - page_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.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_cache_operator().get_page_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.get_cpu_cache_operator().read_req( - cpu_page_buffers=buffer, - cpu_page_indices=[2], - token_indices=torch.arange( - start=8, - end=16, - dtype=torch.long, - device="cuda", - ), - state_token_len=8, - layouts=resources.get_cpu_cache_operator().get_page_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() - - -@pytest.mark.parametrize("hybrid", [False, True]) -def test_production_offload_preserves_ready_pages_without_host_sync( - 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=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) - 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.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.get_cpu_cache_operator(), "write_req", 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_offload_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 - token_indexes = torch.arange(16, device="cuda") - req.token_index = SimpleNamespace(get_gpu_token_indexes=lambda start, end: token_indexes[start:end]) - module.backend = SimpleNamespace( - is_master_in_dp=True, - is_hybrid_att_model=hybrid, - dp_world_size=1, - model=SimpleNamespace( - att_state_manager=resources, - req_manager=SimpleNamespace(req_to_token_indexs=token_indexes.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) - module.get_cpu_offload_stream = lambda: copy_stream - task = module.write_req_to_cpu(req) - assert waited_streams == [] - task.sync_event.synchronize() - assert write_batches == [(pages, [True, False])] - assert req.cpu_cache_offload_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.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_cache_operator().get_page_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( - 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=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.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_cache_operator().get_page_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.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_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 deleted file mode 100644 index 514aa19bf0..0000000000 --- a/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py +++ /dev/null @@ -1,362 +0,0 @@ -from unit_tests.common.att_resources.model_helpers import build_hybrid_resources -from unit_tests.common.req_token_index_test_utils import attach_req_token_index -import dataclasses -import json -from collections import deque -from types import SimpleNamespace - -import pytest -import torch - -from lightllm.common.attention_state.allocators import FullAttKVAllocator - -from lightllm.common.attention_state.full_att import ( - GQAFullAttentionState, - GQAAttentionConfig, -) -from lightllm.common.req_manager import ReqManager -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.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"), - pytest.mark.usefixtures("kv_allocator_environment"), -] - - -@pytest.fixture(params=[False, True], ids=["default_chunks", "whole_mm_blocks"]) -def linear_cache(monkeypatch, tmp_path, request): - args = StartArgs( - tp=1, - page_size=4, - max_req_total_len=512, - linear_att_hash_page_size=64, - linear_att_page_block_num=2, - cpu_cache_token_page_size=128, - disable_chunked_mm_input=request.param, - ) - monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(dataclasses.asdict(args))) - monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_NODE", "0") - monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") - monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_DP", "0") - monkeypatch.setenv("LIGHTLLM_DP_WORLD_SIZE", "1") - monkeypatch.setenv("LIGHTLLM_UNIQUE_SERVICE_NAME_ID", tmp_path.name) - 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 = Qwen3NextResourceConfig( - tp_world_size=1, - full_att_all_num_kv_heads=1, - full_att_head_dim=8, - global_linear_k_heads=1, - global_linear_v_heads=1, - head_linear_k_dim=8, - head_linear_v_dim=8, - 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( - GQAAttentionConfig( - logical_layer_groups=tuple((layer,) for layer in config.get_full_att_layer_indices()), - global_head_num=1, - head_dim=8, - dtype=torch.bfloat16, - ), - page_size=args.page_size, - ) - full.allocate_storage(384) - requests = ReqManager(2, 512) - 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, - ) - 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( - g_infer_context, - "radix_cache", - SimpleNamespace( - get_big_page_ids_by_node=lambda node: [], - free_one_small_page_buffer=lambda: None, - free_radix_cache_to_get_enough_token=lambda need: None, - ), - ) - yield full, resources, requests, config - get_env_start_args.cache_clear() - get_unique_server_name.cache_clear() - - -@pytest.mark.parametrize("token_num", [320, 384]) -@pytest.mark.parametrize("token_index_mode", ["pin_mem_cache", "pin_mem_page_cache", "pure_gpu"]) -def test_cpu_roundtrip_requires_all_linear_pages_and_restores_final_state( - monkeypatch, - linear_cache, - cpu_cache_factory, - token_num, - token_index_mode, -): - 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 - - 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_conv_bf16_ssm_fp32_state - big = linear.big_pages_pool - small = linear.small_pages_pool - big_states = {} - for end in range(128, token_num + 1, 128): - slot = resources.bounded_att_big_page_allocator.alloc(1)[0] - big_states[end] = slot - for tensor in (big.conv_state[slot], big.ssm_state[slot]): - tensor.normal_() - tail_slot = None - if token_num % 128: - 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 = (small.conv_state[tail_slot], small.ssm_state[tail_slot]) - else: - 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.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, resources.get_cpu_cache_operator().get_page_bytes()), - ) - indexes = torch.arange(token_num, dtype=torch.int32, device="cuda") - req = SimpleNamespace( - req_idx=0, - cache_tiers=(CacheTier.CPU,), - cpu_cache_offload_task_status=InferReq._CpuCacheTaskStatus.NOT_STARTED, - cur_kv_len=token_num, - hold_kv_len=token_num, - hybrid_cache_len=token_num, - hybrid_len_to_big_page_id=big_states, - tail_small_page_buffer_id=tail_slot, - shared_kv_node=None, - mm_input_spans=[(96, 160)] if disable_chunked_mm_input else [], - sampling_param=SimpleNamespace(shm_param=SimpleNamespace(prompt_logprobs=-1)), - shm_req=SimpleNamespace( - 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]), - ), - ) - - # 通过接收 InferReq 的生产入口完成页来源查询、CPU 分配与快照写出。 - attach_req_token_index(req, requests, indexes.cpu(), mode=token_index_mode) - module = MultiLevelKvCacheModule.__new__(MultiLevelKvCacheModule) - module.load_tasks = deque() - 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_state_manager=resources), - ) - monkeypatch.setattr(common.dist, "broadcast_object_list", lambda *args, **kwargs: None) - module.get_cpu_offload_stream().wait_stream(torch.cuda.current_stream()) - task = module.write_req_to_cpu(req) - assert task is not None - task.sync_event.synchronize() - pages = task.page_indexes.tolist() - client.update_pages_status_to_ready(pages, deref=False) - resources.bounded_att_big_page_allocator.free(list(big_states.values())) - if tail_slot is not None: - 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 - req.token_index.reset() - full.k_storage.zero_() - full.v_storage.zero_() - big.conv_state.zero_() - big.ssm_state.zero_() - - module.args = get_env_start_args() - module.backend = SimpleNamespace( - is_master_in_dp=True, - dp_world_size=1, - model=SimpleNamespace(req_manager=requests, att_state_manager=resources), - is_hybrid_att_model=True, - ) - module.init_sync_group = None - module.need_sync_compute_stream = lambda: False - monkeypatch.setattr(common.g_infer_context, "get_can_alloc_token_num", lambda: 384) - monkeypatch.setattr(common.dist, "barrier", lambda **kwargs: None) - module.load_cpu_cache_to_reqs([req]) - torch.cuda.synchronize() - - assert req.cur_kv_len == req.hold_kv_len == token_num - 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)) - # 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.conv_state[slot], big.ssm_state[slot]), expected_big_states[end]): - torch.testing.assert_close(actual, expected) - torch.testing.assert_close(linear.conv_state[:, 0].cpu(), expected_states[0]) - torch.testing.assert_close(linear.ssm_state[:, 0].cpu(), expected_states[1]) - - # Model data validity and multimodal resume boundaries are separate checks. - tree = HybridAttPagedRadixCache( - total_token_num=384, - rank_in_node=103, - hash_page_size=64, - big_page_num=2, - att_state_manager=resources, - ) - cache_len = max(req.hybrid_len_to_big_page_id) - tree.insert( - 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), - len_to_big_page_id=req.hybrid_len_to_big_page_id, - ) - for length, expected in [(128, 128 if first_page_has_state else 0), (256, 256)]: - node, matched, _ = tree.match_prefix( - key=torch.arange(length), - block_hashs=list(range(1, length // 64 + 1)), - update_refs=True, - mm_input_spans=req.mm_input_spans, - ) - assert matched == expected - tree.dec_node_ref_counter(node) - req.shared_kv_node, _, _ = tree.match_prefix( - key=torch.arange(cache_len), - block_hashs=list(range(1, cache_len // 64 + 1)), - update_refs=True, - 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, - hold_kv_len=256, - cache_tiers=(CacheTier.CPU,), - shm_req=SimpleNamespace( - input_len=257, - token_hash_list=SimpleNamespace(get_all=lambda: [11, 22]), - token_hash_page_len_list=SimpleNamespace(get_all=lambda: [128, 256]), - ), - ) - attach_req_token_index(req, requests, torch.arange(256, dtype=torch.int32)) - 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_state_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.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) - if allocated_pages == 0: - assert task is None - assert req.cpu_cache_offload_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) - 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 deleted file mode 100644 index eb64c6ed1d..0000000000 --- a/unit_tests/server/router/model_infer/test_multimodal_chunking.py +++ /dev/null @@ -1,404 +0,0 @@ -from types import SimpleNamespace - -import pytest -import torch - - -@pytest.mark.parametrize("hybrid", [False, True]) -@pytest.mark.parametrize("spans", [[], [(450, 600)], [(100, 800)], [(0, 384), (384, 768)]]) -def test_chunk_ids_and_lengths_use_the_same_complete_boundaries(hybrid, spans): - from lightllm.server.router.model_infer.infer_batch import InferReq - - req = InferReq.__new__(InferReq) - req.args = SimpleNamespace( - chunked_prefill_size=256, - batch_max_tokens=1024, - linear_att_hash_page_size=64, - linear_att_page_block_num=8, - ) - req.cur_kv_len = req.cur_output_len = 0 - req.mm_input_spans = spans - req.hybrid_cache_len = 960 - req.shm_req = SimpleNamespace(input_len=1000, shm_prompt_ids=SimpleNamespace(arr=list(range(1000)))) - if hybrid: - req.get_chuncked_input_token_len = req.get_chuncked_input_token_len_for_hybrid_att - req.get_chuncked_input_token_ids = req.get_chuncked_input_token_ids_for_hybrid_att - - while req.cur_kv_len < req.get_cur_total_len(): - end = req.get_chuncked_input_token_len() - assert req.get_chuncked_input_token_ids() == list(range(end)) - assert 0 < end - req.cur_kv_len <= req.args.batch_max_tokens - assert not any(start < end < stop for start, stop in spans) - req.cur_kv_len = end - - -@pytest.mark.usefixtures("kv_allocator_environment") -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for bounded transfers") -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 - - monkeypatch.setattr( - "lightllm.common.attention_state.bounded_att.utils.get_env_start_args", - lambda: SimpleNamespace(enable_dp_prompt_cache_fetch=False), - ) - resources = composition(device="cuda") - slots = init_source(resources) - resources.reset_bounded_att_state([1]) - resources.restore_bounded_att_runtime_from_pages( - req_indices=[1], - state_token_lens=[8], - page_kind="big", - 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_state_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_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 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", - [ - (2049, 8, [], [512, 1024, 1536, 2048], []), - (2049, 8, [(450, 600)], [1024, 1536, 2048], []), - (2049, 8, [(1100, 1700)], [512, 1024, 2048], []), - (2049, 8, [(0, 700)], [1024, 1536, 2048], []), - (2049, 8, [(450, 1700)], [2048], []), - (2049, 8, [(1900, 2049)], [512, 1024, 1536], []), - (2010, 8, [], [512, 1024, 1536], [1984]), - (2010, 8, [(1900, 2000)], [512, 1024, 1536], []), - (2010, 8, [(1900, 1984)], [512, 1024, 1536], [1984]), - (2010, 10000000, [], [], [1984]), - (2010, 10000000, [(1900, 2000)], [], []), - ], -) -def test_hybrid_chunks_save_state_only_at_original_boundaries( - input_len, - big_page_num, - spans, - expected_big, - expected_small, -): - from lightllm.server.router.model_infer.infer_batch import InferReq, InferenceContext - - req = InferReq.__new__(InferReq) - req.args = SimpleNamespace( - chunked_prefill_size=256, - batch_max_tokens=2048, - linear_att_hash_page_size=64, - linear_att_page_block_num=big_page_num, - disable_chunked_prefill=False, - ) - req.cur_kv_len = req.cur_output_len = req.req_idx = 0 - req.mm_input_spans = spans - req.hybrid_cache_len = (input_len - 1) // 64 * 64 - req.hybrid_len_to_big_page_id = {} - req.tail_small_page_buffer_id = None - req.shm_req = SimpleNamespace(input_len=input_len) - req.get_chuncked_input_token_len = req.get_chuncked_input_token_len_for_hybrid_att - big_snapshots, small_snapshots = [], [] - page_ends = list(range(64 * big_page_num, req.hybrid_cache_len + 1, 64 * big_page_num)) - slots = iter(range(len(page_ends))) - context = InferenceContext() - context.args = req.args - context.is_hybrid_att_model = True - context.radix_cache = SimpleNamespace(free_one_small_page_buffer=lambda: None) - 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_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), - ) - - 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() - - assert big_snapshots == expected_big - assert small_snapshots == expected_small - assert sorted(req.hybrid_len_to_big_page_id) == page_ends - assert len(set(req.hybrid_len_to_big_page_id.values())) == len(page_ends) - assert big_snapshots == [end for end in page_ends if not any(start < end < stop for start, stop in spans)] - - -@pytest.mark.parametrize( - "page_size,same_image,gpu_end,cpu_end", - [(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, -): - from lightllm.server.multi_level_kv_cache.manager import MultiLevelKVCacheManager - from lightllm.server.router.dynamic_prompt.radix_cache import RadixCache - from lightllm.server.router.model_infer.infer_batch import InferReq, g_infer_context - from lightllm.common.req_manager import ReqManager - from lightllm.utils import shm_utils - from lightllm.utils.kv_cache_utils import CpuKVCacheMeta, compute_token_list_hash - - # Neo layout: text + + independent virtual image IDs + + text. - prefix = [7] * 699 + [10] - source = prefix + list(range(1000000, 1000500)) + [11] + [20] * 400 - image_start_id = 1000000 if same_image else 2000000 - query = prefix + list(range(image_start_id, image_start_id + 500)) + [11] + [21] * 400 - monkeypatch.setattr(shm_utils, "get_unique_server_name", lambda: tmp_path.name) - radix_cache = RadixCache(4096, 0, page_size=page_size) - radix_cache.insert(torch.tensor(source)) - - monkeypatch.setattr(g_infer_context, "is_hybrid_att_model", False) - monkeypatch.setattr(g_infer_context, "radix_cache", radix_cache) - from lightllm.server.router.model_infer.req_token_index import ( - PinMemPageCacheReqTokenIndex, - create_req_token_index, - ) - - if ( - PinMemPageCacheReqTokenIndex._page_offsets is not None - and PinMemPageCacheReqTokenIndex._page_offsets.numel() != page_size - ): - PinMemPageCacheReqTokenIndex._page_offsets = None - - req_manager = ReqManager.__new__(ReqManager) - req_manager.page_size = page_size - req_manager.req_to_token_indexs = torch.zeros((1, len(query)), dtype=torch.int32) - monkeypatch.setattr(g_infer_context, "req_manager", req_manager) - req = InferReq.__new__(InferReq) - req.req_manager = req_manager - req.args = SimpleNamespace(page_size=page_size, req_token_index_mode="pin_mem_page_cache") - req.sampling_param = SimpleNamespace(disable_prompt_cache=False) - req.cur_kv_len = req.cur_output_len = req.hold_kv_len = req.req_idx = 0 - req.shared_kv_node = None - req.mm_input_spans = [(700, 1200)] - req.shm_req = SimpleNamespace(input_len=len(query), shm_prompt_ids=SimpleNamespace(arr=query)) - req.token_index = create_req_token_index(req, torch.cuda.current_stream()) - req._match_radix_cache() - assert req.cur_kv_len == req.hold_kv_len == gpu_end - if gpu_end: - got = req.token_index.get_cpu_token_indexes(0, gpu_end) - assert got[::page_size].tolist() == query[:gpu_end:page_size] - assert not any(start < gpu_end < end for start, end in req.mm_input_spans) - assert radix_cache.get_refed_tokens_num() == gpu_end - - hashes = compute_token_list_hash(source, 256) - 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) - manager.cpu_cache_client = client - matched = manager._cpu_cache_match(compute_token_list_hash(query, 256)) - assert len(matched) * 256 == cpu_end - # The page containing the image tail also contains the differing text. - assert any(start < cpu_end < end for start, end in req.mm_input_spans) == same_image - client.deref_pages(matched) - radix_cache.dec_node_ref_counter(req.shared_kv_node) - assert radix_cache.get_refed_tokens_num() == 0 - assert radix_cache.root_node.ref_counter == 1 - - -@pytest.mark.parametrize("disable_chunked_mm_input", [False, True]) -def test_mm_chunking_keeps_the_original_state_cache_position(monkeypatch, disable_chunked_mm_input): - from lightllm.server.router.model_infer import infer_batch - - shm_req = SimpleNamespace( - link_prompt_ids_shm_array=lambda: None, - link_logprobs_shm_array=lambda: None, - hybrid_token_hash_list=SimpleNamespace(size=3), - input_len=1000, - shm_prompt_ids=SimpleNamespace(arr=list(range(1000))), - ) - monkeypatch.setattr(infer_batch.g_infer_context, "is_hybrid_att_model", True) - monkeypatch.setattr( - infer_batch.g_infer_context, "shm_req_manager", SimpleNamespace(get_req_obj_by_index=lambda index: shm_req) - ) - monkeypatch.setattr( - infer_batch.g_infer_context, - "req_manager", - SimpleNamespace(req_sampling_params_manager=SimpleNamespace(init_req_sampling_params=lambda req: None)), - ) - monkeypatch.setattr( - infer_batch, - "InferSamplingParams", - lambda *args: SimpleNamespace( - pd_decode_node=None, - shm_param=SimpleNamespace(stop_sequences=SimpleNamespace(to_token_ids_list=lambda: [])), - ), - ) - monkeypatch.setattr(infer_batch, "PromptSelectedLogprobsExt", lambda req: None) - monkeypatch.setattr(infer_batch, "FinalTokenMetadataExt", lambda req: None) - req = infer_batch.InferReq.__new__(infer_batch.InferReq) - req.args = SimpleNamespace( - linear_att_hash_page_size=256, - linear_att_page_block_num=8, - disable_chunked_mm_input=disable_chunked_mm_input, - batch_max_tokens=1024, - ) - req.shm_index, req.vocab_size = 0, 100 - 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}, - ], - "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]) -@pytest.mark.parametrize( - "loaded_start,load_end,spans,disk_tokens,expected,cpu_tokens,remaining_disk_tokens", - [ - (0, 0, [], 0, 0, 0, 0), - (0, 4096, [], 2048, 4096, 2048, 2048), - (0, 2048, [(2048, 2500)], 0, 2048, 2048, 0), - (0, 4096, [(3500, 4096)], 2048, 4096, 2048, 2048), - (0, 4096, [(3500, 4500)], 0, 2048, 2048, 0), - (0, 4096, [(3500, 4500)], 2048, 2048, 2048, 0), - (0, 4096, [(3500, 4500)], 4096, 2048, 0, 2048), - (0, 4096, [(1000, 2500), (3500, 4500)], 2048, 0, 0, 0), - (2304, 4096, [(3000, 4500)], 2048, 2304, 0, 0), - (0, 6144, [(1000, 2500), (3500, 4500)], 2048, 6144, 4096, 2048), - (2304, 6144, [(3000, 4500)], 2048, 6144, 1792, 2048), - ], -) -def test_cpu_cache_match_and_load_use_a_complete_restore_point( - monkeypatch, - cpu_cache_factory, - hybrid, - loaded_start, - load_end, - spans, - disk_tokens, - expected, - cpu_tokens, - remaining_disk_tokens, -): - from lightllm.server.router.model_infer.mode_backend import multi_level_kv_cache as common - 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)), 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 = [] - req = SimpleNamespace( - req_idx=0, - cur_kv_len=loaded_start, - hold_kv_len=loaded_start, - mm_input_spans=spans, - sampling_param=SimpleNamespace(shm_param=SimpleNamespace(prompt_logprobs=-1)), - shm_req=SimpleNamespace( - input_len=8192, - disk_prompt_cache_len=disk_tokens, - cpu_cache_match_page_indexes=SimpleNamespace(get_all=lambda: page_ids), - token_hash_page_len_list=SimpleNamespace(get_all=lambda: ends), - ), - ) - table = torch.arange(8192, dtype=torch.int32).view(1, -1) - - def allocate(alloc_token_num): - if alloc_token_num == 0: - return None - start = req.cur_kv_len - req.hold_kv_len += alloc_token_num - return torch.arange(start, start + alloc_token_num, dtype=torch.int32) - - req.args = SimpleNamespace(page_size=256) - req.token_index = SimpleNamespace(get_gpu_token_indexes=lambda start, end: table[0, start:end]) - req.kv_cache_alloc_need = lambda target: max(0, target - req.hold_kv_len) - req.alloc_kv_mem = allocate - - from collections import deque - - module = common.MultiLevelKvCacheModule.__new__(common.MultiLevelKvCacheModule) - module.load_tasks = deque() - module.args = SimpleNamespace(page_size=256, cpu_cache_token_page_size=256) - module.backend = SimpleNamespace( - is_master_in_dp=True, - dp_world_size=1, - is_hybrid_att_model=hybrid, - model=SimpleNamespace( - req_manager=SimpleNamespace( - req_to_token_indexs=table, alloc_token_indexes=allocate, wait_token_indexes=lambda: None - ), - att_state_manager=SimpleNamespace( - 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))), - full_att_allocator=None, - ), - ), - ) - module.cpu_cache_client = client - module.init_sync_group = None - module.need_sync_compute_stream = lambda: False - monkeypatch.setattr( - common.g_infer_context, - "req_manager", - SimpleNamespace(alloc_token_indexes=allocate), - ) - monkeypatch.setattr(common.g_infer_context, "get_can_alloc_token_num", lambda: 8192) - monkeypatch.setattr(common.g_infer_context, "is_hybrid_att_model", hybrid) - monkeypatch.setattr(common.dist, "barrier", lambda **kwargs: None) - monkeypatch.setattr(torch.Tensor, "cuda", lambda self, **kwargs: self) - req.hybrid_len_to_big_page_id = {} - module.load_cpu_cache_to_reqs([req]) - for task in getattr(module, "load_tasks", ()): - task.event.synchronize() - module.update_cpu_cache_load_task_states() - assert req.cur_kv_len == req.hold_kv_len == expected - assert loaded == ([expected] if expected > loaded_start else []) - assert req.shm_req.cpu_prompt_cache_len == cpu_tokens - assert req.shm_req.disk_prompt_cache_len == remaining_disk_tokens - assert all(client.page_items.linked_items[page].ref_count == 0 for page in page_ids) - - -@pytest.mark.parametrize("state_slot,expected", [(None, 1), (3, 2)]) -def test_cpu_offload_requires_tail_state(monkeypatch, state_slot, expected): - from lightllm.server.router.model_infer.mode_backend import multi_level_kv_cache as common - - monkeypatch.setattr(common.g_infer_context, "is_hybrid_att_model", True) - module = common.MultiLevelKvCacheModule.__new__(common.MultiLevelKvCacheModule) - module.args = SimpleNamespace(cpu_cache_token_page_size=512, disable_linear_att_small_page_cpu_cache=False) - req = SimpleNamespace(tail_small_page_buffer_id=state_slot) - assert module._handle_hybrid_att_last_page(req, 2, [512, 960]) == expected diff --git a/unit_tests/server/router/req_queue/test_pd_queue_selection.py b/unit_tests/server/router/req_queue/test_pd_queue_selection.py deleted file mode 100644 index 0aa1a45c9e..0000000000 --- a/unit_tests/server/router/req_queue/test_pd_queue_selection.py +++ /dev/null @@ -1,108 +0,0 @@ -from types import SimpleNamespace -from unittest.mock import MagicMock - -from lightllm.server.router.req_queue import _get_req_queue_class -from lightllm.server.router.batch import Batch -from lightllm.server.router.req_queue.chunked_prefill.impl_for_pd_decode import PDDecodeQueue -from lightllm.server.router.req_queue.chunked_prefill.impl_for_pd_prefill import PDPrefillQueue - - -def _make_args(run_mode: str): - return SimpleNamespace( - run_mode=run_mode, - diverse_mode=False, - output_constraint_mode="none", - disable_chunked_prefill=False, - ) - - -def test_pd_prefill_uses_prefill_queue(): - args = _make_args("prefill") - - assert _get_req_queue_class(args, router=None, dp_size_in_node=1) is PDPrefillQueue - - -def test_pd_decode_uses_decode_queue(): - args = _make_args("decode") - - assert _get_req_queue_class(args, router=None, dp_size_in_node=1) is PDDecodeQueue - - -def test_pd_prefill_peak_tokens_sum_request_tuples(): - queue = PDPrefillQueue.__new__(PDPrefillQueue) - queue.dp_index = 0 - queue.args = SimpleNamespace(page_size=16) - queue.is_busy = lambda: False - queue.router = SimpleNamespace(router_statics=SimpleNamespace(ema_req_out_len=10)) - reqs = [ - SimpleNamespace( - request_id="req-0", - input_len=10, - sample_params=SimpleNamespace(max_new_tokens=1, suggested_dp_index=0), - get_tuple_tokens=MagicMock(return_value=(11, 34)), - ), - SimpleNamespace( - request_id="req-1", - input_len=20, - sample_params=SimpleNamespace(max_new_tokens=1, suggested_dp_index=0), - get_tuple_tokens=MagicMock(return_value=(21, 34)), - ), - SimpleNamespace( - request_id="req-2", - input_len=100, - sample_params=SimpleNamespace(max_new_tokens=1, suggested_dp_index=1), - get_tuple_tokens=MagicMock(return_value=(101, 34)), - ), - ] - batch = Batch(batch_id=1, reqs=reqs, dp_size_in_node=2) - - assert queue._caclu_batch_estimated_peak_token_num(batch) == 100 - reqs[0].get_tuple_tokens.assert_called_once_with(False, 10) - reqs[1].get_tuple_tokens.assert_called_once_with(False, 10) - reqs[2].get_tuple_tokens.assert_not_called() - - -def test_pd_decode_sums_request_tuples_for_non_decode_requests(): - queue = PDDecodeQueue.__new__(PDDecodeQueue) - queue.args = SimpleNamespace(page_size=16) - queue.dp_index = 0 - queue.is_busy = lambda: False - queue.router = SimpleNamespace(router_statics=SimpleNamespace(ema_req_out_len=10)) - reqs = [ - SimpleNamespace( - request_id="req-0", - input_len=20, - sample_params=SimpleNamespace(max_new_tokens=1, suggested_dp_index=0), - is_infer_decode=lambda: False, - get_tuple_tokens=MagicMock(return_value=(21, 34)), - ), - SimpleNamespace( - request_id="req-1", - input_len=100, - sample_params=SimpleNamespace(max_new_tokens=1, suggested_dp_index=1), - is_infer_decode=lambda: False, - get_tuple_tokens=MagicMock(return_value=(101, 34)), - ), - ] - batch = Batch(batch_id=1, reqs=reqs, dp_size_in_node=2) - - assert queue._caclu_batch_estimated_peak_token_num(batch) == 55 - reqs[0].get_tuple_tokens.assert_called_once_with(False, 10) - reqs[1].get_tuple_tokens.assert_not_called() - - -def test_pd_decode_uses_tuple_estimation(): - queue = PDDecodeQueue.__new__(PDDecodeQueue) - queue.args = SimpleNamespace(page_size=16) - queue.dp_index = 0 - queue.is_busy = lambda: False - queue.router = SimpleNamespace(router_statics=SimpleNamespace(ema_req_out_len=10)) - req = SimpleNamespace( - request_id="req-0", - sample_params=SimpleNamespace(suggested_dp_index=0), - is_infer_decode=lambda: True, - get_tuple_tokens=lambda is_busy, ema_req_out_len: (16, 32), - ) - batch = Batch(batch_id=1, reqs=[req], dp_size_in_node=1) - - assert queue._caclu_batch_estimated_peak_token_num(batch) == 48 diff --git a/unit_tests/server/router/test_kv_capacity.py b/unit_tests/server/router/test_kv_capacity.py deleted file mode 100644 index c243c07a49..0000000000 --- a/unit_tests/server/router/test_kv_capacity.py +++ /dev/null @@ -1,62 +0,0 @@ -from __future__ import annotations - -from types import SimpleNamespace - -import pytest - -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") - - -@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/router/test_profiler_service.py b/unit_tests/server/router/test_profiler_service.py deleted file mode 100644 index 020c495cb8..0000000000 --- a/unit_tests/server/router/test_profiler_service.py +++ /dev/null @@ -1,53 +0,0 @@ -import socket -import types - -import pytest -import rpyc - -from lightllm.server.router.profiler_service import RouterProfilerCmdQueue, start_router_profiler_server - - -def _free_port(): - with socket.socket() as s: - s.bind(("127.0.0.1", 0)) - return s.getsockname()[1] - - -@pytest.fixture -def profiler_server(monkeypatch): - port = _free_port() - from lightllm.utils import shm_port_args - - monkeypatch.setattr(shm_port_args, "get_shm_port_args", lambda: types.SimpleNamespace(router_profiler_port=port)) - queue = RouterProfilerCmdQueue() - server, _ = start_router_profiler_server(types.SimpleNamespace(enable_profiling="torch_profiler"), queue) - yield server, port, queue - server.close() - - -def test_profiler_server_binds_loopback_only(profiler_server): - server, _port, _queue = profiler_server - assert server.listener.getsockname()[0] == "127.0.0.1" - - -def test_profiler_server_accepts_known_cmds(profiler_server): - _server, port, queue = profiler_server - conn = rpyc.connect("127.0.0.1", port) - conn.root.profiler_cmd("start") - assert queue.pop() == "start" - conn.root.profiler_cmd("stop") - assert queue.pop() == "stop" - conn.close() - - -def test_profiler_server_rejects_unknown_cmd(profiler_server): - _server, port, queue = profiler_server - conn = rpyc.connect("127.0.0.1", port) - with pytest.raises(Exception): - conn.root.profiler_cmd("drop_tables") - assert queue.pop() is None - conn.close() - - -def test_disabled_profiling_starts_nothing(): - assert start_router_profiler_server(types.SimpleNamespace(enable_profiling=None), None) == (None, None) diff --git a/unit_tests/server/test_api_responses.py b/unit_tests/server/test_api_responses.py deleted file mode 100644 index 498904b613..0000000000 --- a/unit_tests/server/test_api_responses.py +++ /dev/null @@ -1,313 +0,0 @@ -import asyncio - -import pytest -import ujson as json - -from lightllm.server.api_responses import ( - _chat_response_to_responses, - _openai_sse_to_responses_events, - _responses_to_chat_request, -) -from lightllm.server.api_models import ChatCompletionRequest - - -async def _chunks(*payloads): - for payload in payloads: - yield f"data: {json.dumps(payload)}\n\n" - - -def _collect_events(*payloads): - async def collect(): - return [event async for event in _openai_sse_to_responses_events(_chunks(*payloads), {"input": "hi"})] - - raw_events = asyncio.run(collect()) - return [json.loads(event.decode().split("data: ", 1)[1]) for event in raw_events] - - -def test_function_call_arguments_done_includes_name(): - events = _collect_events( - { - "choices": [ - { - "delta": { - "tool_calls": [ - { - "id": "call_1", - "function": {"name": "get_weather", "arguments": '{"city":"Paris"}'}, - } - ] - } - } - ] - } - ) - - done = next(event for event in events if event["type"] == "response.function_call_arguments.done") - assert done["name"] == "get_weather" - - -def test_content_array_function_output_is_preserved(): - request = _responses_to_chat_request( - { - "input": [ - { - "type": "function_call_output", - "call_id": "call_1", - "output": [ - {"type": "input_text", "text": "sunny"}, - {"type": "input_image", "image_url": "https://example.com/weather.png"}, - ], - } - ] - } - ) - - assert request["messages"] == [ - { - "role": "tool", - "tool_call_id": "call_1", - "content": [ - {"type": "text", "text": "sunny"}, - {"type": "image_url", "image_url": {"url": "https://example.com/weather.png"}}, - ], - } - ] - - -def test_replayed_function_call_preserves_content_after_validation(): - request = _responses_to_chat_request( - { - "input": [ - { - "type": "function_call", - "call_id": "call_1", - "name": "get_weather", - "arguments": '{"city":"Paris"}', - } - ] - } - ) - - chat_request = ChatCompletionRequest(**request) - message = chat_request.messages[0].model_dump(exclude_none=True) - assert message["content"] == "" - - -def test_parallel_function_calls_share_one_assistant_message(): - request = _responses_to_chat_request( - { - "input": [ - { - "type": "function_call", - "call_id": "call_paris", - "name": "get_weather", - "arguments": '{"city":"Paris"}', - }, - { - "type": "function_call", - "call_id": "call_tokyo", - "name": "get_weather", - "arguments": '{"city":"Tokyo"}', - }, - {"type": "function_call_output", "call_id": "call_paris", "output": "sunny"}, - {"type": "function_call_output", "call_id": "call_tokyo", "output": "rainy"}, - ] - } - ) - - assert request["messages"] == [ - { - "role": "assistant", - "content": "", - "tool_calls": [ - { - "id": "call_paris", - "type": "function", - "function": {"name": "get_weather", "arguments": '{"city":"Paris"}'}, - }, - { - "id": "call_tokyo", - "type": "function", - "function": {"name": "get_weather", "arguments": '{"city":"Tokyo"}'}, - }, - ], - }, - {"role": "tool", "tool_call_id": "call_paris", "content": "sunny"}, - {"role": "tool", "tool_call_id": "call_tokyo", "content": "rainy"}, - ] - - -@pytest.mark.parametrize( - ("field", "part_type"), - [("content", "reasoning_text"), ("summary", "summary_text")], -) -def test_replayed_reasoning_is_attached_to_assistant_tool_call(field, part_type): - request = _responses_to_chat_request( - { - "input": [ - { - "type": "reasoning", - field: [{"type": part_type, "text": "I should check the weather."}], - }, - { - "type": "function_call", - "call_id": "call_1", - "name": "get_weather", - "arguments": '{"city":"Paris"}', - }, - {"type": "function_call_output", "call_id": "call_1", "output": "sunny"}, - ] - } - ) - - assert request["messages"][0]["reasoning_content"] == "I should check the weather." - assert request["messages"][0]["tool_calls"][0]["id"] == "call_1" - assert len(request["messages"]) == 2 - - -def test_replayed_reasoning_is_attached_to_assistant_message(): - request = _responses_to_chat_request( - { - "input": [ - { - "type": "reasoning", - "content": [{"type": "reasoning_text", "text": "I found the answer."}], - }, - { - "type": "message", - "role": "assistant", - "content": [{"type": "output_text", "text": "It is sunny."}], - }, - { - "type": "message", - "role": "user", - "content": [{"type": "input_text", "text": "Thanks"}], - }, - ] - } - ) - - assert request["messages"] == [ - { - "role": "assistant", - "content": [{"type": "text", "text": "It is sunny."}], - "reasoning_content": "I found the answer.", - }, - {"role": "user", "content": [{"type": "text", "text": "Thanks"}]}, - ] - - -def test_non_streamed_truncated_function_call_is_incomplete(): - response = _chat_response_to_responses( - { - "choices": [ - { - "finish_reason": "length", - "message": { - "tool_calls": [ - { - "id": "call_1", - "function": {"name": "get_weather", "arguments": '{"city":'}, - } - ] - }, - } - ] - }, - {"input": "hi"}, - ) - - assert response["status"] == "incomplete" - assert response["output"][0]["status"] == "incomplete" - - -def test_streamed_truncated_function_call_is_incomplete(): - events = _collect_events( - { - "choices": [ - { - "delta": { - "tool_calls": [ - { - "id": "call_1", - "function": {"name": "get_weather", "arguments": '{"city":'}, - } - ] - } - } - ] - }, - {"choices": [{"delta": {}, "finish_reason": "length"}]}, - ) - - item_done = next(event for event in events if event["type"] == "response.output_item.done") - response_done = events[-1] - assert item_done["item"]["status"] == "incomplete" - assert response_done["type"] == "response.incomplete" - assert response_done["response"]["output"][0]["status"] == "incomplete" - - -def test_streamed_message_adds_text_content_once(): - events = _collect_events({"choices": [{"delta": {"content": "hello"}, "finish_reason": "stop"}]}) - - item_added = next(event for event in events if event["type"] == "response.output_item.added") - part_added = next(event for event in events if event["type"] == "response.content_part.added") - assert item_added["item"]["content"] == [] - assert part_added["part"] == {"type": "output_text", "text": "", "annotations": []} - - -def test_route_maps_downstream_value_error_to_bad_request(monkeypatch): - from types import SimpleNamespace - - from lightllm.server import api_http, api_responses - - async def raise_value_error(raw_request): - raise ValueError("Unrecognized image input.") - - monkeypatch.setattr(api_http, "get_env_start_args", lambda: SimpleNamespace(run_mode="normal")) - monkeypatch.setattr(api_http.g_objs, "metric_client", SimpleNamespace(counter_inc=lambda *args: None)) - monkeypatch.setattr(api_responses, "responses_impl", raise_value_error) - - response = asyncio.run(api_http.openai_responses(None)) - assert response.status_code == 400 - - -@pytest.mark.parametrize( - "partial_payload", - [ - {"choices": [{"delta": {"content": "partial"}}]}, - { - "choices": [ - { - "delta": { - "tool_calls": [ - { - "id": "call_1", - "function": {"name": "get_weather", "arguments": '{"city":'}, - } - ] - } - } - ] - }, - ], -) -def test_stream_failure_does_not_complete_partial_item(partial_payload): - events = _collect_events(partial_payload, {"error": {"message": "generation failed"}}) - event_types = [event["type"] for event in events] - - assert "response.output_text.done" not in event_types - assert "response.function_call_arguments.done" not in event_types - assert "response.output_item.done" not in event_types - assert event_types[-1] == "response.failed" - - -@pytest.mark.parametrize("effort", ["low", "medium", "high", "none", "minimal", "xhigh"]) -def test_reasoning_effort_is_rejected(effort): - with pytest.raises(ValueError, match="reasoning.effort is not supported"): - _responses_to_chat_request({"input": "hi", "reasoning": {"effort": effort}}) - - -def test_automatic_truncation_is_rejected(): - with pytest.raises(ValueError, match="truncation"): - _responses_to_chat_request({"input": "hi", "truncation": "auto"}) diff --git a/unit_tests/server/test_api_start_deepep.py b/unit_tests/server/test_api_start_deepep.py deleted file mode 100644 index 949021cd34..0000000000 --- a/unit_tests/server/test_api_start_deepep.py +++ /dev/null @@ -1,45 +0,0 @@ -import json - -import pytest - -from lightllm.server.api_start import _launch_subprocesses -from lightllm.server.core.objs.start_args_type import StartArgs - - -def test_decode_node_batch_max_tokens_covers_mtp_rows(monkeypatch, tmp_path): - for name in ( - "_set_envs_and_config", - "auto_set_max_req_total_len", - "auto_set_fused_shared_experts", - "set_unique_server_name", - "auto_set_page_size", - ): - monkeypatch.setattr(f"lightllm.server.api_start.{name}", lambda args: None) - - def stop(args): - raise RuntimeError("startup validation passed") - - monkeypatch.setattr("lightllm.server.api_start.auto_set_response_parsers", stop) - (tmp_path / "config.json").write_text(json.dumps({"model_type": "llama"})) - args = StartArgs( - run_mode="decode", - model_dir=str(tmp_path), - enable_ep_moe=True, - running_max_req_size=256, - mtp_mode="vanilla_no_att", - mtp_step=2, - max_req_total_len=4096, - batch_max_tokens=4096, - chunked_prefill_size=2048, - eos_id=[0], - data_type="float16", - disable_vision=True, - disable_audio=True, - disable_shm_warning=True, - router_token_ratio=0.0, - ) - - with pytest.raises(RuntimeError, match="startup validation passed"): - _launch_subprocesses(args) - - assert args.batch_max_tokens == 768 diff --git a/unit_tests/server/test_application_stop.py b/unit_tests/server/test_application_stop.py deleted file mode 100644 index 0603c29c67..0000000000 --- a/unit_tests/server/test_application_stop.py +++ /dev/null @@ -1,200 +0,0 @@ -import asyncio -import json -from types import SimpleNamespace -from unittest.mock import AsyncMock - -import pytest - -from lightllm.server.api_lightllm import lightllm_generate, lightllm_generate_stream -from lightllm.server.api_tgi import tgi_generate_impl, tgi_generate_stream_impl -from lightllm.server.core.objs import FinishStatus, SamplingParams, StopStrMatchMode -from lightllm.server.httpserver.manager import HttpServerManager -from lightllm.server.pd_io_struct import NodeRole - - -def _manager(tokenizer, args, raw_results): - manager = HttpServerManager.__new__(HttpServerManager) - manager.output_grammar_compiler = None - manager.tokenizer = tokenizer - manager.args = args - manager.pd_mode = NodeRole.NORMAL - manager.is_multinode_tp_slave = False - manager._generate = raw_results - manager.stop_request = AsyncMock() - return manager - - -@pytest.mark.parametrize("stream", [False, True]) -def test_native_api_discards_output_after_application_stop(stream): - async def run(): - closed = [] - exhausted = [] - - async def raw_results(*args, **kwargs): - metadata = {"prompt_tokens": 2, "prompt_cache_len": 0, "logprobs": {}, "logprob": -0.1} - try: - yield 80, "helloE", {**metadata, "id": 10}, FinishStatus() - yield 80, "NDextra", {**metadata, "id": 11}, FinishStatus() - yield 80, "late", {**metadata, "id": 12}, FinishStatus(FinishStatus.FINISHED_LENGTH) - exhausted.append(True) - finally: - closed.append(True) - - manager = _manager( - SimpleNamespace(encode=lambda text, **kwargs: [ord(c) for c in text]), - SimpleNamespace(reasoning_parser=None, enable_rl=False, enable_return_routed_experts=False), - raw_results, - ) - request = SimpleNamespace( - json=AsyncMock( - return_value={ - "inputs": "prompt", - "parameters": {"stop_sequences": ["END"], "max_new_tokens": 10}, - } - ) - ) - response = await (lightllm_generate_stream if stream else lightllm_generate)(request, manager) - if stream: - events = [json.loads(chunk.decode().split(":", 1)[1]) async for chunk in response.body_iterator] - assert "".join(event["token"]["text"] for event in events) == "hello" - assert events[-1]["finished"] - assert events[-1]["finish_reason"] == "stop" - else: - data = json.loads(response.body) - assert data["generated_text"] == ["hello"] - assert data["count_output_tokens"] == 2 - assert data["finish_reason"] == "stop" - manager.stop_request.assert_awaited_once_with(80) - assert closed == [True] - assert exhausted == [True] - - asyncio.run(run()) - - -@pytest.mark.parametrize("stream", [False, True]) -@pytest.mark.parametrize("include", [False, True]) -def test_tgi_api_uses_shared_stop_wrapper(stream, include): - async def run(): - async def raw_results(*args, **kwargs): - metadata = {"prompt_tokens": 2, "logprobs": {}, "logprob": -0.1, "cumlogprob": -0.1} - yield 80, "helloE", {**metadata, "id": 10}, FinishStatus() - yield 80, "NDextra", {**metadata, "id": 11}, FinishStatus() - yield 80, "late", {**metadata, "id": 12}, FinishStatus(FinishStatus.FINISHED_LENGTH) - - manager = _manager( - SimpleNamespace(encode=lambda text, **kwargs: [ord(c) for c in text]), - SimpleNamespace(reasoning_parser=None, enable_rl=False, enable_return_routed_experts=False), - raw_results, - ) - request = SimpleNamespace( - json=AsyncMock( - return_value={ - "inputs": "prompt", - "parameters": {"stop": ["END"], "details": True, "include_stop_str_in_output": include}, - } - ) - ) - response = await (tgi_generate_stream_impl if stream else tgi_generate_impl)(request, manager) - expected = "helloEND" if include else "hello" - if stream: - events = [] - async for chunk in response.body_iterator: - chunk = chunk.decode() if isinstance(chunk, bytes) else chunk - events.append(json.loads(chunk.split(":", 1)[1])) - assert "".join(event["token"]["text"] for event in events) == expected - assert events[-1]["finish_reason"] == "stop" - else: - data = json.loads(response.body) - data = data[0] if isinstance(data, list) else data - assert data["generated_text"] == expected - assert data["details"]["finish_reason"] == "stop" - manager.stop_request.assert_awaited_once_with(80) - - asyncio.run(run()) - - -@pytest.mark.parametrize("mode", [NodeRole.NORMAL, NodeRole.P, NodeRole.D]) -@pytest.mark.parametrize("slave", [False, True]) -def test_http_generate_only_wraps_client_output(mode, slave): - async def run(): - tokenizer = SimpleNamespace(encode=lambda text, **kwargs: [ord(c) for c in text]) - params = SamplingParams() - params.init(tokenizer, stop_sequences="END", max_new_tokens=10) - websocket, pd_event = object(), object() - - async def raw_results(prompt, sampling_params, multimodal, request, upload, event): - assert upload is websocket - assert event is pd_event - yield 80, "helloENDextra", {"id": 10}, FinishStatus(FinishStatus.FINISHED_STOP) - - manager = _manager( - tokenizer, - SimpleNamespace(reasoning_parser=None, enable_rl=False, enable_return_routed_experts=False), - raw_results, - ) - manager.pd_mode = mode - manager.is_multinode_tp_slave = slave - results = [result async for result in manager.generate("prompt", params, None, None, websocket, pd_event)] - expected = "hello" if mode == NodeRole.NORMAL and not slave else "helloENDextra" - assert "".join(result[1] for result in results) == expected - manager.stop_request.assert_not_awaited() - - asyncio.run(run()) - - -@pytest.mark.parametrize("mode", [StopStrMatchMode.ALL_TEXT, StopStrMatchMode.CONTENT_ONLY]) -@pytest.mark.parametrize("feature", ["enable_rl", "enable_return_routed_experts", "prompt_logprobs"]) -@pytest.mark.parametrize("stop", [None, [], "END", [[10, 11]], ["END", [10, 11]]]) -@pytest.mark.parametrize("include", [False, True]) -@pytest.mark.parametrize("split", [False, True]) -def test_final_metadata_stop_forwards_late_tokens_and_backend_finish(mode, feature, stop, include, split): - async def run(): - tokenizer = SimpleNamespace(encode=lambda text, **kwargs: [ord(c) for c in text]) - params = SamplingParams() - params.init(tokenizer, stop_sequences=stop, stop_str_match_mode=mode, include_stop_str_in_output=include) - if feature == "prompt_logprobs": - params.prompt_logprobs = 0 - metadata = {"id": 13, "prompt_token_ids": [1, 2], "prompt_logprobs": [None, {}], "routed_experts": [[2]]} - - async def raw_results(*args, **kwargs): - yield 80, "helloE" if split else "hello", {"id": 10}, FinishStatus() - yield 80, "NDextra" if split else "ENDextra", {"id": 11}, FinishStatus() - yield 80, "lateEND", {"id": 12}, FinishStatus() - yield 80, "tail", metadata, FinishStatus(FinishStatus.FINISHED_ABORTED) - yield 80, "after finish", {"id": 14}, FinishStatus() - - manager = _manager( - tokenizer, - SimpleNamespace( - reasoning_parser=None, - enable_rl=feature == "enable_rl", - enable_return_routed_experts=feature == "enable_return_routed_experts", - ), - raw_results, - ) - results = [result async for result in manager.generate("prompt", params, None, None)] - first_text = "helloE" if split else "hello" - hit_text = "NDextra" if split else "ENDextra" - if stop: - if params.stop_sequences.has_stop_strs(): - first_text = first_text if include else "hello" - hit_text = ("ND" if split else "END") if include else "" - assert [result[1] for result in results] == [ - first_text, - hit_text, - "lateEND", - "tail", - ] - assert [result[2]["id"] for result in results] == [10, 11, 12, 13] - expected_stops = ( - 2 if params.stop_sequences.has_stop_strs() and params.stop_sequences.to_token_ids_list() else 1 - ) - assert [call.args for call in manager.stop_request.await_args_list] == [(80,)] * expected_stops - else: - assert [result[1] for result in results] == [first_text, hit_text, "lateEND", "tail", "after finish"] - manager.stop_request.assert_not_awaited() - assert all(not result[3].is_finished() for result in results[:3]) - assert results[3][2] is metadata - assert results[3][3].status == FinishStatus.FINISHED_ABORTED - - asyncio.run(run()) diff --git a/unit_tests/server/test_deepseek_dsml_stream_args.py b/unit_tests/server/test_deepseek_dsml_stream_args.py deleted file mode 100644 index 0e23fce189..0000000000 --- a/unit_tests/server/test_deepseek_dsml_stream_args.py +++ /dev/null @@ -1,72 +0,0 @@ -import json - -import pytest - -from lightllm.server.api_models import Tool -from lightllm.server.function_call_parser import FunctionCallParser - - -TOOLS = [ - Tool.model_validate( - { - "type": "function", - "function": { - "name": "write_file", - "parameters": { - "type": "object", - "properties": {"path": {"type": "string"}, "content": {"type": "string"}}, - }, - }, - } - ), - Tool.model_validate( - { - "type": "function", - "function": { - "name": "get_weather", - "parameters": { - "type": "object", - "properties": {"location": {"type": "string"}, "days": {"type": "integer"}}, - }, - }, - } - ), -] - -DSML = "｜DSML｜" - - -def dsml_text(wrapper): - return ( - f"<{DSML}{wrapper}>\n" - f'<{DSML}invoke name="write_file">\n' - f'<{DSML}parameter name="path" string="true">src/app.py\n' - f'<{DSML}parameter name="content" string="true">print("hi")\n\n' - f"\n" - f'<{DSML}invoke name="get_weather">\n' - f'<{DSML}parameter name="location" string="true">Hangzhou\n' - f'<{DSML}parameter name="days" string="false">3\n' - f"\n" - f"" - ) - - -def collect_stream(parser_name, text, chunk_size): - parser = FunctionCallParser(TOOLS, parser_name) - calls = {} - for start in range(0, len(text), chunk_size): - _, deltas = parser.parse_stream_chunk(text[start : start + chunk_size]) - for delta in deltas: - if delta.name is not None: - calls[delta.tool_index] = {"name": delta.name, "arguments": ""} - calls[delta.tool_index]["arguments"] += delta.parameters - return calls - - -@pytest.mark.parametrize("parser_name,wrapper", [("deepseekv32", "function_calls"), ("deepseekv4", "tool_calls")]) -@pytest.mark.parametrize("chunk_size", [1, 5, 16]) -def test_streamed_arguments_concatenate_to_valid_json(parser_name, wrapper, chunk_size): - calls = collect_stream(parser_name, dsml_text(wrapper), chunk_size) - assert [call["name"] for call in calls.values()] == ["write_file", "get_weather"] - assert json.loads(calls[0]["arguments"]) == {"path": "src/app.py", "content": 'print("hi")\n'} - assert json.loads(calls[1]["arguments"]) == {"location": "Hangzhou", "days": 3} diff --git a/unit_tests/server/test_deepseek_tool_parser_stream.py b/unit_tests/server/test_deepseek_tool_parser_stream.py deleted file mode 100644 index e18eb9a005..0000000000 --- a/unit_tests/server/test_deepseek_tool_parser_stream.py +++ /dev/null @@ -1,143 +0,0 @@ -import json - -import pytest - -from lightllm.server.api_models import Tool -from lightllm.server.function_call_parser import FunctionCallParser - -TOOLS = [ - Tool.model_validate( - { - "type": "function", - "function": { - "name": "write_file", - "parameters": {"type": "object", "properties": {"content": {"type": "string"}}}, - }, - } - ), - Tool.model_validate( - { - "type": "function", - "function": { - "name": "get_weather", - "parameters": {"type": "object", "properties": {"location": {"type": "string"}}}, - }, - } - ), -] - -DSML = "｜DSML｜" - - -def _dsml_invoke(name, param, value): - return ( - f'<{DSML}invoke name="{name}">\n' - f'<{DSML}parameter name="{param}" string="true">{value}\n' - f"" - ) - - -def _stream(parser, deltas): - names = [] - arguments = {} - for delta in deltas: - _, calls = parser.parse_stream_chunk(delta) - for call in calls: - if call.name: - names.append(call.name) - arguments[call.tool_index] = arguments.get(call.tool_index, "") + (call.parameters or "") - return names, arguments - - -@pytest.mark.parametrize("tool_call_parser,block", [("deepseekv32", "function_calls"), ("deepseekv4", "tool_calls")]) -def test_dsml_streams_long_invokes_in_small_deltas(tool_call_parser, block): - content = "print(1)\n" * 500 - deltas = [f"<{DSML}{block}>", "\n", f"<{DSML}invoke", ' name="write_file">\n', f"<{DSML}parameter"] - deltas += [' name="content" string="true">'] - deltas += [content[i : i + 4] for i in range(0, len(content), 4)] - deltas += [f"", "\n", f"", "\n"] - deltas += [f"<{DSML}invoke", ' name="get_weather">\n', f"<{DSML}parameter", ' name="location" string="true">'] - deltas += ["Hang", "zhou", f"", "\n", f"", "\n", f""] - - parser = FunctionCallParser(TOOLS, tool_call_parser) - names, arguments = _stream(parser, deltas) - - assert names == ["write_file", "get_weather"] - assert json.loads(arguments[0]) == {"content": content} - assert json.loads(arguments[1]) == {"location": "Hangzhou"} - - -def test_dsml_emits_buffered_invoke_on_next_delta(): - # Both invokes arrive in one delta; the second one is still in the buffer - # and must be emitted by the next delta, which has no invoke end tag. - text = f"<{DSML}function_calls>\n" - text += _dsml_invoke("write_file", "content", "x = 1") + "\n" - text += _dsml_invoke("get_weather", "location", "Paris") + "\n" - - parser = FunctionCallParser(TOOLS, "deepseekv32") - names, arguments = _stream(parser, [text, f""]) - - assert names == ["write_file", "get_weather"] - assert json.loads(arguments[0]) == {"content": "x = 1"} - assert json.loads(arguments[1]) == {"location": "Paris"} - - -def _v3_call(name, payload): - return "<｜tool▁call▁begin｜>function<｜tool▁sep｜>" + name + "\n```json\n" + payload + "\n```<｜tool▁call▁end｜>" - - -def test_deepseekv3_streams_long_arguments_in_small_deltas(): - arguments = {"content": "print(1)\n" * 500} - payload = json.dumps(arguments) - call = _v3_call("write_file", payload) - deltas = ["<｜tool▁calls▁begin｜>"] + [call[i : i + 4] for i in range(0, len(call), 4)] - - parser = FunctionCallParser(TOOLS, "deepseekv3") - names, streamed = _stream(parser, deltas) - - assert names == ["write_file"] - assert json.loads(streamed[0]) == arguments - - -def test_deepseekv3_holds_arguments_until_the_json_fence_closes(): - payload = json.dumps({"content": "x" * 200}) - head = "<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>write_file\n```json\n" - parser = FunctionCallParser(TOOLS, "deepseekv3") - names, streamed = _stream(parser, [head] + [payload[i : i + 4] for i in range(0, len(payload), 4)]) - assert names == [] - assert streamed == {} - - # The closing fence is split across deltas. The name is emitted on the delta that - # completes it, and the arguments follow on the next delta, same as one regex search. - names, streamed = _stream(parser, ["\n", "``", "`<｜tool▁call▁end｜>", "<｜tool▁calls▁end｜>"]) - assert names == ["write_file"] - assert json.loads(streamed[0]) == {"content": "x" * 200} - - -def test_deepseekv3_keeps_a_fence_inside_the_argument(): - content = "line1\n```\nline2\n```python\ncode" - payload = json.dumps({"content": content}) - text = "<｜tool▁calls▁begin｜>" + _v3_call("write_file", payload) - deltas = [text[: len("<｜tool▁calls▁begin｜>")]] + [ - text[i : i + 4] for i in range(len("<｜tool▁calls▁begin｜>"), len(text), 4) - ] - - parser = FunctionCallParser(TOOLS, "deepseekv3") - names, streamed = _stream(parser, deltas) - - assert names == ["write_file"] - assert json.loads(streamed[0]) == {"content": content} - - -def test_deepseekv31_streams_long_arguments_in_small_deltas(): - arguments = {"content": "print(1)\n" * 500} - payload = json.dumps(arguments) - deltas = ["<｜tool▁calls▁begin｜>", "<｜tool▁call▁begin｜>", "write_file", "<｜tool▁sep｜>"] - deltas += [payload[i : i + 4] for i in range(0, len(payload), 4)] - deltas += ["<｜tool▁call▁end｜>", "<｜tool▁calls▁end｜>"] - - parser = FunctionCallParser(TOOLS, "deepseekv31") - names, streamed = _stream(parser, deltas) - - assert names == ["write_file"] - assert json.loads(streamed[0]) == arguments diff --git a/unit_tests/server/test_detokenization.py b/unit_tests/server/test_detokenization.py deleted file mode 100644 index 4d12363290..0000000000 --- a/unit_tests/server/test_detokenization.py +++ /dev/null @@ -1,114 +0,0 @@ -"""Incremental detokenization must ignore virtual image/audio cache IDs.""" - -import unittest -from types import SimpleNamespace -from unittest.mock import patch - -import numpy as np -from tokenizers import Tokenizer, decoders, models, pre_tokenizers -from transformers import PreTrainedTokenizerFast - -from lightllm.server.detokenization import decode_req as decode_req_module -from lightllm.server.detokenization.decode import decode_token -from lightllm.server.detokenization.decode_req import DecodeReq -from lightllm.server.embed_cache.utils import MULTIMODAL_TOKEN_ID_START - - -class TestDetokenization(unittest.TestCase): - def setUp(self): - vocab = {char: i for i, char in enumerate(sorted(pre_tokenizers.ByteLevel.alphabet()))} - backend = Tokenizer(models.BPE(vocab=vocab, merges=[])) - backend.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False) - backend.decoder = decoders.ByteLevel() - self.tokenizer = PreTrainedTokenizerFast(tokenizer_object=backend) - self.tokenizer.add_special_tokens({"additional_special_tokens": [""]}) - self.tokenizer.add_tokens(["added_text"]) - prefix_length = patch.object(decode_req_module, "LIGHTLLM_DECODE_PREFIX_LENGTH", 5) - prefix_length.start() - self.addCleanup(prefix_length.stop) - - def make_decode_req(self, prompt_ids, output_ids, **sampling_params): - params = dict(print_eos_token=False, skip_special_tokens=True, add_spaces_between_special_tokens=False) - params.update(sampling_params) - req = SimpleNamespace( - request_id=1, - group_req_id=1, - input_len=len(prompt_ids), - shm_prompt_ids=SimpleNamespace(arr=np.array(prompt_ids + output_ids, dtype=np.int64)), - sample_params=SimpleNamespace(**params), - ) - return DecodeReq(req) - - def decode_stream(self, decode_req, output_ids): - pieces = [] - for token_id in output_ids: - decode_req.output_ids.append(token_id) - pieces.append(decode_token(self.tokenizer, decode_req, token_id, [])) - return pieces - - def test_virtual_prompt_ids_at_boundary_and_large_values(self): - output_ids = self.tokenizer.encode("hello world", add_special_tokens=False) - for virtual_id in (MULTIMODAL_TOKEN_ID_START, 2 ** 31, 2 ** 32, 2 ** 63 - 2): - with self.subTest(virtual_id=virtual_id): - prompt_ids = self.tokenizer.encode("prompt", add_special_tokens=False) + [virtual_id] - decode_req = self.make_decode_req(prompt_ids, output_ids) - self.assertEqual("".join(self.decode_stream(decode_req, output_ids)), "hello world") - - def test_five_token_prompt_tail_and_partial_utf8_keep_original_positions(self): - prompt_ids = self.tokenizer.encode("prompt", add_special_tokens=False) + [2 ** 40 + i for i in range(5)] - output_ids = self.tokenizer.encode("你 world", add_special_tokens=False) - decode_req = self.make_decode_req(prompt_ids, output_ids) - original_tokens = decode_req.req.shm_prompt_ids.arr.copy() - initial_offsets = (len(prompt_ids) - 5, len(prompt_ids)) - self.assertEqual((decode_req.prefix_offset, decode_req.read_offset), initial_offsets) - - pieces = [] - for index, token_id in enumerate(output_ids): - shared_token_id, shared_index = decode_req.get_next_token_id_and_index() - self.assertEqual(shared_token_id, token_id) - self.assertEqual(shared_index, len(prompt_ids) + index) - decode_req.output_ids.append(shared_token_id) - pieces.append(decode_token(self.tokenizer, decode_req, token_id, [])) - if index < 2: - self.assertEqual((decode_req.prefix_offset, decode_req.read_offset), initial_offsets) - elif index == 2: - self.assertEqual( - (decode_req.prefix_offset, decode_req.read_offset), (len(prompt_ids), len(prompt_ids) + 3) - ) - - self.assertEqual(pieces[:2], ["", ""]) - self.assertEqual("".join(pieces), "你 world") - self.assertEqual(decode_req.prompt_ids, prompt_ids) - self.assertEqual(decode_req.output_ids, output_ids) - np.testing.assert_array_equal(decode_req.req.shm_prompt_ids.arr, original_tokens) - self.assertEqual(decode_req.read_offset, len(original_tokens)) - self.assertEqual(decode_req.prefix_offset, len(original_tokens) - 1) - - def test_added_and_special_tokens_keep_tokenizer_options(self): - output_ids = self.tokenizer.encode("hello", add_special_tokens=False) - output_ids += self.tokenizer.convert_tokens_to_ids(["", "added_text"]) - for skip_special_tokens in (False, True): - for add_spaces in (False, True): - with self.subTest(skip_special_tokens=skip_special_tokens, add_spaces=add_spaces): - decode_req = self.make_decode_req( - [2 ** 40], - output_ids, - skip_special_tokens=skip_special_tokens, - add_spaces_between_special_tokens=add_spaces, - ) - expected = self.tokenizer.decode( - output_ids, - skip_special_tokens=skip_special_tokens, - spaces_between_special_tokens=add_spaces, - ) - self.assertEqual("".join(self.decode_stream(decode_req, output_ids)), expected) - - def test_text_only_prompt(self): - prompt_ids = self.tokenizer.encode("prompt", add_special_tokens=False) - output_ids = self.tokenizer.encode("你 world", add_special_tokens=False) - decode_req = self.make_decode_req(prompt_ids, output_ids) - self.assertEqual("".join(self.decode_stream(decode_req, output_ids)), "你 world") - - -if __name__ == "__main__": - unittest.main() diff --git a/unit_tests/server/test_dp_request_capacity.py b/unit_tests/server/test_dp_request_capacity.py deleted file mode 100644 index 800f0cdda3..0000000000 --- a/unit_tests/server/test_dp_request_capacity.py +++ /dev/null @@ -1,109 +0,0 @@ -from types import SimpleNamespace -from unittest.mock import MagicMock - -import pytest - -from lightllm.server.core.objs.shm_req_manager import ShmReqManager -from lightllm.server.core.objs.start_args_type import StartArgs -from lightllm.server.router.batch import Batch -from lightllm.server.router.req_queue.chunked_prefill.impl import ChunkedPrefillQueue -from lightllm.server.router.req_queue.dp_base_queue import DpQueue -from lightllm.utils.config_utils import get_running_max_req_size_per_dp - - -@pytest.mark.parametrize( - "global_capacity,dp,nnodes,expected", - [(32, 8, 1, 4), (33, 8, 1, 5), (4, 8, 1, 1), (32, 8, 2, 8), (32, 1, 2, 32), (32, 1, 1, 32)], -) -def test_hybrid_capacity_uses_local_dp_count(monkeypatch, global_capacity, dp, nnodes, expected): - monkeypatch.setattr("lightllm.utils.config_utils.is_hybrid_att_model", lambda _: True) - args = StartArgs(running_max_req_size=global_capacity, dp=dp, nnodes=nnodes) - assert get_running_max_req_size_per_dp(args) == expected - # HTTP/SHM indexes remain global even when GPU request indexes are local. - monkeypatch.setattr("lightllm.server.core.objs.shm_req_manager.get_env_start_args", lambda: args) - assert ShmReqManager.get_max_req_num(None) == global_capacity - - -@pytest.mark.parametrize("mode", ["full_attention", "diverse", "cache_fetch"]) -def test_modes_requiring_existing_capacity_keep_global_default(monkeypatch, mode): - monkeypatch.setattr("lightllm.utils.config_utils.is_hybrid_att_model", lambda _: mode != "full_attention") - args = StartArgs( - running_max_req_size=32, - dp=8, - diverse_mode=mode == "diverse", - enable_dp_prompt_cache_fetch=mode == "cache_fetch", - ) - assert get_running_max_req_size_per_dp(args) == 32 - - -@pytest.mark.parametrize("capacity", [0, -1]) -def test_reject_invalid_global_capacity(capacity): - args = StartArgs(running_max_req_size=capacity, dp=8) - with pytest.raises(ValueError, match="running_max_req_size"): - get_running_max_req_size_per_dp(args) - - -def _make_dp_queue(monkeypatch, balancer): - monkeypatch.setattr("lightllm.utils.config_utils.is_hybrid_att_model", lambda _: True) - args = StartArgs( - running_max_req_size=32, - dp=8, - max_total_token_num=1048576, - batch_max_tokens=4096, - router_token_ratio=0.85, - dp_balancer=balancer, - ) - router = SimpleNamespace( - router_statics=SimpleNamespace(ema_req_out_len=16), - get_used_tokens=lambda dp: 0, - shared_token_load=MagicMock(), - ) - return DpQueue(args, router, ChunkedPrefillQueue, dp_size_in_node=8) - - -def _req(request_id, dp=-1): - return SimpleNamespace( - request_id=request_id, - sample_params=SimpleNamespace(suggested_dp_index=dp, high_priority_request=False), - is_aborted=False, - get_tuple_tokens=lambda busy, ema: (64, 16), - get_first_router_need_tokens=lambda: 64, - get_decode_need_tokens=lambda: 3, - ) - - -@pytest.mark.parametrize("balancer", ["bs_balancer", "round_robin"]) -def test_global_32_requests_run_as_four_per_rank(monkeypatch, balancer): - queue = _make_dp_queue(monkeypatch, balancer) - for i in range(32): - queue.extend([_req(i)]) - batch = queue.generate_new_batch(None) - assert batch.get_all_dp_req_num() == [4] * 8 - assert queue.get_wait_req_num() == 0 - - -def test_pinned_overflow_waits_until_a_local_slot_is_released(monkeypatch): - queue = _make_dp_queue(monkeypatch, "bs_balancer") - # A skewed client can use any global SHM slot, but only four may enter - # inference on this rank. The fifth must not exhaust its GPU request pool. - pinned = [_req(i, dp=2) for i in range(6)] - for req in pinned: - queue.extend([req]) - queue.extend([_req(6, dp=3)]) - batch = queue.generate_new_batch(None) - assert batch.get_all_dp_req_num() == [0, 0, 4, 1, 0, 0, 0, 0] - assert queue.get_wait_req_num() == 2 - assert queue.generate_new_batch(batch) is None - # Releasing a request on a different rank must not open a slot on rank 2. - batch.pop_req(6) - assert queue.generate_new_batch(batch) is None - batch.pop_req(0) - resumed = queue.generate_new_batch(batch) - assert [req.request_id for req in resumed.reqs] == [4] - batch.merge(resumed) - assert batch.get_all_dp_req_num()[2] == 4 - assert queue.get_wait_req_num() == 1 - assert queue.generate_new_batch(batch) is None - rest = queue.generate_new_batch(Batch(-1, [], dp_size_in_node=8)) - assert [req.request_id for req in rest.reqs] == [5] - assert queue.get_wait_req_num() == 0 diff --git a/unit_tests/server/test_function_call_common_prefix.py b/unit_tests/server/test_function_call_common_prefix.py deleted file mode 100644 index 498f5d45a2..0000000000 --- a/unit_tests/server/test_function_call_common_prefix.py +++ /dev/null @@ -1,67 +0,0 @@ -import json - -import pytest - -from lightllm.server.api_models import Tool -from lightllm.server.function_call_parser import FunctionCallParser, _find_common_prefix - - -def _reference_common_prefix(s1, s2): - prefix = "" - for a, b in zip(s1, s2): - if a != b: - break - prefix += a - return prefix - - -@pytest.mark.parametrize( - "s1,s2", - [ - ("", ""), - ("", "abc"), - ("abc", ""), - ("abc", "abc"), - ("abc", "abcd"), - ("abcd", "abc"), - ("abc", "abd"), - ("xbc", "abc"), - ('{"a": "x"}', '{"a": "xy"}'), - ('{"a": "é😀"}', '{"a": "é😀b"}'), - ("a" * 1000 + "b", "a" * 1000 + "c"), - ], -) -def test_find_common_prefix_matches_reference(s1, s2): - assert _find_common_prefix(s1, s2) == _reference_common_prefix(s1, s2) - - -def test_qwen25_streams_long_arguments_in_small_deltas(): - tools = [ - Tool.model_validate( - { - "type": "function", - "function": { - "name": "write_file", - "parameters": {"type": "object", "properties": {"content": {"type": "string"}}}, - }, - } - ) - ] - arguments = {"content": "print(1)\n" * 500} - payload = json.dumps(arguments) - deltas = ["", '\n{"name": "write_file", "arguments": '] - deltas += [payload[i : i + 4] for i in range(0, len(payload), 4)] - deltas += ["}\n", ""] - - parser = FunctionCallParser(tools, "qwen25") - streamed = "" - names = [] - for delta in deltas: - _, calls = parser.parse_stream_chunk(delta) - for call in calls: - if call.name: - names.append(call.name) - streamed += call.parameters or "" - - assert names == ["write_file"] - assert json.loads(streamed) == arguments diff --git a/unit_tests/server/test_generation_error_response.py b/unit_tests/server/test_generation_error_response.py deleted file mode 100644 index 87670d5bf0..0000000000 --- a/unit_tests/server/test_generation_error_response.py +++ /dev/null @@ -1,91 +0,0 @@ -import json -from types import SimpleNamespace -from unittest.mock import AsyncMock, Mock - -import pytest - -from fastapi import FastAPI -from fastapi.testclient import TestClient - -from lightllm.server import api_openai, api_stream_obj -from lightllm.server.api_models import ChatCompletionRequest -from lightllm.server.core.objs import FinishStatus, SamplingParams, StartArgs -from lightllm.server.api_http import g_objs -from lightllm.server.api_http import generation_exception_handler -from lightllm.utils.error_utils import GenerationError - - -def test_generation_failure_before_first_chunk_returns_http_500(monkeypatch): - monkeypatch.setattr(g_objs, "metric_client", Mock()) - monkeypatch.setattr(api_stream_obj, "get_env_start_args", lambda: StartArgs()) - app = FastAPI() - app.add_exception_handler(GenerationError, generation_exception_handler) - - @app.get("/generate") - async def generate(): - async def failed_stream(): - raise GenerationError("Generation failed before producing output") - yield - - return api_stream_obj.CustomStreamingResponse(failed_stream(), media_type="text/event-stream") - - response = TestClient(app).get("/generate") - assert response.status_code == 500 - assert "Generation failed" in response.json()["error"]["message"] - - -@pytest.mark.parametrize( - "early_headers, finish_status, expected_status", - [ - (False, FinishStatus.FINISHED_ERROR, 500), - (True, FinishStatus.FINISHED_ERROR, 200), - (False, FinishStatus.FINISHED_STOP, 200), - (True, FinishStatus.FINISHED_STOP, 200), - ], -) -def test_chat_first_result_sets_http_status(monkeypatch, early_headers, finish_status, expected_status): - args = StartArgs(disable_delay_response_start=early_headers) - - async def generate(*_args, **_kwargs): - yield 0, "answer", {"prompt_tokens": 2, "prompt_cache_len": 0}, FinishStatus(finish_status) - - monkeypatch.setattr(g_objs, "metric_client", Mock()) - monkeypatch.setattr(g_objs, "args", args) - monkeypatch.setattr( - g_objs, - "httpserver_manager", - SimpleNamespace(tokenizer=Mock(), output_grammar_compiler=None, generate=generate), - ) - monkeypatch.setattr(api_openai, "get_env_start_args", lambda: args) - monkeypatch.setattr(api_stream_obj, "get_env_start_args", lambda: args) - monkeypatch.setattr(api_openai, "build_prompt", AsyncMock(return_value="prompt")) - monkeypatch.setattr(SamplingParams, "init", lambda *_args, **_kwargs: None) - app = FastAPI() - app.add_exception_handler(GenerationError, generation_exception_handler) - - @app.post("/v1/chat/completions") - async def chat(request: ChatCompletionRequest): - return await api_openai.chat_completions_impl(request, Mock()) - - response = TestClient(app).post( - "/v1/chat/completions", - json={"model": "test", "messages": [{"role": "user", "content": "hi"}], "stream": True}, - ) - assert response.status_code == expected_status - if expected_status == 500: - assert "Generation failed" in response.json()["error"]["message"] - elif finish_status == FinishStatus.FINISHED_ERROR: - assert response.headers["content-type"].startswith("text/event-stream") - events = response.text.strip().split("\n\n") - assert len(events) == 2 - assert json.loads(events[0].removeprefix("data: ")) == { - "error": { - "message": "Generation failed before producing output", - "type": "server_error", - "code": "stream_error", - } - } - assert events[1] == "data: [DONE]" - else: - assert "answer" in response.text - assert "data: [DONE]" in response.text diff --git a/unit_tests/server/test_glm47_tool_parser.py b/unit_tests/server/test_glm47_tool_parser.py deleted file mode 100644 index 8f626d6343..0000000000 --- a/unit_tests/server/test_glm47_tool_parser.py +++ /dev/null @@ -1,127 +0,0 @@ -import json - -import pytest - -from lightllm.server.api_models import Tool -from lightllm.server.function_call_parser import FunctionCallParser - - -TOOLS = [ - Tool.model_validate( - { - "type": "function", - "function": { - "name": "Read", - "parameters": {"type": "object", "properties": {"file_path": {"type": "string"}}}, - }, - } - ), - Tool.model_validate( - { - "type": "function", - "function": { - "name": "Write", - "parameters": { - "type": "object", - "properties": {"content": {"type": "string"}, "overwrite": {"type": "boolean"}}, - }, - }, - } - ), -] - - -def collect_stream(text, chunk_size): - parser = FunctionCallParser(TOOLS, "glm47") - content = "" - calls = {} - for start in range(0, len(text), chunk_size): - normal, deltas = parser.parse_stream_chunk(text[start : start + chunk_size]) - content += normal - for delta in deltas: - if delta.name is not None: - assert delta.tool_index not in calls, "a tool call must have exactly one name event" - calls[delta.tool_index] = {"name": delta.name, "arguments": ""} - calls[delta.tool_index]["arguments"] += delta.parameters - return content, calls - - -@pytest.mark.parametrize("separator", ["", "\n"]) -@pytest.mark.parametrize("chunk_size", [1, 7, 39, 4096]) -def test_repeated_tool_calls_have_separate_indices(separator, chunk_size): - text = ( - f'Write{separator}content{{"ready": true}}' - "overwritefalse" - f"Read{separator}file_path/tmp/a" - f"Read{separator}file_path/tmp/b" - ) - _, calls = collect_stream(text, chunk_size) - assert list(calls) == [0, 1, 2] - assert [call["name"] for call in calls.values()] == ["Write", "Read", "Read"] - assert [json.loads(call["arguments"]) for call in calls.values()] == [ - {"content": '{"ready": true}', "overwrite": False}, - {"file_path": "/tmp/a"}, - {"file_path": "/tmp/b"}, - ] - _, nonstream = FunctionCallParser(TOOLS, "glm47").parse_non_stream(text) - assert [call.tool_index for call in nonstream] == [0, 1, 2] - assert [call.parameters for call in nonstream] == [call["arguments"] for call in calls.values()] - - -@pytest.mark.parametrize("value", ['{"count": 1}', "123", "false", " return 1\n", '"quoted"']) -def test_string_arguments_preserve_type_and_whitespace(value): - text = f"Writecontent{value}" - _, calls = collect_stream(text, 1) - assert json.loads(calls[0]["arguments"])["content"] == value - - -@pytest.mark.parametrize("chunk_size", [1, 7, 4096]) -def test_text_around_tool_calls_survives_chunk_boundaries(chunk_size): - text = ( - "Before Readfile_path/tmp/a" - " between Readfile_path/tmp/b after" - ) - content, calls = collect_stream(text, chunk_size) - assert content == "Before between after" - assert list(calls) == [0, 1] - - -def test_buffered_arguments_keep_the_stream_alive_without_repeating_the_name(): - parser = FunctionCallParser(TOOLS, "glm47") - normal, calls = parser.parse_stream_chunk("Writecontent") - assert normal == "" - assert len(calls) == 1 - assert calls[0].tool_index == 0 - assert calls[0].name == "Write" - assert calls[0].parameters == "" - - value = ' line 1\n{"literal": true}\n' + "x" * 100_000 - for chunk in (value[:13], value[13:27], value[27:]): - normal, calls = parser.parse_stream_chunk(chunk) - assert normal == "" - assert len(calls) == 1 - assert calls[0].tool_index == 0 - assert calls[0].name is None - assert calls[0].parameters == "" - - normal, calls = parser.parse_stream_chunk("") - assert normal == "" - assert len(calls) == 1 - assert calls[0].tool_index == 0 - assert calls[0].name is None - assert json.loads(calls[0].parameters) == {"content": value} - - _, calls = parser.parse_stream_chunk("Write\n") - assert len(calls) == 1 - assert calls[0].tool_index == 1 - assert calls[0].name == "Write" - - -def test_undefined_buffered_tool_does_not_consume_a_call_index(): - parser = FunctionCallParser(TOOLS, "glm47") - _, calls = parser.parse_stream_chunk("unknown\ncontentx") - assert calls == [] - _, calls = parser.parse_stream_chunk("Read\n") - assert len(calls) == 1 - assert calls[0].name == "Read" - assert calls[0].tool_index == 0 diff --git a/unit_tests/server/test_hybrid_start_args.py b/unit_tests/server/test_hybrid_start_args.py deleted file mode 100644 index 1c391c09cd..0000000000 --- a/unit_tests/server/test_hybrid_start_args.py +++ /dev/null @@ -1,138 +0,0 @@ -import json - -import pytest - -from lightllm.server.api_start import _launch_subprocesses -from lightllm.server.core.objs.start_args_type import StartArgs - - -@pytest.fixture -def launch_args(monkeypatch, tmp_path): - for name in ( - "_set_envs_and_config", - "auto_set_max_req_total_len", - "auto_set_fused_shared_experts", - "set_unique_server_name", - ): - monkeypatch.setattr(f"lightllm.server.api_start.{name}", lambda args: None) - - def validation_finished(args): - raise RuntimeError("startup validation passed") - - monkeypatch.setattr("lightllm.server.api_start.auto_set_response_parsers", validation_finished) - - def make_args(model_type, **kwargs): - (tmp_path / "config.json").write_text(json.dumps({"model_type": model_type})) - return StartArgs( - model_dir=str(tmp_path), - max_req_total_len=2048, - eos_id=[2], - router_token_ratio=0.0, - disable_vision=True, - disable_audio=True, - disable_shm_warning=True, - **kwargs, - ) - - return make_args - - -@pytest.mark.parametrize("run_mode", ["normal", "prefill", "decode"]) -@pytest.mark.parametrize("output_constraint_mode", ["none", "xgrammar"]) -@pytest.mark.parametrize( - "options, message", - [ - ({"disable_dynamic_prompt_cache": True}, "hybrid attention models require dynamic prompt cache"), - ({"disable_chunked_prefill": True}, "hybrid attention models require chunked prefill"), - ({"diverse_mode": True}, "--diverse_mode does not support hybrid attention models"), - ], -) -def test_hybrid_rejects_unsupported_modes(launch_args, run_mode, output_constraint_mode, options, message): - args = launch_args("qwen3_5_text", run_mode=run_mode, output_constraint_mode=output_constraint_mode, **options) - with pytest.raises(AssertionError, match=message): - _launch_subprocesses(args) - - -@pytest.mark.parametrize("output_constraint_mode", ["none", "xgrammar"]) -@pytest.mark.parametrize( - "model_type, options", - [ - ("qwen3_5_text", {}), - ("qwen3", {"disable_dynamic_prompt_cache": True}), - ("qwen3", {"disable_chunked_prefill": True}), - ("qwen3", {"diverse_mode": True}), - ], -) -def test_supported_models_pass_startup_validation(launch_args, output_constraint_mode, model_type, options): - args = launch_args(model_type, output_constraint_mode=output_constraint_mode, **options) - with pytest.raises(RuntimeError, match="startup validation passed"): - _launch_subprocesses(args) - - -@pytest.mark.parametrize( - "model_type,disable_mm,disable_vision,batch_tokens,passes", - [ - ("qwen3", True, False, 512, False), - ("qwen3", True, False, 1024, True), - ("qwen3", True, False, 2048, True), - ("qwen3", True, False, None, True), - ("qwen3", False, False, 512, True), - ("qwen3", True, True, 512, True), - ("neo_chat", False, False, 512, False), - ("deepseek_v4", False, False, 512, False), - ("gemma4", False, False, 512, False), - ("neo_chat", False, False, 1024, True), - ("qwen3_vl", False, False, 512, True), - ("qwen3_5", False, False, 512, True), - ], -) -def test_whole_image_blocks_fit_batch_at_startup( - launch_args, model_type, disable_mm, disable_vision, batch_tokens, passes -): - args = launch_args( - model_type, - disable_chunked_mm_input=disable_mm, - batch_max_tokens=batch_tokens, - max_image_token_count=1024, - ) - args.disable_vision = disable_vision - if passes: - with pytest.raises(RuntimeError, match="startup validation passed"): - _launch_subprocesses(args) - else: - with pytest.raises(AssertionError, match="batch_max_tokens >= max_image_token_count"): - _launch_subprocesses(args) - - -@pytest.mark.parametrize("mode,step", [("native", 0), ("replay", 0), ("replay", 4)]) -def test_s3_configuration_is_prepared_before_ssm_validation(launch_args, monkeypatch, tmp_path, mode, step): - from pathlib import Path - from lightllm.utils.petrel_helper import PetrelHelper - - monkeypatch.chdir(tmp_path) - model_dir = "s3://test-bucket/qwen" - config = {"model_type": "qwen3_5_text"} - downloads = [] - - def list_files(path, extension): - assert path == model_dir - return ["config.json"] if extension == "json" else [] - - def download_json(remote, local): - downloads.append(remote) - Path(local).write_text(json.dumps(config)) - - monkeypatch.setattr(PetrelHelper, "list", list_files) - monkeypatch.setattr(PetrelHelper, "download_json", download_json) - args = launch_args("qwen3_5_text", ssm_state_mode=mode, mtp_step=step, mtp_mode="eagle_with_att" if step else None) - args.model_dir = model_dir - assert not Path(model_dir, "config.json").exists() - with pytest.raises(RuntimeError, match="startup validation passed"): - _launch_subprocesses(args) - assert downloads == [model_dir + "/config.json"] - if mode == "replay": - expected_compute_mode = "precompute" if step > 0 else "inline" - assert args.ssm_replay_config == { - "cache_len": 16, - "compute_mode": expected_compute_mode, - } diff --git a/unit_tests/server/test_hypercorn_config.py b/unit_tests/server/test_hypercorn_config.py deleted file mode 100644 index 46ee8b7e7a..0000000000 --- a/unit_tests/server/test_hypercorn_config.py +++ /dev/null @@ -1,26 +0,0 @@ -from lightllm.server.api_cli import make_argument_parser -from lightllm.server.api_start import _hypercorn_config_args - - -def test_hypercorn_config_defaults_to_none(): - args = make_argument_parser().parse_args([]) - - assert args.hypercorn_config is None - - -def test_hypercorn_config_is_parsed(): - args = make_argument_parser().parse_args(["--hypercorn_config", "hypercorn.toml"]) - - assert args.hypercorn_config == "hypercorn.toml" - - -def test_hypercorn_args_use_default_keep_alive_without_config(): - args = make_argument_parser().parse_args([]) - - assert _hypercorn_config_args(args) == ["--keep-alive", "10"] - - -def test_hypercorn_args_do_not_set_keep_alive_with_config(): - args = make_argument_parser().parse_args(["--hypercorn_config", "hypercorn.toml"]) - - assert _hypercorn_config_args(args) == ["--config", "hypercorn.toml"] diff --git a/unit_tests/server/test_mtp_start_args.py b/unit_tests/server/test_mtp_start_args.py deleted file mode 100644 index 3aa269bb0e..0000000000 --- a/unit_tests/server/test_mtp_start_args.py +++ /dev/null @@ -1,55 +0,0 @@ -import pytest - -from lightllm.server.api_start import _launch_subprocesses -from lightllm.server.core.objs.start_args_type import StartArgs - - -class _StartupStopped(Exception): - pass - - -def test_mtp_requires_cuda_graph(monkeypatch): - monkeypatch.setattr("lightllm.server.api_start._set_envs_and_config", lambda args: None) - args = StartArgs(mtp_mode="vanilla_no_att", disable_cudagraph=True, output_constraint_mode="none") - - with pytest.raises(AssertionError, match="only supported on Prefill nodes"): - _launch_subprocesses(args) - - -def test_mtp_prefill_still_requires_positive_step(monkeypatch): - monkeypatch.setattr("lightllm.server.api_start._set_envs_and_config", lambda args: None) - monkeypatch.setattr("lightllm.server.api_start.auto_set_max_req_total_len", lambda args: None) - monkeypatch.setattr("lightllm.server.api_start.auto_set_fused_shared_experts", lambda args: None) - monkeypatch.setattr("lightllm.server.api_start.set_unique_server_name", lambda args: None) - args = StartArgs( - run_mode="prefill", - mtp_mode="dspark", - output_constraint_mode="none", - mtp_draft_model_dir=["draft"], - mtp_step=0, - disable_cudagraph=True, - disable_vision=True, - disable_audio=True, - disable_shm_warning=True, - ) - - with pytest.raises(AssertionError): - _launch_subprocesses(args) - - -def test_dsv4_cpu_cache_uses_common_adaptive_placement(monkeypatch): - monkeypatch.setattr("lightllm.server.api_start._set_envs_and_config", lambda args: None) - monkeypatch.setattr("lightllm.server.api_start.auto_set_max_req_total_len", lambda args: None) - monkeypatch.setattr("lightllm.server.api_start.auto_set_fused_shared_experts", lambda args: None) - monkeypatch.setattr("lightllm.server.api_start.set_unique_server_name", lambda args: None) - monkeypatch.setattr("lightllm.utils.config_utils.get_model_type", lambda model_dir: "deepseek_v4") - monkeypatch.setattr( - "lightllm.server.api_start.has_vision_module", - lambda model_dir: (_ for _ in ()).throw(_StartupStopped), - ) - args = StartArgs(enable_cpu_cache=True, disable_audio=True) - - with pytest.raises(_StartupStopped): - _launch_subprocesses(args) - - assert args.cache_placement_strategy == "adaptive" diff --git a/unit_tests/server/test_multimodal_chunk_args.py b/unit_tests/server/test_multimodal_chunk_args.py deleted file mode 100644 index c799235904..0000000000 --- a/unit_tests/server/test_multimodal_chunk_args.py +++ /dev/null @@ -1,15 +0,0 @@ -import argparse - -from lightllm.server.api_cli import add_cli_args -from lightllm.server.core.objs.start_args_type import StartArgs - - -def test_multimodal_chunk_flag(): - parser = add_cli_args(argparse.ArgumentParser()) - assert parser.parse_args(["--disable_chunked_mm_input"]).disable_chunked_mm_input is True - - -def test_multimodal_chunking_keeps_existing_defaults(): - parser = add_cli_args(argparse.ArgumentParser()) - assert parser.parse_args([]).disable_chunked_mm_input is False - assert StartArgs().disable_chunked_mm_input is False diff --git a/unit_tests/server/test_pd_cache_aware.py b/unit_tests/server/test_pd_cache_aware.py deleted file mode 100644 index 1b111122fb..0000000000 --- a/unit_tests/server/test_pd_cache_aware.py +++ /dev/null @@ -1,321 +0,0 @@ -import asyncio -from types import SimpleNamespace -from unittest.mock import patch - -import pytest - -from lightllm.server.httpserver_for_pd_master.pd_selector.prompt_cache_tree_client import ( - create_prompt_cache_tree_client, -) -from lightllm.server.httpserver_for_pd_master.pd_selector.cache_aware import ( - BalanceRelThresholdController, - CacheAwareConfig, - CacheAwarePolicy, -) -from lightllm.server.httpserver_for_pd_master.pd_selector.prompt_cache_tree import PromptCacheTree -from lightllm.server.httpserver_for_pd_master.pd_selector.pd_selector import ( - LoadBalancedCacheAwareSelector, - PDSelectionExtraInfo, -) - -pytestmark = pytest.mark.usefixtures("single_http_worker_start_args") - - -def _worker(address: str, dispatched_prompt_chars: int = 0, dispatched_req_num: int = 0): - return SimpleNamespace( - client_ip_port=address, - connection_key=address, - dispatched_prompt_chars=dispatched_prompt_chars, - dispatched_req_num=dispatched_req_num, - ) - - -def _local_policy(config): - return CacheAwarePolicy(prompt_cache_tree=create_prompt_cache_tree_client(config), config=config) - - -def _insert_prompt(policy, prompt, worker): - sample_stride = policy.config.sample_stride - asyncio.run(policy.prompt_cache_tree.insert_sampled(prompt[::sample_stride], worker.connection_key, sample_stride)) - - -def _prefix_match(policy, prompt): - sample_stride = policy.config.sample_stride - return asyncio.run( - policy.prompt_cache_tree.prefix_match_sampled(prompt[::sample_stride], len(prompt), sample_stride) - ) - - -def test_balance_threshold_controller_adjusts_threshold_each_window(): - config = CacheAwareConfig(cache_hit_rate_window_size=3, balance_rel_threshold_step=0.1) - controller = BalanceRelThresholdController() - - for cache_hit_rate, expected_threshold in ( - (0.5, 1.5), - (0.6, 1.6), - (0.4, 1.5), - (0.4, 1.5), - ): - for _ in range(config.cache_hit_rate_window_size): - controller.append(cache_hit_rate) - controller.update_config(config) - assert config.balance_rel_threshold == pytest.approx(expected_threshold) - - -@pytest.mark.parametrize( - ("initial_threshold", "first_hit_rate", "second_hit_rate", "expected_threshold"), - ( - (1.95, 0.5, 0.6, 2.0), - (1.05, 0.5, 0.4, 1.0), - ), -) -def test_balance_threshold_controller_limits_threshold_range( - initial_threshold, - first_hit_rate, - second_hit_rate, - expected_threshold, -): - config = CacheAwareConfig( - balance_rel_threshold=initial_threshold, - cache_hit_rate_window_size=1, - balance_rel_threshold_step=0.1, - ) - controller = BalanceRelThresholdController() - - controller.append(first_hit_rate) - controller.update_config(config) - controller.append(second_hit_rate) - controller.update_config(config) - - assert config.balance_rel_threshold == pytest.approx(expected_threshold) - - -def test_cache_aware_updates_threshold_from_inference_cache_hit_rate(): - policy = _local_policy(CacheAwareConfig(cache_hit_rate_window_size=2)) - - for _ in range(2): - policy.record_prompt_cache_hit_rate(0.25) - assert policy.config.balance_rel_threshold == 1.5 - for _ in range(2): - policy.record_prompt_cache_hit_rate(0.75) - - assert policy.config.balance_rel_threshold == pytest.approx(1.55) - - -@pytest.mark.parametrize("operation", ["insert", "prefix_match"]) -def test_prompt_cache_tree_rejects_mismatched_sample_stride(operation): - tree = PromptCacheTree(sample_stride=4) - - with pytest.raises(AssertionError, match="sample_stride mismatch"): - if operation == "insert": - tree.insert_sampled("po", "prefill", 2) - else: - tree.prefix_match_sampled("po", len("prompt"), 2) - - -def test_prompt_cache_tree_handles_key_deeper_than_python_recursion_limit(): - tree = PromptCacheTree(sample_stride=1) - key = "x" * 5000 - - tree.insert_sampled(key, "prefill", sample_stride=1) - result = tree.prefix_match_sampled(key, input_char_count=len(key), sample_stride=1) - - assert result.prefill_node == "prefill" - assert result.matched_char_count == len(key) - - -def test_prompt_cache_tree_empty_key_does_not_create_nodes(): - tree = PromptCacheTree(sample_stride=1) - - tree.insert_sampled("", "prefill", sample_stride=1) - result = tree.prefix_match_sampled("", input_char_count=0, sample_stride=1) - - assert tree._node_count == 0 - assert not tree.root.children - assert result.prefill_node is None - assert result.matched_char_count == 0 - assert result.input_char_count == 0 - assert result.last_insert_time is None - - -def test_prompt_cache_tree_zero_capacity_evicts_the_complete_path(): - tree = PromptCacheTree(sample_stride=1, max_node_count=0, evict_node_batch=1) - - tree.insert_sampled("你好🙂", "prefill", sample_stride=1) - result = tree.prefix_match_sampled("你好🙂", input_char_count=3, sample_stride=1) - - assert tree._node_count == 0 - assert not tree.root.children - assert result.prefill_node is None - assert result.matched_char_count == 0 - assert result.input_char_count == 3 - assert result.last_insert_time is None - - -def test_prompt_cache_tree_matches_unicode_partial_prefix(): - tree = PromptCacheTree(sample_stride=1) - tree.insert_sampled("你好🙂", "prefill", sample_stride=1) - - result = tree.prefix_match_sampled("你好🙃", input_char_count=3, sample_stride=1) - - assert result.prefill_node == "prefill" - assert result.matched_char_count == 2 - assert result.input_char_count == 3 - assert result.last_insert_time is not None - - -def test_prompt_cache_tree_iterative_insert_keeps_lru_eviction_order(): - tree = PromptCacheTree(sample_stride=1, max_node_count=3, evict_node_batch=1) - tree.insert_sampled("ab", "prefill-ab", sample_stride=1) - tree.insert_sampled("ac", "prefill-ac", sample_stride=1) - - # 再次访问 ab，使 ac 成为最早未使用的叶节点。 - tree.insert_sampled("ab", "prefill-ab", sample_stride=1) - # 插入 d 后超出容量，依次驱逐 ac 和 ab；共同父节点 a 随后成为叶节点。 - tree.insert_sampled("d", "prefill-d", sample_stride=1) - - assert tree._node_count == 2 - assert set(tree.root.children) == {"a", "d"} - result = tree.prefix_match_sampled("ab", input_char_count=2, sample_stride=1) - assert result.prefill_node == "prefill-ab" - assert result.matched_char_count == 1 - - -def test_cache_aware_selector_returns_selected_nodes_and_estimated_cache_hit_rate(): - selector = LoadBalancedCacheAwareSelector(pd_manager=None) - p_node = _worker("10.0.0.1:8000") - d_node = SimpleNamespace( - client_ip_port="10.0.0.2:8000", - run_status=SimpleNamespace(total_token_usage_rate=0.0), - ) - selector.update_nodes([p_node], [d_node]) - prompt = "x" * 1025 - with patch( - "lightllm.server.httpserver_for_pd_master.pd_selector.prompt_cache_tree.time.monotonic", - return_value=123.0, - ): - _insert_prompt(selector.policy, prompt, p_node) - - selected_p_node, selected_d_node, selection_extra_info = asyncio.run(selector.select_p_d_node(prompt, None, None)) - - assert selected_p_node is p_node - assert selected_d_node is d_node - assert selection_extra_info == PDSelectionExtraInfo( - estimated_cache_hit_rate=pytest.approx(1.0), - cache_last_insert_time=123.0, - ) - - -def test_cache_aware_reinsert_refreshes_matched_node_time(): - policy = _local_policy(CacheAwareConfig()) - worker = _worker("10.0.0.1:8000") - prompt = "shared prefix " * 100 - - with patch( - "lightllm.server.httpserver_for_pd_master.pd_selector.prompt_cache_tree.time.monotonic", - return_value=100.0, - ): - _insert_prompt(policy, prompt, worker) - first_result = _prefix_match(policy, prompt) - with patch( - "lightllm.server.httpserver_for_pd_master.pd_selector.prompt_cache_tree.time.monotonic", - return_value=160.0, - ): - _insert_prompt(policy, prompt, worker) - second_result = _prefix_match(policy, prompt) - - assert first_result.last_insert_time == 100.0 - assert second_result.last_insert_time == 160.0 - - -def test_cache_aware_inserts_prompt_only_when_explicitly_recorded(): - policy = _local_policy(CacheAwareConfig()) - worker = _worker("10.0.0.1:8000") - prompt = "new prompt" - - selected_worker, cache_info = asyncio.run(policy.select_worker([worker], prompt)) - assert selected_worker is worker - assert cache_info == PDSelectionExtraInfo() - assert _prefix_match(policy, prompt).prefill_node is None - - _insert_prompt(policy, prompt, worker) - - assert _prefix_match(policy, prompt).prefill_node == worker.connection_key - - -def test_cache_aware_keeps_cache_worker_when_inflight_load_is_balanced(): - policy = _local_policy(CacheAwareConfig()) - cache_worker = _worker("10.0.0.1:8000", dispatched_prompt_chars=110, dispatched_req_num=2) - least_loaded_worker = _worker("10.0.0.2:8000", dispatched_prompt_chars=100, dispatched_req_num=2) - prompt = "shared prefix " * 100 - _insert_prompt(policy, prompt, cache_worker) - - selected_worker, cache_info = asyncio.run(policy.select_worker([cache_worker, least_loaded_worker], prompt)) - - assert selected_worker is cache_worker - # 前缀树每 512 个字符抽样一次，命中率保持原有的保守估算方式。 - assert cache_info.estimated_cache_hit_rate == pytest.approx(1025 / len(prompt)) - - -def test_cache_aware_uses_least_loaded_worker_when_cache_worker_is_overloaded(): - policy = _local_policy(CacheAwareConfig()) - cache_worker = _worker("10.0.0.1:8000", dispatched_prompt_chars=2000, dispatched_req_num=2) - least_loaded_worker = _worker("10.0.0.2:8000", dispatched_prompt_chars=100, dispatched_req_num=2) - prompt = "shared prefix " * 100 - _insert_prompt(policy, prompt, cache_worker) - - selected_worker, cache_info = asyncio.run(policy.select_worker([cache_worker, least_loaded_worker], prompt)) - - assert selected_worker is least_loaded_worker - assert cache_info.estimated_cache_hit_rate == 0.0 - - -def test_cache_aware_keeps_cache_worker_when_both_workers_are_idle(): - policy = _local_policy(CacheAwareConfig()) - cache_worker = _worker("10.0.0.1:8000") - other_worker = _worker("10.0.0.2:8000") - prompt = "shared prefix " * 100 - _insert_prompt(policy, prompt, cache_worker) - - selected_worker, _ = asyncio.run(policy.select_worker([other_worker, cache_worker], prompt)) - - assert selected_worker is cache_worker - - -def test_cache_aware_forces_idle_worker_over_busy_cache_worker(): - policy = _local_policy(CacheAwareConfig()) - cache_worker = _worker("10.0.0.1:8000", dispatched_prompt_chars=100, dispatched_req_num=1) - idle_worker = _worker("10.0.0.2:8000") - prompt = "shared prefix " * 100 - _insert_prompt(policy, prompt, cache_worker) - - selected_worker, _ = asyncio.run(policy.select_worker([cache_worker, idle_worker], prompt)) - - assert selected_worker is idle_worker - - -def test_cache_aware_randomly_selects_an_idle_worker_without_cache_match(): - policy = _local_policy(CacheAwareConfig()) - first = _worker("10.0.0.1:8000") - second = _worker("10.0.0.2:8000") - - with patch( - "lightllm.server.httpserver_for_pd_master.pd_selector.cache_aware.random.choice", return_value=second - ) as choice: - selected_worker, _ = asyncio.run(policy.select_worker([first, second], "new prompt")) - - assert selected_worker is second - choice.assert_called_once_with([first, second]) - - -def test_cache_aware_matches_cache_only_within_idle_workers(): - policy = _local_policy(CacheAwareConfig()) - busy_worker = _worker("10.0.0.1:8000", dispatched_prompt_chars=100, dispatched_req_num=1) - cache_idle_worker = _worker("10.0.0.2:8000") - other_idle_worker = _worker("10.0.0.3:8000") - prompt = "shared prefix " * 100 - _insert_prompt(policy, prompt, cache_idle_worker) - - selected_worker, _ = asyncio.run(policy.select_worker([other_idle_worker, busy_worker, cache_idle_worker], prompt)) - - assert selected_worker is cache_idle_worker diff --git a/unit_tests/server/test_pd_local_selection.py b/unit_tests/server/test_pd_local_selection.py deleted file mode 100644 index 669be609b6..0000000000 --- a/unit_tests/server/test_pd_local_selection.py +++ /dev/null @@ -1,919 +0,0 @@ -"""Local PD policies use shared cache/load snapshots without remote selection.""" - -import asyncio -import random -from types import SimpleNamespace -from unittest.mock import AsyncMock, Mock, call - -import pytest - -from lightllm.server.core.objs import FinishStatus, SamplingParams -from lightllm.server.httpserver_for_pd_master.manager import ( - HttpServerManagerForPDMaster, - PDManager, -) -from lightllm.server.httpserver_for_pd_master.pd_selector import pd_node_stats_client -from lightllm.server.httpserver_for_pd_master.pd_selector import pd_selector -from lightllm.server.httpserver_for_pd_master.pd_selector import prompt_cache_tree_client -from lightllm.server.httpserver_for_pd_master.pd_selector.prompt_cache_tree_client import ( - LocalPromptCacheTreeClient, - RemotePromptCacheTreeClient, -) -from lightllm.server.httpserver_for_pd_master.pd_selector.pd_selector import ( - AdaptiveLoadSelector, - LoadBalancedCacheAwareSelector, - PDSelectionExtraInfo, - RandomSelector, - RoundRobinSelector, -) -from lightllm.server.httpserver_for_pd_master.pd_selector.prompt_cache_tree import ( - PromptCacheMatchResult, -) -from lightllm.server.httpserver_for_pd_master.pd_selector.pd_node_stats_server import ( - PDNodeStatsState, -) -from lightllm.server.httpserver_for_pd_master.pd_selector.prompt_cache_tree_server import ( - PromptCacheTreeState, -) -from lightllm.server.pd_io_struct import PD_Client_Obj -from lightllm.utils.error_utils import ServerBusyError - - -PROMPT = "x" * 2048 -COUNTER_METHODS = ("add_dispatched_info", "sub_dispatched_info") - - -@pytest.fixture(autouse=True) -def local_pd_master_mode(monkeypatch): - start_args = SimpleNamespace(httpserver_workers=1) - monkeypatch.setattr(pd_node_stats_client, "get_env_start_args", lambda: start_args) - monkeypatch.setattr(prompt_cache_tree_client, "get_env_start_args", lambda: start_args) - - -def _node( - node_id, - address, - *, - mode="prefill", - chars=0, - requests=0, - load=0.0, - connection_id=None, -): - node = PD_Client_Obj( - node_id=node_id, - client_ip_port=address, - mode=mode, - start_args={}, - connection_id=connection_id or f"connection-{node_id}", - dispatched_prompt_chars=chars, - dispatched_req_num=requests, - ) - node.run_status.total_token_usage_rate = load - return node - - -def _match(node=None, *, matched=1500): - return PromptCacheMatchResult( - prefill_node=None if node is None else node.connection_key, - matched_char_count=0 if node is None else matched, - input_char_count=len(PROMPT), - last_insert_time=None if node is None else 123.0, - ) - - -def _assert_backend_client(selector): - assert isinstance( - selector.policy.prompt_cache_tree, - (LocalPromptCacheTreeClient, RemotePromptCacheTreeClient), - ) - - -def _set_nodes(manager, prefill_nodes, decode_nodes): - manager.pd_manager.prefill_nodes = prefill_nodes - manager.pd_manager.decode_nodes = decode_nodes - manager.pd_manager.url_to_pd_nodes = {node.connection_key: node for node in [*prefill_nodes, *decode_nodes]} - manager.pd_manager.selector.update_nodes(prefill_nodes, decode_nodes) - - -def _manager(strategy="cache_aware", *, prefill_nodes=None, decode_nodes=None, mock_prefix=True): - manager = HttpServerManagerForPDMaster.__new__(HttpServerManagerForPDMaster) - manager.args = SimpleNamespace(select_p_d_node_strategy=strategy, pd_master_mode="elastic") - manager.pd_manager = PDManager(manager.args) - _set_nodes( - manager, - prefill_nodes if prefill_nodes is not None else [_node(100, "p1:8000"), _node(101, "p2:8000")], - decode_nodes if decode_nodes is not None else [_node(200, "d:8000", mode="decode")], - ) - if mock_prefix: - selector = manager.pd_manager.selector - selector.pd_node_stats.update_dispatched_info = AsyncMock(wraps=selector.pd_node_stats.update_dispatched_info) - if strategy == "cache_aware": - selector.policy.prompt_cache_tree.prefix_match_sampled = AsyncMock(return_value=_match()) - manager.pd_node_stats = manager.pd_manager.selector.pd_node_stats - return manager - - -@pytest.mark.parametrize("busy", [False, True]) -def test_shared_tree_is_used_for_selection_and_cache_metadata(monkeypatch, busy): - selector = LoadBalancedCacheAwareSelector(None) - cached = _node(100, "p2:8000", chars=110, requests=int(busy)) - other = _node(101, "p1:8000", chars=100, requests=int(busy)) - decode = _node(200, "d:8000", mode="decode") - selector.update_nodes([other, cached], [decode]) - _assert_backend_client(selector) - selector.policy.prompt_cache_tree.prefix_match_sampled = AsyncMock(return_value=_match(cached)) - - p_node, d_node, extra = asyncio.run(selector.select_p_d_node(PROMPT, None, None)) - - assert p_node is cached - assert d_node is decode - assert extra == PDSelectionExtraInfo(1500 / len(PROMPT), 123.0) - - -def test_cache_metadata_comes_from_the_selection_query(monkeypatch): - selector = LoadBalancedCacheAwareSelector(None) - cached, other = _node(100, "p2:8000"), _node(101, "p1:8000") - selector.update_nodes([other, cached], [_node(200, "d:8000", mode="decode")]) - _assert_backend_client(selector) - selector.policy.prompt_cache_tree.prefix_match_sampled = AsyncMock(return_value=_match(cached)) - - selected, _, extra = asyncio.run(selector.select_p_d_node(PROMPT, None, None)) - - assert selected is cached - assert extra == PDSelectionExtraInfo(1500 / len(PROMPT), 123.0) - selector.policy.prompt_cache_tree.prefix_match_sampled.assert_awaited_once_with( - PROMPT[:: selector.policy.config.sample_stride], - len(PROMPT), - selector.policy.config.sample_stride, - ) - - -@pytest.mark.parametrize("restarted", [False, True]) -def test_cache_hint_survives_endpoint_reconnect_or_process_restart(monkeypatch, restarted): - original = _node(100, "p2:8000", connection_id="old-socket") - current = _node(300 if restarted else 100, "p2:8000", connection_id="new-socket") - other = _node(101, "p1:8000") - selector = LoadBalancedCacheAwareSelector(None) - selector.update_nodes([other, current], [_node(200, "d:8000", mode="decode")]) - _assert_backend_client(selector) - selector.policy.prompt_cache_tree.prefix_match_sampled = AsyncMock(return_value=_match(original)) - - selected, _, extra = asyncio.run(selector.select_p_d_node(PROMPT, None, None)) - - assert original.connection_key == current.connection_key - assert selected is current - assert extra == PDSelectionExtraInfo(1500 / len(PROMPT), 123.0) - - -@pytest.mark.parametrize("match_kind", ["miss", "remote_worker"]) -def test_shared_match_never_falls_back_to_local_tree_or_remote_candidates(monkeypatch, match_kind): - busy = _node(100, "p1:8000", chars=2000, requests=1) - available = _node(101, "p2:8000", chars=100, requests=1) - selector = LoadBalancedCacheAwareSelector(None) - selector.update_nodes([busy, available], [_node(200, "d:8000", mode="decode")]) - asyncio.run( - selector.policy.prompt_cache_tree.insert_sampled( - PROMPT[:: selector.policy.config.sample_stride], - busy.connection_key, - selector.policy.config.sample_stride, - ) - ) - _assert_backend_client(selector) - shared_match = _match() if match_kind == "miss" else _match(_node(999, "remote:8000")) - selector.policy.prompt_cache_tree.prefix_match_sampled = AsyncMock(return_value=shared_match) - - selected, _, extra = asyncio.run(selector.select_p_d_node(PROMPT, None, None)) - - assert selected is available - assert extra == PDSelectionExtraInfo() - - -def test_default_cache_selection_uses_connection_keys(): - selector = LoadBalancedCacheAwareSelector(None) - cached, other = _node(100, "p2:8000"), _node(101, "p1:8000") - selector.update_nodes([other, cached], [_node(200, "d:8000", mode="decode")]) - asyncio.run( - selector.policy.prompt_cache_tree.insert_sampled( - PROMPT[:: selector.policy.config.sample_stride], - cached.connection_key, - selector.policy.config.sample_stride, - ) - ) - - selected, _, extra = asyncio.run(selector.select_p_d_node(PROMPT, None, None)) - - assert selected is cached - assert extra.estimated_cache_hit_rate > 0.5 - assert ( - asyncio.run( - selector.policy.prompt_cache_tree.prefix_match_sampled( - PROMPT[:: selector.policy.config.sample_stride], - len(PROMPT), - selector.policy.config.sample_stride, - ) - ).prefill_node - == cached.connection_key - ) - - -def test_round_robin_progress_is_local_to_each_worker(): - nodes = [_node(100, "p1:8000"), _node(101, "p2:8000")] - decode = [ - _node(200, "d1:8000", mode="decode"), - _node(201, "d2:8000", mode="decode"), - ] - first, second = RoundRobinSelector(None), RoundRobinSelector(None) - for selector in (first, second): - selector.update_nodes(nodes, decode) - - async def run(): - assert (await first.select_p_d_node(PROMPT, None, None))[:2] == ( - nodes[0], - decode[0], - ) - assert (await first.select_p_d_node(PROMPT, None, None))[:2] == ( - nodes[1], - decode[1], - ) - assert (await second.select_p_d_node(PROMPT, None, None))[:2] == ( - nodes[0], - decode[0], - ) - - asyncio.run(run()) - - -@pytest.mark.parametrize("selector_type", [RandomSelector, AdaptiveLoadSelector]) -def test_random_and_adaptive_keep_their_local_load_preferences(monkeypatch, selector_type): - monkeypatch.setattr(pd_selector, "random", random.Random(23)) - busy, available = _node(100, "p1:8000", load=1.0), _node(101, "p2:8000", load=0.0) - selector = selector_type(None) - selector.update_nodes([busy, available], [_node(200, "d:8000", mode="decode")]) - - async def select_many(): - return [(await selector.select_p_d_node(PROMPT, None, None))[0] for _ in range(100)] - - selections = asyncio.run(select_many()) - - assert all(node is busy or node is available for node in selections) - available_count = sum(node is available for node in selections) - if selector_type is AdaptiveLoadSelector: - assert available_count > 90 - else: - assert 20 < available_count < 80 - - -@pytest.mark.parametrize("strategy", ["random", "round_robin", "adaptive_load", "cache_aware"]) -def test_only_cache_aware_policy_queries_prefix_data_and_all_policies_choose_local_nodes(monkeypatch, strategy): - manager = _manager(strategy) - local_nodes = manager.pd_manager.prefill_nodes - if strategy == "cache_aware": - _assert_backend_client(manager.pd_manager.selector) - manager.pd_manager.selector.policy.prompt_cache_tree.prefix_match_sampled.return_value = _match(local_nodes[-1]) - - selected, decode, _ = asyncio.run(manager.select_p_d_node(PROMPT, None, None)) - - assert any(selected is node for node in local_nodes) - assert decode is manager.pd_manager.decode_nodes[0] - if strategy == "cache_aware": - assert manager.pd_manager.selector.policy.prompt_cache_tree.prefix_match_sampled.await_count == 1 - manager.pd_manager.selector.policy.prompt_cache_tree.prefix_match_sampled.assert_awaited_with( - PROMPT[:: manager.pd_manager.selector.policy.config.sample_stride], - len(PROMPT), - manager.pd_manager.selector.policy.config.sample_stride, - ) - manager.pd_manager.selector.pd_node_stats.update_dispatched_info.assert_awaited_once_with(p_nodes=local_nodes) - else: - manager.pd_manager.selector.pd_node_stats.update_dispatched_info.assert_not_awaited() - - -def test_manager_uses_shared_node_load_and_local_balancing_threshold_for_selection( - monkeypatch, -): - cached, available = _node(100, "p1:8000"), _node(101, "p2:8000") - manager = _manager(prefill_nodes=[cached, available]) - _assert_backend_client(manager.pd_manager.selector) - - async def update_dispatched_info(p_nodes): - assert p_nodes == [cached, available] - cached.dispatched_prompt_chars, cached.dispatched_req_num = 500, 1 - available.dispatched_prompt_chars, available.dispatched_req_num = 100, 1 - - manager.pd_manager.selector.pd_node_stats.update_dispatched_info.side_effect = update_dispatched_info - manager.pd_manager.selector.policy.prompt_cache_tree.prefix_match_sampled.return_value = _match(cached) - - manager.pd_manager.selector.policy.config.balance_rel_threshold = 1.1 - first = asyncio.run(manager.select_p_d_node(PROMPT, None, None)) - manager.pd_manager.selector.policy.config.balance_rel_threshold = 2.0 - second = asyncio.run(manager.select_p_d_node(PROMPT, None, None)) - - assert first[0] is available - assert first[2] == PDSelectionExtraInfo() - assert second[0] is cached - assert second[2].estimated_cache_hit_rate == 1500 / len(PROMPT) - - -@pytest.mark.parametrize("replacement_load", [(0, 0), (999, 9)]) -def test_query_completion_uses_current_connections_without_overwriting_their_load(monkeypatch, replacement_load): - async def run(): - old_p, old_d = _node(100, "p:8000"), _node(200, "d:8000", mode="decode") - new_p = _node(300, "p:8000", chars=replacement_load[0], requests=replacement_load[1]) - new_d = _node(400, "d:8000", mode="decode") - manager = _manager(prefill_nodes=[old_p], decode_nodes=[old_d]) - _assert_backend_client(manager.pd_manager.selector) - started, complete = asyncio.Event(), asyncio.Event() - - async def update_dispatched_info(p_nodes): - assert len(p_nodes) == 1 - if not started.is_set(): - assert p_nodes[0] is old_p - started.set() - await complete.wait() - old_p.dispatched_prompt_chars, old_p.dispatched_req_num = 800, 2 - - manager.pd_manager.selector.pd_node_stats.update_dispatched_info.side_effect = update_dispatched_info - manager.pd_manager.selector.policy.prompt_cache_tree.prefix_match_sampled.return_value = _match(old_p) - task = asyncio.create_task(manager.select_p_d_node(PROMPT, None, None)) - await asyncio.wait_for(started.wait(), timeout=2) - _set_nodes(manager, [new_p], [new_d]) - complete.set() - selected, decode, extra = await task - - assert selected is new_p and decode is new_d - assert ( - new_p.dispatched_prompt_chars, - new_p.dispatched_req_num, - ) == replacement_load - assert extra == PDSelectionExtraInfo(1500 / len(PROMPT), 123.0) - assert manager.pd_manager.selector.policy.prompt_cache_tree.prefix_match_sampled.await_count == 1 - manager.pd_manager.selector.policy.prompt_cache_tree.prefix_match_sampled.assert_awaited_with( - PROMPT[:: manager.pd_manager.selector.policy.config.sample_stride], - len(PROMPT), - manager.pd_manager.selector.policy.config.sample_stride, - ) - - asyncio.run(run()) - - -@pytest.mark.parametrize("replaced", ["prefill", "decode"]) -def test_selection_query_uses_current_decode_and_rejects_a_replaced_prefill(replaced): - async def run(): - manager = _manager() - selector = manager.pd_manager.selector - prefill, decode = ( - manager.pd_manager.prefill_nodes, - manager.pd_manager.decode_nodes, - ) - selected_prefill = prefill[0] - replacement = _node(300, "p1:8000") if replaced == "prefill" else _node(400, "d:8000", mode="decode") - started, complete = asyncio.Event(), asyncio.Event() - - async def query(*_args): - started.set() - await complete.wait() - return _match(selected_prefill) - - selector.policy.prompt_cache_tree.prefix_match_sampled.side_effect = query - task = asyncio.create_task(manager.select_p_d_node(PROMPT, None, None)) - await asyncio.wait_for(started.wait(), timeout=2) - _set_nodes( - manager, - [replacement, *prefill[1:]] if replaced == "prefill" else prefill, - [replacement] if replaced == "decode" else decode, - ) - complete.set() - if replaced == "prefill": - with pytest.raises(ServerBusyError, match="connection changed"): - await task - else: - selected, current_decode, _ = await task - assert selected is selected_prefill - assert current_decode is replacement - - asyncio.run(run()) - - -@pytest.fixture(params=["local", "rpyc"]) -def backend(request, monkeypatch): - start_args = SimpleNamespace(httpserver_workers=2 if request.param == "rpyc" else 1) - monkeypatch.setattr(pd_node_stats_client, "get_env_start_args", lambda: start_args) - monkeypatch.setattr(prompt_cache_tree_client, "get_env_start_args", lambda: start_args) - monkeypatch.setattr( - pd_node_stats_client, - "get_pd_master_unix_socket_address", - lambda service_name: f"\0lightllm_test-pd-info_{service_name}".encode(), - ) - monkeypatch.setattr( - prompt_cache_tree_client, - "get_pd_master_unix_socket_address", - lambda service_name: f"\0lightllm_test-pd-info_{service_name}".encode(), - ) - if request.param == "rpyc": - monkeypatch.setattr(pd_node_stats_client, "unix_connect", Mock()) - monkeypatch.setattr(prompt_cache_tree_client, "unix_connect", Mock()) - return request.param - - -def _generator_manager(backend, strategy="round_robin"): - manager = _manager( - strategy, - prefill_nodes=[_node(100, "p:8000")], - mock_prefix=False, - ) - states = SimpleNamespace(stats=None, tree=None) - if backend != "local": - states.stats = PDNodeStatsState() - stats = manager.pd_manager.selector.pd_node_stats - stats._get_dispatched_info = states.stats.get_dispatched_info - stats._add_dispatched_info = states.stats.add_dispatched_info - stats._sub_dispatched_info = states.stats.sub_dispatched_info - if strategy == "cache_aware": - states.tree = PromptCacheTreeState() - tree = manager.pd_manager.selector.policy.prompt_cache_tree - - def prefix_match_sampled(*args): - return PromptCacheMatchResult(*states.tree.prefix_match_sampled(*args)) - - tree._prefix_match_sampled = prefix_match_sampled - tree._insert_sampled = states.tree.insert_sampled - for name in COUNTER_METHODS: - stats = manager.pd_manager.selector.pd_node_stats - setattr(stats, name, AsyncMock(wraps=getattr(stats, name))) - manager.pd_manager.selector.record_prompt_cache_hit_rate = Mock( - wraps=manager.pd_manager.selector.record_prompt_cache_hit_rate - ) - if strategy == "cache_aware": - tree = manager.pd_manager.selector.policy.prompt_cache_tree - tree.insert_sampled = AsyncMock(wraps=tree.insert_sampled) - manager.active_req_ids = set() - manager.disable_pd_cache_high_priority = True - manager.enable_pd_node_self_request_limit = False - manager.remove_req = AsyncMock() - manager.abort = AsyncMock() - ids = iter(range(8, 80, 8)) - manager.id_gen = SimpleNamespace(generate_id=lambda: next(ids)) - return manager, states - - -def _pending(state, node): - if state.stats is None: - return node.dispatched_prompt_chars, node.dispatched_req_num - return state.stats.get_dispatched_info((node.connection_key,))[0] - - -def test_backend_tree_queries_preserve_cache_hint_for_restarted_endpoint(monkeypatch, backend): - async def run(): - manager, state = _generator_manager(backend, "cache_aware") - cached, other = _node(100, "p2:8000"), _node(101, "p1:8000") - decode = manager.pd_manager.decode_nodes - _set_nodes(manager, [other, cached], decode) - selector = manager.pd_manager.selector - await selector.policy.prompt_cache_tree.insert_sampled( - PROMPT[:: selector.policy.config.sample_stride], - cached.connection_key, - selector.policy.config.sample_stride, - ) - injected_tree = selector.policy.prompt_cache_tree - tree = ( - selector.policy.prompt_cache_tree.prompt_cache_tree if state.tree is None else state.tree.prompt_cache_tree - ) - prefix_match = Mock(wraps=tree.prefix_match_sampled) - monkeypatch.setattr(tree, "prefix_match_sampled", prefix_match) - assert isinstance( - injected_tree, - (LocalPromptCacheTreeClient, RemotePromptCacheTreeClient), - ) - - selected, _, extra = await manager.select_p_d_node(PROMPT, None, None) - - assert selected is cached - assert extra.estimated_cache_hit_rate > 0.5 - assert extra.cache_last_insert_time is not None - original_extra = extra - assert prefix_match.call_count == 1 - - restarted = _node(300, cached.client_ip_port) - _set_nodes(manager, [other, restarted], decode) - selected, _, extra = await manager.select_p_d_node(PROMPT, None, None) - assert selected is restarted - assert extra == original_extra - assert prefix_match.call_count == 2 - assert selector.policy.prompt_cache_tree is injected_tree - - asyncio.run(run()) - - -def test_cache_feedback_and_dynamic_thresholds_stay_private_to_each_worker(backend): - first, state = _generator_manager(backend, "cache_aware") - second, _ = _generator_manager(backend, "cache_aware") - stats_calls = tree_calls = None - if state.stats is not None: - stats_calls = Mock(wraps=state.stats.get_dispatched_info) - - def prefix_match_sampled(*args): - return PromptCacheMatchResult(*state.tree.prefix_match_sampled(*args)) - - tree_calls = Mock(wraps=prefix_match_sampled) - for manager in (first, second): - selector = manager.pd_manager.selector - selector.pd_node_stats._get_dispatched_info = stats_calls - selector.policy.prompt_cache_tree._prefix_match_sampled = tree_calls - - for manager, rates in ( - (first, (0.2, 0.2, 0.8, 0.8)), - (second, (0.8, 0.8, 0.2, 0.2)), - ): - selector = manager.pd_manager.selector - selector.policy.config.cache_hit_rate_window_size = 2 - for rate in rates: - assert selector.record_prompt_cache_hit_rate(rate) is None - - first_policy, second_policy = ( - first.pd_manager.selector.policy, - second.pd_manager.selector.policy, - ) - assert first_policy.config.balance_rel_threshold == pytest.approx(1.55) - assert second_policy.config.balance_rel_threshold == pytest.approx(1.45) - assert first_policy.balance_rel_threshold_controller is not second_policy.balance_rel_threshold_controller - if stats_calls is not None: - stats_calls.assert_not_called() - tree_calls.assert_not_called() - - async def query_both_workers(): - for manager in (first, second): - await manager.select_p_d_node(PROMPT, None, None) - - asyncio.run(query_both_workers()) - - assert first_policy.config.balance_rel_threshold == pytest.approx(1.55) - assert second_policy.config.balance_rel_threshold == pytest.approx(1.45) - if stats_calls is not None: - assert stats_calls.call_count == 2 - assert tree_calls.call_count == 2 - - -def _params(max_new_tokens=3): - params = SamplingParams() - params.init(None, max_new_tokens=max_new_tokens) - return params - - -async def _consume(manager, params): - return [item async for item in manager._generate_one_attempt("prompt", params, None, None, 0, 80, 2)] - - -def test_cache_aware_selection_and_initial_reservation_are_serialized_per_worker(backend): - async def run(): - manager, state = _generator_manager(backend, "cache_aware") - cached = _node(100, "p1:8000") - other = _node(101, "p2:8000") - _set_nodes(manager, [cached, other], manager.pd_manager.decode_nodes) - manager.pd_manager.selector.policy.prompt_cache_tree.prefix_match_sampled = AsyncMock( - return_value=_match(cached) - ) - - active_selections = 0 - max_active_selections = 0 - original_select = manager.select_p_d_node - - async def select_p_d_node(*args): - nonlocal active_selections, max_active_selections - active_selections += 1 - max_active_selections = max(max_active_selections, active_selections) - try: - # 主动让出事件循环，验证并发请求仍无法同时进入选点区间。 - await asyncio.sleep(0) - return await original_select(*args) - finally: - active_selections -= 1 - - manager.select_p_d_node = select_p_d_node - selected_nodes = [] - all_reserved = asyncio.Event() - release_requests = asyncio.Event() - request_count = 16 - request_ids = iter(range(8, 8 * (request_count + 1), 8)) - manager.id_gen = SimpleNamespace(generate_id=lambda: next(request_ids)) - - async def generate_segment(p_node, _d_node, _start_time, _prompt, params, _multimodal, _request): - selected_nodes.append(p_node) - if len(selected_nodes) == request_count: - all_reserved.set() - await release_requests.wait() - yield ( - params.group_request_id, - "x", - { - "prompt_tokens": len(PROMPT), - "prompt_cache_len": 0, - "needs_prefill_first_token": False, - }, - FinishStatus(FinishStatus.FINISHED_STOP), - ) - - manager._wait_to_token_package = generate_segment - - async def consume(origin_request_id): - return [ - item - async for item in manager._generate_one_attempt( - PROMPT, _params(max_new_tokens=1), None, None, 0, origin_request_id, 2 - ) - ] - - tasks = [asyncio.create_task(consume(1000 + index * 8)) for index in range(request_count)] - try: - await asyncio.wait_for(all_reserved.wait(), timeout=5) - assert max_active_selections == 1 - assert cached in selected_nodes - assert other in selected_nodes - finally: - release_requests.set() - await asyncio.gather(*tasks) - - assert _pending(state, cached) == (0, 0) - assert _pending(state, other) == (0, 0) - - asyncio.run(run()) - - -@pytest.mark.parametrize("strategy", ["random", "round_robin", "adaptive_load", "cache_aware"]) -@pytest.mark.parametrize("needs_prefill", [True, False]) -def test_segment_counters_clear_at_first_token_and_restart_for_continuation(backend, strategy, needs_prefill): - manager, state = _generator_manager(backend, strategy) - prefill = manager.pd_manager.prefill_nodes[0] - segments = [] - - async def generate_segment(p_node, d_node, start_time, prompt, params, multimodal, request): - segments.append(prompt) - assert _pending(state, p_node) == (len(prompt), 1) - yield ( - params.group_request_id, - "x", - { - "prompt_tokens": 100, - "prompt_cache_len": 20, - "needs_prefill_first_token": needs_prefill, - }, - FinishStatus(), - ) - assert _pending(state, p_node) == (0, 0) - finish = FinishStatus.FINISHED_PD_DECODE_CAPACITY if len(segments) == 1 else FinishStatus.FINISHED_STOP - yield params.group_request_id, "", {"prompt_tokens": 100}, FinishStatus(finish) - - manager._wait_to_token_package = generate_segment - results = asyncio.run(_consume(manager, _params())) - - assert segments == ["prompt", "promptx"] - assert len(results) == 3 - assert _pending(state, prefill) == (0, 0) - assert not manager.active_req_ids - stats = manager.pd_manager.selector.pd_node_stats - assert stats.add_dispatched_info.await_args_list == [ - call(p_node=prefill, chars=len("prompt"), req_num=1), - call(p_node=prefill, chars=len("promptx"), req_num=1), - ] - assert stats.sub_dispatched_info.await_args_list == [ - call(p_node=prefill, chars=len("prompt"), req_num=1), - call(p_node=prefill, chars=len("promptx"), req_num=1), - ] - selector = manager.pd_manager.selector - selector.record_prompt_cache_hit_rate.assert_called_once_with(0.2) - if strategy == "cache_aware": - assert selector.policy.balance_rel_threshold_controller._cache_hit_rates == [0.2] - if needs_prefill: - sample_stride = manager.pd_manager.selector.policy.config.sample_stride - selector.policy.prompt_cache_tree.insert_sampled.assert_awaited_once_with( - "prompt"[::sample_stride], prefill.connection_key, sample_stride - ) - else: - selector.policy.prompt_cache_tree.insert_sampled.assert_not_awaited() - - -def test_continuation_tokenization_failure_clears_reserved_counter(backend): - manager, state = _generator_manager(backend) - prefill = manager.pd_manager.prefill_nodes[0] - params = SamplingParams() - params.init(None, max_new_tokens=3, regular_constraint="ab") - - async def generate_first_segment(_p_node, _d_node, _start_time, _prompt, params, _multimodal, _request): - yield ( - params.group_request_id, - "x", - { - "prompt_tokens": 100, - "prompt_cache_len": 0, - "needs_prefill_first_token": False, - }, - FinishStatus(), - ) - yield ( - params.group_request_id, - "", - {"prompt_tokens": 100}, - FinishStatus(FinishStatus.FINISHED_PD_DECODE_CAPACITY), - ) - - manager._wait_to_token_package = generate_first_segment - manager.tokens = Mock(side_effect=RuntimeError("continuation tokenization failed")) - - with pytest.raises(RuntimeError, match="continuation tokenization failed"): - asyncio.run(_consume(manager, params)) - - assert manager.pd_node_stats.add_dispatched_info.await_args_list == [ - call(p_node=prefill, chars=len("prompt"), req_num=1), - call(p_node=prefill, chars=len("promptx"), req_num=1), - ] - assert manager.pd_node_stats.sub_dispatched_info.await_args_list == [ - call(p_node=prefill, chars=len("prompt"), req_num=1), - call(p_node=prefill, chars=len("promptx"), req_num=1), - ] - assert _pending(state, prefill) == (0, 0) - - -def test_error_output_records_hit_rate_without_inserting_cache(backend): - manager, state = _generator_manager(backend, "cache_aware") - prefill = manager.pd_manager.prefill_nodes[0] - - async def generate_segment(p_node, d_node, start_time, prompt, params, multimodal, request): - yield ( - params.group_request_id, - "", - { - "prompt_tokens": 100, - "prompt_cache_len": 20, - "needs_prefill_first_token": True, - }, - FinishStatus(FinishStatus.FINISHED_ERROR), - ) - - manager._wait_to_token_package = generate_segment - results = asyncio.run(_consume(manager, _params())) - - assert len(results) == 1 and results[0][3].is_error_finished() - assert _pending(state, prefill) == (0, 0) - selector = manager.pd_manager.selector - selector.record_prompt_cache_hit_rate.assert_called_once_with(0.2) - assert selector.policy.balance_rel_threshold_controller._cache_hit_rates == [0.2] - selector.policy.prompt_cache_tree.insert_sampled.assert_not_awaited() - - -def test_error_before_first_token_clears_segment_counters(backend): - manager, state = _generator_manager(backend) - prefill = manager.pd_manager.prefill_nodes[0] - - async def generate_segment(*_args): - assert _pending(state, prefill) == (len("prompt"), 1) - raise RuntimeError("before first token") - yield - - manager._wait_to_token_package = generate_segment - with pytest.raises(RuntimeError, match="before first token"): - asyncio.run(_consume(manager, _params())) - - assert _pending(state, prefill) == (0, 0) - assert not manager.active_req_ids - manager.abort.assert_awaited_once() - manager.remove_req.assert_awaited_once() - manager.pd_node_stats.sub_dispatched_info.assert_awaited_once_with(p_node=prefill, chars=len("prompt"), req_num=1) - - -def test_sampling_params_copy_failure_clears_initial_reservation(monkeypatch, backend): - manager, state = _generator_manager(backend) - prefill = manager.pd_manager.prefill_nodes[0] - manager._wait_to_token_package = Mock() - params = _params() - - def fail_copy(*_args, **_kwargs): - raise RuntimeError("segment initialization failed") - - monkeypatch.setattr(SamplingParams, "from_buffer_copy", fail_copy) - with pytest.raises(RuntimeError, match="segment initialization"): - asyncio.run(_consume(manager, params)) - manager.pd_node_stats.add_dispatched_info.assert_awaited_once_with(p_node=prefill, chars=len("prompt"), req_num=1) - manager.pd_node_stats.sub_dispatched_info.assert_awaited_once_with(p_node=prefill, chars=len("prompt"), req_num=1) - assert _pending(state, prefill) == (0, 0) - manager._wait_to_token_package.assert_not_called() - manager.remove_req.assert_awaited_once() - assert not manager.active_req_ids - - -@pytest.mark.parametrize("replaced", ["prefill", "decode"]) -def test_connection_closed_during_counter_update_fails_at_dispatch_and_clears_load(backend, replaced): - async def run(): - manager, state = _generator_manager(backend) - prefill, decode = ( - manager.pd_manager.prefill_nodes[0], - manager.pd_manager.decode_nodes[0], - ) - created, complete = asyncio.Event(), asyncio.Event() - original_add_info = manager.pd_node_stats.add_dispatched_info - - async def add_info(**kwargs): - await original_add_info(**kwargs) - created.set() - await complete.wait() - - async def generate_segment(p_node, d_node, *_args): - assert p_node is prefill and d_node is decode - raise ConnectionError("selected PD connection closed") - yield - - manager.pd_node_stats.add_dispatched_info = add_info - manager._wait_to_token_package = Mock(side_effect=generate_segment) - task = asyncio.create_task(_consume(manager, _params())) - await asyncio.wait_for(created.wait(), timeout=2) - previous = prefill if replaced == "prefill" else decode - replacement = _node( - previous.node_id, - previous.client_ip_port, - mode=previous.mode, - chars=23, - requests=2, - ) - _set_nodes( - manager, - [replacement] if replaced == "prefill" else [prefill], - [replacement] if replaced == "decode" else [decode], - ) - complete.set() - with pytest.raises(ConnectionError, match="selected PD connection closed"): - await task - - manager._wait_to_token_package.assert_called_once() - manager.abort.assert_awaited_once() - manager.remove_req.assert_awaited_once() - assert _pending(state, prefill) == (0, 0) - assert ( - replacement.dispatched_prompt_chars, - replacement.dispatched_req_num, - ) == (23, 2) - assert not manager.active_req_ids - - asyncio.run(run()) - - -def test_cancelled_request_clears_counters_before_removing_request(backend): - manager, state = _generator_manager(backend) - prefill = manager.pd_manager.prefill_nodes[0] - - async def run(): - started = asyncio.Event() - - async def remove_req(*_args): - assert _pending(state, prefill) == (0, 0) - - manager.remove_req.side_effect = remove_req - - async def generate_segment(*_args): - started.set() - await asyncio.Event().wait() - yield - - manager._wait_to_token_package = generate_segment - task = asyncio.create_task(_consume(manager, _params())) - try: - await asyncio.wait_for(started.wait(), timeout=2) - assert _pending(state, prefill) == (len("prompt"), 1) - task.cancel() - with pytest.raises(asyncio.CancelledError): - await task - finally: - task.cancel() - await asyncio.gather(task, return_exceptions=True) - - asyncio.run(run()) - assert _pending(state, prefill) == (0, 0) - assert not manager.active_req_ids - manager.remove_req.assert_awaited_once() - - -def test_cancellation_after_counter_release_does_not_release_twice(backend): - manager, state = _generator_manager(backend) - prefill = manager.pd_manager.prefill_nodes[0] - - async def run(): - released = asyncio.Event() - original_sub_dispatched_info = manager.pd_node_stats.sub_dispatched_info - - async def sub_dispatched_info(**kwargs): - await original_sub_dispatched_info(**kwargs) - released.set() - await asyncio.Event().wait() - - manager.pd_node_stats.sub_dispatched_info = AsyncMock(side_effect=sub_dispatched_info) - - async def generate_segment(_p_node, _d_node, _start_time, _prompt, params, _multimodal, _request): - yield params.group_request_id, "x", {"prompt_tokens": 1}, FinishStatus() - - manager._wait_to_token_package = generate_segment - task = asyncio.create_task(_consume(manager, _params())) - await asyncio.wait_for(released.wait(), timeout=2) - task.cancel() - with pytest.raises(asyncio.CancelledError): - await task - - asyncio.run(run()) - assert _pending(state, prefill) == (0, 0) - manager.pd_node_stats.sub_dispatched_info.assert_awaited_once() diff --git a/unit_tests/server/test_pd_master_haproxy_proxy.py b/unit_tests/server/test_pd_master_haproxy_proxy.py deleted file mode 100644 index 3642bc69e3..0000000000 --- a/unit_tests/server/test_pd_master_haproxy_proxy.py +++ /dev/null @@ -1,434 +0,0 @@ -"""Exercise the real streaming proxy bundled by ray-haproxy.""" - -from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer -import json -import os -from pathlib import Path -import signal -import socket -import subprocess -import threading -import time -from urllib.error import HTTPError, URLError -from urllib.request import ProxyHandler, Request, build_opener - -import pytest -import psutil - -from lightllm.server.core.objs.start_args_type import StartArgs -from lightllm.server.httpserver_for_pd_master.multiprocess import ( - _build_haproxy_config, - _build_initial_health_state, -) -from lightllm.utils import start_utils -from lightllm.utils.start_utils import kill_recursive -from unit_tests.server.pd_master_test_utils import get_test_haproxy_binary - - -def test_real_haproxy_starts_with_a_low_fd_hard_limit(tmp_path): - binary = str(Path(get_test_haproxy_binary()).resolve()) - with socket.socket() as listener: - listener.bind(("127.0.0.1", 0)) - public_port = listener.getsockname()[1] - - args = StartArgs( - run_mode="pd_master", - host="127.0.0.1", - port=public_port, - httpserver_workers=2, - ) - state_file = tmp_path / "haproxy-low-fd.state" - state_file.write_text(_build_initial_health_state([1, 2]), encoding="utf-8") - config = tmp_path / "haproxy-low-fd.cfg" - config.write_text(_build_haproxy_config(args, [1, 2], state_file), encoding="utf-8") - error_log_path = tmp_path / "haproxy-low-fd.stderr" - # 同时降低 soft/hard limit，验证 HAProxy 会按实际 FD 上限自动缩小 maxconn。 - command = [ - "/bin/bash", - "-c", - 'ulimit -n 1024 && exec "$@"', - "lightllm-haproxy-low-fd", - binary, - "-db", - "-f", - str(config), - ] - with error_log_path.open("wb") as error_log: - haproxy = subprocess.Popen(command, stdout=subprocess.DEVNULL, stderr=error_log) - try: - deadline = time.monotonic() + 5 - while True: - assert haproxy.poll() is None, error_log_path.read_text() - try: - with socket.create_connection(("127.0.0.1", public_port), timeout=0.2): - break - except OSError: - assert time.monotonic() < deadline, error_log_path.read_text() - time.sleep(0.02) - finally: - if haproxy.poll() is None: - haproxy.send_signal(signal.SIGUSR1) - haproxy.wait(timeout=5) - - assert haproxy.returncode == 0 - - -@pytest.fixture -def haproxy_proxy(tmp_path, request): - binary = get_test_haproxy_binary() - state = { - "requests": [], - "lock": threading.Lock(), - "release_stream": threading.Event(), - "release_initial_probes": threading.Event(), - "stream_disconnected": threading.Event(), - "runtime": tmp_path, - "probe_status": {worker: {path: 200 for path in ("/liveness", "/readiness", "/health")} for worker in range(2)}, - "probe_delays": {}, - "probe_observations": [], - } - if getattr(request, "param", None) != "block_initial_health_checks": - state["release_initial_probes"].set() - - class Handler(BaseHTTPRequestHandler): - protocol_version = "HTTP/1.1" - - def log_message(self, *args): - pass - - def do_GET(self): - if self.path in ("/stream", "/stream_until_disconnect"): - self.send_response(200) - self.send_header("Content-Type", "text/event-stream") - self.send_header("Connection", "close") - self.end_headers() - self.close_connection = True - try: - self.wfile.write(b"data: first\n\n") - self.wfile.flush() - if self.path == "/stream_until_disconnect": - for _ in range(500): - self.wfile.write(b"data: " + b"x" * (64 * 1024) + b"\n\n") - self.wfile.flush() - time.sleep(0.01) - else: - state["release_stream"].wait(timeout=5) - self.wfile.write(b"data: last\n\n") - self.wfile.flush() - except (BrokenPipeError, ConnectionResetError): - state["stream_disconnected"].set() - return - with state["lock"]: - status = state["probe_status"][self.server.worker_id].get(self.path, 200) - delay = state["probe_delays"].get((self.server.worker_id, self.path), 0) - if self.path in ("/liveness", "/readiness", "/health"): - state["release_initial_probes"].wait(timeout=15) - if delay: - time.sleep(delay) - with state["lock"]: - if self.path in ("/liveness", "/readiness", "/health"): - state["probe_observations"].append((self.server.worker_id, self.path, status)) - body = json.dumps({"worker_id": self.server.worker_id}).encode() - self.send_response(status) - self.send_header("Content-Type", "application/json") - self.send_header("Content-Length", str(len(body))) - self.end_headers() - try: - self.wfile.write(body) - except (BrokenPipeError, ConnectionResetError): - pass - - def do_POST(self): - body = self.rfile.read(int(self.headers["Content-Length"])) - with state["lock"]: - state["requests"].append((self.server.worker_id, body)) - # Simulate a worker dying after accepting a non-idempotent request. - # A proxy retry would submit this same body to the other worker. - self.close_connection = True - self.connection.shutdown(socket.SHUT_RDWR) - self.connection.close() - - servers = [] - threads = [] - haproxy = None - try: - for worker_id in range(2): - server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) - server.daemon_threads = True - server.worker_id = worker_id - thread = threading.Thread(target=server.serve_forever, kwargs={"poll_interval": 0.05}, daemon=True) - thread.start() - servers.append(server) - threads.append(thread) - with socket.socket() as listener: - listener.bind(("127.0.0.1", 0)) - public_port = listener.getsockname()[1] - args = StartArgs( - run_mode="pd_master", - host="127.0.0.1", - port=public_port, - httpserver_workers=2, - ) - worker_ports = [server.server_port for server in servers] - state_file = tmp_path / "haproxy-health.state" - state_file.write_text(_build_initial_health_state(worker_ports), encoding="utf-8") - config = tmp_path / "haproxy.cfg" - config.write_text(_build_haproxy_config(args, worker_ports, state_file), encoding="utf-8") - executable = str(Path(binary).resolve()) - check = subprocess.run([executable, "-c", "-f", str(config)], capture_output=True, text=True, timeout=5) - assert check.returncode == 0, check.stderr - command = [executable, "-db", "-f", str(config)] - with (tmp_path / "haproxy.stderr").open("wb") as error_log: - haproxy = subprocess.Popen(command, stdout=subprocess.DEVNULL, stderr=error_log, start_new_session=True) - state["haproxy"] = haproxy - opener = build_opener(ProxyHandler({})) - base_url = f"http://127.0.0.1:{public_port}" - deadline = time.monotonic() + 5 - while True: - assert haproxy.poll() is None, (tmp_path / "haproxy.stderr").read_text() - try: - with opener.open(base_url + "/worker", timeout=0.5) as response: - assert response.status == 200 - break - except (OSError, URLError): - if time.monotonic() >= deadline: - pytest.fail("HAProxy did not start: " + (tmp_path / "haproxy.stderr").read_text()) - time.sleep(0.02) - yield base_url, opener, state - finally: - state["release_stream"].set() - state["release_initial_probes"].set() - if haproxy is not None: - if haproxy.poll() is None: - haproxy.send_signal(signal.SIGUSR1) - try: - haproxy.wait(timeout=5) - except subprocess.TimeoutExpired: - kill_recursive(haproxy) - haproxy.wait(timeout=5) - try: - os.killpg(haproxy.pid, signal.SIGKILL) - except ProcessLookupError: - pass - for server in servers: - server.shutdown() - server.server_close() - for thread in threads: - thread.join(timeout=1) - - -def test_real_haproxy_balances_workers_and_delivers_sse_before_stream_completion(haproxy_proxy): - base_url, opener, state = haproxy_proxy - proxy_process = psutil.Process(state["haproxy"].pid) - assert proxy_process.children(recursive=True) == [] - assert proxy_process.num_threads() >= 2 - workers = set() - for _ in range(8): - with opener.open(base_url + "/worker", timeout=3) as response: - workers.add(json.load(response)["worker_id"]) - assert workers == {0, 1} - - # The backend waits for this client to read the first event before it sends - # the last one. Buffering the complete stream would deadlock until timeout. - with opener.open(base_url + "/stream", timeout=3) as response: - assert response.headers["Content-Type"] == "text/event-stream" - assert response.readline() == b"data: first\n" - assert response.readline() == b"\n" - assert not state["release_stream"].is_set() - state["release_stream"].set() - assert response.read() == b"data: last\n\n" - - -def test_real_haproxy_never_replays_a_post_and_keeps_pd_websockets_direct(haproxy_proxy): - base_url, opener, state = haproxy_proxy - body = b'{"prompt":"submit once"}' - with pytest.raises(HTTPError) as response: - opener.open(Request(base_url + "/generate", data=body, method="POST"), timeout=3) - assert response.value.code == 502 - response.value.close() - assert [request_body for _worker_id, request_body in state["requests"]] == [body] - - websocket_headers = { - "Connection": "Upgrade", - "Upgrade": "websocket", - "Sec-WebSocket-Version": "13", - "Sec-WebSocket-Key": "dGhlIHNhbXBsZSBub25jZQ==", - } - for path in ("/pd_register", "/kv_move_status"): - for headers in ({}, websocket_headers): - with pytest.raises(HTTPError) as response: - opener.open(Request(base_url + path, headers=headers), timeout=3) - assert response.value.code == 426 - response.value.close() - - -def test_real_haproxy_propagates_client_disconnect_to_the_request_worker(haproxy_proxy): - base_url, opener, state = haproxy_proxy - response = opener.open(base_url + "/stream_until_disconnect", timeout=3) - assert response.readline() == b"data: first\n" - response.close() - assert state["stream_disconnected"].wait(timeout=3), "Backend continued producing after the client disconnected" - - -def test_launcher_sigterm_drains_an_active_sse_response(haproxy_proxy, monkeypatch): - base_url, opener, state = haproxy_proxy - registered_handlers = {} - shutdown_result = [] - process_manager = start_utils.SubmoduleManager() - monkeypatch.setattr( - start_utils.signal, - "signal", - lambda sig, handler: registered_handlers.__setitem__(sig, handler), - ) - monkeypatch.setattr(process_manager, "terminate_all_processes", lambda: None) - process_manager.setup_signal_handlers( - state["haproxy"], - graceful_shutdown_signal=signal.SIGUSR1, - ) - - def stop_launcher(): - try: - registered_handlers[signal.SIGTERM](signal.SIGTERM, None) - except SystemExit as error: - shutdown_result.append(error.code) - - with opener.open(base_url + "/stream", timeout=3) as response: - assert response.readline() == b"data: first\n" - assert response.readline() == b"\n" - - shutdown_thread = threading.Thread(target=stop_launcher) - shutdown_thread.start() - time.sleep(0.1) - assert state["haproxy"].poll() is None, "HAProxy exited before the active stream completed" - assert state["runtime"].is_dir() - - state["release_stream"].set() - assert response.read() == b"data: last\n\n" - shutdown_thread.join(timeout=5) - assert not shutdown_thread.is_alive() - assert shutdown_result == [0] - assert state["haproxy"].returncode == 0 - - -def _probe_response(opener, base_url, path, method="GET"): - request = Request(base_url + path, data=b"" if method == "POST" else None, method=method) - try: - response = opener.open(request, timeout=3) - except HTTPError as error: - response = error - with response: - return response.status, response.read(), response.headers.get_content_type() - - -def _wait_for_probe_response(opener, base_url, path, status, timeout=15): - deadline = time.monotonic() + timeout - while True: - response = _probe_response(opener, base_url, path) - if response[0] == status: - return response - assert time.monotonic() < deadline, response - time.sleep(0.05) - - -@pytest.mark.parametrize("haproxy_proxy", ["block_initial_health_checks"], indirect=True) -def test_real_haproxy_reports_unready_until_every_initial_probe_succeeds(haproxy_proxy): - base_url, opener, state = haproxy_proxy - - # 后端首次探测仍在处理中时，三类健康接口及 healthz 别名必须全部返回 503。 - for path in ("/liveness", "/readiness", "/health", "/healthz"): - assert _probe_response(opener, base_url, path)[0] == 503 - with opener.open(base_url + "/worker", timeout=3) as response: - assert response.status == 200 - - state["release_initial_probes"].set() - for path in ("/liveness", "/readiness", "/health", "/healthz"): - assert _wait_for_probe_response(opener, base_url, path, 200)[0] == 200 - - -@pytest.mark.parametrize( - "path,failure,success_body,error_body", - [ - ("/liveness", 503, {"status": "ok"}, {"status": "not alive"}), - ("/readiness", 302, {"status": "ok"}, {"status": "not ready"}), - ("/health", 500, {"message": "Ok"}, {"message": "Error"}), - ], -) -def test_real_haproxy_requires_every_worker_for_each_probe(haproxy_proxy, path, failure, success_body, error_body): - base_url, opener, state = haproxy_proxy - for initial_path in ("/liveness", "/readiness", "/health"): - _wait_for_probe_response(opener, base_url, initial_path, 200) - assert json.loads(_probe_response(opener, base_url, path)[1]) == success_body - - def change_worker_status(worker, status): - with state["lock"]: - before = len(state["probe_observations"]) - state["probe_status"][worker][path] = status - # Wait for a new probe of this worker, rather than treating HAProxy's - # initial UP state or an old response as evidence of recovery. - deadline = time.monotonic() + 15 - while True: - with state["lock"]: - observed = (worker, path, status) in state["probe_observations"][before:] - if observed: - return - assert time.monotonic() < deadline, (worker, path, status) - time.sleep(0.05) - - change_worker_status(1, failure) - status, body, content_type = _wait_for_probe_response(opener, base_url, path, 503) - assert content_type == "application/json" - assert json.loads(body) == error_body - for other in ("/liveness", "/readiness", "/health"): - if other != path: - assert _probe_response(opener, base_url, other)[0] == 200 - - # Probe failure must not filter ordinary traffic. - for ordinary_path in ("/worker", "/metrics"): - workers = set() - for _ in range(6): - with opener.open(base_url + ordinary_path, timeout=3) as response: - workers.add(json.load(response)["worker_id"]) - assert workers == {0, 1} - aliases = ("/health", "/healthz") if path == "/health" else (path,) - for alias in aliases: - assert _probe_response(opener, base_url, alias + "?detail=1")[0] == 503 - method = "HEAD" if path == "/health" else "POST" - method_status, method_body, _ = _probe_response(opener, base_url, path, method) - assert method_status == 503 - if method == "HEAD": - assert method_body == b"" - else: - assert json.loads(method_body) == error_body - # Unsupported methods still reach the original backend handler. - assert _probe_response(opener, base_url, path, "HEAD")[0] == 501 - - change_worker_status(0, failure) - assert _wait_for_probe_response(opener, base_url, path, 503)[0] == 503 - change_worker_status(1, 200) - assert _wait_for_probe_response(opener, base_url, path, 503)[0] == 503 - change_worker_status(0, 200) - assert json.loads(_wait_for_probe_response(opener, base_url, path, 200)[1]) == success_body - if method == "HEAD": - assert _probe_response(opener, base_url, path, method)[:2] == (200, b"") - else: - assert json.loads(_probe_response(opener, base_url, path, method)[1]) == success_body - - -def test_real_haproxy_allows_slow_health_checks_and_rejects_a_timed_out_worker(haproxy_proxy): - base_url, opener, state = haproxy_proxy - with state["lock"]: - state["probe_status"][1]["/health"] = 503 - _wait_for_probe_response(opener, base_url, "/health", 503) - - # Recover from a known DOWN state so the initial UP default cannot hide a - # check timeout that is too short for the fixed-topology eight-second probe. - with state["lock"]: - state["probe_status"][1]["/health"] = 200 - state["probe_delays"][(1, "/health")] = 8.5 - _wait_for_probe_response(opener, base_url, "/health", 200, timeout=25) - - with state["lock"]: - state["probe_delays"][(1, "/health")] = 13 - _wait_for_probe_response(opener, base_url, "/health", 503, timeout=30) - assert _probe_response(opener, base_url, "/liveness")[0] == 200 - assert _probe_response(opener, base_url, "/readiness")[0] == 200 diff --git a/unit_tests/server/test_pd_master_launcher_integration.py b/unit_tests/server/test_pd_master_launcher_integration.py deleted file mode 100644 index 8833be9202..0000000000 --- a/unit_tests/server/test_pd_master_launcher_integration.py +++ /dev/null @@ -1,97 +0,0 @@ -"""Exercise real PD master startup, supervision, and cleanup.""" - -import json -from pathlib import Path -import time -from urllib.error import HTTPError - -import pytest - -from unit_tests.server.pd_master_test_utils import assert_processes_exit, running_pd_master - - -@pytest.mark.parametrize( - "failed_process", - [ - "api_server", - "pd_node_stats_server", - "prompt_cache_tree_server", - "req_id_generator_server", - "metric_manager", - "haproxy", - ], -) -def test_real_launcher_discovers_workers_before_pd_readiness_and_fails_as_a_group(tmp_path, failed_process): - with running_pd_master(tmp_path) as master: - assert len({worker["node_id"] for worker in master.workers}) == 2 - assert len({worker["port"] for worker in master.workers}) == 2 - public_port = int(master.public_url.rsplit(":", 1)[1]) - for endpoint in master.workers: - info = master.get_json("/get_server_info", "http://" + endpoint["host_ip_port"]) - assert info["port"] == public_port - assert info["pd_node_id"] == endpoint["node_id"] - assert info["httpserver_workers"] == len(master.workers) - assert set(info["pd_master_multi_ports"]) == {worker["port"] for worker in master.workers} - assert info["config_server_host"] is None - assert info["config_server_port"] is None - assert "pd_master_ipc_name" not in info - assert "metric_port" not in info - - deadline = time.monotonic() + 15 - while True: - try: - master.get_json("/readiness") - except HTTPError as readiness: - assert readiness.code == 503 - assert json.load(readiness) == {"status": "not ready"} - break - assert time.monotonic() < deadline, "Public readiness checks did not converge" - time.sleep(0.1) - with master.opener.open(master.public_url + "/metrics", timeout=3) as response: - assert b"lightllm_request_count" in response.read() - - runtime = master.haproxy_runtime_dir() - assert runtime.is_dir() - descendants = master.descendants() - if failed_process == "haproxy": - candidates = [ - process - for process in descendants - if "haproxy" in process.name() and process.ppid() == master.process.pid - ] - else: - candidates = [process for process in descendants if "::" + failed_process in process.name()] - assert candidates, [(process.pid, process.name()) for process in descendants] - candidates[0].kill() - assert master.process.wait(timeout=20) != 0 - assert_processes_exit(descendants) - assert not runtime.exists() - - -@pytest.mark.parametrize( - "failure_stage", - ["backend_start", "proxy_start"], -) -def test_partial_startup_failure_stops_every_started_process_and_removes_runtime(tmp_path, failure_stage): - launcher_script = Path(__file__).with_name("pd_master_fault_launcher.py") - observations = tmp_path / "cleanup.json" - with running_pd_master( - tmp_path, - launcher_script=launcher_script, - wait_for_startup=False, - extra_args=("--failure-stage", failure_stage, "--cleanup-observations", str(observations)), - ) as master: - descendants = {} - deadline = time.monotonic() + 60 - while master.process.poll() is None: - descendants.update({process.pid: process for process in master.descendants()}) - assert time.monotonic() < deadline, master.log_path.read_text() - time.sleep(0.05) - assert master.process.returncode != 0 - assert "injected" in master.log_path.read_text() or failure_stage == "backend_start" - assert_processes_exit(descendants.values()) - cleanup = json.loads(observations.read_text()) - assert all(not Path(runtime).exists() for runtime in cleanup["runtime_paths"]) - assert cleanup["runtime_before_stop"] == [True] - assert cleanup["runtime_after_stop"] == [True] - assert cleanup["backends_stopped"] is True diff --git a/unit_tests/server/test_pd_master_mode.py b/unit_tests/server/test_pd_master_mode.py deleted file mode 100644 index d9163366f0..0000000000 --- a/unit_tests/server/test_pd_master_mode.py +++ /dev/null @@ -1,612 +0,0 @@ -import asyncio -import json -from types import SimpleNamespace -from unittest.mock import Mock - -import pytest -from easydict import EasyDict - -from lightllm.server.api_cli import make_argument_parser -from lightllm.server.core.objs.start_args_type import StartArgs -from lightllm.server.core.objs import StopStrMatchMode, FinishStatus, SamplingParams -from lightllm.server.httpserver_for_pd_master.pd_selector import pd_node_stats_client -from lightllm.server.httpserver_for_pd_master.manager import HttpServerManagerForPDMaster, PDManager -from lightllm.server.httpserver_for_pd_master.pd_selector import prompt_cache_tree_client - - -@pytest.fixture(autouse=True) -def local_pd_master_clients(monkeypatch): - start_args = SimpleNamespace(httpserver_workers=1) - monkeypatch.setattr(pd_node_stats_client, "get_env_start_args", lambda: start_args) - monkeypatch.setattr( - prompt_cache_tree_client, - "get_env_start_args", - lambda: start_args, - ) - - -def test_pd_node_self_request_limit_cli_defaults_to_enabled_and_can_be_disabled(): - parser = make_argument_parser() - - assert parser.parse_args([]).disable_pd_node_self_request_limit is False - assert parser.parse_args(["--disable_pd_node_self_request_limit"]).disable_pd_node_self_request_limit is True - assert StartArgs().disable_pd_node_self_request_limit is False - - -def test_pd_cache_high_priority_cli_defaults_to_enabled_and_can_be_disabled(): - parser = make_argument_parser() - - assert parser.parse_args([]).disable_pd_cache_high_priority is False - assert parser.parse_args(["--disable_pd_cache_high_priority"]).disable_pd_cache_high_priority is True - assert StartArgs().disable_pd_cache_high_priority is False - - -def test_pd_master_request_slot_is_reserved_noop(): - manager = HttpServerManagerForPDMaster.__new__(HttpServerManagerForPDMaster) - - asyncio.run(manager._wait_for_pd_master_request_slot()) - - -def test_auto_set_response_parsers_from_qwen35_model_config(tmp_path): - from lightllm.utils.config_utils import auto_set_response_parsers - - (tmp_path / "config.json").write_text('{"model_type": "qwen3_5"}', encoding="utf-8") - args = StartArgs(model_dir=str(tmp_path)) - - auto_set_response_parsers(args) - - assert args.tool_call_parser == "qwen3_coder" - assert args.reasoning_parser == "qwen3" - - -def test_auto_set_response_parsers_preserves_explicit_values(tmp_path): - from lightllm.utils.config_utils import auto_set_response_parsers - - (tmp_path / "config.json").write_text('{"model_type": "qwen3_5"}', encoding="utf-8") - args = StartArgs( - model_dir=str(tmp_path), - tool_call_parser="llama3", - reasoning_parser="deepseek-r1", - ) - - auto_set_response_parsers(args) - - assert args.tool_call_parser == "llama3" - assert args.reasoning_parser == "deepseek-r1" - - -def test_get_server_info_serializes_runtime_easydict(monkeypatch): - from lightllm.server import api_http - - runtime_args = EasyDict( - model_name="qwen35_0.8b", - tool_call_parser="qwen3_coder", - reasoning_parser="qwen3", - ) - monkeypatch.setattr(api_http.g_objs, "args", runtime_args) - - assert api_http.get_server_info() == dict(runtime_args) - - -@pytest.mark.parametrize( - "runtime,worker_ports", - [(False, None), (True, None), (True, [8100, 8101])], -) -def test_pd_master_server_info_reports_worker_ports_without_changing_start_args(monkeypatch, runtime, worker_ports): - from lightllm.server import api_http - - worker_count = 1 if worker_ports is None else len(worker_ports) - start_args = StartArgs( - run_mode="pd_master", - httpserver_workers=worker_count, - pd_node_id=100, - ) - original = vars(start_args).copy() - args = EasyDict(original) if runtime else start_args - monkeypatch.setattr(api_http.g_objs, "args", args) - if worker_ports is not None: - monkeypatch.setattr( - api_http, - "get_shm_port_args", - lambda: SimpleNamespace(pd_master_multi_ports=worker_ports), - ) - - expected_ports = worker_ports or [start_args.port] - assert api_http.get_server_info() == { - **original, - "pd_master_multi_ports": expected_ports, - } - assert (dict(args) if runtime else vars(args)) == original - - -def test_get_server_info_serializes_start_args_dataclass(monkeypatch): - from lightllm.server import api_http - - start_args = StartArgs(model_name="qwen35_0.8b") - monkeypatch.setattr(api_http.g_objs, "args", start_args) - - server_info = api_http.get_server_info() - assert server_info["model_name"] == "qwen35_0.8b" - assert server_info["run_mode"] == "normal" - - -def test_pd_master_models_endpoint_has_created_timestamp(monkeypatch): - from lightllm.server import api_http - - args = StartArgs(run_mode="pd_master", model_dir="/tmp/test-model", model_name="test-model") - manager = SimpleNamespace(get_real_supported_max_req_total_len=lambda: 1024, metric_client=object()) - monkeypatch.setattr(api_http, "init_tokenizer", lambda *_args, **_kwargs: None) - monkeypatch.setattr(api_http.SamplingParams, "load_generation_cfg", lambda _model_dir: None) - monkeypatch.setattr(api_http.CompletionRequest, "load_generation_cfg", lambda _model_dir: None) - monkeypatch.setattr(api_http.ChatCompletionRequest, "load_generation_cfg", lambda _model_dir: None) - monkeypatch.setattr(api_http, "MetricClient", lambda _port: object()) - monkeypatch.setattr(api_http, "get_shm_port_args", lambda: SimpleNamespace(metric_port=1234)) - monkeypatch.setattr(api_http, "HttpServerManagerForPDMaster", lambda args: manager) - monkeypatch.setattr(api_http, "get_unique_server_name", lambda: "test-server") - monkeypatch.setattr(api_http.setproctitle, "setproctitle", lambda _title: None) - monkeypatch.setattr(api_http.time, "time", lambda: 1234.9) - - global_objs = api_http.G_Objs() - global_objs.set_args(args) - monkeypatch.setattr(api_http, "g_objs", global_objs) - - response = asyncio.run(api_http.get_models(None)) - assert response.data[0].created == 1234 - - -def test_elastic_pd_nodes_are_ready_with_at_least_one_node_of_each_role(): - manager = PDManager(StartArgs(pd_master_mode="elastic")) - assert manager.is_pd_nodes_ready() is False - - manager.prefill_nodes = [object()] - assert manager.is_pd_nodes_ready() is False - - manager.decode_nodes = [object()] - assert manager.is_pd_nodes_ready() is True - - manager.prefill_nodes.append(object()) - manager.decode_nodes.append(object()) - assert manager.is_pd_nodes_ready() is True - - -def test_fixed_pd_nodes_are_ready_only_with_exact_node_counts(): - manager = PDManager(StartArgs(pd_master_mode="2p1d")) - assert manager.is_pd_nodes_ready() is False - - manager.prefill_nodes = [SimpleNamespace(client_ip_port="p1:8000"), SimpleNamespace(client_ip_port="p2:8000")] - manager.decode_nodes = [object()] - assert manager.is_pd_nodes_ready() is True - - manager.decode_nodes.append(object()) - assert manager.is_pd_nodes_ready() is False - - -@pytest.mark.parametrize("mode", ["flex", "xpxd", "1p1"]) -def test_malformed_pd_master_mode_is_not_ready(mode): - manager = PDManager(StartArgs(pd_master_mode=mode)) - assert manager.is_pd_nodes_ready() is False - - -def test_pd_master_readiness_endpoint_uses_pd_node_readiness(monkeypatch): - from lightllm.server import api_http - - args = StartArgs(run_mode="pd_master", pd_master_mode="elastic") - pd_manager = PDManager(args) - monkeypatch.setattr(api_http.g_objs, "args", args) - monkeypatch.setattr(api_http.g_objs, "httpserver_manager", SimpleNamespace(pd_manager=pd_manager)) - - response = api_http.readiness() - assert response.status_code == 503 - assert json.loads(response.body) == {"status": "not ready"} - - pd_manager.prefill_nodes = [object()] - pd_manager.decode_nodes = [object()] - response = api_http.readiness() - assert response.status_code == 200 - assert json.loads(response.body) == {"status": "ok"} - - -def test_pd_manager_checks_all_connected_node_health(monkeypatch): - from lightllm.server.httpserver_for_pd_master import manager as manager_module - - requested_urls = [] - responses = { - "http://10.0.0.1:8000/health": SimpleNamespace(status_code=200), - "http://10.0.0.2:8000/health": SimpleNamespace(status_code=200), - } - - class FakeAsyncClient: - def __init__(self, **kwargs): - assert kwargs == {"timeout": 8, "trust_env": False} - - async def __aenter__(self): - return self - - async def __aexit__(self, exc_type, exc_value, traceback): - return False - - async def get(self, url): - requested_urls.append(url) - result = responses[url] - if isinstance(result, BaseException): - raise result - return result - - monkeypatch.setattr(manager_module.httpx, "AsyncClient", FakeAsyncClient) - - manager = PDManager(StartArgs()) - manager.prefill_nodes = [SimpleNamespace(mode="prefill", client_ip_port="10.0.0.1:8000")] - manager.decode_nodes = [SimpleNamespace(mode="decode", client_ip_port="10.0.0.2:8000")] - - assert asyncio.run(manager.check_pd_nodes_health()) is True - assert set(requested_urls) == set(responses) - - responses["http://10.0.0.2:8000/health"] = SimpleNamespace(status_code=503) - assert asyncio.run(manager.check_pd_nodes_health()) is False - - responses["http://10.0.0.2:8000/health"] = RuntimeError("connection failed") - assert asyncio.run(manager.check_pd_nodes_health()) is False - - -def test_pd_manager_without_connected_nodes_is_healthy(): - manager = PDManager(StartArgs()) - assert asyncio.run(manager.check_pd_nodes_health()) is True - - -def test_prefill_registration_preserves_existing_inflight_prompt_chars(): - args = StartArgs() - manager = PDManager(args) - - def register_prefill(node_id, client_ip_port): - manager.register_pd( - { - "node_id": node_id, - "client_ip_port": client_ip_port, - "mode": "prefill", - "start_args": { - "max_req_total_len": args.max_req_total_len, - "max_image_pixels": args.max_image_pixels, - "disable_image_resize": args.disable_image_resize, - }, - }, - websocket=object(), - ) - - register_prefill(1, "10.0.0.1:8000") - manager.prefill_nodes[0].dispatched_prompt_chars = 1234 - manager.prefill_nodes[0].dispatched_req_num = 12 - - register_prefill(2, "10.0.0.2:8000") - - assert [node.dispatched_prompt_chars for node in manager.prefill_nodes] == [1234, 0] - assert [node.dispatched_req_num for node in manager.prefill_nodes] == [12, 0] - - -def test_prefill_reconnection_preserves_other_nodes_inflight_prompt_chars(): - args = StartArgs() - manager = PDManager(args) - - def pd_info(node_id, client_ip_port): - return { - "node_id": node_id, - "client_ip_port": client_ip_port, - "mode": "prefill", - "start_args": { - "max_req_total_len": args.max_req_total_len, - "max_image_pixels": args.max_image_pixels, - "disable_image_resize": args.disable_image_resize, - }, - } - - manager.register_pd(pd_info(1, "10.0.0.1:8000"), websocket=object()) - manager.register_pd(pd_info(2, "10.0.0.2:8000"), websocket=object()) - manager.prefill_nodes[0].dispatched_prompt_chars = 100 - manager.prefill_nodes[1].dispatched_prompt_chars = 200 - manager.prefill_nodes[0].dispatched_req_num = 1 - manager.prefill_nodes[1].dispatched_req_num = 2 - - manager.register_pd(pd_info(3, "10.0.0.2:8000"), websocket=object()) - - assert [node.client_ip_port for node in manager.prefill_nodes] == ["10.0.0.1:8000", "10.0.0.2:8000"] - assert [node.dispatched_prompt_chars for node in manager.prefill_nodes] == [100, 0] - assert [node.dispatched_req_num for node in manager.prefill_nodes] == [1, 0] - - -def test_pd_master_inference_health_matches_normal_node_semantics(monkeypatch): - monkeypatch.setattr("lightllm.server.httpserver_for_pd_master.manager.time.time", lambda: 1000) - manager = SimpleNamespace( - health_timeout=200, - latest_success_infer_time=900, - running_request_count=1, - req_id_to_out_inf={1: object()}, - ) - - assert HttpServerManagerForPDMaster.is_healthy(manager) is True - - manager.latest_success_infer_time = 700 - assert HttpServerManagerForPDMaster.is_healthy(manager) is False - - manager.running_request_count = 0 - manager.req_id_to_out_inf.clear() - assert HttpServerManagerForPDMaster.is_healthy(manager) is True - - -def test_pd_master_restores_request_count_when_preload_fails(): - class FailingMultimodalParams: - async def verify_and_preload(self, request): - raise RuntimeError("preload failed") - - manager = HttpServerManagerForPDMaster.__new__(HttpServerManagerForPDMaster) - manager.args = SimpleNamespace(reasoning_parser=None, enable_rl=False, enable_return_routed_experts=False) - manager.running_request_count = 0 - - params = SamplingParams() - params.init(None) - - async def consume_generate(): - async for _ in manager.generate("prompt", params, FailingMultimodalParams(), None): - pass - - with pytest.raises(RuntimeError, match="preload failed"): - asyncio.run(consume_generate()) - - assert manager.running_request_count == 0 - - -@pytest.mark.parametrize("parser", [None, "qwen3"]) -@pytest.mark.parametrize("enabled", [False, True]) -def test_pd_master_request_count_covers_async_generator_lifecycle(monkeypatch, parser, enabled): - monkeypatch.setenv("LIGHTLLM_STOP_IN_REASONING", "1" if enabled else "0") - manager = HttpServerManagerForPDMaster.__new__(HttpServerManagerForPDMaster) - manager.args = SimpleNamespace(reasoning_parser=None, enable_rl=False, enable_return_routed_experts=False) - manager.running_request_count = 0 - manager.args.reasoning_parser = parser - inner_generator_closed = False - - async def fake_generate(prompt, sampling_params, multimodal_params, request): - nonlocal inner_generator_closed - try: - assert manager.running_request_count == 1 - yield "result" - finally: - inner_generator_closed = True - - manager._generate_choices = fake_generate - - async def consume_one_result_and_close(): - params = SamplingParams() - params.init(None) - results_generator = manager.generate(None, params, None, None) - assert await results_generator.__anext__() == "result" - assert manager.running_request_count == 1 - await results_generator.aclose() - - asyncio.run(consume_one_result_and_close()) - assert manager.running_request_count == 0 - assert inner_generator_closed is True - - -@pytest.mark.parametrize("include", [False, True]) -def test_pd_master_forwards_filtered_choices(include): - manager = HttpServerManagerForPDMaster.__new__(HttpServerManagerForPDMaster) - manager.args = SimpleNamespace(reasoning_parser=None, enable_rl=False, enable_return_routed_experts=False) - manager.running_request_count = 0 - manager.tokenizer = SimpleNamespace(encode=lambda text, **kwargs: [ord(c) for c in text]) - params = SamplingParams() - params.init(manager.tokenizer, stop_sequences="END", n=2, best_of=2, include_stop_str_in_output=include) - - async def fake_generate(*args): - # Prefill's first token and subsequent decode tokens arrive in one stream. - yield 80, "first", {"id": 1}, FinishStatus() - yield 81, "second", {"id": 2}, FinishStatus() - yield 80, "END" if include else "", {"id": 3}, FinishStatus(FinishStatus.FINISHED_STOP) - yield 81, "EN", {"id": 4}, FinishStatus(FinishStatus.FINISHED_LENGTH) - - manager._generate_choices = fake_generate - - async def collect(): - return [result async for result in manager.generate("prompt", params, None, None)] - - results = asyncio.run(collect()) - assert "".join(result[1] for result in results if result[0] == 80) == ("firstEND" if include else "first") - assert "".join(result[1] for result in results if result[0] == 81) == "secondEN" - assert sorted(result[2]["id"] for result in results) == [1, 2, 3, 4] - assert manager.running_request_count == 0 - - -def test_fixed_pd_master_health_endpoint_combines_ready_and_health_status(monkeypatch): - from lightllm.server import api_http - - args = StartArgs(run_mode="pd_master", pd_master_mode="2p1d") - pd_manager = PDManager(args) - pd_node_health_results = iter([True, False]) - pd_node_health_check_count = 0 - - async def check_pd_nodes_health(): - nonlocal pd_node_health_check_count - pd_node_health_check_count += 1 - return next(pd_node_health_results) - - pd_manager.check_pd_nodes_health = check_pd_nodes_health - httpserver_manager = SimpleNamespace(pd_manager=pd_manager, is_healthy=lambda: True) - monkeypatch.setattr(api_http.g_objs, "args", args) - monkeypatch.setattr(api_http.g_objs, "httpserver_manager", httpserver_manager) - - response = asyncio.run(api_http.healthcheck(None)) - assert response.status_code == 503 - assert json.loads(response.body) == { - "message": "Error", - "pd_master_mode": "2p1d", - "inference_healthy": True, - "pd_nodes_ready": False, - "pd_nodes_healthy": False, - "registered_prefill_nodes": 0, - "registered_decode_nodes": 0, - } - assert pd_node_health_check_count == 0 - - pd_manager.prefill_nodes = [SimpleNamespace(client_ip_port="p1:8000"), SimpleNamespace(client_ip_port="p2:8000")] - pd_manager.decode_nodes = [object()] - response = asyncio.run(api_http.healthcheck(None)) - assert response.status_code == 200 - assert json.loads(response.body)["message"] == "Ok" - assert json.loads(response.body)["pd_nodes_ready"] is True - assert pd_node_health_check_count == 1 - - response = asyncio.run(api_http.healthcheck(None)) - assert response.status_code == 503 - assert json.loads(response.body)["pd_nodes_healthy"] is False - assert pd_node_health_check_count == 2 - - httpserver_manager.is_healthy = lambda: False - response = asyncio.run(api_http.healthcheck(None)) - assert response.status_code == 503 - assert json.loads(response.body)["inference_healthy"] is False - assert pd_node_health_check_count == 2 - - -def test_elastic_pd_master_health_does_not_check_individual_pd_nodes(monkeypatch): - from lightllm.server import api_http - - args = StartArgs(run_mode="pd_master", pd_master_mode="elastic") - pd_manager = PDManager(args) - pd_manager.prefill_nodes = [object()] - pd_manager.decode_nodes = [object()] - - async def unexpected_pd_node_health_check(): - raise AssertionError("elastic mode must not check individual PD node health") - - pd_manager.check_pd_nodes_health = unexpected_pd_node_health_check - httpserver_manager = SimpleNamespace(pd_manager=pd_manager, is_healthy=lambda: True) - monkeypatch.setattr(api_http.g_objs, "args", args) - monkeypatch.setattr(api_http.g_objs, "httpserver_manager", httpserver_manager) - - response = asyncio.run(api_http.healthcheck(None)) - response_body = json.loads(response.body) - assert response.status_code == 200 - assert response_body["message"] == "Ok" - assert "pd_nodes_healthy" not in response_body - - -@pytest.mark.parametrize( - "run_mode,start_name,workers", - [ - ("pd_master", "pd_master_start", 1), - ("pd_master", "pd_master_multi_worker_start", 2), - ("prefill", "_launch_subprocesses", 1), - ("decode", "_launch_subprocesses", 1), - ], -) -@pytest.mark.parametrize("feature", ["enable_rl", "enable_return_routed_experts", "enable_prompt_logprobs"]) -def test_pd_rejects_final_metadata_features_before_startup(monkeypatch, run_mode, start_name, workers, feature): - from lightllm.server import api_start - - def unexpected_startup(*args, **kwargs): - pytest.fail("Unsupported PD features must be rejected before process setup") - - monkeypatch.setattr(api_start.mp, "set_start_method", unexpected_startup) - start = getattr(api_start, start_name) - with pytest.raises(ValueError, match="PD mode does not support"): - start(StartArgs(run_mode=run_mode, httpserver_workers=workers, **{feature: True})) - - -@pytest.mark.parametrize("workers", [1, 2, 0]) -def test_pd_master_cli_dispatches_to_the_requested_startup_mode(monkeypatch, workers): - from lightllm.server import api_server, api_start - - single, multiple = Mock(), Mock() - monkeypatch.setattr(api_start, "pd_master_start", single) - monkeypatch.setattr(api_start, "pd_master_multi_worker_start", multiple) - monkeypatch.setattr(api_server, "set_log_node_role", Mock()) - monkeypatch.setattr(api_server.torch.multiprocessing, "set_start_method", Mock()) - args = StartArgs(run_mode="pd_master", httpserver_workers=workers) - - if workers == 0: - with pytest.raises(ValueError, match="httpserver_workers.*at least 1"): - api_server.launch_server(args) - single.assert_not_called() - multiple.assert_not_called() - else: - api_server.launch_server(args) - expected, unexpected = (single, multiple) if workers == 1 else (multiple, single) - expected.assert_called_once_with(args) - unexpected.assert_not_called() - - -def test_multiworker_preflight_failure_precedes_environment_and_process_initialization(monkeypatch): - from lightllm.server import api_start - from lightllm.server.httpserver_for_pd_master import multiprocess - - preflight = Mock(side_effect=ValueError("preflight rejected")) - configure = Mock(side_effect=AssertionError("Environment setup must follow preflight")) - launch = Mock(side_effect=AssertionError("Process startup must follow preflight")) - monkeypatch.setattr(multiprocess, "validate_multiprocess_pd_master_args", preflight) - monkeypatch.setattr(multiprocess, "start_multiprocess_pd_master", launch) - monkeypatch.setattr(api_start, "_set_envs_and_config", configure) - args = StartArgs(run_mode="pd_master", httpserver_workers=2) - - with pytest.raises(ValueError, match="preflight rejected"): - api_start.pd_master_multi_worker_start(args) - - preflight.assert_called_once_with(args) - configure.assert_not_called() - launch.assert_not_called() - - -def test_multiworker_starts_required_services_in_order(monkeypatch): - from lightllm.server import api_start - from lightllm.server.httpserver_for_pd_master import multiprocess - from lightllm.server.httpserver_for_pd_master.pd_selector import pd_node_stats_server - from lightllm.server.httpserver_for_pd_master.pd_selector import prompt_cache_tree_server - - args = StartArgs(run_mode="pd_master", port=8000, httpserver_workers=2) - ports = SimpleNamespace( - close=Mock(), - metric_port=8104, - pd_master_multi_ports=[8100, 8101], - ) - metric_server = Mock(name="metric_server") - request_id_server = Mock(name="request_id_server") - stats_server = Mock(name="stats_server") - tree_server = Mock(name="tree_server") - started_services = [] - manager = Mock() - manager.start_submodule_processes.side_effect = lambda *, start_funcs, start_args: started_services.append( - start_funcs[0] - ) - launch = Mock(return_value="completed") - publish = Mock() - read_start_args = Mock() - validate_ports = Mock() - - monkeypatch.setattr(api_start, "process_manager", manager) - monkeypatch.setattr(api_start, "start_metric_manager", metric_server) - monkeypatch.setattr(api_start, "start_req_id_generator_server", request_id_server) - monkeypatch.setattr(api_start, "_set_envs_and_config", Mock()) - monkeypatch.setattr(api_start, "set_unique_server_name", Mock()) - monkeypatch.setattr(api_start, "auto_set_max_req_total_len", Mock()) - monkeypatch.setattr(api_start, "auto_set_response_parsers", Mock()) - monkeypatch.setattr(api_start, "set_env_start_args", publish) - monkeypatch.setattr(api_start, "get_env_start_args", read_start_args) - monkeypatch.setattr(api_start, "get_shm_port_args", Mock(return_value=ports)) - monkeypatch.setattr(api_start, "validate_ports", validate_ports) - monkeypatch.setattr(multiprocess, "validate_multiprocess_pd_master_args", Mock(return_value="haproxy")) - monkeypatch.setattr(multiprocess, "start_multiprocess_pd_master", launch) - monkeypatch.setattr(pd_node_stats_server, "start_pd_node_stats_server", stats_server) - monkeypatch.setattr(prompt_cache_tree_server, "start_prompt_cache_tree_server", tree_server) - - assert api_start.pd_master_multi_worker_start(args) == "completed" - assert started_services == [metric_server, request_id_server, stats_server, tree_server] - manager.setup_exit_controller.assert_called_once_with() - manager.terminate_all_processes.assert_not_called() - ports.close.assert_not_called() - multiprocess.validate_multiprocess_pd_master_args.assert_called_once_with(args) - publish.assert_called_once_with(args) - read_start_args.cache_clear.assert_called_once_with() - api_start.get_shm_port_args.assert_called_once_with(create=True) - validate_ports.assert_called_once_with([8000, 8104, 8100, 8101]) - launch.assert_called_once_with( - args, - ports=ports, - haproxy_binary="haproxy", - ) diff --git a/unit_tests/server/test_pd_master_multiprocess_e2e.py b/unit_tests/server/test_pd_master_multiprocess_e2e.py deleted file mode 100644 index cc4037a90b..0000000000 --- a/unit_tests/server/test_pd_master_multiprocess_e2e.py +++ /dev/null @@ -1,496 +0,0 @@ -"""Real PD master service tests with CPU-only P/D protocol peers. - -The functional cases cover worker discovery, HTTP generation, SSE, busy replies, -cancellation, shared IDs and metrics, health aggregation, and private channel -isolation. The opt-in performance case compares the same end-to-end protocol path -with one worker and multiple workers without loading model weights. -""" - -import asyncio -from contextlib import AsyncExitStack -import json -import os -import pickle -import statistics -import time -from types import SimpleNamespace -from urllib.parse import urlsplit - -import httpx -import psutil -import pytest -import websockets -from prometheus_client.parser import text_string_to_metric_families - -from lightllm.server.core.objs import FinishStatus -from lightllm.server.httpserver import pd_loop -from lightllm.server.pd_io_struct import ObjType, PDDecodeNodeInfo, PDUpKVStatus -from unit_tests.server.pd_master_test_utils import running_pd_master - - -class FakePDPeers: - """Drive the actual registration, prompt-ID, KV status, and token channels.""" - - def __init__(self, workers, decode_delay=0.02): - self.workers = workers - self.decode_delay = decode_delay - self.stack = AsyncExitStack() - self.tasks = [] - self.requests = [] - self.prefill_done = {} - self.generations = {} - self.generation_errors = [] - self.aborted = set() - self.blocked_prefills = [] - self.hold_prompt_ids = False - self.hold_decode = False - self.busy_prefill_requests = 0 - self.node_addresses = {"prefill": "prefill.internal:18001", "decode": "decode.internal:18002"} - - async def __aenter__(self): - async with httpx.AsyncClient(trust_env=False, timeout=5) as client: - for endpoint in self.workers: - owner = endpoint["node_id"] - response = await client.get("http://" + endpoint["host_ip_port"] + "/get_server_info") - response.raise_for_status() - start_args = response.json() - kv = await self.stack.enter_async_context( - websockets.connect("ws://" + endpoint["host_ip_port"] + "/kv_move_status", proxy=None) - ) - for mode, node_id in (("prefill", 501), ("decode", 502)): - websocket = await self.stack.enter_async_context( - websockets.connect("ws://" + endpoint["host_ip_port"] + "/pd_register", proxy=None) - ) - await websocket.send( - json.dumps( - { - "node_id": node_id, - "client_ip_port": self.node_addresses[mode], - "mode": mode, - "start_args": dict(start_args, run_mode=mode), - "connection_id": f"{mode}-{owner}", - } - ) - ) - self.tasks.append(asyncio.create_task(self._run(mode, owner, websocket, kv))) - return self - - async def __aexit__(self, *exc): - tasks = [*self.tasks, *self.generations.values()] - for task in tasks: - task.cancel() - await asyncio.gather(*tasks, return_exceptions=True) - await self.stack.aclose() - - def _event(self, request_id): - return self.prefill_done.setdefault(request_id, asyncio.Event()) - - async def _token(self, websocket, mode, owner, request_id, index, text, finish): - metadata = { - "id": 2 if index == 1 else 3, - "logprob": -0.1, - "logprobs": None, - "prompt_tokens": 3, - "count_output_tokens": index, - "node_mode": mode, - "prompt_cache_len": 0, - } - await websocket.send( - pickle.dumps( - ( - ObjType.TOKEN_PACKS, - [(request_id, text, metadata, FinishStatus(finish))], - { - "connection_key": self.node_addresses[mode] + "/dpNone", - "total_token_usage_rate": 0.1, - }, - ) - ) - ) - - async def _run(self, mode, owner, websocket, kv): - async for data in websocket: - message = pickle.loads(data) - if message[0] == ObjType.ABORT: - request_id = message[1] - self.aborted.add((mode, owner, request_id)) - task = self.generations.get((owner, request_id)) - if mode == "decode" and task is not None: - task.cancel() - continue - if mode == "prefill": - if message[0] == ObjType.REQ: - _, params, _ = message[1] - assert params.pd_master_node_id.get() == owner - if self.busy_prefill_requests: - self.busy_prefill_requests -= 1 - await websocket.send( - pickle.dumps( - (ObjType.PD_UPLOAD_SERVER_BUSY, params.group_request_id, "test capacity exhausted") - ) - ) - continue - if self.hold_prompt_ids: - self.blocked_prefills.append((owner, params.group_request_id)) - continue - await websocket.send( - pickle.dumps((ObjType.PD_UPLOAD_PREFILL_PROMPT_IDS, params.group_request_id, [2, 3, 2])) - ) - elif message[0] == ObjType.PD_REQ_DECODE_NODE_INFO: - request_id = message[1] - await self._token(websocket, mode, owner, request_id, 1, "hello", FinishStatus.FINISHED_LENGTH) - self._event(request_id).set() - elif message[0] == ObjType.REQ: - task = asyncio.create_task(self._decode(owner, websocket, kv, message[1])) - self.generations[(owner, message[1][1].group_request_id)] = task - task.add_done_callback(self._generation_finished) - - def _generation_finished(self, task): - if not task.cancelled() and task.exception() is not None: - self.generation_errors.append(task.exception()) - - async def _decode(self, owner, websocket, kv, request): - prompt_ids, params, _ = request - request_id = params.group_request_id - assert params.pd_master_node_id.get() == owner - self.requests.append((owner, request_id)) - info = PDDecodeNodeInfo( - decode_node_id=502, - pd_master_node_id=owner, - agent_name="test", - agent_metadata=b"", - num_pages=0, - page_reg_desc=b"", - request_id=request_id, - ready_kv_len=0, - ) - assert len(prompt_ids) == 3 - await kv.send(pickle.dumps(PDUpKVStatus(request_id, owner, pickle.dumps(info)))) - await asyncio.wait_for(self._event(request_id).wait(), timeout=5) - # Decode also reports the transferred first token. P/D WebSocket delivery - # order is independent; the master deduplicates their index-one outputs. - await self._token(websocket, "decode", owner, request_id, 1, "hello", FinishStatus.NO_FINISH) - if self.hold_decode: - await asyncio.Event().wait() - # Keep enough requests in flight to exercise load balancing. - if self.decode_delay: - await asyncio.sleep(self.decode_delay) - await self._token(websocket, "decode", owner, request_id, 2, " world", FinishStatus.FINISHED_STOP) - - -async def _wait_until(predicate, timeout=5): - async def poll(): - while not predicate(): - await asyncio.sleep(0.01) - - await asyncio.wait_for(poll(), timeout=timeout) - - -async def _wait_http_status(client, url, expected_status, timeout=20): - deadline = asyncio.get_running_loop().time() + timeout - while True: - response = await client.get(url) - if response.status_code == expected_status: - return response - assert asyncio.get_running_loop().time() < deadline, response.text - await asyncio.sleep(0.1) - - -@pytest.mark.parametrize("worker_count,strategy", [(1, "cache_aware"), (2, "round_robin"), (4, "cache_aware")]) -def test_real_pd_master_generation_sse_busy_and_cancellation(tmp_path, monkeypatch, worker_count, strategy): - monkeypatch.setenv("LIGHTLLM_PD_NODE_BUSY_RETRY_TIMEOUT_SECONDS", "0") - - async def run(master): - public = urlsplit(master.public_url) - monkeypatch.setattr( - pd_loop, - "get_shm_port_args", - lambda: SimpleNamespace(config_server_port=public.port, pd_master_port=public.port), - ) - discovery_args = SimpleNamespace( - config_server_host=None, - config_server_port=None, - pd_master_ip=public.hostname, - ) - discovered = await pd_loop._get_pd_master_objs(discovery_args) - assert discovered is not None - assert {(node.node_id, node.host_ip_port) for node in discovered.values()} == { - (worker["node_id"], worker["host_ip_port"]) for worker in master.workers - } - async with FakePDPeers(master.workers) as peers: - async with httpx.AsyncClient(trust_env=False, timeout=15) as client: - deadline = asyncio.get_running_loop().time() + 10 - while True: - response = await client.get(master.public_url + "/readiness") - if response.status_code == 200: - break - assert asyncio.get_running_loop().time() < deadline, response.text - await asyncio.sleep(0.1) - - async def generate(): - response = await client.post( - master.public_url + "/generate", - json={ - "inputs": "hello world", - "parameters": {"max_new_tokens": 2, "n": 2, "best_of": 2}, - }, - ) - assert response.status_code == 200, response.text - assert response.json()["generated_text"] == ["hello world", "hello world"] - - await asyncio.gather(*(generate() for _ in range(32))) - assert len(peers.requests) == 64 - assert {owner for owner, _ in peers.requests} == {worker["node_id"] for worker in master.workers} - ids = [request_id for _, request_id in peers.requests] - assert len(set(ids)) == len(ids) - assert all(request_id % 8 == 0 for request_id in ids) - - # Scraping through each worker shares its manager's metrics client - # and exposes the same service-wide request counter. - for worker in master.workers: - deadline = asyncio.get_running_loop().time() + 5 - while True: - response = await client.get("http://" + worker["host_ip_port"] + "/metrics") - assert response.status_code == 200 - counts = [ - sample.value - for family in text_string_to_metric_families(response.text) - for sample in family.samples - if sample.name == "lightllm_request_count_total" - ] - if counts == [32]: - break - assert asyncio.get_running_loop().time() < deadline, response.text - await asyncio.sleep(0.05) - - async with client.stream( - "POST", - master.public_url + "/generate_stream", - json={ - "inputs": "hello world", - "parameters": {"max_new_tokens": 2}, - }, - ) as response: - assert response.status_code == 200 - events = [json.loads(line[5:]) async for line in response.aiter_lines() if line.startswith("data:")] - assert [event["token"]["text"] for event in events] == ["hello", " world"] - assert [event["finished"] for event in events] == [False, True] - - # A rejected Prefill must release its reservation and preserve the - # public busy envelope, without killing a control connection. - peers.busy_prefill_requests = 1 - response = await client.post( - master.public_url + "/generate", - json={ - "inputs": "hello world", - "parameters": {"max_new_tokens": 2}, - }, - ) - assert response.status_code == 429, response.text - assert response.json()["error"]["type"] == "RateLimitError" - - # Closing SSE after the first token must reach both P/D actors. - peers.hold_decode = True - async with client.stream( - "POST", - master.public_url + "/generate_stream", - json={ - "inputs": "hello world", - "parameters": {"max_new_tokens": 2}, - }, - ) as response: - assert response.status_code == 200 - async for line in response.aiter_lines(): - if line.startswith("data:"): - assert json.loads(line[5:])["token"]["text"] == "hello" - break - owner, request_id = peers.requests[-1] - await _wait_until( - lambda: all((mode, owner, request_id) in peers.aborted for mode in ("prefill", "decode")) - ) - peers.hold_decode = False - - # Cancellation before prompt IDs must clean up a still-held - # Prefill reservation, rather than only the output iterator. - peers.hold_prompt_ids = True - pending = asyncio.create_task( - client.post( - master.public_url + "/generate", - json={ - "inputs": "hello world", - "parameters": {"max_new_tokens": 2}, - }, - ) - ) - await _wait_until(lambda: bool(peers.blocked_prefills)) - owner, request_id = peers.blocked_prefills[-1] - pending.cancel() - with pytest.raises(asyncio.CancelledError): - await pending - await _wait_until( - lambda: all((mode, owner, request_id) in peers.aborted for mode in ("prefill", "decode")) - ) - peers.hold_prompt_ids = False - - await generate() - assert not peers.generation_errors - for task in peers.tasks: - assert not task.done(), task.exception() - - with running_pd_master( - tmp_path, extra_args=["--httpserver_workers", str(worker_count), "--select_p_d_node_strategy", strategy] - ) as master: - assert len(master.workers) == worker_count - try: - asyncio.run(run(master)) - except BaseException: - print(master.log_path.read_text()[-12000:]) - raise - - -def test_real_health_aggregation_and_private_channel_isolation(tmp_path): - """Public health succeeds only while every real worker has usable P/D peers.""" - - async def run(master): - async with httpx.AsyncClient(trust_env=False, timeout=15) as client: - assert (await client.get(master.public_url + "/liveness")).status_code == 200 - assert (await client.get(master.public_url + "/pd_register")).status_code == 426 - await _wait_http_status(client, master.public_url + "/readiness", 503) - await _wait_http_status(client, master.public_url + "/health", 503) - - async with FakePDPeers(master.workers): - readiness = await _wait_http_status(client, master.public_url + "/readiness", 200) - health = await _wait_http_status(client, master.public_url + "/health", 200) - assert readiness.json() == {"status": "ok"} - assert health.json() == {"message": "Ok"} - - readiness = await _wait_http_status(client, master.public_url + "/readiness", 503) - health = await _wait_http_status(client, master.public_url + "/health", 503) - assert readiness.json() == {"status": "not ready"} - assert health.json() == {"message": "Error"} - - with running_pd_master(tmp_path, extra_args=["--httpserver_workers", "2"]) as master: - try: - asyncio.run(run(master)) - except BaseException: - print(master.log_path.read_text()[-12000:]) - raise - - -def _service_cpu_seconds(master): - processes = [psutil.Process(master.process.pid), *master.descendants()] - cpu_seconds = 0.0 - for process in {process.pid: process for process in processes}.values(): - try: - cpu_time = process.cpu_times() - except (psutil.NoSuchProcess, psutil.AccessDenied): - continue - cpu_seconds += cpu_time.user + cpu_time.system - return cpu_seconds - - -async def _run_high_load(master, request_count, concurrency): - prompt = "hello world " * 48 - payload = { - "inputs": prompt, - "parameters": {"max_new_tokens": 2}, - } - limits = httpx.Limits(max_connections=concurrency, max_keepalive_connections=concurrency) - timeout = httpx.Timeout(60) - - async with FakePDPeers(master.workers, decode_delay=0) as peers, httpx.AsyncClient( - trust_env=False, - timeout=timeout, - limits=limits, - ) as client: - await _wait_http_status(client, master.public_url + "/readiness", 200) - semaphore = asyncio.Semaphore(concurrency) - - async def generate(latencies=None): - async with semaphore: - start = time.perf_counter() - response = await client.post(master.public_url + "/generate", json=payload) - if latencies is not None: - latencies.append(time.perf_counter() - start) - assert response.status_code == 200, response.text - assert response.json()["generated_text"] == ["hello world"] - - # 预热连接池、tokenizer 和请求 ID 客户端，避免一次性初始化污染结果。 - await asyncio.gather(*(generate() for _ in range(concurrency))) - peers.requests.clear() - cpu_start = _service_cpu_seconds(master) - latencies = [] - start = time.perf_counter() - await asyncio.gather(*(generate(latencies) for _ in range(request_count))) - elapsed = time.perf_counter() - start - cpu_seconds = _service_cpu_seconds(master) - cpu_start - - assert len(peers.requests) == request_count - request_ids = [request_id for _, request_id in peers.requests] - assert len(set(request_ids)) == request_count - assert {owner for owner, _ in peers.requests} == {worker["node_id"] for worker in master.workers} - - sorted_latencies = sorted(latencies) - return { - "requests": request_count, - "concurrency": concurrency, - "elapsed_seconds": elapsed, - "requests_per_second": request_count / elapsed, - "p50_ms": statistics.median(sorted_latencies) * 1000, - "p95_ms": sorted_latencies[int((len(sorted_latencies) - 1) * 0.95)] * 1000, - "service_cpu_cores": cpu_seconds / elapsed, - } - - -@pytest.mark.skipif( - os.getenv("LIGHTLLM_RUN_PD_MASTER_PERF") != "1", - reason="set LIGHTLLM_RUN_PD_MASTER_PERF=1 to run the real PD master load benchmark", -) -def test_multiprocess_pd_master_high_load_speedup(tmp_path): - """Compare the complete HTTP -> PD master -> fake P/D protocol path.""" - available_cpus = len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else (os.cpu_count() or 1) - if available_cpus < 2: - pytest.skip("the multiprocess benchmark requires at least two available CPUs") - - request_count = int(os.getenv("LIGHTLLM_PD_MASTER_PERF_REQUESTS", "2000")) - concurrency = int(os.getenv("LIGHTLLM_PD_MASTER_PERF_CONCURRENCY", "128")) - multi_worker_count = int(os.getenv("LIGHTLLM_PD_MASTER_PERF_WORKERS", str(min(4, available_cpus)))) - if not 2 <= multi_worker_count <= available_cpus: - raise ValueError("LIGHTLLM_PD_MASTER_PERF_WORKERS must be between 2 and the available CPU count") - results = {} - - for worker_count in (1, multi_worker_count): - runtime_dir = tmp_path / f"workers-{worker_count}" - with running_pd_master( - runtime_dir, - extra_args=[ - "--httpserver_workers", - str(worker_count), - "--select_p_d_node_strategy", - "round_robin", - ], - ) as master: - try: - results[worker_count] = asyncio.run(_run_high_load(master, request_count, concurrency)) - except BaseException: - print(master.log_path.read_text()[-12000:]) - raise - - single_qps = results[1]["requests_per_second"] - multi_qps = results[multi_worker_count]["requests_per_second"] - speedup = multi_qps / single_qps - print( - "PD master high-load benchmark: " - + json.dumps( - { - "single_worker": results[1], - f"{multi_worker_count}_workers": results[multi_worker_count], - "throughput_speedup": speedup, - }, - sort_keys=True, - ) - ) - - minimum_speedup = float(os.getenv("LIGHTLLM_PD_MASTER_MIN_SPEEDUP", "1.0")) - if minimum_speedup: - assert speedup >= minimum_speedup diff --git a/unit_tests/server/test_pd_master_multiprocess_launcher.py b/unit_tests/server/test_pd_master_multiprocess_launcher.py deleted file mode 100644 index 18282463c7..0000000000 --- a/unit_tests/server/test_pd_master_multiprocess_launcher.py +++ /dev/null @@ -1,399 +0,0 @@ -import argparse -import json -from pathlib import Path -import signal -import sys -import uuid -from types import SimpleNamespace -from unittest.mock import Mock - -import pytest - -from lightllm.server.core.objs.start_args_type import StartArgs -from lightllm.server.httpserver_for_pd_master import multiprocess -from lightllm.utils.envs_utils import get_unique_server_name -from lightllm.utils.net_utils import get_pd_master_unix_socket_address - - -def _args(**kwargs): - options = dict( - run_mode="pd_master", - host="10.0.0.1", - port=8000, - httpserver_workers=2, - ) - options.update(kwargs) - return StartArgs(**options) - - -def _worker_ports(): - return [8100, 8101] - - -def _ports(**kwargs): - values = dict( - port=8000, - metric_port=8104, - pd_master_multi_ports=[8100, 8101], - ) - values.update(kwargs) - return SimpleNamespace(**values) - - -def test_workers_use_the_shared_service_ipc_address(monkeypatch): - monkeypatch.setenv("LIGHTLLM_UNIQUE_SERVICE_NAME_ID", "shared-service") - get_unique_server_name.cache_clear() - launcher_address = get_pd_master_unix_socket_address("pd_node_stats") - - worker_address = get_pd_master_unix_socket_address("pd_node_stats") - - assert worker_address == launcher_address - get_unique_server_name.cache_clear() - - -@pytest.mark.parametrize("exited_worker", [None, 1]) -def test_backend_launch_tracks_processes_and_passes_worker_index(monkeypatch, exited_worker): - children = [Mock(pid=2101), Mock(pid=2102)] - managed = [object(), object()] - manager = SimpleNamespace(processes=[], process_names={}) - process_factory = Mock(side_effect=children) - process_context = Mock(Process=process_factory) - monkeypatch.setattr( - multiprocess.multiprocessing, - "get_context", - Mock(return_value=process_context), - ) - monkeypatch.setattr(multiprocess, "process_manager", manager) - - def lookup(pid): - index = [child.pid for child in children].index(pid) - if index == exited_worker: - raise multiprocess.psutil.NoSuchProcess(pid) - return managed[index] - - monkeypatch.setattr(multiprocess.psutil, "Process", lookup) - if exited_worker is None: - multiprocess._start_backends(2) - else: - with pytest.raises(multiprocess.psutil.NoSuchProcess): - multiprocess._start_backends(2) - - registered_count = len(children) if exited_worker is None else exited_worker - assert manager.processes == managed[:registered_count] - assert manager.process_names == { - process: f"pd_master_worker_{index}" for index, process in enumerate(managed[:registered_count]) - } - multiprocess.multiprocessing.get_context.assert_called_once_with("spawn") - for index, call in enumerate(process_factory.call_args_list): - assert call.kwargs == { - "target": multiprocess._run_pd_master_worker, - "args": (index,), - } - children[index].start.assert_called_once_with() - - -def test_worker_bootstrap_uses_its_assigned_port(monkeypatch): - import hypercorn.__main__ - - args = _args(port=8000, httpserver_workers=2) - ports = SimpleNamespace(pd_master_multi_ports=[8100, 8101]) - create_ports = Mock(return_value=ports) - monkeypatch.setenv("LIGHTLLM_PD_MASTER_WORKER_INDEX", "1") - monkeypatch.setattr(multiprocess, "get_env_start_args", lambda: args) - monkeypatch.setattr(multiprocess, "get_shm_port_args", create_ports) - - def start_hypercorn(argv): - assert args.port == 8000 - assert args.pd_node_id == 8101 - assert argv[argv.index("--workers") + 1] == "0" - assert argv[argv.index("--bind") + 1] == "10.0.0.1:8101" - assert argv[-1] == "lightllm.server.api_http:app" - - monkeypatch.setattr(hypercorn.__main__, "main", start_hypercorn) - multiprocess._start_pd_master_worker() - create_ports.assert_called_once_with() - - -def test_worker_port_table_preserves_dynamically_allocated_internal_ports(monkeypatch): - from lightllm.utils.envs_utils import get_env_start_args, get_unique_server_name - from lightllm.utils.shm_port_args import ShmPortArgs - - args = _args(port=8000, pd_master_port=8001, config_server_port=8103) - monkeypatch.setenv("LIGHTLLM_UNIQUE_SERVICE_NAME_ID", "test_pd_worker_ports_" + uuid.uuid4().hex) - monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(vars(args))) - get_unique_server_name.cache_clear() - get_env_start_args.cache_clear() - ports = ShmPortArgs(create=True) - linked = None - try: - allocator = Mock(side_effect=[8120, 8104]) - monkeypatch.setattr(ports, "_alloc_free_port", allocator) - assert ports.router_port == 8120 - assert ports.metric_port == 8104 - assert ports.port == 8000 - assert ports.config_server_port == 8103 - linked = ShmPortArgs() - assert linked.router_port == 8120 - assert linked.metric_port == 8104 - assert linked.port == 8000 - assert linked.config_server_port == 8103 - assert allocator.call_count == 2 - finally: - if linked is not None: - linked.close() - ports.shm.unlink() - ports.close() - get_unique_server_name.cache_clear() - get_env_start_args.cache_clear() - - -def test_haproxy_config_streams_without_request_replay(): - state_path = Path("/tmp/lightllm-test-health.state") - config = multiprocess._build_haproxy_config(_args(), _worker_ports(), state_path) - - assert "balance leastconn" in config - assert "nbthread 2" in config - assert "fd-hard-limit 131072" in config - assert f"server-state-file {state_path}" in config - assert config.count("load-server-state-from-file global") == 3 - assert "maxconn 65536" not in config - assert "server worker_0 10.0.0.1:8100" in config - assert "server worker_1 10.0.0.1:8101" in config - assert "bind 10.0.0.1:8000" in config - assert "no option http-buffer-request" in config - assert "option http-server-close" in config - assert "retries 0" in config - assert "retry-on none" in config - assert "timeout client 3600s" in config - assert "timeout server 3600s" in config - assert "acl private_pd_channel path -i /pd_register /kv_move_status" in config - assert "http-request return status 426 if private_pd_channel" in config - assert "option forwardfor" in config - assert "registered_objects" not in config - assert "allocate_global_unique_multimodal_id_range" not in config - assert "pd_master_config_server" not in config - - -@pytest.mark.parametrize( - ("host", "expected_address"), - [("0.0.0.0", "0.0.0.0:8100"), ("::", "[::]:8100")], -) -def test_haproxy_worker_addresses_follow_the_configured_listen_host(host, expected_address): - config = multiprocess._build_haproxy_config( - _args(host=host), _worker_ports(), Path("/tmp/lightllm-test-health.state") - ) - - assert f"server worker_0 {expected_address}" in config - - -def test_initial_health_state_only_contains_down_health_backends(): - state = multiprocess._build_initial_health_state(_worker_ports()) - state_lines = state.splitlines() - - assert state_lines[0] == "1" - assert "pd_master_workers" not in state - assert len(state_lines) == 2 + 3 * len(_worker_ports()) - for backend_name in ("pd_master_liveness", "pd_master_readiness", "pd_master_health"): - backend_lines = [line for line in state_lines if f" {backend_name} " in line] - assert len(backend_lines) == len(_worker_ports()) - for line in backend_lines: - fields = line.split() - assert fields[4] == "-" # srv_addr: keep the configured server address - assert fields[5] == "0" # srv_op_state: DOWN - assert fields[12] == "0" # srv_check_health: no successful check - assert fields[13] == "6" # srv_check_state: configured and enabled - - -@pytest.mark.parametrize("config", [{"config_server_host": "10.0.0.2"}, {"config_server_port": 9000}]) -def test_multiworker_rejects_config_server_before_haproxy_lookup(monkeypatch, config): - lookup = Mock() - monkeypatch.setitem( - sys.modules, - "ray_haproxy", - SimpleNamespace(get_haproxy_binary=lookup), - ) - - with pytest.raises(ValueError, match="does not yet support"): - multiprocess.validate_multiprocess_pd_master_args(_args(**config)) - lookup.assert_not_called() - - -@pytest.mark.parametrize("port", [None, 0, 65536]) -def test_multiworker_requires_a_valid_public_port_before_haproxy_lookup(monkeypatch, port): - lookup = Mock() - monkeypatch.setitem( - sys.modules, - "ray_haproxy", - SimpleNamespace(get_haproxy_binary=lookup), - ) - - with pytest.raises(ValueError, match="between 1 and 65535"): - multiprocess.validate_multiprocess_pd_master_args(_args(port=port)) - lookup.assert_not_called() - - -def test_cli_rejects_the_removed_worker_ports_option(capsys): - from lightllm.server.api_cli import add_cli_args - - parser = add_cli_args(argparse.ArgumentParser()) - with pytest.raises(SystemExit) as error: - parser.parse_args(["--pd_master_worker_ports", "8100", "8101"]) - assert error.value.code == 2 - assert "unrecognized arguments: --pd_master_worker_ports" in capsys.readouterr().err - - -def test_cli_rejects_the_removed_multiprocess_ports_option(capsys): - from lightllm.server.api_cli import add_cli_args - - parser = add_cli_args(argparse.ArgumentParser()) - with pytest.raises(SystemExit) as error: - parser.parse_args(["--pd_master_multi_ports", "8100", "8101"]) - assert error.value.code == 2 - assert "unrecognized arguments: --pd_master_multi_ports" in capsys.readouterr().err - - -def test_missing_haproxy_package_reports_the_install_command(monkeypatch): - monkeypatch.setitem(sys.modules, "ray_haproxy", None) - with pytest.raises(RuntimeError, match=r"pip install ray-haproxy==2\.8\.25"): - multiprocess.validate_multiprocess_pd_master_args(_args()) - - -def test_haproxy_uses_the_native_package(monkeypatch): - package = SimpleNamespace(get_haproxy_binary=Mock(return_value="/tmp/native-wheel/haproxy")) - monkeypatch.setitem(sys.modules, "ray_haproxy", package) - args = _args() - assert multiprocess.validate_multiprocess_pd_master_args(args) == "/tmp/native-wheel/haproxy" - package.get_haproxy_binary.assert_called_once_with() - - -def test_package_resolution_failure_remains_an_error(monkeypatch): - error = FileNotFoundError("bundled executable unavailable") - package = SimpleNamespace(get_haproxy_binary=Mock(side_effect=error)) - monkeypatch.setitem(sys.modules, "ray_haproxy", package) - with pytest.raises(RuntimeError, match="Cannot locate the packaged HAProxy executable"): - multiprocess.validate_multiprocess_pd_master_args(_args()) - - -def test_worker_ports_are_allocated_as_one_launcher_owned_list(monkeypatch): - from lightllm.utils.envs_utils import get_env_start_args, get_unique_server_name - from lightllm.utils.shm_port_args import ShmPortArgs - - args = _args() - monkeypatch.setenv( - "LIGHTLLM_UNIQUE_SERVICE_NAME_ID", - "test_pd_master_multi_ports_" + uuid.uuid4().hex, - ) - monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(vars(args))) - get_unique_server_name.cache_clear() - get_env_start_args.cache_clear() - ports = ShmPortArgs(create=True) - linked = None - try: - allocator = Mock(side_effect=[18100, 18101]) - monkeypatch.setattr(ports, "_alloc_free_port", allocator) - - assert ports.pd_master_multi_ports == [18100, 18101] - linked = ShmPortArgs() - assert linked.pd_master_multi_ports == [18100, 18101] - assert allocator.call_count == 2 - finally: - if linked is not None: - linked.close() - ports.shm.unlink() - ports.close() - get_unique_server_name.cache_clear() - get_env_start_args.cache_clear() - - -@pytest.mark.parametrize("run_mode", ["pd_master", "prefill", "config_server"]) -def test_config_server_port_is_never_allocated_dynamically(monkeypatch, run_mode): - from lightllm.utils.envs_utils import get_env_start_args, get_unique_server_name - from lightllm.utils.shm_port_args import ShmPortArgs - - args = _args(run_mode=run_mode, config_server_port=None) - monkeypatch.setenv("LIGHTLLM_UNIQUE_SERVICE_NAME_ID", "test_pd_config_port_" + uuid.uuid4().hex) - monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(vars(args))) - get_unique_server_name.cache_clear() - get_env_start_args.cache_clear() - ports = ShmPortArgs(create=True) - try: - with pytest.raises(RuntimeError, match="config_server_port.*None"): - _ = ports.config_server_port - finally: - ports.shm.unlink() - ports.close() - get_unique_server_name.cache_clear() - get_env_start_args.cache_clear() - - -@pytest.mark.parametrize("failure_stage", [None, "backend", "proxy"]) -def test_launcher_uses_process_manager_before_releasing_runtime_and_ports(monkeypatch, tmp_path, failure_stage): - args, ports = _args(), _ports(close=Mock()) - proxy = Mock(pid=3103) - managed_proxy = object() - runtime_paths, events = [], [] - create_runtime = multiprocess.tempfile.TemporaryDirectory - - def runtime_directory(**kwargs): - runtime = create_runtime(dir=tmp_path, **kwargs) - runtime_paths.append(multiprocess.Path(runtime.name)) - return runtime - - def start_backends(worker_count): - assert worker_count == 2 - if failure_stage == "backend": - raise RuntimeError("injected backend startup failure") - - def terminate(): - assert runtime_paths[0].is_dir() - ports.close.assert_not_called() - events.append("stop") - - manager = SimpleNamespace( - processes=[], - process_names={}, - terminate_all_processes=Mock(side_effect=terminate), - setup_signal_handlers=Mock(), - supervise_processes=Mock(), - ) - ports.close.side_effect = lambda: events.append("close ports") - monkeypatch.setattr(multiprocess, "process_manager", manager) - monkeypatch.setattr(multiprocess, "get_unique_server_name", lambda: "test-cleanup-order") - monkeypatch.setattr(multiprocess.tempfile, "TemporaryDirectory", runtime_directory) - monkeypatch.setattr(multiprocess, "_start_backends", start_backends) - config_check = Mock() - monkeypatch.setattr(multiprocess.subprocess, "run", config_check) - monkeypatch.setattr(multiprocess.psutil, "Process", Mock(return_value=managed_proxy)) - launch_proxy = Mock(return_value=proxy) - if failure_stage == "proxy": - launch_proxy.side_effect = OSError("injected proxy startup failure") - monkeypatch.setattr(multiprocess.subprocess, "Popen", launch_proxy) - - if failure_stage is None: - multiprocess.start_multiprocess_pd_master( - args, - ports=ports, - haproxy_binary="/bundled/haproxy", - ) - manager.supervise_processes.assert_called_once_with(proxy) - manager.setup_signal_handlers.assert_called_once_with( - proxy, - graceful_shutdown_signal=signal.SIGUSR1, - ) - assert manager.processes == [managed_proxy] - assert manager.process_names == {managed_proxy: "pd_master_haproxy"} - else: - with pytest.raises((RuntimeError, OSError), match=f"injected {failure_stage} startup failure"): - multiprocess.start_multiprocess_pd_master( - args, - ports=ports, - haproxy_binary="/bundled/haproxy", - ) - manager.supervise_processes.assert_not_called() - - manager.terminate_all_processes.assert_called_once_with() - assert events == ["stop", "close ports"] - assert not runtime_paths[0].exists() - ports.close.assert_called_once_with() - if failure_stage != "backend": - assert config_check.call_args.args[0][0:3] == ["/bundled/haproxy", "-c", "-f"] diff --git a/unit_tests/server/test_pd_master_req_ids.py b/unit_tests/server/test_pd_master_req_ids.py deleted file mode 100644 index d4f80acb03..0000000000 --- a/unit_tests/server/test_pd_master_req_ids.py +++ /dev/null @@ -1,149 +0,0 @@ -"""Request IDs are owned by one RPyC service instead of shared memory.""" - -import socket -import threading -import uuid -from types import SimpleNamespace -from unittest.mock import Mock, patch - -from rpyc.utils.server import ThreadedServer - -from lightllm.server.httpserver import req_id_generator_client -from lightllm.server.httpserver.req_id_generator_server import ( - CLIENT_REQUEST_ID_RANGE_SIZE, - CONFIG_SERVER_REQUEST_TIMEOUT_SECONDS, - MAX_BEST_OF, - get_req_id_generator_socket_address, -) -from lightllm.utils.net_utils import get_abstract_unix_socket_address -from lightllm.server.httpserver.req_id_generator_client import ReqIDGenerator -from lightllm.server.httpserver.req_id_generator_server import ( - ConfigServerReqIDGeneratorState, - StandaloneReqIDGeneratorState, - _create_req_id_generator_state, -) - - -def _args(**overrides): - values = { - "run_mode": "normal", - "config_server_host": None, - "config_server_port": None, - } - values.update(overrides) - return SimpleNamespace(**values) - - -def _remote_client(ipc_name): - socket_address = get_abstract_unix_socket_address("req_id_generator", ipc_name) - with patch.object( - req_id_generator_client, - "get_req_id_generator_socket_address", - return_value=socket_address, - ): - return ReqIDGenerator() - - -def test_server_allocates_aligned_request_id_ranges(): - state = StandaloneReqIDGeneratorState() - - first = state.allocate_id_range(4) - second = state.allocate_id_range(2) - - assert first == (8, 40) - assert second == (40, 56) - - -def test_client_generates_locally_until_its_range_is_exhausted(monkeypatch): - connection = Mock() - connection.root.allocate_id_range.side_effect = [(8, 8200), (8200, 16392)] - monkeypatch.setattr(req_id_generator_client, "unix_connect", Mock(return_value=connection)) - client = _remote_client("request-id-client") - - first_range = [client.generate_id() for _ in range(CLIENT_REQUEST_ID_RANGE_SIZE)] - - assert first_range == list(range(8, 8200, MAX_BEST_OF)) - connection.root.allocate_id_range.assert_called_once_with(CLIENT_REQUEST_ID_RANGE_SIZE) - assert client.generate_id() == 8200 - assert connection.root.allocate_id_range.call_count == 2 - - -def test_config_server_range_is_used_and_renewed(monkeypatch): - responses = [ - Mock(status_code=200, json=Mock(return_value={"start_id": 0, "end_id": 800000})), - Mock( - status_code=200, - json=Mock(return_value={"start_id": 800000, "end_id": 1600000}), - ), - ] - get = Mock(side_effect=responses) - monkeypatch.setattr("lightllm.server.httpserver.req_id_generator_server.requests.get", get) - state = ConfigServerReqIDGeneratorState("10.0.0.1", 9000) - - assert state.allocate_id_range(1) == (8, 16) - state.next_id = state.range_end - MAX_BEST_OF - assert state.allocate_id_range(1) == (799992, 800000) - assert state.allocate_id_range(2) == (800000, 800016) - assert [call.args for call in get.call_args_list] == [ - ("http://10.0.0.1:9000/allocate_global_unique_id_range",), - ("http://10.0.0.1:9000/allocate_global_unique_id_range",), - ] - assert [call.kwargs for call in get.call_args_list] == [ - {"timeout": CONFIG_SERVER_REQUEST_TIMEOUT_SECONDS}, - {"timeout": CONFIG_SERVER_REQUEST_TIMEOUT_SECONDS}, - ] - - -def test_request_id_socket_uses_configured_ipc_name(): - with patch( - "lightllm.server.httpserver.req_id_generator_server.get_pd_master_unix_socket_address", - return_value=get_abstract_unix_socket_address("req_id_generator", "launcher-name"), - ): - assert get_req_id_generator_socket_address().endswith(b"launcher-name_req_id_generator") - - -def test_state_type_is_selected_once_during_server_startup(): - assert isinstance(_create_req_id_generator_state(_args()), StandaloneReqIDGeneratorState) - assert isinstance( - _create_req_id_generator_state(_args(config_server_host="127.0.0.1")), - StandaloneReqIDGeneratorState, - ) - assert isinstance( - _create_req_id_generator_state(_args(config_server_port=9000)), - StandaloneReqIDGeneratorState, - ) - assert isinstance( - _create_req_id_generator_state( - _args( - run_mode="normal", - config_server_host="127.0.0.1", - config_server_port=9000, - ) - ), - ConfigServerReqIDGeneratorState, - ) - - -def test_real_rpc_clients_share_one_sequence(): - ipc_name = uuid.uuid4().hex - state = StandaloneReqIDGeneratorState() - server = ThreadedServer( - state, - socket_path=get_abstract_unix_socket_address("req_id_generator", ipc_name), - protocol_config={"allow_pickle": True}, - ) - thread = threading.Thread(target=server.start, daemon=True) - thread.start() - clients = [_remote_client(ipc_name) for _ in range(4)] - try: - ids = [clients[index % len(clients)].generate_id() for index in range(200)] - assert len(set(ids)) == len(ids) - assert all(request_id % MAX_BEST_OF == 0 for request_id in ids) - assert state.next_id == 8 + 4 * CLIENT_REQUEST_ID_RANGE_SIZE * MAX_BEST_OF - assert clients[0].connection._channel.stream.sock.family == socket.AF_UNIX - finally: - for client in clients: - client.connection.close() - server.close() - thread.join(timeout=5) - assert not thread.is_alive() diff --git a/unit_tests/server/test_pd_master_topology_api.py b/unit_tests/server/test_pd_master_topology_api.py deleted file mode 100644 index de92995ddc..0000000000 --- a/unit_tests/server/test_pd_master_topology_api.py +++ /dev/null @@ -1,96 +0,0 @@ -import asyncio -import json -import pickle -from types import SimpleNamespace -from unittest.mock import AsyncMock, Mock - -import pytest -from fastapi.testclient import TestClient - -from lightllm.server import api_http -from lightllm.server.core.objs import StartArgs -from lightllm.server.httpserver_for_pd_master.manager import HttpServerManagerForPDMaster -from lightllm.server.pd_io_struct import PDUpKVStatus - - -def _install(monkeypatch, *, ready=False, healthy=False): - # Exercise the real HTTP/WS handlers without loading a tokenizer/model. - monkeypatch.setattr(api_http.app.router, "on_startup", []) - monkeypatch.setattr(api_http.app.router, "on_shutdown", []) - args = StartArgs( - run_mode="pd_master", - host="10.0.0.1", - port=18090, - pd_node_id=100, - ) - manager = HttpServerManagerForPDMaster.__new__(HttpServerManagerForPDMaster) - manager.__dict__.update( - args=args, - pd_manager=SimpleNamespace( - prefill_nodes=[SimpleNamespace(node_id=199)], - decode_nodes=[SimpleNamespace(node_id=200)], - is_pd_nodes_ready=Mock(return_value=ready), - check_pd_nodes_health=AsyncMock(return_value=True), - ), - is_healthy=Mock(return_value=healthy), - update_req_status=AsyncMock(), - ) - monkeypatch.setattr(api_http, "g_objs", SimpleNamespace(args=args, httpserver_manager=manager)) - return args, manager - - -@pytest.mark.parametrize("ready,healthy", [(False, False), (True, False), (True, True)]) -def test_local_readiness_and_health_preserve_distinct_checks(monkeypatch, ready, healthy): - _, manager = _install(monkeypatch, ready=ready, healthy=healthy) - readiness = api_http.readiness() - assert readiness.status_code == (200 if ready else 503) - assert json.loads(readiness.body) == {"status": "ok" if ready else "not ready"} - response = asyncio.run(api_http.healthcheck(None)) - assert response.status_code == (200 if ready and healthy else 503) - assert json.loads(response.body) == { - "inference_healthy": healthy, - "pd_nodes_ready": ready, - "message": "Ok" if ready and healthy else "Error", - "pd_master_mode": "elastic", - "registered_prefill_nodes": 1, - "registered_decode_nodes": 1, - } - manager.pd_manager.check_pd_nodes_health.assert_not_awaited() - - -async def _wait_until(predicate): - async def poll(): - while not predicate(): - await asyncio.sleep(0) - - await asyncio.wait_for(poll(), timeout=1) - - -def test_local_disconnect_makes_readiness_and_health_unavailable(monkeypatch): - _, manager = _install(monkeypatch, ready=True, healthy=True) - manager.pd_manager.is_pd_nodes_ready.return_value = False - readiness = api_http.readiness() - health = asyncio.run(api_http.healthcheck(None)) - assert readiness.status_code == health.status_code == 503 - assert json.loads(readiness.body) == {"status": "not ready"} - assert json.loads(health.body)["pd_nodes_ready"] is False - assert json.loads(health.body)["message"] == "Error" - - -def test_fixed_topology_health_failure_has_consistent_payload(monkeypatch): - args, manager = _install(monkeypatch, ready=True, healthy=True) - args.pd_master_mode = "1p1d" - manager.pd_manager.check_pd_nodes_health.return_value = False - response = asyncio.run(api_http.healthcheck(None)) - assert response.status_code == 503 - assert json.loads(response.body)["pd_nodes_healthy"] is False - assert json.loads(response.body)["message"] == "Error" - manager.pd_manager.check_pd_nodes_health.assert_awaited_once() - - -def test_kv_status_keeps_existing_wire_protocol(monkeypatch): - _, manager = _install(monkeypatch) - with TestClient(api_http.app) as client: - with client.websocket_connect("/kv_move_status") as websocket: - websocket.send_bytes(pickle.dumps(PDUpKVStatus(8, 100, b"metadata"))) - client.portal.call(_wait_until, lambda: manager.update_req_status.await_count == 1) diff --git a/unit_tests/server/test_pd_node_stats.py b/unit_tests/server/test_pd_node_stats.py deleted file mode 100644 index d8a536ce4b..0000000000 --- a/unit_tests/server/test_pd_node_stats.py +++ /dev/null @@ -1,321 +0,0 @@ -"""PD node dispatch statistics stay independent from prompt-cache state.""" - -import asyncio -import errno -import os -import signal -import socket -import threading -import uuid -from unittest.mock import Mock, patch - -import pytest -from rpyc.utils.server import ThreadedServer - -from lightllm.server.httpserver_for_pd_master.pd_selector import pd_node_stats_client -from lightllm.server.httpserver_for_pd_master.pd_selector import pd_node_stats_server -from lightllm.server.httpserver_for_pd_master.pd_selector.pd_node_stats_client import ( - LocalPDNodeStats, - RemotePDNodeStats, - create_pd_node_stats_client, -) -from lightllm.server.httpserver_for_pd_master.pd_selector.pd_node_stats_server import PDNodeStatsState -from lightllm.server.pd_io_struct import PD_Client_Obj -from lightllm.utils.net_utils import get_abstract_unix_socket_address - - -def _node(address="p", node_id=1): - return PD_Client_Obj( - node_id=node_id, - client_ip_port=address, - mode="prefill", - start_args={}, - dp_index=0, - ) - - -def _counts(node): - return node.dispatched_prompt_chars, node.dispatched_req_num - - -def _query(state, *keys): - return dict(zip(keys, state.get_dispatched_info(tuple(keys)))) - - -def _remote_client(ipc_name="test-pd-node-stats"): - with patch.object( - pd_node_stats_client, - "get_pd_master_unix_socket_address", - return_value=get_abstract_unix_socket_address("pd_node_stats", ipc_name), - ): - return RemotePDNodeStats() - - -def test_ipc_address_uses_the_launcher_ipc_name(monkeypatch): - monkeypatch.setattr(pd_node_stats_client, "unix_connect", Mock()) - client = _remote_client("deployment-a") - assert client.socket_path == get_abstract_unix_socket_address("pd_node_stats", "deployment-a") - assert client.socket_path != get_abstract_unix_socket_address("pd_node_stats", "deployment-b") - - -def test_local_client_updates_node_objects_without_rpc(monkeypatch): - monkeypatch.setattr( - pd_node_stats_client, - "get_env_start_args", - lambda: Mock(httpserver_workers=1), - ) - client = create_pd_node_stats_client() - assert isinstance(client, LocalPDNodeStats) - node = _node() - - async def run(): - await client.add_dispatched_info(p_node=node, chars=10, req_num=2) - assert _counts(node) == (10, 2) - await client.sub_dispatched_info(p_node=node, chars=3, req_num=1) - assert _counts(node) == (7, 1) - await client.sub_dispatched_info(p_node=node, chars=100, req_num=100) - assert _counts(node) == (0, 0) - assert await client.update_dispatched_info([node]) is None - - asyncio.run(run()) - - -def test_remote_client_sends_only_connection_keys(monkeypatch): - monkeypatch.setattr(pd_node_stats_client, "unix_connect", Mock()) - client = _remote_client() - client._add_dispatched_info = Mock() - client._sub_dispatched_info = Mock() - client._get_dispatched_info = Mock(return_value=((17, 1),)) - node = _node() - - async def run(): - await client.add_dispatched_info(p_node=node, chars=17, req_num=1) - await client.sub_dispatched_info(p_node=node, chars=3, req_num=1) - await client.update_dispatched_info([node]) - - asyncio.run(run()) - client._add_dispatched_info.assert_called_once_with(node.connection_key, 17, 1) - client._sub_dispatched_info.assert_called_once_with(node.connection_key, 3, 1) - client._get_dispatched_info.assert_called_once_with((node.connection_key,)) - assert _counts(node) == (17, 1) - - -def test_state_creates_keys_lazily_and_clamps_subtraction(): - state = PDNodeStatsState() - assert _query(state, "a", "b") == {"a": (0, 0), "b": (0, 0)} - state.add_dispatched_info("a", 10, 2) - state.sub_dispatched_info("a", 100, 1) - assert _query(state, "a")["a"] == (0, 1) - state.sub_dispatched_info("a", 0, 100) - assert _query(state, "a")["a"] == (0, 0) - - -@pytest.mark.parametrize("key", ["", None, 1]) -def test_invalid_key_does_not_change_statistics(key): - state = PDNodeStatsState() - with pytest.raises(ValueError, match="nonempty string"): - state.add_dispatched_info(key, 1, 1) - assert not state.stats - - -def test_idle_and_capacity_eviction_are_lru(monkeypatch): - now = [1000.0] - monkeypatch.setattr(pd_node_stats_server.time, "monotonic", lambda: now[0]) - monkeypatch.setattr(pd_node_stats_server, "MAX_NODE_STATS", 2) - state = PDNodeStatsState() - _query(state, "a", "b") - _query(state, "a") - _query(state, "c") - assert list(state.stats) == ["a", "c"] - now[0] += pd_node_stats_server.NODE_STATS_IDLE_SECONDS - _query(state, "new") - assert list(state.stats) == ["new"] - - -def test_cancelled_update_finishes_before_following_update(monkeypatch): - monkeypatch.setattr(pd_node_stats_client, "unix_connect", Mock()) - state = PDNodeStatsState() - client = _remote_client() - node = _node() - started, finish = threading.Event(), threading.Event() - operations = [] - - def add_dispatched_info(*args): - started.set() - assert finish.wait(5) - state.add_dispatched_info(*args) - operations.append("add_dispatched_info") - - def sub_dispatched_info(*args): - state.sub_dispatched_info(*args) - operations.append("sub_dispatched_info") - - client._add_dispatched_info = add_dispatched_info - client._sub_dispatched_info = sub_dispatched_info - - async def run(): - first = asyncio.create_task(client.add_dispatched_info(p_node=node, chars=10, req_num=1)) - assert await asyncio.to_thread(started.wait, 5) - second = asyncio.create_task(client.sub_dispatched_info(p_node=node, chars=10, req_num=1)) - first.cancel() - await asyncio.sleep(0) - assert not first.done() and not second.done() - finish.set() - with pytest.raises(asyncio.CancelledError): - await first - await second - - asyncio.run(run()) - assert operations == ["add_dispatched_info", "sub_dispatched_info"] - assert _query(state, node.connection_key)[node.connection_key] == (0, 0) - - -def test_cancelled_update_waits_for_the_client_lock_and_still_executes(monkeypatch): - monkeypatch.setattr(pd_node_stats_client, "unix_connect", Mock()) - state = PDNodeStatsState() - client = _remote_client() - client._add_dispatched_info = state.add_dispatched_info - node = _node() - - async def run(): - await client.async_lock.acquire() - task = asyncio.create_task(client.add_dispatched_info(p_node=node, chars=10, req_num=1)) - await asyncio.sleep(0) - task.cancel() - await asyncio.sleep(0) - assert not task.done() - client.async_lock.release() - with pytest.raises(asyncio.CancelledError): - await task - - asyncio.run(run()) - assert _query(state, node.connection_key)[node.connection_key] == (10, 1) - - -def test_rpc_failure_terminates_the_worker(monkeypatch): - monkeypatch.setattr(pd_node_stats_client, "unix_connect", Mock()) - terminate = Mock() - monkeypatch.setattr(pd_node_stats_client.os, "kill", terminate) - client = _remote_client() - client._add_dispatched_info = Mock(side_effect=ConnectionError("connection lost")) - - with pytest.raises(ConnectionError, match="connection lost"): - asyncio.run(client.add_dispatched_info(p_node=_node(), chars=10, req_num=1)) - - terminate.assert_called_once_with(os.getpid(), signal.SIGTERM) - - -def test_rpc_failure_after_cancellation_still_terminates_the_worker(monkeypatch): - monkeypatch.setattr(pd_node_stats_client, "unix_connect", Mock()) - terminate = Mock() - monkeypatch.setattr(pd_node_stats_client.os, "kill", terminate) - client = _remote_client() - client._add_dispatched_info = Mock(side_effect=ConnectionError("connection lost")) - - async def run(): - await client.async_lock.acquire() - task = asyncio.create_task(client.add_dispatched_info(p_node=_node(), chars=10, req_num=1)) - await asyncio.sleep(0) - task.cancel() - await asyncio.sleep(0) - client.async_lock.release() - with pytest.raises(ConnectionError, match="connection lost"): - await task - - asyncio.run(run()) - terminate.assert_called_once_with(os.getpid(), signal.SIGTERM) - - -@pytest.fixture -def rpc_stats_cluster(): - state = PDNodeStatsState() - ipc_name = uuid.uuid4().hex - server = ThreadedServer( - state, - socket_path=get_abstract_unix_socket_address("pd_node_stats", ipc_name), - protocol_config={"allow_pickle": True}, - ) - thread = threading.Thread(target=server.start, daemon=True) - thread.start() - clients = [_remote_client(ipc_name) for _ in range(4)] - try: - yield state, clients - finally: - for client in clients: - connection = client.connection - if connection is not None: - connection.close() - server.close() - thread.join(timeout=5) - assert not thread.is_alive() - - -def test_real_rpc_clients_share_statistics(rpc_stats_cluster): - _, clients = rpc_stats_cluster - node = _node() - node.websocket = threading.Lock() - - async def run(): - await clients[0].add_dispatched_info(p_node=node, chars=17, req_num=1) - await clients[1].update_dispatched_info([node]) - assert _counts(node) == (17, 1) - transport = clients[0].connection._channel.stream.sock - assert transport.family == socket.AF_UNIX - await clients[1].sub_dispatched_info(p_node=node, chars=17, req_num=1) - await clients[0].update_dispatched_info([node]) - assert _counts(node) == (0, 0) - - asyncio.run(run()) - - -def test_abstract_socket_is_released_after_server_close(): - address = get_abstract_unix_socket_address("pd_node_stats", uuid.uuid4().hex) - server = ThreadedServer( - PDNodeStatsState(), - socket_path=address, - protocol_config={"allow_pickle": True}, - ) - try: - with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as duplicate: - with pytest.raises(OSError) as conflict: - duplicate.bind(address) - assert conflict.value.errno == errno.EADDRINUSE - finally: - server.close() - ThreadedServer( - PDNodeStatsState(), - socket_path=address, - protocol_config={"allow_pickle": True}, - ).close() - - -def test_initial_connection_failure_is_reported_during_construction(monkeypatch): - attempts = [] - - def connect(*args, **kwargs): - attempts.append(1) - raise ConnectionRefusedError("not listening") - - monkeypatch.setattr(pd_node_stats_client, "unix_connect", connect) - - for _ in range(2): - with pytest.raises(ConnectionRefusedError): - _remote_client() - - assert len(attempts) == 2 - - -def test_real_rpc_combined_updates_are_atomic(rpc_stats_cluster): - state, clients = rpc_stats_cluster - node = _node() - - async def run(): - await asyncio.gather( - *( - clients[index % len(clients)].add_dispatched_info(p_node=node, chars=1, req_num=1) - for index in range(200) - ) - ) - - asyncio.run(run()) - assert _query(state, node.connection_key)[node.connection_key] == (200, 200) diff --git a/unit_tests/server/test_pd_start_args.py b/unit_tests/server/test_pd_start_args.py deleted file mode 100644 index 3109fe11b9..0000000000 --- a/unit_tests/server/test_pd_start_args.py +++ /dev/null @@ -1,311 +0,0 @@ -import json - -import pytest - -from lightllm.server.api_start import _launch_subprocesses -from lightllm.server.core.objs.start_args_type import StartArgs - - -@pytest.fixture -def startup_args(monkeypatch): - for name in ( - "_set_envs_and_config", - "auto_set_max_req_total_len", - "auto_set_fused_shared_experts", - "set_unique_server_name", - ): - monkeypatch.setattr("lightllm.server.api_start." + name, lambda args: None) - - def validation_finished(args): - raise RuntimeError("startup validation complete") - - monkeypatch.setattr("lightllm.server.api_start.auto_set_response_parsers", validation_finished) - return StartArgs( - model_dir="unused", - max_req_total_len=8192, - eos_id=[2], - disable_vision=True, - disable_audio=True, - disable_shm_warning=True, - ) - - -@pytest.mark.parametrize( - "model_type,block_num,cpu_page_size,expected", - [ - ("deepseek_v4", 8, None, 4096), - ("deepseek_v4", 8, 4096, 4096), - ("deepseek_v4", 10000000, None, 512 * 10000000), - ("deepseek_v4", 10000000, 4096, 512 * 10000000), - ("glm5_next", 8, None, 4096), - ("glm5_next", 8, 256, 4096), - ("llama", 8, None, 256), - ("llama", 8, 1024, 1024), - ], -) -def test_cpu_cache_page_defaults_preserve_model_contracts( - startup_args, monkeypatch, model_type, block_num, cpu_page_size, expected -): - monkeypatch.setattr("lightllm.utils.config_utils.get_model_type", lambda _: model_type) - monkeypatch.setattr( - "lightllm.server.api_start.is_hybrid_att_model", lambda _: model_type in ("deepseek_v4", "glm5_next") - ) - args = startup_args - args.enable_cpu_cache = True - args.linear_att_page_block_num = block_num - args.cpu_cache_token_page_size = cpu_page_size - with pytest.raises(RuntimeError, match="startup validation complete"): - _launch_subprocesses(args) - assert args.cpu_cache_token_page_size == expected - assert args.linear_att_page_block_num == block_num - - -@pytest.mark.parametrize("model_type", ["deepseek_v4", "llama"]) -@pytest.mark.parametrize("enable_cpu_cache", [False, True]) -@pytest.mark.parametrize("disable_chunked_prefill", [False, True]) -def test_model_kv_defaults_preserve_hash_pages_kv_selection_and_explicit_pd_tuning( - startup_args, monkeypatch, model_type, enable_cpu_cache, disable_chunked_prefill -): - monkeypatch.setattr("lightllm.utils.config_utils.get_model_type", lambda _: model_type) - monkeypatch.setattr("lightllm.server.api_start.is_hybrid_att_model", lambda _: model_type == "deepseek_v4") - args = startup_args - args.enable_cpu_cache = enable_cpu_cache - args.disable_chunked_prefill = disable_chunked_prefill - args.pd_kv_page_num, args.pd_kv_page_size = 8, 2048 - if model_type == "deepseek_v4" and disable_chunked_prefill: - with pytest.raises(AssertionError, match="hybrid attention models require chunked prefill"): - _launch_subprocesses(args) - return - with pytest.raises(RuntimeError, match="startup validation complete"): - _launch_subprocesses(args) - assert args.page_size == (256 if model_type == "deepseek_v4" else 1) - assert args.linear_att_hash_page_size == 512 - assert args.llm_kv_type == "None" - assert (args.pd_kv_page_num, args.pd_kv_page_size) == (8, 2048) - assert args.cache_placement_strategy == "adaptive" - - -def test_disabled_cpu_cache_keeps_page_size_unset(startup_args): - with pytest.raises(RuntimeError, match="startup validation complete"): - _launch_subprocesses(startup_args) - assert startup_args.cpu_cache_token_page_size is None - - -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", - ) - assert isinstance(state, GQAFullAttentionState) - - -def test_dsv4_cpu_cache_factory_resolves_packed_layout_without_overriding_kv_type(startup_args, monkeypatch, tmp_path): - 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 - - (tmp_path / "config.json").write_text( - json.dumps( - { - "model_type": "deepseek_v4", - "num_hidden_layers": 2, - "head_dim": 512, - "index_head_dim": 128, - "qk_rope_head_dim": 64, - "sliding_window": 128, - "compress_ratios": [4, 128], - } - ) - ) - args = startup_args - args.model_dir = str(tmp_path) - args.enable_cpu_cache = True - args.linear_att_page_block_num = 8 - args.data_type = "bf16" - with pytest.raises(RuntimeError, match="startup validation complete"): - _launch_subprocesses(args) - monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(vars(args))) - cached_functions = ( - get_env_start_args, - get_added_mtp_kv_layer_num, - get_llm_data_type, - get_llm_model_class, - calcu_cpu_cache_meta, - ) - for function in cached_functions: - function.cache_clear() - try: - 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 args.cpu_cache_token_page_size == args.linear_att_hash_page_size * args.linear_att_page_block_num - 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: - function.cache_clear() - - -@pytest.mark.parametrize( - "page_size,hash_page_size,cpu_page_size", - [ - (1, 256, None), - (256, 512, None), - (256, 256, 4096), - (256, 256, None), - (256, 512, 4096), - (512, 512, None), - (768, 768, None), - (1024, 1024, None), - ], -) -def test_dsv4_page_and_checkpoint_configuration(startup_args, monkeypatch, page_size, hash_page_size, cpu_page_size): - monkeypatch.setattr("lightllm.utils.config_utils.get_model_type", lambda model_dir: "deepseek_v4") - monkeypatch.setattr("lightllm.server.api_start.is_hybrid_att_model", lambda model_dir: True) - - def validation_finished(args): - assert args.page_size == (256 if page_size == 1 else page_size) - assert args.linear_att_hash_page_size == hash_page_size - assert args.cpu_cache_token_page_size == hash_page_size * 8 - assert args.llm_kv_type == "None" - assert (args.pd_kv_page_num, args.pd_kv_page_size) == (16, 1024) - raise RuntimeError("DSV4 page-size validation passed") - - monkeypatch.setattr("lightllm.server.api_start.auto_set_response_parsers", validation_finished) - args = startup_args - args.page_size = page_size - args.linear_att_hash_page_size = hash_page_size - args.linear_att_page_block_num = 8 - args.enable_cpu_cache = True - args.cpu_cache_token_page_size = cpu_page_size - with pytest.raises(RuntimeError, match="DSV4 page-size validation passed"): - _launch_subprocesses(args) - - -@pytest.mark.parametrize( - "page_size,hash_page_size,error", - [ - (256, 128, "--linear_att_hash_page_size must be divisible by --page_size"), - (256, 384, "--linear_att_hash_page_size must be divisible by --page_size"), - (512, 256, "--linear_att_hash_page_size must be divisible by --page_size"), - (0, 512, "--page_size must be >= 1"), - ], -) -def test_dsv4_rejects_unaligned_cache_pages(startup_args, monkeypatch, page_size, hash_page_size, error): - monkeypatch.setattr("lightllm.utils.config_utils.get_model_type", lambda _: "deepseek_v4") - monkeypatch.setattr("lightllm.server.api_start.is_hybrid_att_model", lambda _: True) - args = startup_args - args.page_size = page_size - args.linear_att_hash_page_size = hash_page_size - with pytest.raises(ValueError, match=error): - _launch_subprocesses(args) - - -def test_pd_cli_and_dataclass_keep_main_defaults(): - import argparse - from lightllm.server.api_cli import add_cli_args - - parsed = add_cli_args(argparse.ArgumentParser()).parse_args([]) - args = StartArgs() - assert (parsed.pd_kv_page_num, parsed.pd_kv_page_size) == (16, 1024) - assert (args.pd_kv_page_num, args.pd_kv_page_size) == (16, 1024) - - -def test_pd_master_does_not_apply_model_local_cpu_cache_restrictions(monkeypatch): - from lightllm.server.api_start import pd_master_start - - monkeypatch.setattr("lightllm.server.api_start._set_envs_and_config", lambda args: None) - monkeypatch.setattr("lightllm.server.api_start.set_unique_server_name", lambda args: None) - - def validation_finished(args): - raise RuntimeError("PD master uses no local CPU cache") - - monkeypatch.setattr("lightllm.server.api_start.auto_set_max_req_total_len", validation_finished) - pd_master_args = StartArgs(run_mode="pd_master", model_dir="unused", enable_cpu_cache=True) - with pytest.raises(RuntimeError, match="PD master uses no local CPU cache"): - pd_master_start(pd_master_args) - - -def test_pd_kv_page_size_must_be_divisible_by_model_page_size(monkeypatch): - monkeypatch.setattr("lightllm.server.api_start._set_envs_and_config", lambda args: None) - monkeypatch.setattr("lightllm.server.api_start.auto_set_max_req_total_len", lambda args: None) - monkeypatch.setattr("lightllm.server.api_start.auto_set_fused_shared_experts", lambda args: None) - monkeypatch.setattr("lightllm.server.api_start.set_unique_server_name", lambda args: None) - args = StartArgs( - run_mode="decode", - page_size=4, - pd_kv_page_size=6, - disable_vision=True, - disable_audio=True, - disable_shm_warning=True, - ) - - with pytest.raises(AssertionError, match="--pd_kv_page_size must be divisible by --page_size"): - _launch_subprocesses(args) - - -def test_normal_mode_does_not_validate_pd_kv_page_size(monkeypatch): - monkeypatch.setattr("lightllm.server.api_start._set_envs_and_config", lambda args: None) - monkeypatch.setattr("lightllm.server.api_start.auto_set_max_req_total_len", lambda args: None) - monkeypatch.setattr("lightllm.server.api_start.auto_set_fused_shared_experts", lambda args: None) - monkeypatch.setattr("lightllm.server.api_start.set_unique_server_name", lambda args: None) - monkeypatch.setattr("lightllm.server.api_start.is_hybrid_att_model", lambda model_dir: False) - - def validation_finished(args): - raise RuntimeError("PD page-size validation passed") - - monkeypatch.setattr("lightllm.server.api_start.auto_set_response_parsers", validation_finished) - args = StartArgs( - run_mode="normal", - model_dir="unused", - page_size=4, - pd_kv_page_size=6, - eos_id=[2], - disable_vision=True, - disable_audio=True, - disable_shm_warning=True, - ) - - with pytest.raises(RuntimeError, match="PD page-size validation passed"): - _launch_subprocesses(args) - - -@pytest.mark.parametrize("llm_kv_type", ["fp8kv_sph", "fp8kv_spt"]) -def test_fp8_kv_cache_allows_multi_token_model_pages(monkeypatch, llm_kv_type): - monkeypatch.setattr("lightllm.server.api_start._set_envs_and_config", lambda args: None) - monkeypatch.setattr("lightllm.server.api_start.auto_set_max_req_total_len", lambda args: None) - monkeypatch.setattr("lightllm.server.api_start.auto_set_fused_shared_experts", lambda args: None) - monkeypatch.setattr("lightllm.server.api_start.set_unique_server_name", lambda args: None) - monkeypatch.setattr("lightllm.server.api_start.is_hybrid_att_model", lambda model_dir: False) - - def validation_finished(args): - raise RuntimeError("FP8 page-size validation passed") - - monkeypatch.setattr("lightllm.server.api_start.auto_set_response_parsers", validation_finished) - args = StartArgs( - llm_kv_type=llm_kv_type, - kv_quant_calibration_config_path="unused", - model_dir="unused", - page_size=16, - eos_id=[2], - disable_vision=True, - disable_audio=True, - disable_shm_warning=True, - ) - - with pytest.raises(RuntimeError, match="FP8 page-size validation passed"): - _launch_subprocesses(args) diff --git a/unit_tests/server/test_prefill_queue_start_args.py b/unit_tests/server/test_prefill_queue_start_args.py deleted file mode 100644 index 6a802fd1ae..0000000000 --- a/unit_tests/server/test_prefill_queue_start_args.py +++ /dev/null @@ -1,27 +0,0 @@ -import pytest - -from lightllm.server.api_cli import make_argument_parser -from lightllm.server.api_start import _auto_set_prefill_queue_strategy -from lightllm.server.core.objs.start_args_type import StartArgs - - -@pytest.mark.parametrize( - "run_mode, automatic_strategy", - [("normal", "default"), ("prefill", "hrrn"), ("decode", "default"), ("visual_only", "default")], -) -@pytest.mark.parametrize("configured_strategy", [None, "auto", "default", "promote_shortest", "hrrn"]) -def test_auto_set_prefill_queue_strategy(run_mode, automatic_strategy, configured_strategy): - cli_args = ["--run_mode", run_mode] - if configured_strategy is not None: - cli_args.extend(["--prefill_queue_strategy", configured_strategy]) - args = StartArgs(**vars(make_argument_parser().parse_args(cli_args))) - assert args.prefill_queue_strategy == (configured_strategy or "auto") - expected = automatic_strategy if configured_strategy in (None, "auto") else configured_strategy - _auto_set_prefill_queue_strategy(args) - assert args.prefill_queue_strategy == expected - - -def test_start_args_default_prefill_queue_strategy_matches_cli(): - assert StartArgs().prefill_queue_strategy == make_argument_parser().parse_args([]).prefill_queue_strategy == "auto" - cli_option = next(action for action in make_argument_parser()._actions if action.dest == "prefill_queue_strategy") - assert set(cli_option.choices) == set(StartArgs.__dataclass_fields__["prefill_queue_strategy"].metadata["choices"]) diff --git a/unit_tests/server/test_prompt_cache_tree_service.py b/unit_tests/server/test_prompt_cache_tree_service.py deleted file mode 100644 index b85cd67902..0000000000 --- a/unit_tests/server/test_prompt_cache_tree_service.py +++ /dev/null @@ -1,228 +0,0 @@ -"""Prompt cache tree local and RPyC component behavior.""" - -import asyncio -import os -import signal -import socket -import threading -import uuid -from unittest.mock import Mock, patch - -import pytest -from rpyc.utils.server import ThreadedServer - -from lightllm.server.httpserver_for_pd_master.pd_selector import prompt_cache_tree_client -from lightllm.server.httpserver_for_pd_master.pd_selector.pd_node_stats_server import PDNodeStatsState -from lightllm.server.httpserver_for_pd_master.pd_selector.cache_aware import ( - CacheAwareConfig, -) -from lightllm.server.httpserver_for_pd_master.pd_selector.prompt_cache_tree import ( - PromptCacheMatchResult, -) -from lightllm.server.httpserver_for_pd_master.pd_selector.prompt_cache_tree_client import ( - LocalPromptCacheTreeClient, - RemotePromptCacheTreeClient, - create_prompt_cache_tree_client, -) -from lightllm.server.httpserver_for_pd_master.pd_selector.prompt_cache_tree_server import PromptCacheTreeState -from lightllm.utils.net_utils import get_abstract_unix_socket_address - -SAMPLE_STRIDE = CacheAwareConfig().sample_stride - - -def _key(prompt): - return prompt[::SAMPLE_STRIDE] - - -def _remote_tree(ipc_name): - with patch.object( - prompt_cache_tree_client, - "get_pd_master_unix_socket_address", - return_value=get_abstract_unix_socket_address("prompt_cache_tree", ipc_name), - ): - return RemotePromptCacheTreeClient() - - -def test_local_client_owns_an_independent_tree(monkeypatch): - monkeypatch.setattr( - prompt_cache_tree_client, - "get_env_start_args", - lambda: Mock(httpserver_workers=1), - ) - config = CacheAwareConfig() - first = create_prompt_cache_tree_client(config) - second = create_prompt_cache_tree_client(config) - assert isinstance(first, LocalPromptCacheTreeClient) - assert isinstance(second, LocalPromptCacheTreeClient) - prompt = "shared prefix" * 200 - - async def run(): - await first.insert_sampled(_key(prompt), "prefill", SAMPLE_STRIDE) - match = await first.prefix_match_sampled(_key(prompt), len(prompt), SAMPLE_STRIDE) - missing = await second.prefix_match_sampled(_key(prompt), len(prompt), SAMPLE_STRIDE) - assert match.prefill_node == "prefill" - assert missing.prefill_node is None - - asyncio.run(run()) - - -def test_remote_client_uses_symmetric_sampled_operations(monkeypatch): - monkeypatch.setattr(prompt_cache_tree_client, "unix_connect", Mock()) - tree = _remote_tree("test-prompt-cache") - prompt = "0123456789abcdef" * 128 - result = PromptCacheMatchResult(None, 0, len(prompt), None) - tree._prefix_match_sampled = Mock(return_value=result) - tree._insert_sampled = Mock() - - async def run(): - assert await tree.prefix_match_sampled(_key(prompt), len(prompt), SAMPLE_STRIDE) is result - await tree.insert_sampled(_key(prompt), "prefill", SAMPLE_STRIDE) - - asyncio.run(run()) - tree._prefix_match_sampled.assert_called_once_with(_key(prompt), len(prompt), SAMPLE_STRIDE) - tree._insert_sampled.assert_called_once_with(_key(prompt), "prefill", SAMPLE_STRIDE) - - -def test_cancelled_tree_operation_waits_for_the_client_lock_and_still_executes(monkeypatch): - monkeypatch.setattr(prompt_cache_tree_client, "unix_connect", Mock()) - tree = _remote_tree("test-prompt-cache-cancellation") - tree._insert_sampled = Mock() - - async def run(): - await tree.async_lock.acquire() - task = asyncio.create_task(tree.insert_sampled("key", "prefill", SAMPLE_STRIDE)) - await asyncio.sleep(0) - task.cancel() - await asyncio.sleep(0) - assert not task.done() - tree.async_lock.release() - with pytest.raises(asyncio.CancelledError): - await task - - asyncio.run(run()) - tree._insert_sampled.assert_called_once_with("key", "prefill", SAMPLE_STRIDE) - - -def test_tree_rpc_failure_terminates_the_worker(monkeypatch): - monkeypatch.setattr(prompt_cache_tree_client, "unix_connect", Mock()) - terminate = Mock() - monkeypatch.setattr(prompt_cache_tree_client.os, "kill", terminate) - tree = _remote_tree("test-prompt-cache-failure") - tree._insert_sampled = Mock(side_effect=ConnectionError("connection lost")) - - with pytest.raises(ConnectionError, match="connection lost"): - asyncio.run(tree.insert_sampled("key", "prefill", SAMPLE_STRIDE)) - - terminate.assert_called_once_with(os.getpid(), signal.SIGTERM) - - -def test_tree_state_does_not_create_node_statistics(): - tree_state = PromptCacheTreeState() - stats_state = PDNodeStatsState() - prompt = "prefix" * 200 - tree_state.insert_sampled(_key(prompt), "prefill", SAMPLE_STRIDE) - match = tree_state.prefix_match_sampled(_key(prompt), len(prompt), SAMPLE_STRIDE) - assert match[0] == "prefill" - assert not stats_state.stats - - -@pytest.fixture -def rpc_tree_cluster(): - state = PromptCacheTreeState() - ipc_name = uuid.uuid4().hex - server = ThreadedServer( - state, - socket_path=get_abstract_unix_socket_address("prompt_cache_tree", ipc_name), - protocol_config={"allow_pickle": True}, - ) - thread = threading.Thread(target=server.start, daemon=True) - thread.start() - clients = [_remote_tree(ipc_name) for _ in range(4)] - try: - yield state, clients - finally: - for client in clients: - connection = client.connection - if connection is not None: - connection.close() - server.close() - thread.join(timeout=5) - assert not thread.is_alive() - - -def test_real_rpc_clients_share_prompt_cache(rpc_tree_cluster): - _, clients = rpc_tree_cluster - prompt = "shared prefix" * 200 - - async def run(): - await clients[0].insert_sampled(_key(prompt), "prefill", SAMPLE_STRIDE) - match = await clients[1].prefix_match_sampled(_key(prompt), len(prompt), SAMPLE_STRIDE) - assert match.prefill_node == "prefill" - assert match.matched_char_count > len(prompt) // 2 - transport = clients[0].connection._channel.stream.sock - assert transport.family == socket.AF_UNIX - - asyncio.run(run()) - - -def test_concurrent_prompts_remain_separate(rpc_tree_cluster): - _, clients = rpc_tree_cluster - prompts = ["a" * 2048, "b" * 2048] - - async def run(): - await asyncio.gather( - *( - clients[index].insert_sampled(_key(prompt), f"prefill-{index}", SAMPLE_STRIDE) - for index, prompt in enumerate(prompts) - ) - ) - matches = await asyncio.gather( - *( - clients[index + 2].prefix_match_sampled(_key(prompt), len(prompt), SAMPLE_STRIDE) - for index, prompt in enumerate(prompts) - ) - ) - assert [match.prefill_node for match in matches] == [ - "prefill-0", - "prefill-1", - ] - - asyncio.run(run()) - - -def test_tree_and_statistics_servers_use_independent_sockets_and_state(): - ipc_name = uuid.uuid4().hex - tree_address = get_abstract_unix_socket_address("prompt_cache_tree", ipc_name) - stats_address = get_abstract_unix_socket_address("pd_node_stats", ipc_name) - tree_state, stats_state = PromptCacheTreeState(), PDNodeStatsState() - tree_server = ThreadedServer( - tree_state, - socket_path=tree_address, - protocol_config={"allow_pickle": True}, - ) - stats_server = ThreadedServer( - stats_state, - socket_path=stats_address, - protocol_config={"allow_pickle": True}, - ) - tree_thread = threading.Thread(target=tree_server.start, daemon=True) - stats_thread = threading.Thread(target=stats_server.start, daemon=True) - tree_thread.start() - stats_thread.start() - tree = _remote_tree(ipc_name) - try: - prompt = "prefix" * 200 - asyncio.run(tree.insert_sampled(_key(prompt), "prefill", SAMPLE_STRIDE)) - assert not stats_state.stats - tree_server.close() - tree_thread.join(timeout=5) - assert not tree_thread.is_alive() - assert stats_thread.is_alive() - assert stats_state.get_dispatched_info(("prefill",)) == ((0, 0),) - finally: - if tree.connection is not None: - tree.connection.close() - tree_server.close() - stats_server.close() - tree_thread.join(timeout=5) - stats_thread.join(timeout=5) diff --git a/unit_tests/server/test_qwen3_coder_stream_fc.py b/unit_tests/server/test_qwen3_coder_stream_fc.py deleted file mode 100644 index bd8ff6f08c..0000000000 --- a/unit_tests/server/test_qwen3_coder_stream_fc.py +++ /dev/null @@ -1,151 +0,0 @@ -"""Coverage for buffered Qwen3-Coder streaming keepalive deltas.""" - -import json -from types import SimpleNamespace - -from lightllm.server.api_cli import make_argument_parser -from lightllm.server.api_models import Function, Tool -from lightllm.server.function_call_parser import FunctionCallParser, Qwen3CoderDetector - - -def _tool() -> Tool: - return Tool( - type="function", - function=Function( - name="write_file", - description="", - parameters={ - "type": "object", - "properties": {"content": {"type": "string"}}, - }, - ), - ) - - -def test_qwen3_coder_selects_buffered_parser(): - assert FunctionCallParser.ToolCallParserEnum["qwen3_coder"] is Qwen3CoderDetector - assert "qwen3_coder_legacy" not in FunctionCallParser.ToolCallParserEnum - - args = make_argument_parser().parse_args(["--tool_call_parser", "qwen3_coder"]) - assert args.tool_call_parser == "qwen3_coder" - - -def test_qwen3_coder_emits_one_tool_delta_for_every_decode_after_the_name(): - parser = FunctionCallParser([_tool()], "qwen3_coder") - - normal_text, calls = parser.parse_stream_chunk("\n\n\n") - assert normal_text == "" - assert len(calls) == 1 - assert calls[0].tool_index == 0 - assert calls[0].name == "write_file" - assert calls[0].parameters == "" - - streamed_calls = list(calls) - for chunk in ("a", "b", "c" * 100_000): - normal_text, calls = parser.parse_stream_chunk(chunk) - assert normal_text == "" - assert len(calls) == 1 - assert calls[0].tool_index == 0 - assert calls[0].name is None - assert calls[0].parameters == "" - streamed_calls.extend(calls) - - normal_text, calls = parser.parse_stream_chunk("\n\n\n") - assert normal_text == "" - assert len(calls) == 1 - assert calls[0].tool_index == 0 - assert calls[0].name is None - assert json.loads(calls[0].parameters) == {"content": "ab" + "c" * 100_000} - streamed_calls.extend(calls) - - assert [call.name for call in streamed_calls].count("write_file") == 1 - assert json.loads("".join(call.parameters for call in streamed_calls)) == {"content": "ab" + "c" * 100_000} - - -def test_qwen3_coder_accepts_an_undefined_tool_when_name_check_is_disabled(monkeypatch): - import lightllm.server.function_call_parser as parser_module - - monkeypatch.setattr(parser_module, "ENABLE_TOOL_NAME_CHECK", False) - parser = FunctionCallParser([_tool()], "qwen3_coder") - - _, calls = parser.parse_stream_chunk("\n\n") - assert len(calls) == 1 - assert calls[0].name == "unknown" - assert calls[0].parameters == "" - - _, calls = parser.parse_stream_chunk("x\n\n") - assert len(calls) == 1 - assert calls[0].name is None - assert json.loads(calls[0].parameters) == {"value": "x"} - - -def test_qwen3_coder_rejects_an_undefined_tool_when_name_check_is_enabled(monkeypatch): - import lightllm.server.function_call_parser as parser_module - - monkeypatch.setattr(parser_module, "ENABLE_TOOL_NAME_CHECK", True) - parser = FunctionCallParser([_tool()], "qwen3_coder") - - _, calls = parser.parse_stream_chunk("\n\n") - assert calls == [] - - _, calls = parser.parse_stream_chunk("x\n\n") - assert calls == [] - - -def test_api_tool_parser_keeps_its_two_item_return_contract(): - from lightllm.server.api_openai import _process_tools_stream - - request = SimpleNamespace(tools=[_tool()]) - parser_dict = {0: FunctionCallParser(request.tools, "qwen3_coder")} - - result = _process_tools_stream( - 0, - "\n\n\n", - parser_dict, - request, - ) - - assert len(result) == 2 - normal_text, calls = result - assert normal_text == "" - assert len(calls) == 1 - assert calls[0].name == "write_file" - - -def test_empty_arguments_are_preserved_in_the_sse_payload(): - from lightllm.server.api_models import ( - ChatCompletionStreamResponse, - ChatCompletionStreamResponseChoice, - DeltaMessage, - FunctionResponse, - ToolCall, - ) - from lightllm.server.api_openai import _serialize_sse_chunk - - chunk = ChatCompletionStreamResponse( - id="chatcmpl-test", - created=0, - model="test", - choices=[ - ChatCompletionStreamResponseChoice( - index=0, - delta=DeltaMessage( - tool_calls=[ - ToolCall( - index=0, - function=FunctionResponse(arguments=""), - ) - ] - ), - finish_reason=None, - ) - ], - ) - - payload = json.loads(_serialize_sse_chunk(chunk, ("logprobs", "token_ids", "finish_reason"))) - function_delta = payload["choices"][0]["delta"]["tool_calls"][0]["function"] - assert function_delta["arguments"] == "" diff --git a/unit_tests/server/test_reasoning_token_usage.py b/unit_tests/server/test_reasoning_token_usage.py deleted file mode 100644 index 9be5611922..0000000000 --- a/unit_tests/server/test_reasoning_token_usage.py +++ /dev/null @@ -1,222 +0,0 @@ -from unittest.mock import patch - -import pytest - -from lightllm.server.reasoning_parser import ReasoningParser -from lightllm.server.httpserver.stop_seq_match import ReasoningStateMachine - - -def _create_parser(model_type: str, force_reasoning: bool) -> ReasoningParser: - with patch("lightllm.server.reasoning_parser.get_token_id", return_value=99): - return ReasoningParser(model_type, force_reasoning=force_reasoning) - - -def _count_tokens(parser: ReasoningParser, token_ids: list[int]) -> int: - for token_id in token_ids: - parser.update_reasoning_token_count(token_id) - return parser.reasoning_tokens - - -def test_counts_tokens_until_single_token_closing_marker(): - parser = _create_parser("qwen3", force_reasoning=True) - - reasoning_tokens = _count_tokens(parser, [1, 2, 3, 99, 4]) - - assert reasoning_tokens == 3 - - -def test_counts_all_tokens_when_generation_is_truncated_before_closing_marker(): - parser = _create_parser("qwen3", force_reasoning=True) - - reasoning_tokens = _count_tokens(parser, [1, 2]) - - assert reasoning_tokens == 2 - - -def test_does_not_count_when_reasoning_is_disabled(): - parser = _create_parser("qwen3", force_reasoning=False) - - reasoning_tokens = _count_tokens(parser, [1, 2, 99]) - - assert reasoning_tokens == 0 - - -def test_counts_for_always_reasoning_detector(): - parser = _create_parser("deepseek-r1", force_reasoning=False) - - reasoning_tokens = _count_tokens(parser, [1, 2, 99]) - - assert reasoning_tokens == 2 - - -def test_minimax_append_think_output_is_not_counted_as_reasoning(): - parser = _create_parser("minimax-append-think", force_reasoning=True) - - reasoning_tokens = _count_tokens(parser, [1, 2, 99]) - - assert reasoning_tokens == 0 - - -class _StopTokenizer: - def __init__(self, end_ids=(92, 93), start_ids=(90, 91)): - self.start_ids = list(start_ids) - self.end_ids = list(end_ids) - - def encode(self, text, **kwargs): - return {"": self.start_ids, "": self.end_ids}[text] - - -def test_stop_state_tracks_multi_token_delimiters(): - state = ReasoningStateMachine("qwen3", _StopTokenizer(), 0, []) - tokens = [1, 90, 91, 2, 92, 93, 3] - assert [state.handle(token) for token in tokens] == [True, True, False, False, False, False, True] - - -@pytest.mark.parametrize( - "tail_prompt_ids,in_reasoning", - [([90, 91], True), ([1, 90, 91], True), ([90, 91, 1], False), ([90], False), ([], False)], -) -def test_prompt_complete_start_initializes_reasoning(tail_prompt_ids, in_reasoning): - state = ReasoningStateMachine("qwen3", _StopTokenizer(), 0, tail_prompt_ids) - assert state.in_reasoning is in_reasoning - assert state.handle(2) is (not in_reasoning) - - -@pytest.mark.parametrize("model,force", [("qwen3", 1), ("deepseek-r1", -1)]) -def test_prompt_end_marker_preserves_initial_reasoning_hint(model, force): - state = ReasoningStateMachine(model, _StopTokenizer(), force, [90, 91, 1, 92, 93]) - assert state.state is state.State.REASONING - assert state.handle(2) is False - - -@pytest.mark.parametrize( - "force,tail_prompt_ids,in_reasoning", - [(0, [90, 91, 1, 2], False), (1, [92, 93, 1, 2], True), (0, [90, 91, 92, 93, 1, 2], False)], -) -def test_prompt_marker_before_trailing_whitespace_does_not_pretransition(force, tail_prompt_ids, in_reasoning): - state = ReasoningStateMachine("qwen3", _StopTokenizer(), force, tail_prompt_ids) - assert state.in_reasoning is in_reasoning - assert state.handle(3) is (not in_reasoning) - - -@pytest.mark.parametrize("force,marker_ids", [(0, [90, 91]), (1, [92, 93])]) -def test_prompt_phase_hint_is_preserved_when_marker_is_outside_tail(force, marker_ids): - state = ReasoningStateMachine("qwen3", _StopTokenizer(), force, marker_ids + [1] * 20) - assert state.in_reasoning is bool(force) - assert state.handle(2) is (not bool(force)) - - -def test_short_prompt_retains_complete_tail_for_long_delimiters(): - tokenizer = _StopTokenizer(end_ids=[94, 95, 96, 97], start_ids=[90, 91, 92]) - state = ReasoningStateMachine("qwen3", tokenizer, 0, [90, 91]) - assert not state.in_reasoning - assert state.handle(92) is False - assert state.in_reasoning - for token in [94, 95, 96, 97]: - assert state.handle(token) is False - assert not state.in_reasoning - - -@pytest.mark.parametrize("model,force", [("qwen3", 1), ("deepseek-r1", -1)]) -def test_stop_state_starts_in_forced_reasoning(model, force): - state = ReasoningStateMachine(model, _StopTokenizer(), force, []) - assert state.state is state.State.REASONING - assert state.handle(1) is False - assert state.handle(92) is False - assert state.state is state.State.REASONING - assert state.in_reasoning - assert state.handle(93) is False - assert state.handle(2) is True - - -def test_stop_state_carries_a_partial_delimiter_from_pd_prompt(): - state = ReasoningStateMachine("qwen3", _StopTokenizer(), 1, [1, 92]) - assert state.state is state.State.REASONING - assert state.in_reasoning - assert state.handle(93) is False - assert state.in_reasoning is False - assert state.handle(2) is True - - -def test_stop_state_can_reenter_reasoning_after_answer_content(): - state = ReasoningStateMachine("qwen3", _StopTokenizer(), 0, []) - assert state.state is state.State.CONTENT - assert state.handle(90) is True - assert state.state is state.State.CONTENT - assert state.handle(1) is True # 尚未完整命中思维链 start 标记，仍处于正文。 - assert state.state is state.State.CONTENT - for _ in range(2): - assert state.handle(90) is True - assert state.handle(91) is False - assert state.in_reasoning is True - assert state.handle(1) is False - assert state.handle(92) is False - assert state.state is state.State.REASONING - assert state.in_reasoning - assert state.handle(93) is False - assert state.in_reasoning is False - assert state.handle(2) is True - - -@pytest.mark.parametrize( - "token,can_match,next_state", - [(91, False, "REASONING"), (90, True, "CONTENT"), (1, True, "CONTENT")], -) -def test_stop_state_requires_complete_start_from_prompt(token, can_match, next_state): - state = ReasoningStateMachine("qwen3", _StopTokenizer(), 0, [1, 90]) - assert state.state is state.State.CONTENT - assert not state.in_reasoning - assert state.handle(token) is can_match - assert state.state is state.State[next_state] - - -@pytest.mark.parametrize("end_ids,can_match,next_state", [([90, 92], True, "CONTENT"), ([90, 91], False, "REASONING")]) -def test_start_matching_only_recognizes_start_marker(end_ids, can_match, next_state): - state = ReasoningStateMachine("qwen3", _StopTokenizer(end_ids), 0, []) - assert state.handle(90) is True - assert state.state is state.State.CONTENT - assert state.handle(end_ids[-1]) is can_match - assert state.state is state.State[next_state] - - -def test_end_marker_outside_reasoning_is_content(): - state = ReasoningStateMachine("qwen3", _StopTokenizer(), 0, []) - assert state.handle(92) is True - assert state.handle(93) is True - assert state.state is state.State.CONTENT - - -@pytest.mark.parametrize("token,next_state", [(93, "CONTENT"), (92, "REASONING"), (1, "REASONING")]) -def test_stop_state_requires_complete_end_from_prompt(token, next_state): - state = ReasoningStateMachine("qwen3", _StopTokenizer(), 1, [1, 92]) - assert state.state is state.State.REASONING - assert state.handle(token) is False - assert state.state is state.State[next_state] - assert state.in_reasoning is (next_state != "CONTENT") - assert state.handle(2) is (next_state == "CONTENT") - - -def test_start_marker_inside_reasoning_is_reasoning_content(): - state = ReasoningStateMachine("qwen3", _StopTokenizer([90, 92]), 1, []) - assert state.handle(90) is False - assert state.state is state.State.REASONING - assert state.handle(91) is False - assert state.state is state.State.REASONING - assert state.in_reasoning - - -@pytest.mark.parametrize( - "end_ids,tokens,states", - [ - ([90, 91], [90, 91], ["REASONING", "CONTENT"]), - ([91], [90, 91], ["REASONING", "CONTENT"]), - ([91, 92], [90, 91, 92], ["REASONING", "REASONING", "CONTENT"]), - ], -) -def test_reasoning_matches_end_even_when_it_overlaps_start(end_ids, tokens, states): - state = ReasoningStateMachine("qwen3", _StopTokenizer(end_ids), 1, []) - for token, expected_state in zip(tokens, states): - assert state.handle(token) is False - assert state.state is state.State[expected_state] - assert not state.in_reasoning - assert state.handle(1) is True diff --git a/unit_tests/server/test_req_id_generator_performance.py b/unit_tests/server/test_req_id_generator_performance.py deleted file mode 100644 index da0ab4fcaa..0000000000 --- a/unit_tests/server/test_req_id_generator_performance.py +++ /dev/null @@ -1,131 +0,0 @@ -"""Opt-in microbenchmark for request ID generation through RPyC. - -Run with:: - - LIGHTLLM_RUN_REQ_ID_PERF=1 pytest -q -s \ - unit_tests/server/test_req_id_generator_performance.py -""" - -import os -import statistics -import threading -import time -import uuid -from concurrent.futures import ThreadPoolExecutor -from unittest.mock import patch - -import pytest -from rpyc.utils.server import ThreadedServer - -from lightllm.server.httpserver import req_id_generator_client -from lightllm.server.httpserver.req_id_generator_client import ReqIDGenerator -from lightllm.server.httpserver.req_id_generator_server import ( - CLIENT_REQUEST_ID_RANGE_SIZE, - MAX_BEST_OF, - StandaloneReqIDGeneratorState, -) -from lightllm.utils.net_utils import get_abstract_unix_socket_address - -RUN_PERFORMANCE_TEST = os.getenv("LIGHTLLM_RUN_REQ_ID_PERF") == "1" - - -class _LocalIDGenerator: - def __init__(self): - self.current_id = MAX_BEST_OF - self.lock = threading.Lock() - - def generate_id(self): - with self.lock: - request_id = self.current_id - self.current_id += MAX_BEST_OF - return request_id - - -def _client(ipc_name): - socket_address = get_abstract_unix_socket_address("req_id_generator", ipc_name) - with patch.object( - req_id_generator_client, - "get_req_id_generator_socket_address", - return_value=socket_address, - ): - return ReqIDGenerator() - - -def _measure_calls(generate_id, count: int): - latencies_ns = [] - request_ids = [] - start = time.perf_counter_ns() - for _ in range(count): - call_start = time.perf_counter_ns() - request_ids.append(generate_id()) - latencies_ns.append(time.perf_counter_ns() - call_start) - elapsed_ns = time.perf_counter_ns() - start - return request_ids, latencies_ns, elapsed_ns - - -@pytest.mark.skipif( - not RUN_PERFORMANCE_TEST, - reason="set LIGHTLLM_RUN_REQ_ID_PERF=1 to run the request-ID microbenchmark", -) -def test_req_id_generator_rpc_performance(): - iterations = int(os.getenv("LIGHTLLM_REQ_ID_PERF_ITERATIONS", "10000")) - client_count = 4 - assert iterations >= client_count - - ipc_name = uuid.uuid4().hex - state = StandaloneReqIDGeneratorState() - server = ThreadedServer( - state, - socket_path=get_abstract_unix_socket_address("req_id_generator", ipc_name), - protocol_config={"allow_pickle": True}, - ) - server_thread = threading.Thread(target=server.start, daemon=True) - server_thread.start() - clients = [_client(ipc_name) for _ in range(client_count)] - try: - for _ in range(100): - clients[0].generate_id() - - local_generator = _LocalIDGenerator() - _, _, local_elapsed_ns = _measure_calls(local_generator.generate_id, iterations) - request_ids, rpc_latencies_ns, rpc_elapsed_ns = _measure_calls(clients[0].generate_id, iterations) - - calls_per_client = iterations // client_count - concurrent_call_count = calls_per_client * client_count - concurrent_start_ns = time.perf_counter_ns() - with ThreadPoolExecutor(max_workers=client_count) as executor: - futures = [executor.submit(_measure_calls, client.generate_id, calls_per_client) for client in clients] - concurrent_ids = [request_id for future in futures for request_id in future.result()[0]] - concurrent_elapsed_ns = time.perf_counter_ns() - concurrent_start_ns - - assert len(set(request_ids)) == len(request_ids) - assert len(set(concurrent_ids)) == len(concurrent_ids) - assert set(request_ids).isdisjoint(concurrent_ids) - assert all(request_id % MAX_BEST_OF == 0 for request_id in request_ids) - - local_ns_per_call = local_elapsed_ns / iterations - rpc_ns_per_call = rpc_elapsed_ns / iterations - rpc_calls_per_second = iterations * 1e9 / rpc_elapsed_ns - concurrent_calls_per_second = concurrent_call_count * 1e9 / concurrent_elapsed_ns - rpc_latency_us = sorted(rpc_latencies_ns) - rpc_p50_us = statistics.median(rpc_latency_us) / 1e3 - rpc_p95_us = rpc_latency_us[int(len(rpc_latency_us) * 0.95)] / 1e3 - print( - "request ID performance: " - f"local={local_ns_per_call:.0f} ns/call, " - f"ranged_client={rpc_calls_per_second:.0f} calls/s, " - f"client_p50={rpc_p50_us:.1f} us, client_p95={rpc_p95_us:.1f} us, " - f"client_overhead={rpc_ns_per_call / local_ns_per_call:.1f}x, " - f"ranged_{client_count}_clients={concurrent_calls_per_second:.0f} calls/s, " - f"range_size={CLIENT_REQUEST_ID_RANGE_SIZE}" - ) - - minimum_calls_per_second = float(os.getenv("LIGHTLLM_REQ_ID_MIN_CALLS_PER_SECOND", "0")) - if minimum_calls_per_second: - assert rpc_calls_per_second >= minimum_calls_per_second - finally: - for client in clients: - client.connection.close() - server.close() - server_thread.join(timeout=5) - assert not server_thread.is_alive() diff --git a/unit_tests/server/test_ssm_state_mode.py b/unit_tests/server/test_ssm_state_mode.py deleted file mode 100644 index e60e71028d..0000000000 --- a/unit_tests/server/test_ssm_state_mode.py +++ /dev/null @@ -1,165 +0,0 @@ -import json - -import pytest - -from lightllm.server.api_cli import make_argument_parser -from lightllm.server.api_start import _validate_and_normalize_ssm_config -from lightllm.server.core.objs.start_args_type import StartArgs - - -@pytest.mark.parametrize( - "mode,dtype,step,model,capacity,error", - [ - ("native", "bfloat16", 0, "llama", 16, None), - ("native", "bfloat16", 32, "llama", 16, None), - ("replay", "float32", 0, "sensenova_loop", 16, None), - ("replay", "bfloat16", 2, "sensenova_loop", 8, None), - ("replay", "float32", 0, "glm5_next", 8, None), - ("replay", "bfloat16", 2, "glm5_next_text", 8, None), - ("replay", "float32", 0, "qwen3_next", 16, None), - ("replay", "float32", 15, "qwen3_5", 16, None), - ("replay", "bfloat16", 2, "qwen3_5", 16, None), - ("replay", "bfloat16", 2, "qwen3_5", 4, None), - ("replay", "bfloat16", 0, "qwen3_5", 16, None), - ("replay", "float32", 4, "qwen3_5", 4, "capacity"), - ("replay", "float32", 2, "llama", 16, "require a GDN"), - ], -) -def test_ssm_mode_validation(monkeypatch, mode, dtype, step, model, capacity, error): - monkeypatch.setattr("lightllm.utils.config_utils.get_model_type", lambda _: model) - args = make_argument_parser().parse_args( - [ - "--ssm_state_mode", - mode, - "--linear_att_ssm_data_type", - dtype, - "--mtp_step", - str(step), - "--ssm_replay_config", - json.dumps({"cache_len": capacity} if mode == "replay" else {}), - ] - ) - if error: - with pytest.raises(AssertionError, match=error): - _validate_and_normalize_ssm_config(args) - else: - _validate_and_normalize_ssm_config(args) - - -@pytest.mark.parametrize("model", ["qwen3_5", "glm5_next"]) -@pytest.mark.parametrize("mode", ["native", "replay"]) -@pytest.mark.parametrize("step", [15, 16]) -def test_gdn_kda_mtp_width_limit(monkeypatch, model, mode, step): - monkeypatch.setattr("lightllm.utils.config_utils.get_model_type", lambda _: model) - args = StartArgs(ssm_state_mode=mode, mtp_step=step) - if step <= 15: - _validate_and_normalize_ssm_config(args) - else: - with pytest.raises(AssertionError, match="GDN/KDA requires mtp_step <= 15"): - _validate_and_normalize_ssm_config(args) - - -@pytest.mark.parametrize("model", ["qwen3_5", "glm5_next"]) -@pytest.mark.parametrize("compute_mode", ["inline", "precompute"]) -@pytest.mark.parametrize("cache_len", [32, 64]) -def test_replay_large_cache_does_not_bypass_mtp_width_limit(monkeypatch, model, compute_mode, cache_len): - monkeypatch.setattr("lightllm.utils.config_utils.get_model_type", lambda _: model) - args = StartArgs( - ssm_state_mode="replay", - ssm_replay_config={"cache_len": cache_len, "compute_mode": compute_mode}, - linear_att_ssm_data_type="bfloat16", - mtp_step=16, - ) - with pytest.raises(AssertionError, match="GDN/KDA requires mtp_step <= 15"): - _validate_and_normalize_ssm_config(args) - - -def test_ssm_mode_defaults_and_removed_flag(): - parser = make_argument_parser() - assert parser.parse_args([]).ssm_state_mode == StartArgs().ssm_state_mode == "native" - assert parser.parse_args([]).ssm_replay_config == StartArgs().ssm_replay_config == {} - with pytest.raises(SystemExit): - parser.parse_args(["--enable_replayssm"]) - with pytest.raises(SystemExit): - parser.parse_args(["--ssm_state_mode", "auto"]) - - -@pytest.mark.parametrize( - "overrides,error", - [ - ({"ssm_replay_config": {"cache_len": 2}}, "power of two and at least 4"), - ({"ssm_replay_config": {"cache_len": 6}}, "power of two and at least 4"), - ({"mtp_step": -1}, "mtp_step >= 0"), - ({"linear_att_ssm_data_type": "float16"}, "FP32 or BF16"), - ], -) -def test_programmatic_replay_config_rejected_before_model_loading(monkeypatch, overrides, error): - monkeypatch.setattr("lightllm.utils.config_utils.get_model_type", lambda _: "qwen3_5") - args = StartArgs(ssm_state_mode="replay") - for key, value in overrides.items(): - setattr(args, key, value) - with pytest.raises(AssertionError, match=error): - _validate_and_normalize_ssm_config(args) - - -def test_replay_compute_mode_requires_replay(monkeypatch): - monkeypatch.setattr("lightllm.utils.config_utils.get_model_type", lambda _: "qwen3_5") - parser = make_argument_parser() - args = parser.parse_args(["--ssm_state_mode", "replay", "--ssm_replay_config", '{"compute_mode":"precompute"}']) - _validate_and_normalize_ssm_config(args) - - -@pytest.mark.parametrize( - "config,error", - [ - ([], "JSON object"), - (None, "JSON object"), - ({"projection_mode": "inline"}, "Unknown"), - ({"cache_len": True}, "capacity"), - ({"cache_len": 16.0}, "capacity"), - ({"cache_len": "16"}, "capacity"), - ({"cache_len": 128}, "capacity"), - ({"compute_mode": "auto"}, "compute_mode"), - ({"fold_method": "raw"}, "Unknown"), - ], -) -def test_invalid_replay_config(config, error): - args = StartArgs(ssm_state_mode="replay", ssm_replay_config=config) - with pytest.raises(AssertionError, match=error): - _validate_and_normalize_ssm_config(args) - - -@pytest.mark.parametrize("mode", ["native"]) -def test_replay_config_rejects_other_modes(mode): - args = StartArgs(ssm_state_mode=mode, ssm_replay_config={"cache_len": 16}) - with pytest.raises(AssertionError, match="requires ssm_state_mode=replay"): - _validate_and_normalize_ssm_config(args) - - -def test_partial_replay_config_defaults_and_serialization(monkeypatch): - from dataclasses import asdict - - monkeypatch.setattr("lightllm.utils.config_utils.get_model_type", lambda _: "qwen3_5") - args = make_argument_parser().parse_args( - ["--ssm_state_mode", "replay", "--ssm_replay_config", '{"compute_mode":"precompute"}'] - ) - _validate_and_normalize_ssm_config(args) - assert args.ssm_replay_config == {"cache_len": 16, "compute_mode": "precompute"} - original = StartArgs(ssm_state_mode="replay", ssm_replay_config=args.ssm_replay_config) - restored = StartArgs(**json.loads(json.dumps(asdict(original)))) - _validate_and_normalize_ssm_config(restored) - assert restored.ssm_replay_config == args.ssm_replay_config - assert StartArgs().ssm_replay_config == {} - - -@pytest.mark.parametrize( - "option,value", - [ - ("--replayssm_cache_len", "16"), - ("--replayssm_projection_mode", "inline"), - ("--ssm_replay_config", "{invalid}"), - ], -) -def test_removed_flags_and_malformed_json(option, value): - with pytest.raises(SystemExit): - make_argument_parser().parse_args([option, value]) diff --git a/unit_tests/server/test_tokenizer.py b/unit_tests/server/test_tokenizer.py deleted file mode 100644 index aa1c6d6273..0000000000 --- a/unit_tests/server/test_tokenizer.py +++ /dev/null @@ -1,177 +0,0 @@ -"""Tokenizer entry points must not eagerly load unrelated model implementations.""" - -import json -import os -from pathlib import Path -import subprocess -import sys -from types import ModuleType, SimpleNamespace -from unittest.mock import Mock - -import pytest - - -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( - [ - sys.executable, - "-c", - """ -import importlib.abc -import json -from pathlib import Path -import sys -import tempfile -from tokenizers import Tokenizer -from tokenizers.models import WordLevel -from transformers import AutoConfig, PreTrainedTokenizerFast - -model_type = sys.argv[1] -class RejectUnrelatedModels(importlib.abc.MetaPathFinder): - def find_spec(self, fullname, path=None, target=None): - if fullname.startswith("lightllm.models."): - allowed = {"lightllm.models.registry", "lightllm.models.draft_registry", "lightllm.models.builtin"} - if model_type == "deepseek_v32": - allowed.update({ - "lightllm.models.deepseek3_2", - "lightllm.models.deepseek3_2.tokenizer", - "lightllm.models.deepseek3_2.encoding_dsv32", - }) - if model_type == "deepseek_v4": - allowed.update({"lightllm.models.deepseek_v4", "lightllm.models.deepseek_v4.tokenizer"}) - if fullname not in allowed: - raise AssertionError("Unexpected model import: " + fullname) -sys.meta_path.insert(0, RejectUnrelatedModels()) -from lightllm.server import tokenizer as module -assert not any(name.startswith("lightllm.models.") for name in sys.modules) - -with tempfile.TemporaryDirectory() as directory: - fast = PreTrainedTokenizerFast(tokenizer_object=Tokenizer(WordLevel({"hello": 0, "[UNK]": 1}, unk_token="[UNK]"))) - fast.save_pretrained(directory) - Path(directory, "config.json").write_text(json.dumps({ - "model_type": model_type, "architectures": ["LlamaForCausalLM"], - })) - 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 - thinking_prompt = tokenizer.apply_chat_template( - messages=[{"role": "user", "content": "hello"}], thinking=True - ) - assert "" in thinking_prompt - tool_prompt = tokenizer.apply_chat_template( - messages=[{"role": "user", "content": "hello"}], - tools=[{"name": "echo", "description": "Echo text", "parameters": {"type": "object"}}], - ) - assert "" in tool_prompt and "echo" in tool_prompt - assert "lightllm.models.deepseek3_2.model" not in sys.modules -""", - model_type, - ], - cwd=ROOT, - env={**os.environ, "CUDA_VISIBLE_DEVICES": "", "HF_HUB_OFFLINE": "1"}, - capture_output=True, - text=True, - timeout=60, - ) - assert result.returncode == 0, result.stdout + result.stderr - - -@pytest.mark.parametrize( - "config,module_name,class_name", - [ - ( - {"model_type": "llava", "architectures": ["TarsierForConditionalGeneration"]}, - "tarsier2.model", - "Tarsier2Tokenizer", - ), - ({"model_type": "llava"}, "llava.model", "LlavaTokenizer"), - ({"model_type": "internlmxcomposer2"}, "llava.model", "LlavaTokenizer"), - ({"model_type": "qwen", "visual": {}}, "qwen_vl.model", "QWenVLTokenizer"), - ({"model_type": "qwen2_vl", "vision_config": {}}, "qwen2_vl.model", "QWen2VLTokenizer"), - ({"model_type": "qwen2_5_vl", "vision_config": {}}, "qwen2_vl.model", "QWen2VLTokenizer"), - ({"model_type": "qwen3_vl", "vision_config": {}}, "qwen3_vl.model", "QWen3VLTokenizer"), - ({"model_type": "qwen3_vl_moe", "vision_config": {}}, "qwen3_vl.model", "QWen3VLTokenizer"), - ({"model_type": "qwen3_5", "vision_config": {}}, "qwen3_5.model", "QWen3_5Tokenizer"), - ({"model_type": "sensenova_loop", "vision_config": {}}, "qwen3_5.model", "QWen3_5Tokenizer"), - ({"model_type": "qwen3_5_moe", "vision_config": {}}, "qwen3_5.model", "QWen3_5Tokenizer"), - ({"thinker_config": {}}, "qwen3_omni_moe_thinker.model", "QWen3OmniTokenizer"), - ({"model_type": "internvl_chat"}, "internvl.model", "InternvlTokenizer"), - ({"model_type": "gemma3"}, "gemma3.model", "Gemma3Tokenizer"), - ({"model_type": "gemma4"}, "gemma4.tokenizer", "Gemma4Tokenizer"), - ({"model_type": "gemma4", "vision_config": {}}, "gemma4.tokenizer", "Gemma4Tokenizer"), - ({"model_type": "neo_chat"}, "neo_chat_moe.model", "NeoChatTokenizer"), - ], -) -def test_selected_multimodal_tokenizer_is_loaded(monkeypatch, config, module_name, class_name): - import transformers - from lightllm.server import tokenizer as module - - config = {"architectures": ["Example"], **config} - base_tokenizer = object() - monkeypatch.setattr(module.PretrainedConfig, "get_config_dict", lambda *a, **kw: (config, {})) - monkeypatch.setattr(module.AutoTokenizer, "from_pretrained", lambda *a, **kw: base_tokenizer) - processor = SimpleNamespace(image_processor=object()) - monkeypatch.setattr(transformers.AutoProcessor, "from_pretrained", lambda *a, **kw: processor) - monkeypatch.setattr(transformers.AutoImageProcessor, "from_pretrained", lambda *a, **kw: processor.image_processor) - vision = ModuleType("lightllm.models.qwen2_vl.vision_process") - vision.Qwen2VLImageProcessor = SimpleNamespace(from_pretrained=lambda *a, **kw: processor.image_processor) - monkeypatch.setitem(sys.modules, vision.__name__, vision) - implementation = ModuleType("lightllm.models." + module_name) - constructor = Mock() - setattr(implementation, class_name, constructor) - monkeypatch.setitem(sys.modules, implementation.__name__, implementation) - - assert module.get_tokenizer("local-checkpoint") is constructor.return_value - constructor.assert_called_once() - args, kwargs = constructor.call_args - assert (args[0] if args else kwargs["tokenizer"]) is base_tokenizer - assert (args[1] if len(args) > 1 else kwargs["model_cfg"]) is config.get("thinker_config", config) diff --git a/unit_tests/server/test_vocab_sampling_start_args.py b/unit_tests/server/test_vocab_sampling_start_args.py deleted file mode 100644 index f76dd0f7ec..0000000000 --- a/unit_tests/server/test_vocab_sampling_start_args.py +++ /dev/null @@ -1,57 +0,0 @@ -import pytest - -from lightllm.server.api_start import _launch_subprocesses, pd_master_start -from lightllm.server.core.objs.start_args_type import StartArgs - - -@pytest.mark.parametrize("run_mode", ["normal", "prefill", "decode"]) -@pytest.mark.parametrize( - "constraint_args", - [ - {}, - {"output_constraint_mode": "xgrammar"}, - ], -) -def test_target_candidates_reject_constraints_before_model_loading(monkeypatch, run_mode, constraint_args): - monkeypatch.setattr("lightllm.server.api_start._set_envs_and_config", lambda args: None) - - def unexpected_model_loading(args): - pytest.fail("Incompatible candidate sampling must be rejected before loading model configuration") - - monkeypatch.setattr("lightllm.server.api_start.auto_set_max_req_total_len", unexpected_model_loading) - args = StartArgs(run_mode=run_mode, target_vocab_topk_sampling=2, **constraint_args) - - with pytest.raises(AssertionError, match="Disable --target_vocab_topk_sampling when using output constraints"): - _launch_subprocesses(args) - - -@pytest.mark.parametrize("run_mode", ["normal", "prefill", "decode", "pd_master"]) -@pytest.mark.parametrize( - "sampling_args", - [ - {}, - {"target_vocab_topk_sampling": 2, "output_constraint_mode": "none"}, - {"draft_vocab_topk_sampling": 2}, - {"diverse_mode": True, "output_constraint_mode": "xgrammar"}, - {"use_reward_model": True, "output_constraint_mode": "xgrammar"}, - { - "dp": 2, - "mtp_mode": "vanilla_no_att", - "enable_prefill_microbatch_overlap": True, - "enable_decode_microbatch_overlap": True, - "output_constraint_mode": "xgrammar", - }, - ], -) -def test_compatible_sampling_options_reach_model_configuration(monkeypatch, run_mode, sampling_args): - monkeypatch.setattr("lightllm.server.api_start._set_envs_and_config", lambda args: None) - monkeypatch.setattr("lightllm.server.api_start.set_unique_server_name", lambda args: None) - - def stop_at_model_configuration(args): - raise RuntimeError("model configuration reached") - - monkeypatch.setattr("lightllm.server.api_start.auto_set_max_req_total_len", stop_at_model_configuration) - start = pd_master_start if run_mode == "pd_master" else _launch_subprocesses - - with pytest.raises(RuntimeError, match="model configuration reached"): - start(StartArgs(run_mode=run_mode, **sampling_args)) diff --git a/unit_tests/test_format_out.py b/unit_tests/test_format_out.py deleted file mode 100644 index ec42be9ef7..0000000000 --- a/unit_tests/test_format_out.py +++ /dev/null @@ -1,47 +0,0 @@ -import json -from unittest.mock import Mock - -import pytest -import xgrammar as xgr -from pydantic import BaseModel - -from format_out.impl import ChatSession, SamplingParams -from lightllm.server.core.objs.sampling_params import SamplingParams as ServerSamplingParams - - -class HelperResult(BaseModel): - answer: str - - -@pytest.mark.parametrize("mode", ["plain", "regex", "grammar", "json", "inherited_json"]) -def test_helper_payload_initializes_server_constraints(monkeypatch, mode): - payloads = [] - - def post(url, headers, data): - payload = json.loads(data) - server_params = ServerSamplingParams() - server_params.init(None, **payload["parameters"]) - assert server_params.constraint_out_params.has_constraint() == (mode != "plain") - assert all(value is not None for value in payload["parameters"].values()) - payloads.append(payload) - return Mock(status_code=200, json=lambda: {"generated_text": ['{"answer":"OK"}']}) - - monkeypatch.setattr("format_out.impl.requests.post", post) - params = SamplingParams(do_sample=False, guided_json={} if mode == "inherited_json" else None) - session = ChatSession("Answer briefly.", params, disable_log=True) - if mode == "regex": - session.generate(regex="OK") - assert payloads[0]["parameters"]["regular_constraint"] == "OK" - elif mode == "grammar": - session.generate(grammar=xgr.Grammar.from_ebnf('root ::= "OK"')) - assert payloads[0]["parameters"]["guided_grammar"] - elif mode == "json": - output = session.gen_json_object(HelperResult, prefix_regex=r"\s{0,20}") - assert HelperResult.model_validate_json(output).answer == "OK" - else: - session.generate() - if mode == "inherited_json": - assert payloads[0]["parameters"]["guided_json"] == {} - assert params.regular_constraint is None - assert params.guided_grammar is None - assert params.guided_json == ({} if mode == "inherited_json" else None) diff --git a/unit_tests/test_static_benchmark_paged_kv.py b/unit_tests/test_static_benchmark_paged_kv.py deleted file mode 100644 index 1f4efc4baa..0000000000 --- a/unit_tests/test_static_benchmark_paged_kv.py +++ /dev/null @@ -1,155 +0,0 @@ -import sys -from pathlib import Path -from types import SimpleNamespace - -import numpy as np -import torch - - -REPO_ROOT = Path(__file__).resolve().parents[1] -STATIC_BENCHMARK_DIR = REPO_ROOT / "test" / "benchmark" / "static_inference" -for path in (REPO_ROOT, STATIC_BENCHMARK_DIR): - if str(path) not in sys.path: - sys.path.append(str(path)) - -import static_benchmark -from static_benchmark import ( - BenchmarkCase, - StaticBenchmarkExecutor, - decode_profile_batch_divisor, - resolve_batch_max_prefill_cases, -) - - -class FakeKVAllocator: - def __init__(self, page_size: int): - self.page_size = page_size - self.next_index = 0 - self.alloc_sizes = [] - - def alloc(self, size: int) -> torch.Tensor: - assert size % self.page_size == 0 - self.alloc_sizes.append(size) - indexes = torch.arange(self.next_index, self.next_index + size, dtype=torch.int32) - self.next_index += size - return indexes - - -def make_executor(page_size: int = 4): - allocator = FakeKVAllocator(page_size=page_size) - req_manager = SimpleNamespace( - req_to_token_indexs=torch.full((4, 32), -1, dtype=torch.int32), - ) - req_manager.page_size = page_size - - 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 - - -def assert_complete_physical_pages(indexes: torch.Tensor, page_size: int): - for start in range(0, indexes.numel(), page_size): - page = indexes[start : start + page_size] - assert torch.all(page // page_size == page[0] // page_size) - - -def test_prefill_allocates_and_reuses_complete_pages_per_request(): - executor, allocator, req_manager = make_executor(page_size=4) - req_idx = torch.tensor([0, 1], dtype=torch.int32) - - model_input = executor._make_prefill_input( - token_chunk=np.arange(6, dtype=np.int64).reshape(2, 3), - req_idx=req_idx, - ready_cache_len=0, - ) - - 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) - assert_complete_physical_pages(req_manager.req_to_token_indexs[1, :4], page_size=4) - - # 两个请求都还在已有页内继续 decode，不应产生新的分配。 - executor._make_decode_input( - batch_size=2, - req_idx=req_idx, - mtp_index=torch.zeros(2, dtype=torch.int32), - seq_len=torch.tensor([4, 4], dtype=torch.int32), - input_ids=torch.tensor([1, 2], dtype=torch.int64), - max_kv_seq_len=4, - ) - assert allocator.alloc_sizes == [4, 4] - - # 跨过页边界后，每个请求分别追加一整页。 - executor._make_decode_input( - batch_size=2, - req_idx=req_idx, - mtp_index=torch.zeros(2, dtype=torch.int32), - seq_len=torch.tensor([5, 5], dtype=torch.int32), - input_ids=torch.tensor([3, 4], dtype=torch.int64), - max_kv_seq_len=5, - ) - 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, allocator, req_manager = make_executor(page_size=4) - executor.args.mtp_mode = "eagle_with_att" - executor.args.mtp_step = 3 - - model_input = executor._make_decode_input( - batch_size=4, - req_idx=torch.tensor([0, 0, 1, 1], dtype=torch.int32), - mtp_index=torch.tensor([0, 1, 0, 1], dtype=torch.int32), - seq_len=torch.tensor([3, 4, 3, 4], dtype=torch.int32), - input_ids=torch.tensor([1, 2, 3, 4], dtype=torch.int64), - max_kv_seq_len=4, - extra_kv_len=2, - ) - - assert model_input.batch_size == 4 - 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) - - -def test_profile_capacity_uses_page_aligned_per_request_cost(): - args = SimpleNamespace( - mtp_mode=None, - mtp_step=0, - page_size=4, - max_batch_size=0, - batch_max_tokens=64, - ) - prefill_case = BenchmarkCase( - name="prefill", - stage="prefill", - batch_size=8, - context_len=5, - output_len=0, - prefill_uncached_len=5, - prefill_batch_size_by_batch_max_tokens=8, - ) - - resolved = resolve_batch_max_prefill_cases(args, [prefill_case], profiled_max_total_token_num=24) - - assert resolved[0].batch_size == 3 - assert resolved[0].profiled_batch_divisor == 8 - - decode_case = BenchmarkCase( - name="decode", - stage="decode", - batch_size=1, - context_len=6, - output_len=2, - ) - assert decode_profile_batch_divisor(args, decode_case) == 20 - - args.mtp_mode = "eagle_with_att" - args.mtp_step = 3 - assert decode_profile_batch_divisor(args, decode_case) == 24 diff --git a/unit_tests/utils/test_atexit_exit_modes.py b/unit_tests/utils/test_atexit_exit_modes.py deleted file mode 100644 index 0883f0efb7..0000000000 --- a/unit_tests/utils/test_atexit_exit_modes.py +++ /dev/null @@ -1,188 +0,0 @@ -"""用独立进程验证 atexit；所有信号只发送给本测试创建的进程。""" - -import selectors -import signal -import subprocess -import sys -import textwrap - -import pytest - - -pytestmark = pytest.mark.skipif(sys.platform != "linux", reason="Exit codes and signals below target Linux") - -PROBE_SCRIPT = textwrap.dedent( - """ - import atexit - import multiprocessing as mp - import os - import resource - import signal - import sys - - - def record(marker_path, event): - with open(marker_path, "a", encoding="utf-8") as output: - output.write(event + "\\n") - - - def cleanup(marker_path, interrupt=False): - record(marker_path, "started") - if interrupt: - print("READY", flush=True) - while True: - signal.pause() - record(marker_path, "finished") - - - def worker(marker_path): - atexit.register(cleanup, marker_path) - record(marker_path, "worker_registered") - - - def main(): - # SIGABRT/SIGSEGV 实验不生成 core dump。 - resource.setrlimit(resource.RLIMIT_CORE, (0, 0)) - mode, marker_path, signal_number = sys.argv[1:] - signal_number = int(signal_number) - signal.signal(signal.SIGINT, signal.default_int_handler) - - if mode.startswith("multiprocessing_"): - context = mp.get_context(mode.removeprefix("multiprocessing_")) - process = context.Process(target=worker, args=(marker_path,)) - process.start() - process.join(timeout=10) - if process.is_alive(): - process.kill() - process.join(timeout=10) - raise RuntimeError("multiprocessing worker failed to exit") - assert process.exitcode == 0, process.exitcode - return - - atexit.register(cleanup, marker_path, mode == "interrupt_cleanup") - - if mode in ("normal_return", "interrupt_cleanup"): - return - if mode == "sys_exit_0": - sys.exit(0) - if mode == "sys_exit_7": - sys.exit(7) - if mode == "unhandled_exception": - raise RuntimeError("intentional test exception") - if mode == "keyboard_interrupt": - raise KeyboardInterrupt - if mode == "os_exit_0": - os._exit(0) - if mode == "os_exit_7": - os._exit(7) - if mode == "abort": - os.abort() - if mode == "exec_replace": - os.execv(sys.executable, [sys.executable, "-I", "-c", "pass"]) - - if mode == "signal_sys_exit": - signal.signal(signal_number, lambda signum, frame: sys.exit(0)) - elif mode == "signal_os_exit": - signal.signal(signal_number, lambda signum, frame: os._exit(7)) - elif mode == "signal_os_default" or (mode == "signal_default" and signal_number != signal.SIGINT): - if signal_number != signal.SIGKILL: - signal.signal(signal_number, signal.SIG_DFL) - elif mode != "signal_default": - raise ValueError(mode) - - # 父进程收到 READY 后才发信号，确保回调和信号处理方式已注册。 - print("READY", flush=True) - while True: - signal.pause() - - - if __name__ == "__main__": - main() - """ -) - - -@pytest.fixture -def probe_script(tmp_path): - script_path = tmp_path / "atexit_probe.py" - script_path.write_text(PROBE_SCRIPT, encoding="utf-8") - return script_path - - -def run_probe(probe_script, tmp_path, mode, signum=0): - marker_path = tmp_path / "callback_events.txt" - with subprocess.Popen( - [sys.executable, "-I", str(probe_script), mode, str(marker_path), str(int(signum))], - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - text=True, - start_new_session=True, - ) as process: - try: - if signum: - with selectors.DefaultSelector() as selector: - selector.register(process.stdout, selectors.EVENT_READ) - assert selector.select(timeout=10), "Probe did not become ready" - assert process.stdout.readline().strip() == "READY", "Probe exited before registering its callback" - process.send_signal(signum) - stdout, stderr = process.communicate(timeout=30) - finally: - if process.poll() is None: - process.kill() - process.communicate(timeout=10) - - events = marker_path.read_text(encoding="utf-8").splitlines() if marker_path.exists() else [] - print(f"{mode} signal={int(signum)} returncode={process.returncode} events={events}") - return process.returncode, events, stdout + stderr - - -@pytest.mark.parametrize( - "mode,signum,expected_returncode,callback_completed", - [ - pytest.param("normal_return", 0, 0, True, id="normal-return"), - pytest.param("sys_exit_0", 0, 0, True, id="sys-exit-0"), - pytest.param("sys_exit_7", 0, 7, True, id="sys-exit-nonzero"), - pytest.param("unhandled_exception", 0, 1, True, id="unhandled-exception"), - pytest.param("keyboard_interrupt", 0, -signal.SIGINT, True, id="unhandled-keyboard-interrupt"), - pytest.param("signal_default", signal.SIGINT, -signal.SIGINT, True, id="sigint-python-default"), - pytest.param("signal_os_default", signal.SIGINT, -signal.SIGINT, False, id="sigint-os-default"), - pytest.param("signal_default", signal.SIGTERM, -signal.SIGTERM, False, id="sigterm-default"), - pytest.param("signal_default", signal.SIGHUP, -signal.SIGHUP, False, id="sighup-default"), - pytest.param("signal_default", signal.SIGQUIT, -signal.SIGQUIT, False, id="sigquit-default"), - pytest.param("signal_default", signal.SIGSEGV, -signal.SIGSEGV, False, id="sigsegv-default"), - pytest.param("signal_default", signal.SIGKILL, -signal.SIGKILL, False, id="sigkill"), - pytest.param("signal_sys_exit", signal.SIGINT, 0, True, id="sigint-handler-sys-exit"), - pytest.param("signal_sys_exit", signal.SIGTERM, 0, True, id="sigterm-handler-sys-exit"), - pytest.param("signal_sys_exit", signal.SIGHUP, 0, True, id="sighup-handler-sys-exit"), - pytest.param("signal_os_exit", signal.SIGTERM, 7, False, id="sigterm-handler-os-exit"), - pytest.param("os_exit_0", 0, 0, False, id="os-exit-0"), - pytest.param("os_exit_7", 0, 7, False, id="os-exit-nonzero"), - pytest.param("abort", 0, -signal.SIGABRT, False, id="abort"), - pytest.param("exec_replace", 0, 0, False, id="exec-replaces-interpreter"), - ], -) -def test_atexit_on_process_exit(probe_script, tmp_path, mode, signum, expected_returncode, callback_completed): - returncode, events, output = run_probe(probe_script, tmp_path, mode, signum) - - assert returncode == expected_returncode, output - assert events == (["started", "finished"] if callback_completed else []), output - - -@pytest.mark.skipif( - sys.version_info[:2] != (3, 10), reason="This experiment records Python 3.10 multiprocessing behavior" -) -@pytest.mark.parametrize("start_method,callback_completed", [("spawn", True), ("fork", False)]) -def test_atexit_in_multiprocessing_worker(probe_script, tmp_path, start_method, callback_completed): - returncode, events, output = run_probe(probe_script, tmp_path, f"multiprocessing_{start_method}") - - assert returncode == 0, output - # 本机 Python 3.10：spawn_main() 使用 sys.exit()，fork 的 _launch() 使用 os._exit()。 - expected_events = ["worker_registered"] + (["started", "finished"] if callback_completed else []) - assert events == expected_events, output - - -@pytest.mark.parametrize("signum", [signal.SIGINT, signal.SIGKILL]) -def test_signal_can_interrupt_atexit_callback(probe_script, tmp_path, signum): - _returncode, events, output = run_probe(probe_script, tmp_path, "interrupt_cleanup", signum) - - assert events == ["started"], output diff --git a/unit_tests/utils/test_config_utils.py b/unit_tests/utils/test_config_utils.py deleted file mode 100644 index ae3e093ca0..0000000000 --- a/unit_tests/utils/test_config_utils.py +++ /dev/null @@ -1,53 +0,0 @@ -import pytest - -from lightllm.utils import config_utils - - -@pytest.fixture -def fake_config(monkeypatch): - def install(cfg_json): - monkeypatch.setattr(config_utils, "get_config_json", lambda model_path: cfg_json) - return "dummy/model/dir" - - return install - - -def test_get_keyvalue_reads_top_level_key(fake_config): - model_path = fake_config({"model_type": "llama", "hidden_size": 4096}) - - assert config_utils._get_config_llm_keyvalue(model_path, ["hidden_size"]) == 4096 - - -def test_get_keyvalue_does_not_crash_when_thinker_config_lacks_text_config(fake_config): - cfg = {"model_type": "qwen3_omni_moe", "hidden_size": 4096, "thinker_config": {"model_type": "text"}} - model_path = fake_config(cfg) - - assert config_utils._get_config_llm_keyvalue(model_path, ["hidden_size"]) == 4096 - - -def test_get_keyvalue_keeps_top_level_value_when_thinker_text_config_lacks_key(fake_config): - cfg = { - "model_type": "qwen3_omni_moe", - "hidden_size": 4096, - "thinker_config": {"text_config": {"model_type": "text"}}, - } - model_path = fake_config(cfg) - - assert config_utils._get_config_llm_keyvalue(model_path, ["hidden_size"]) == 4096 - - -def test_get_keyvalue_prefers_thinker_text_config_value_when_present(fake_config): - cfg = { - "model_type": "qwen3_omni_moe", - "hidden_size": 4096, - "thinker_config": {"text_config": {"hidden_size": 2048}}, - } - model_path = fake_config(cfg) - - assert config_utils._get_config_llm_keyvalue(model_path, ["hidden_size"]) == 2048 - - -def test_get_keyvalue_returns_none_when_key_missing_everywhere(fake_config): - model_path = fake_config({"model_type": "llama", "hidden_size": 4096}) - - assert config_utils._get_config_llm_keyvalue(model_path, ["head_dim"]) is None diff --git a/unit_tests/utils/test_custom_kernel_utils.py b/unit_tests/utils/test_custom_kernel_utils.py deleted file mode 100644 index 4f9d63388d..0000000000 --- a/unit_tests/utils/test_custom_kernel_utils.py +++ /dev/null @@ -1,35 +0,0 @@ -import torch -import time -import pytest -from lightllm.utils.custom_kernel_utis import pad2dim_tensor_to_new_batch, torch_cat_3 - - -def test_torch_cat(): - a = torch.tensor([[[1, 2], [3, 4]]], device="cuda") - b = torch.tensor([[[5, 6], [7, 8]]], device="cuda") - c = torch_cat_3([a, b], dim=0) - torch.equal(torch.cat((a, b), dim=0), c) - - d = torch_cat_3([a, b], dim=1) - torch.equal(torch.cat((a, b), dim=1), d) - - e = torch_cat_3([a, b], dim=-1) - torch.equal(torch.cat((a, b), dim=-1), e) - - empty = torch.empty((0, 2), device="cuda") - torch_cat_3([a, empty, b], dim=0) - return - - -def test_pad2dim_tensor_to_new_batch(): - input_tensor = torch.tensor([[1.0, 2.0]]) - padded_tensor = pad2dim_tensor_to_new_batch(input=input_tensor, new_batch_size=3) - assert torch.equal(padded_tensor, torch.tensor([[1.0, 2.0], [1.0, 2.0], [1.0, 2.0]])) - - empty_input = torch.empty((0, 2)) - padded_empty_input = pad2dim_tensor_to_new_batch(input=empty_input, new_batch_size=2) - assert torch.equal(padded_empty_input, torch.zeros((2, 2))) - - -if __name__ == "__main__": - pytest.main() diff --git a/unit_tests/utils/test_envs_utils.py b/unit_tests/utils/test_envs_utils.py deleted file mode 100644 index 75f48cb174..0000000000 --- a/unit_tests/utils/test_envs_utils.py +++ /dev/null @@ -1,179 +0,0 @@ -from types import SimpleNamespace - -from lightllm.utils import envs_utils -from lightllm.utils.envs_utils import ( - get_deepep_num_max_dispatch_tokens_per_rank, - get_deepep_num_max_dispatch_tokens_per_rank_decode, - get_deepep_num_max_dispatch_tokens_per_rank_prefill, - is_stop_matching_in_reasoning_enabled, - get_pd_cache_high_priority_max_age_seconds, - get_pd_cache_high_priority_min_prompt_tokens, - get_pd_node_busy_retry_timeout_seconds, - get_pd_node_continuation_resource_wait_timeout_seconds, - get_pd_node_resource_wait_timeout_seconds, -) - - -def test_stop_in_reasoning_defaults_to_false(monkeypatch): - monkeypatch.delenv("LIGHTLLM_STOP_IN_REASONING", raising=False) - is_stop_matching_in_reasoning_enabled.cache_clear() - assert is_stop_matching_in_reasoning_enabled() is False - is_stop_matching_in_reasoning_enabled.cache_clear() - - -def test_stop_in_reasoning_reads_environment(monkeypatch): - for value, expected in [("1", True), ("true", True), ("ON", True), ("0", False), ("false", False)]: - monkeypatch.setenv("LIGHTLLM_STOP_IN_REASONING", value) - is_stop_matching_in_reasoning_enabled.cache_clear() - assert is_stop_matching_in_reasoning_enabled() is expected - is_stop_matching_in_reasoning_enabled.cache_clear() - - -def test_pd_cache_high_priority_max_age_defaults_to_180_seconds(monkeypatch): - monkeypatch.delenv("LIGHTLLM_PD_CACHE_HIGH_PRIORITY_MAX_AGE_SECONDS", raising=False) - get_pd_cache_high_priority_max_age_seconds.cache_clear() - - assert get_pd_cache_high_priority_max_age_seconds() == 180 - - get_pd_cache_high_priority_max_age_seconds.cache_clear() - - -def test_pd_cache_high_priority_max_age_reads_environment_variable(monkeypatch): - monkeypatch.setenv("LIGHTLLM_PD_CACHE_HIGH_PRIORITY_MAX_AGE_SECONDS", "30") - get_pd_cache_high_priority_max_age_seconds.cache_clear() - - assert get_pd_cache_high_priority_max_age_seconds() == 30 - - get_pd_cache_high_priority_max_age_seconds.cache_clear() - - -def test_pd_cache_high_priority_min_prompt_tokens_defaults_to_2048(monkeypatch): - monkeypatch.delenv("LIGHTLLM_PD_CACHE_HIGH_PRIORITY_MIN_PROMPT_TOKENS", raising=False) - get_pd_cache_high_priority_min_prompt_tokens.cache_clear() - - assert get_pd_cache_high_priority_min_prompt_tokens() == 2048 - - get_pd_cache_high_priority_min_prompt_tokens.cache_clear() - - -def test_pd_cache_high_priority_min_prompt_tokens_reads_environment_variable(monkeypatch): - monkeypatch.setenv("LIGHTLLM_PD_CACHE_HIGH_PRIORITY_MIN_PROMPT_TOKENS", "2048") - get_pd_cache_high_priority_min_prompt_tokens.cache_clear() - - assert get_pd_cache_high_priority_min_prompt_tokens() == 2048 - - get_pd_cache_high_priority_min_prompt_tokens.cache_clear() - - -def test_pd_node_resource_wait_timeout_defaults_to_20_seconds(monkeypatch): - monkeypatch.delenv("LIGHTLLM_PD_NODE_RESOURCE_WAIT_TIMEOUT_SECONDS", raising=False) - get_pd_node_resource_wait_timeout_seconds.cache_clear() - - assert get_pd_node_resource_wait_timeout_seconds() == 20 - - get_pd_node_resource_wait_timeout_seconds.cache_clear() - - -def test_pd_node_resource_wait_timeout_reads_environment_variable(monkeypatch): - monkeypatch.setenv("LIGHTLLM_PD_NODE_RESOURCE_WAIT_TIMEOUT_SECONDS", "30") - get_pd_node_resource_wait_timeout_seconds.cache_clear() - - assert get_pd_node_resource_wait_timeout_seconds() == 30 - - get_pd_node_resource_wait_timeout_seconds.cache_clear() - - -def test_pd_node_continuation_resource_wait_timeout_defaults_to_60_seconds(monkeypatch): - monkeypatch.delenv("LIGHTLLM_PD_NODE_CONTINUATION_RESOURCE_WAIT_TIMEOUT_SECONDS", raising=False) - get_pd_node_continuation_resource_wait_timeout_seconds.cache_clear() - - assert get_pd_node_continuation_resource_wait_timeout_seconds() == 60 - - get_pd_node_continuation_resource_wait_timeout_seconds.cache_clear() - - -def test_pd_node_continuation_resource_wait_timeout_reads_environment_variable(monkeypatch): - monkeypatch.setenv("LIGHTLLM_PD_NODE_CONTINUATION_RESOURCE_WAIT_TIMEOUT_SECONDS", "90") - get_pd_node_continuation_resource_wait_timeout_seconds.cache_clear() - - assert get_pd_node_continuation_resource_wait_timeout_seconds() == 90 - - get_pd_node_continuation_resource_wait_timeout_seconds.cache_clear() - - -def test_pd_node_busy_retry_timeout_defaults_to_zero(monkeypatch): - monkeypatch.delenv("LIGHTLLM_PD_NODE_BUSY_RETRY_TIMEOUT_SECONDS", raising=False) - get_pd_node_busy_retry_timeout_seconds.cache_clear() - - assert get_pd_node_busy_retry_timeout_seconds() == 0 - - get_pd_node_busy_retry_timeout_seconds.cache_clear() - - -def test_pd_node_busy_retry_timeout_reads_environment_variable(monkeypatch): - monkeypatch.setenv("LIGHTLLM_PD_NODE_BUSY_RETRY_TIMEOUT_SECONDS", "60") - get_pd_node_busy_retry_timeout_seconds.cache_clear() - - assert get_pd_node_busy_retry_timeout_seconds() == 60 - - get_pd_node_busy_retry_timeout_seconds.cache_clear() - - -def test_deepep_decode_buffer_capacity_covers_mtp_rows_and_safety_margin(monkeypatch): - monkeypatch.delenv("NUM_MAX_DISPATCH_TOKENS_PER_RANK_DECODE", raising=False) - monkeypatch.setattr( - envs_utils, - "get_env_start_args", - lambda: SimpleNamespace(running_max_req_size=256, mtp_step=2), - ) - get_deepep_num_max_dispatch_tokens_per_rank_decode.cache_clear() - - assert get_deepep_num_max_dispatch_tokens_per_rank_decode() == 784 - - get_deepep_num_max_dispatch_tokens_per_rank_decode.cache_clear() - - -def test_deepep_prefill_buffer_capacity_adds_safety_margin(monkeypatch): - monkeypatch.delenv("NUM_MAX_DISPATCH_TOKENS_PER_RANK_PREFILL", raising=False) - monkeypatch.setattr(envs_utils, "get_env_start_args", lambda: SimpleNamespace(batch_max_tokens=6000)) - get_deepep_num_max_dispatch_tokens_per_rank_prefill.cache_clear() - - assert get_deepep_num_max_dispatch_tokens_per_rank_prefill() == 6128 - - get_deepep_num_max_dispatch_tokens_per_rank_prefill.cache_clear() - - -def test_deepep_buffer_capacity_uses_larger_phase_capacity(monkeypatch): - cases = [ - (4096, 256, 4096), - (256, 512, 512), - ] - - for prefill_capacity, decode_capacity, expected_capacity in cases: - monkeypatch.setattr( - envs_utils, - "get_deepep_num_max_dispatch_tokens_per_rank_prefill", - lambda prefill_capacity=prefill_capacity: prefill_capacity, - ) - monkeypatch.setattr( - envs_utils, - "get_deepep_num_max_dispatch_tokens_per_rank_decode", - lambda decode_capacity=decode_capacity: decode_capacity, - ) - get_deepep_num_max_dispatch_tokens_per_rank.cache_clear() - assert get_deepep_num_max_dispatch_tokens_per_rank() == expected_capacity - - get_deepep_num_max_dispatch_tokens_per_rank.cache_clear() - - -def test_deepep_buffer_capacity_keeps_explicit_environment_override(monkeypatch): - monkeypatch.setenv("NUM_MAX_DISPATCH_TOKENS_PER_RANK_PREFILL", "1024") - monkeypatch.setenv("NUM_MAX_DISPATCH_TOKENS_PER_RANK_DECODE", "128") - get_deepep_num_max_dispatch_tokens_per_rank_prefill.cache_clear() - get_deepep_num_max_dispatch_tokens_per_rank_decode.cache_clear() - - assert get_deepep_num_max_dispatch_tokens_per_rank_prefill() == 1024 - assert get_deepep_num_max_dispatch_tokens_per_rank_decode() == 128 - - get_deepep_num_max_dispatch_tokens_per_rank_prefill.cache_clear() - get_deepep_num_max_dispatch_tokens_per_rank_decode.cache_clear() diff --git a/unit_tests/utils/test_log_utils.py b/unit_tests/utils/test_log_utils.py deleted file mode 100644 index 3534656835..0000000000 --- a/unit_tests/utils/test_log_utils.py +++ /dev/null @@ -1,42 +0,0 @@ -import logging - -import pytest - -from lightllm.utils import log_utils - - -def _format_message(message: str) -> str: - formatter = log_utils.NewLineFormatter("%(node_role)s%(message)s") - record = logging.LogRecord("test", logging.INFO, __file__, 1, message, (), None) - return formatter.format(record) - - -@pytest.mark.parametrize( - ("run_mode", "expected_marker"), - [ - ("prefill", "P"), - ("decode", "D"), - ("normal", ""), - ("pd_master", ""), - ], -) -def test_log_node_role_marker_is_limited_to_pd_workers(monkeypatch, run_mode, expected_marker): - monkeypatch.delenv("LIGHTLLM_LOG_NODE_ROLE", raising=False) - monkeypatch.setattr(log_utils, "_LOG_NODE_ROLE", "") - - log_utils.set_log_node_role(run_mode) - - formatted_marker = f"[{expected_marker}] " if expected_marker else "" - assert _format_message("ready") == f"{formatted_marker}ready" - if expected_marker: - assert log_utils.os.environ["LIGHTLLM_LOG_NODE_ROLE"] == expected_marker - else: - assert "LIGHTLLM_LOG_NODE_ROLE" not in log_utils.os.environ - - -def test_log_node_role_marker_is_applied_to_every_line(monkeypatch): - monkeypatch.delenv("LIGHTLLM_LOG_NODE_ROLE", raising=False) - monkeypatch.setattr(log_utils, "_LOG_NODE_ROLE", "") - log_utils.set_log_node_role("prefill") - - assert _format_message("first\nsecond") == "[P] first\r\n[P] second" diff --git a/unit_tests/utils/test_multimodal_chunk_utils.py b/unit_tests/utils/test_multimodal_chunk_utils.py deleted file mode 100644 index d553a62994..0000000000 --- a/unit_tests/utils/test_multimodal_chunk_utils.py +++ /dev/null @@ -1,68 +0,0 @@ -import random -import unittest - -from lightllm.utils.multimodal_chunk_utils import ( - adjust_mm_chunk_end, - get_mm_input_spans, -) - - -class MultimodalChunkTests(unittest.TestCase): - def test_embedding_position_locates_a_qwen_image(self): - params = {"images": [{"start_idx": 3, "token_id": 100, "token_num": 4}], "audios": []} - input_ids = [7, 8, 9, 100, 101, 102, 103] - self.assertEqual(get_mm_input_spans(input_ids, params), [(3, 7)]) - - def test_adjacent_images_audio_and_reused_embeddings_remain_separate(self): - params = { - "images": [ - {"start_idx": 1, "token_id": 100, "token_num": 3}, - {"start_idx": 6, "token_id": 100, "token_num": 3}, - ], - "audios": [{"start_idx": 4, "token_id": 200, "token_num": 2}], - } - input_ids = [7, 100, 101, 102, 200, 201, 100, 101, 102] - self.assertCountEqual(get_mm_input_spans(input_ids, params), [(1, 4), (4, 6), (6, 9)]) - - def test_text_before_and_after_image_can_still_be_chunked(self): - spans = [(1800, 2100)] - chunks = [] - start = 0 - while start < 4000: - end = adjust_mm_chunk_end(start, min(start + 2048, 4000), spans) - chunks.append((start, end)) - start = end - self.assertEqual(chunks, [(0, 1800), (1800, 3848), (3848, 4000)]) - - def test_small_chunk_target_can_yield_to_a_complete_item(self): - spans = [(100, 800)] - self.assertEqual(adjust_mm_chunk_end(0, 256, spans), 100) - self.assertEqual(adjust_mm_chunk_end(100, 356, spans), 800) - - def test_checkpoint_boundary_can_be_crossed_without_splitting_an_item(self): - spans = [(1900, 2200)] - self.assertEqual(adjust_mm_chunk_end(0, 2048, spans), 1900) - self.assertEqual(adjust_mm_chunk_end(1900, 2048, spans), 2200) - - def test_random_chunk_partitions_always_make_progress_within_budget(self): - rng = random.Random(0) - for _ in range(100): - spans = [] - position = 0 - for _ in range(8): - start = position + rng.randrange(0, 50) - position = start + rng.randrange(1, 129) - spans.append((start, position)) - total = position + 50 - target = rng.randrange(1, 129) - position = 0 - while position < total: - end = adjust_mm_chunk_end(position, min(position + target, total), spans) - self.assertLess(position, end) - self.assertLessEqual(end - position, 128) - self.assertFalse(any(start < end < stop for start, stop in spans)) - position = end - - -if __name__ == "__main__": - unittest.main() diff --git a/unit_tests/utils/test_service_shm_cleanup.py b/unit_tests/utils/test_service_shm_cleanup.py deleted file mode 100644 index d28434b995..0000000000 --- a/unit_tests/utils/test_service_shm_cleanup.py +++ /dev/null @@ -1,280 +0,0 @@ -import ctypes -import json -import os -from pathlib import Path -import signal -import socket -import subprocess -import sys -import time -from types import SimpleNamespace -import uuid - -import psutil -import pytest - -from lightllm.utils import service_shm_cleanup - - -def test_cleanup_service_shm_only_removes_matching_service(monkeypatch, tmp_path): - shm_dir = tmp_path / "shm" - shm_dir.mkdir() - matching_names = ["service_0_req_pool", "service_0_token_load"] - for name in [*matching_names, "service_1_req_pool", "other_service_0_value"]: - (shm_dir / name).touch() - monkeypatch.setattr(service_shm_cleanup, "SHM_DIR", shm_dir) - - service_shm_cleanup.ServiceShmCleanup.cleanup_posix_shm("service_0") - - assert all(not (shm_dir / name).exists() for name in matching_names) - assert (shm_dir / "service_1_req_pool").exists() - assert (shm_dir / "other_service_0_value").exists() - - -def test_cleanup_ipc_files_only_removes_matching_service(monkeypatch, tmp_path): - service_socket = tmp_path / "_service_0_router" - other_socket = tmp_path / "_service_1_router" - sockets = [] - for path in [service_socket, other_socket]: - sock = socket.socket(socket.AF_UNIX) - sock.bind(str(path)) - sockets.append(sock) - service_lock = tmp_path / "service_0_shm_port_args.lock" - other_lock = tmp_path / "service_1_shm_port_args.lock" - service_lock.touch() - other_lock.touch() - monkeypatch.setattr(service_shm_cleanup, "IPC_DIR", tmp_path) - - try: - assert service_shm_cleanup.ServiceShmCleanup.cleanup_ipc_files("service_0") == 2 - assert not service_socket.exists() - assert not service_lock.exists() - assert other_socket.exists() - assert other_lock.exists() - finally: - for sock in sockets: - sock.close() - - -def test_system_v_shm_keys_follow_feature_switches(monkeypatch): - monkeypatch.setenv( - "LIGHTLLM_START_ARGS", - json.dumps( - { - "run_mode": "prefill", - "enable_cpu_cache": True, - "enable_multimodal": False, - "cpu_kv_cache_shm_id": 21, - "multi_modal_cache_shm_id": 22, - } - ), - ) - removed_system_v_keys = [] - monkeypatch.setattr( - service_shm_cleanup.ServiceShmCleanup, - "cleanup_posix_shm", - staticmethod(lambda service_name: 0), - ) - monkeypatch.setattr( - service_shm_cleanup.ServiceShmCleanup, - "cleanup_ipc_files", - staticmethod(lambda service_name: 0), - ) - monkeypatch.setattr( - service_shm_cleanup.ServiceShmCleanup, - "cleanup_system_v_shm", - staticmethod(lambda keys: removed_system_v_keys.extend(keys) or len(keys)), - ) - - service_shm_cleanup.ServiceShmCleanup("current_service_0").cleanup_service_resources() - - assert removed_system_v_keys == [21] - - -def test_non_inference_mode_skips_system_v_shm_cleanup(monkeypatch): - monkeypatch.setenv( - "LIGHTLLM_START_ARGS", - json.dumps( - { - "run_mode": "visual_only", - "enable_cpu_cache": True, - "enable_multimodal": True, - "cpu_kv_cache_shm_id": 21, - "multi_modal_cache_shm_id": 22, - } - ), - ) - system_v_cleanup_calls = [] - monkeypatch.setattr( - service_shm_cleanup.ServiceShmCleanup, - "cleanup_posix_shm", - staticmethod(lambda service_name: 0), - ) - monkeypatch.setattr( - service_shm_cleanup.ServiceShmCleanup, - "cleanup_ipc_files", - staticmethod(lambda service_name: 0), - ) - monkeypatch.setattr( - service_shm_cleanup.ServiceShmCleanup, - "cleanup_system_v_shm", - staticmethod(lambda keys: system_v_cleanup_calls.append(keys) or 0), - ) - - service_shm_cleanup.ServiceShmCleanup("current_service_0").cleanup_service_resources() - - assert system_v_cleanup_calls == [] - - -def test_active_process_rejects_missing_pid_and_zombie(monkeypatch): - monkeypatch.setattr( - service_shm_cleanup.psutil, - "Process", - lambda pid: SimpleNamespace(status=lambda: psutil.STATUS_RUNNING), - ) - assert service_shm_cleanup.is_process_active(123) - monkeypatch.setattr( - service_shm_cleanup.psutil, - "Process", - lambda pid: SimpleNamespace(status=lambda: psutil.STATUS_ZOMBIE), - ) - assert not service_shm_cleanup.is_process_active(123) - - -def test_cleanup_process_checks_parent_every_two_seconds(monkeypatch): - states = iter([object(), object(), None]) - sleeps = [] - cleanup_calls = [] - signal_calls = [] - monkeypatch.setattr(service_shm_cleanup, "is_process_active", lambda pid: next(states)) - monkeypatch.setattr(service_shm_cleanup.time, "sleep", sleeps.append) - monkeypatch.setattr(service_shm_cleanup.signal, "signal", lambda sig, handler: signal_calls.append((sig, handler))) - monkeypatch.setattr( - service_shm_cleanup.ServiceShmCleanup, - "cleanup_service_resources", - lambda self: cleanup_calls.append(self.service_name), - ) - - service_shm_cleanup.run_launcher_shm_cleanup_process("service_0", 123) - - assert sleeps == [2.0, 2.0] - assert cleanup_calls == ["service_0"] - assert signal_calls == [ - (signal.SIGINT, signal.SIG_IGN), - (signal.SIGTERM, signal.SIG_IGN), - (signal.SIGHUP, signal.SIG_IGN), - ] - - -def test_start_cleanup_process_is_independent_session(monkeypatch): - popen_calls = [] - monkeypatch.setattr(service_shm_cleanup.os, "getpid", lambda: 123) - monkeypatch.setattr( - service_shm_cleanup.subprocess, "Popen", lambda *args, **kwargs: popen_calls.append((args, kwargs)) - ) - - assert service_shm_cleanup.start_launcher_shm_cleanup_process("service_0") is None - - command = popen_calls[0][0][0] - assert command[-2:] == ["service_0", "123"] - assert popen_calls[0][1]["start_new_session"] is True - assert "env" not in popen_calls[0][1] - - -LAUNCHER_CODE = r""" -import json, os, psutil, sys, time -from pathlib import Path -from lightllm.utils.service_shm_cleanup import start_launcher_shm_cleanup_process -service, directory, exit_mode = sys.argv[1:] -directory = Path(directory) -start_launcher_shm_cleanup_process(service) -watcher = psutil.Process().children()[0] -(Path('/dev/shm')/(service+'_req_pool')).write_bytes(b'test') -lock = Path('/tmp')/(service+'_shm_port_args.lock') -lock.touch() -(directory/'launcher.json').write_text(json.dumps({'watcher':watcher.pid})) -while not (directory/'exit').exists(): time.sleep(.01) -if exit_mode == 'normal': sys.exit(0) -if exit_mode == 'exception': raise RuntimeError('test launcher failure') -if exit_mode == 'os_exit': os._exit(7) -while True: time.sleep(1) -""" - - -def _wait_until(predicate, timeout=15): - deadline = time.monotonic() + timeout - while time.monotonic() < deadline: - if predicate(): - return - time.sleep(0.05) - assert predicate(), "Timed out waiting for resource cleanup" - - -def _process_stopped(pid): - try: - return psutil.Process(pid).status() == psutil.STATUS_ZOMBIE - except psutil.NoSuchProcess: - return True - - -def _cleanup_signals_are_ignored(pid): - status = Path(f"/proc/{pid}/status").read_text() - ignored = int(next(line.split()[1] for line in status.splitlines() if line.startswith("SigIgn:")), 16) - return all(ignored & (1 << (sig.value - 1)) for sig in [signal.SIGINT, signal.SIGTERM, signal.SIGHUP]) - - -@pytest.mark.parametrize("exit_mode", ["normal", "exception", "os_exit", "SIGKILL", "SIGINT", "SIGTERM", "SIGHUP"]) -def test_cleanup_process_after_real_launcher_exit(tmp_path, exit_mode): - service = "cleanup_test_" + uuid.uuid4().hex - libc = ctypes.CDLL(None) - while True: - key = uuid.uuid4().int % (2 ** 30) + 1 - shmid = libc.shmget(key, 4096, 0o1000 | 0o2000 | 0o600) - if shmid >= 0: - break - environment = os.environ.copy() - environment["LIGHTLLM_START_ARGS"] = json.dumps( - {"run_mode": "normal", "enable_cpu_cache": True, "cpu_kv_cache_shm_id": key} - ) - process = None - watcher_pid = None - try: - process = subprocess.Popen( - [sys.executable, "-c", LAUNCHER_CODE, service, str(tmp_path), exit_mode], - env=environment, - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - start_new_session=True, - ) - _wait_until(lambda: (tmp_path / "launcher.json").exists()) - watcher_pid = json.loads((tmp_path / "launcher.json").read_text())["watcher"] - assert os.getsid(watcher_pid) == watcher_pid - assert os.getsid(watcher_pid) != os.getsid(process.pid) - _wait_until(lambda: _cleanup_signals_are_ignored(watcher_pid)) - assert (Path("/dev/shm") / (service + "_req_pool")).exists() - assert libc.shmget(key, 0, 0) == shmid - if exit_mode == "normal": - os.kill(watcher_pid, signal.SIGTERM) - time.sleep(0.1) - assert not _process_stopped(watcher_pid) - if exit_mode.startswith("SIG"): - if exit_mode == "SIGKILL": - process.kill() - else: - os.killpg(process.pid, getattr(signal, exit_mode)) - else: - (tmp_path / "exit").touch() - process.wait(timeout=10) - _wait_until(lambda: _process_stopped(watcher_pid)) - assert not (Path("/dev/shm") / (service + "_req_pool")).exists() - assert not (Path("/tmp") / (service + "_shm_port_args.lock")).exists() - assert libc.shmget(key, 0, 0) == -1 - finally: - if process is not None and process.poll() is None: - process.kill() - process.wait() - if watcher_pid and not _process_stopped(watcher_pid): - psutil.Process(watcher_pid).kill() - libc.shmctl(shmid, 0, None) - (Path("/dev/shm") / (service + "_req_pool")).unlink(missing_ok=True) - (Path("/tmp") / (service + "_shm_port_args.lock")).unlink(missing_ok=True) diff --git a/unit_tests/utils/test_sgl_utils.py b/unit_tests/utils/test_sgl_utils.py deleted file mode 100644 index 85bb041e68..0000000000 --- a/unit_tests/utils/test_sgl_utils.py +++ /dev/null @@ -1,288 +0,0 @@ -import weakref -import collections -import inspect -import math -from types import SimpleNamespace - -import pytest -import torch -from frozendict import frozendict - -from lightllm.common.basemodel.infer_struct import InferStateInfo -from lightllm.common.basemodel.attention.fa3 import fp as fa3_module -from lightllm.common.triton_utils import autotuner as autotuner_module -from lightllm.common.triton_utils.autotuner import AutotuneKernelType, AutotuneLevel, Autotuner -from lightllm.utils import sgl_utils -from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len - - -@pytest.fixture(autouse=True) -def autotune_seq_len_environment(monkeypatch): - monkeypatch.delenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", raising=False) - monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", None) - get_decode_attn_autotune_seq_len.cache_clear() - yield - get_decode_attn_autotune_seq_len.cache_clear() - - -@pytest.mark.parametrize("seq_len, expected_len", [(None, 32768), ("8192", 8192), ("16384", 16384), ("8193", 8704)]) -@pytest.mark.parametrize("level", [AutotuneLevel.ADAPTIVE_AUTOTUNE, AutotuneLevel.FORCE_AUTOTUNE]) -def test_fa3_run_key_uses_configured_length_during_tuning(monkeypatch, seq_len, expected_len, level): - if seq_len is not None: - monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", seq_len) - monkeypatch.setattr(sgl_utils, "get_triton_autotune_level", lambda: level) - with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - for graph_capacity in [16, 32768]: - page_table = torch.empty(2, graph_capacity, dtype=torch.int32) - assert ( - sgl_utils._flash_attn_kvcache_run_key(page_table, 3, 2) - == 2 * 10_000_000_000_000 + 3 * 10_000_000 + expected_len - ) - - -@pytest.mark.parametrize( - "phase, level", - [ - (None, AutotuneLevel.USE_AUTOTUNE_HIS_CONFIG), - (None, AutotuneLevel.ADAPTIVE_AUTOTUNE), - (None, AutotuneLevel.FORCE_AUTOTUNE), - (AutotuneKernelType.GENERAL, AutotuneLevel.ADAPTIVE_AUTOTUNE), - (AutotuneKernelType.GENERAL, AutotuneLevel.FORCE_AUTOTUNE), - (AutotuneKernelType.DECODE_ATTENTION, AutotuneLevel.USE_AUTOTUNE_HIS_CONFIG), - ], -) -def test_fa3_run_key_uses_actual_length_without_tuning(monkeypatch, phase, level): - # 未调优时不应读取该环境变量，即使它无效也不影响正常配置查找。 - monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", "invalid") - monkeypatch.setattr(sgl_utils, "get_triton_autotune_level", lambda: level) - monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", phase) - page_table = torch.empty(2, 32768, dtype=torch.int32) - for actual_len, expected_len in [(1, 512), (512, 512), (513, 1024), (8192, 8192), (16384, 16384)]: - assert ( - sgl_utils._flash_attn_kvcache_run_key(page_table, 3, actual_len) - == 2 * 10_000_000_000_000 + 3 * 10_000_000 + expected_len - ) - - -def test_fa3_runtime_matches_cached_config_by_actual_length(monkeypatch): - kernel = sgl_utils.flash_attn_with_kvcache_autotune - monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: AutotuneLevel.USE_AUTOTUNE_HIS_CONFIG) - monkeypatch.setattr(autotuner_module.dist, "is_initialized", lambda: False) - monkeypatch.setattr(kernel, "fn", lambda **kwargs: kwargs["run_config"]) - monkeypatch.setattr(kernel, "fast_match_configs", collections.defaultdict(dict)) - inputs = dict( - q=torch.empty(2, 4, 8), - k_cache=torch.empty(1, 1, 2, 8), - v_cache=torch.empty(1, 1, 2, 8), - page_table=torch.empty(2, 32768, dtype=torch.int32), - max_seqlen_q=1, - causal=True, - window_size=(-1, -1), - softcap=0.0, - sinks=None, - k_descale=None, - v_descale=None, - ) - static_key = frozendict(kernel._static_key(**inputs)) - base_key = 2 * 10_000_000_000_000 + 10_000_000 - monkeypatch.setattr( - kernel, - "cached_configs", - {static_key: {str(base_key + 8192): {"num_splits": 16}, str(base_key + 16384): {"num_splits": 32}}}, - ) - - assert kernel(**inputs, max_seqlen_k=8192) == {"num_splits": 16} - assert kernel(**inputs, max_seqlen_k=8191) == {"num_splits": 16} - assert kernel(**inputs, max_seqlen_k=15000) == {"num_splits": 32} - assert set(kernel.fast_match_configs[static_key]) == {str(base_key + 8192), str(base_key + 15360)} - - -@pytest.mark.parametrize("page_size, actual_kv_len", [(1, 8192), (16, 8191), (256, 8192)]) -def test_fa3_decode_preserves_actual_length_before_graph_capture(monkeypatch, page_size, actual_kv_len): - graph = SimpleNamespace(can_run=lambda **kwargs: True, graph_max_len_in_batch=32768) - model = SimpleNamespace( - graph=graph, - req_manager=SimpleNamespace(req_to_token_indexs=torch.zeros(2, 8192, dtype=torch.int32)), - is_mtp_draft_model=False, - mtp_manager=SimpleNamespace(get_decode_draft_step=lambda _: 0), - ) - infer_state = InferStateInfo() - infer_state.batch_size = 2 - infer_state.max_kv_seq_len = actual_kv_len - infer_state.microbatch_index = 0 - infer_state.b_req_idx = torch.tensor([0, 1], dtype=torch.int32) - infer_state.b_seq_len = torch.tensor([4096, actual_kv_len], dtype=torch.int32) - infer_state.b1_cu_q_seq_len = torch.tensor([0, 1, 2], dtype=torch.int32) - infer_state.b1_cu_kv_seq_len = torch.tensor([0, 4096, 4096 + actual_kv_len], dtype=torch.int32) - state = fa3_module.Fa3DecodeAttState( - backend=SimpleNamespace( - model=model, - infer_page_size=page_size, - uses_causal_attention=lambda: True, - uses_dynamic_spec_verify_layout=lambda: False, - get_page_table_view=lambda att_batch_size, max_kv_len, microbatch_index: torch.zeros( - att_batch_size, math.ceil(max_kv_len / page_size), dtype=torch.int32 - ), - ), - _infer_state_weak_ref=weakref.ref(infer_state), - ) - copy_calls = [] - monkeypatch.setattr(fa3_module, "page_table_copy", lambda **kwargs: copy_calls.append(kwargs)) - state.init_state() - # 模拟 CudaGraph._capture_decode 对 infer_state 的修改。 - state.infer_state.max_kv_seq_len = graph.graph_max_len_in_batch - calls = [] - - def attention(**kwargs): - calls.append(kwargs) - return kwargs["q"] - - monkeypatch.setattr(fa3_module, "flash_attn_with_kvcache_autotune", attention) - state.decode_att(torch.empty(2, 4, 8), torch.empty(page_size, 2, 8), torch.empty(page_size, 2, 8)) - assert calls[0]["max_seqlen_k"] == actual_kv_len - assert calls[0]["page_table"].shape == (2, 32768 // page_size) - - assert copy_calls[0]["page_table"].shape == (2, math.ceil(actual_kv_len / page_size)) - assert copy_calls[0]["page_size"] == page_size - assert calls[0]["k_cache"].shape == (1, page_size, 2, 8) - assert calls[0]["v_cache"].shape == (1, page_size, 2, 8) - - -@pytest.mark.parametrize("seq_len", ["0", "-1", "invalid"]) -def test_invalid_decode_autotune_seq_len_is_rejected(monkeypatch, seq_len): - monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", seq_len) - with pytest.raises(ValueError): - get_decode_attn_autotune_seq_len() - - -@pytest.mark.parametrize("page_size, num_pages", [(1, 64), (1, 3), (256, 3)]) -@pytest.mark.parametrize("query_lengths", [[1, 1], [3, 3], [0, 3]]) -def test_fa3_rebuilds_valid_kv_metadata_without_changing_originals(monkeypatch, page_size, num_pages, query_lengths): - kv_len = 17 - monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", str(kv_len)) - max_pages = (kv_len + page_size - 1) // page_size - q = torch.randn(sum(query_lengths), 4, 8) - k = torch.randn(num_pages, page_size, 2, 8) - v = torch.randn_like(k) - page_table = torch.full((2, 8), -1, dtype=torch.int32) - seq_lens = torch.full((2,), 2, dtype=torch.int32) - cu_q = torch.tensor([0, query_lengths[0], sum(query_lengths)], dtype=torch.int32) - original = (q, k, v, seq_lens, page_table, cu_q, max(query_lengths), 2) - snapshots = [tensor.clone() for tensor in original[:6]] - options = { - "cu_seqlens_k_new": None, - "causal": True, - "window_size": (16, 0), - "softmax_scale": 0.25, - "return_softmax_lse": False, - } - - args, kwargs = sgl_utils._flash_attn_kvcache_rebuild_inputs(*original, **options) - - assert all(args[i] is original[i] for i in [0, 1, 2, 5]) - assert args[6:] == (max(query_lengths), kv_len) - assert kwargs == options - assert args[4].shape == (2, max_pages) - assert args[4].min() >= 0 and args[4].max() < num_pages - if num_pages >= 2 * max_pages: - assert args[4].unique().numel() == 2 * max_pages - torch.testing.assert_close(args[3], torch.full_like(seq_lens, kv_len)) - for tensor, snapshot in zip(original[:6], snapshots): - torch.testing.assert_close(tensor, snapshot) - - # Batched Q does not require cumulative query or KV lengths. - batched_q = torch.randn(2, 1, 4, 8) - args, kwargs = sgl_utils._flash_attn_kvcache_rebuild_inputs( - q=batched_q, - k_cache=k, - v_cache=v, - page_table=page_table, - cache_seqlens=seq_lens, - cu_seqlens_q=None, - max_seqlen_q=1, - max_seqlen_k=2, - ) - assert args[0] is batched_q - assert args[5:] == (None, 1, kv_len) - - -@pytest.mark.parametrize("kv_len", [8192, 16384]) -@pytest.mark.parametrize("query_lengths", [[1, 1], [3, 3], [0, 3]]) -def test_fa3_autotunes_long_kv_then_captures_original_inputs(tmp_path, monkeypatch, kv_len, query_lengths): - if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 9: - pytest.skip("FA3 requires a Hopper GPU") - if sgl_utils.flash_attn_with_kvcache is None: - pytest.skip("sgl_kernel FA3 is unavailable") - - monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", str(kv_len)) - kernel = sgl_utils.flash_attn_with_kvcache_autotune - monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", None) - monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: AutotuneLevel.ADAPTIVE_AUTOTUNE) - monkeypatch.setattr(sgl_utils, "get_triton_autotune_level", lambda: AutotuneLevel.ADAPTIVE_AUTOTUNE) - monkeypatch.setattr(autotuner_module.dist, "is_initialized", lambda: False) - monkeypatch.setattr(kernel, "_cache_dir", str(tmp_path), raising=False) - monkeypatch.setattr(kernel, "cached_configs", {}) - monkeypatch.setattr(kernel, "fast_match_configs", collections.defaultdict(dict)) - monkeypatch.setattr(kernel, "warmuped_configs_set", set()) - - q = torch.randn(sum(query_lengths), 4, 64, device="cuda", dtype=torch.bfloat16) - k = torch.randn(2 * kv_len, 1, 2, 64, device="cuda", dtype=q.dtype) - v = torch.randn_like(k) - inputs = dict( - q=q, - k_cache=k, - v_cache=v, - page_table=torch.zeros(2, 32768, device="cuda", dtype=torch.int32), - cache_seqlens=torch.full((2,), 2, device="cuda", dtype=torch.int32), - cu_seqlens_q=torch.tensor([0, query_lengths[0], sum(query_lengths)], device="cuda", dtype=torch.int32), - cu_seqlens_k_new=None, - max_seqlen_q=max(query_lengths), - max_seqlen_k=2, - causal=True, - window_size=(-1, -1), - softcap=0.0, - sinks=None, - k_descale=None, - v_descale=None, - ) - snapshots = {name: value.clone() for name, value in inputs.items() if isinstance(value, torch.Tensor)} - reference = kernel.fn(**inputs) - benchmark = kernel._bench - timings = [] - - def checked_bench(*args, **kwargs): - bound = inspect.signature(kernel.fn).bind(*args, **kwargs).arguments - assert bound["cache_seqlens"].tolist() == [kv_len, kv_len] - assert bound["page_table"].shape == (2, kv_len) - assert bound["cu_seqlens_k_new"] is None - assert bound["max_seqlen_k"] == kv_len - assert bound["q"] is q - elapsed = benchmark(*args, **kwargs) - assert math.isfinite(elapsed), f"FA3 benchmark failed for {kwargs['run_config']}" - timings.append(elapsed) - return elapsed - - monkeypatch.setattr(kernel, "_bench", checked_bench) - with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - output = kernel(**inputs) - assert len(timings) == 3 - expected_key = str(2 * 10_000_000_000_000 + max(query_lengths) * 10_000_000 + kv_len) - assert kernel._run_key(**inputs) == 2 * 10_000_000_000_000 + max(query_lengths) * 10_000_000 + 512 - assert all(list(configs) == [expected_key] for configs in kernel.cached_configs.values()) - torch.testing.assert_close(output, reference) - - def unexpected_rebuild(*args, **kwargs): - pytest.fail("Cached execution and CUDA Graph capture must use the original inputs") - - monkeypatch.setattr(kernel, "rebuild_input_func", unexpected_rebuild) - with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): - kernel(**inputs) - torch.cuda.synchronize() - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - captured_output = kernel(**inputs) - graph.replay() - torch.cuda.synchronize() - torch.testing.assert_close(captured_output, reference) - for name, snapshot in snapshots.items(): - torch.testing.assert_close(inputs[name], snapshot) diff --git a/unit_tests/utils/test_shm_utils.py b/unit_tests/utils/test_shm_utils.py deleted file mode 100644 index 00203aba3d..0000000000 --- a/unit_tests/utils/test_shm_utils.py +++ /dev/null @@ -1,88 +0,0 @@ -from concurrent.futures import ThreadPoolExecutor -from multiprocessing import resource_tracker, shared_memory -from threading import Event - -import pytest - -from lightllm.utils import shm_utils - - -def test_get_service_shm_name_requires_service_name(monkeypatch): - monkeypatch.setattr(shm_utils, "get_unique_server_name", lambda: None) - - with pytest.raises(RuntimeError, match="LIGHTLLM_UNIQUE_SERVICE_NAME_ID is unset"): - shm_utils.get_service_shm_name("req_pool") - - -def test_get_service_shm_name_adds_prefix_once(monkeypatch): - monkeypatch.setattr(shm_utils, "get_unique_server_name", lambda: "service_uuid_0") - - assert shm_utils.get_service_shm_name("req_pool") == "service_uuid_0_req_pool" - assert shm_utils.get_service_shm_name("service_uuid_0_req_pool") == "service_uuid_0_req_pool" - - -def test_create_or_link_shm_passes_scoped_name_to_shared_memory_layer(monkeypatch): - created_names = [] - monkeypatch.setattr(shm_utils, "get_unique_server_name", lambda: "service_uuid_1") - monkeypatch.setattr( - shm_utils, - "_force_create_shm", - lambda name, expected_size: created_names.append((name, expected_size)) or object(), - ) - - shm_utils.create_or_link_shm("token_load", 128, force_mode="create") - - assert created_names == [("service_uuid_1_token_load", 128)] - - -@pytest.mark.parametrize("method, tracker_method", [("__init__", "register"), ("unlink", "unregister")]) -def test_service_shm_concurrent_calls_restore_tracker(monkeypatch, method, tracker_method): - original_tracker_method = getattr(resource_tracker, tracker_method) - # 即使回归导致 patch 恢复错误，也不污染后续测试。 - monkeypatch.setattr(resource_tracker, tracker_method, original_tracker_method) - first_entered = Event() - second_started = Event() - second_entered = Event() - release_first = Event() - first_finished = Event() - - def ordered_operation(self, *args, **kwargs): - if self._test_index == 0: - first_entered.set() - assert release_first.wait(5) - else: - second_entered.set() - # 无锁时强制先退出第一个 patch，再退出第二个，复现恢复顺序错误。 - assert first_finished.wait(5) - - monkeypatch.setattr(shared_memory.SharedMemory, method, ordered_operation) - - def worker(index): - # 只测试 patch 生命周期，不申请真实共享内存。 - shm = shm_utils.ServiceSharedMemory.__new__(shm_utils.ServiceSharedMemory) - shm._test_index = index - if index == 1: - second_started.set() - try: - if method == "__init__": - shm.__init__(name="test_tracker") - else: - shm.unlink() - finally: - if index == 0: - first_finished.set() - - with ThreadPoolExecutor(max_workers=2) as executor: - first = executor.submit(worker, 0) - try: - assert first_entered.wait(5) - second = executor.submit(worker, 1) - assert second_started.wait(5) - overlapped = second_entered.wait(0.2) - finally: - release_first.set() - first.result(timeout=5) - second.result(timeout=5) - - assert not overlapped, "Concurrent calls entered the patched operation together" - assert getattr(resource_tracker, tracker_method) is original_tracker_method diff --git a/unit_tests/utils/test_speculative_utils.py b/unit_tests/utils/test_speculative_utils.py deleted file mode 100644 index a58b35198d..0000000000 --- a/unit_tests/utils/test_speculative_utils.py +++ /dev/null @@ -1,413 +0,0 @@ -import weakref -import json -from importlib import import_module -from types import SimpleNamespace - -import pytest -import torch - -from lightllm.common.basemodel.infer_struct import InferStateInfo -import lightllm.common.basemodel.attention.base_att as base_att_module -import lightllm.common.basemodel.hidden_collector as hidden_collector_module -from lightllm.common.basemodel.attention.base_att import BaseAttBackend -from lightllm.common.basemodel.attention.fa3.fp import Fa3DecodeAttState, Fa3PrefillAttState -from lightllm.common.basemodel.attention.fa3.mla import MlaFa3DecodeAttState, MlaFa3PrefillAttState -from lightllm.common.basemodel.attention.triton.fp import TritonDecodeAttState -from lightllm.models import get_draft_model_class -from lightllm.models.qwen3_eagle.layer_weights.transformer_layer_weight import Qwen3EagleTransformerLayerWeight -from lightllm.utils import envs_utils - - -@pytest.mark.parametrize( - "module_name, class_name", - [ - ("lightllm.models.deepseek_mtp.model", "Deepseek3MTPModel"), - ("lightllm.models.glm4_moe_lite_mtp.model", "Glm4MoeLiteMTPModel"), - ("lightllm.models.mistral_mtp.model", "MistralMTPModel"), - ("lightllm.models.qwen3_moe_mtp.model", "Qwen3MOEMTPModel"), - ("lightllm.models.qwen3_5_mtp.model", "Qwen3_5MTPModel"), - ("lightllm.models.qwen3_5_moe_mtp.model", "Qwen3_5MoeMTPModel"), - ("lightllm.models.qwen3_eagle.model", "Qwen3EagleModel"), - ("lightllm.models.qwen3_dflash.model", "Qwen3DFlashModel"), - ("lightllm.models.qwen3_dflash2.model", "Qwen3DFlash2Model"), - ("lightllm.models.qwen3_5_dflash.model", "Qwen3_5DFlashModel"), - ("lightllm.models.qwen3_dspark.model", "Qwen3DSparkModel"), - ("lightllm.models.qwen3_5_dspark.model", "Qwen3_5DSparkModel"), - ], -) -def test_spec_draft_model_class_is_marked(module_name, class_name): - model_class = getattr(import_module(module_name), class_name) - assert model_class.is_mtp_draft_model is True - - -def test_qwen3_eagle_uses_layers_checkpoint_prefix(): - layer_weight = Qwen3EagleTransformerLayerWeight.__new__(Qwen3EagleTransformerLayerWeight) - layer_weight.layer_num_ = 2 - - layer_weight._init_weight_names() - - assert layer_weight._q_weight_name == "layers.2.self_attn.q_proj.weight" - assert layer_weight._hidden_norm_weight_name == "layers.2.hidden_norm.weight" - - -@pytest.mark.parametrize( - "mtp_mode, is_draft_model, draft_step, dynamic_spec, expected", - [ - ("dspark", False, 7, True, True), - ("dspark", True, 7, True, False), - ("dflash", False, 7, True, True), - ("dflash", True, 7, True, False), - ("dflash2", False, 7, False, False), - ("dflash2", True, 7, False, False), - ("dflash2", False, 7, True, True), - ("dflash2", True, 7, True, False), - ("vanilla_with_att", True, 7, True, False), - ("vanilla_with_att", True, 0, True, False), - ("eagle3", True, 0, True, False), - ("eagle3", False, 7, True, True), - ("eagle3", False, 7, False, False), - ], -) -def test_attention_backend_selects_dynamic_spec_layout( - monkeypatch, - mtp_mode, - is_draft_model, - draft_step, - dynamic_spec, - expected, -): - monkeypatch.setattr( - base_att_module, - "get_env_start_args", - lambda: SimpleNamespace(mtp_mode=mtp_mode, mtp_dynamic_verify=dynamic_spec), - ) - backend = SimpleNamespace( - model=SimpleNamespace( - is_mtp_draft_model=is_draft_model, - mtp_manager=SimpleNamespace(get_decode_draft_step=lambda _: draft_step), - ) - ) - - assert BaseAttBackend.uses_dynamic_spec_verify_layout(backend) is expected - - -@pytest.mark.parametrize( - "mtp_mode, is_draft_model, expected", - [ - (None, False, True), - ("dflash", False, True), - ("dflash", True, False), - ("dflash2", True, False), - ("dspark", True, False), - ("eagle3", True, True), - ("vanilla_with_att", True, True), - ], -) -def test_attention_backend_selects_causality(monkeypatch, mtp_mode, is_draft_model, expected): - monkeypatch.setattr( - base_att_module, - "get_env_start_args", - lambda: SimpleNamespace(mtp_mode=mtp_mode), - ) - backend = SimpleNamespace(model=SimpleNamespace(is_mtp_draft_model=is_draft_model)) - - assert BaseAttBackend.uses_causal_attention(backend) is expected - - -@pytest.mark.parametrize("state_class", [Fa3PrefillAttState, MlaFa3PrefillAttState]) -def test_fa3_prefill_state_owns_causality(monkeypatch, state_class): - monkeypatch.setattr(import_module(Fa3PrefillAttState.__module__), "page_table_copy", lambda **kwargs: None) - infer_state = InferStateInfo() - infer_state.b1_cu_q_seq_len = torch.tensor([0, 1, 2], dtype=torch.int32) - infer_state.b1_cu_kv_seq_len = torch.tensor([0, 3, 7], dtype=torch.int32) - infer_state.b_req_idx = torch.tensor([0, 1], dtype=torch.int32) - infer_state.batch_size = 2 - infer_state.max_kv_seq_len = 4 - infer_state.input_ids = torch.empty(2, dtype=torch.int64) - infer_state.req_manager = SimpleNamespace(req_to_token_indexs=torch.arange(8, dtype=torch.int32).reshape(2, 4)) - state = state_class( - backend=SimpleNamespace(uses_causal_attention=lambda: False, page_size=1, infer_page_size=1), - _infer_state_weak_ref=weakref.ref(infer_state), - ) - - state.init_state() - - assert state.causal is False - - -@pytest.mark.parametrize("state_class", [Fa3DecodeAttState, MlaFa3DecodeAttState]) -def test_fa3_decode_state_owns_causality(state_class): - infer_state = InferStateInfo() - infer_state.b1_cu_q_seq_len = torch.tensor([0, 1, 2], dtype=torch.int32) - infer_state.b1_cu_kv_seq_len = torch.tensor([0, 3, 7], dtype=torch.int32) - infer_state.b_req_idx = torch.tensor([0, 1], dtype=torch.int32) - infer_state.b_seq_len = torch.tensor([3, 4], dtype=torch.int32) - model = SimpleNamespace( - is_mtp_draft_model=False, - mtp_manager=SimpleNamespace(get_decode_draft_step=lambda _: 0), - ) - state = state_class( - backend=SimpleNamespace( - model=model, - uses_causal_attention=lambda: False, - uses_dynamic_spec_verify_layout=lambda: False, - ), - _infer_state_weak_ref=weakref.ref(infer_state), - ) - state._init_page_table = lambda _: None - - state.init_state() - - assert state.causal is False - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") -def test_triton_mtp_decode_state_builds_group_markers(): - model = SimpleNamespace( - is_mtp_draft_model=False, - mtp_manager=SimpleNamespace(get_decode_draft_step=lambda _: 2), - req_manager=SimpleNamespace(HOLD_REQUEST_ID=-1), - ) - infer_state = InferStateInfo() - infer_state.b_req_idx = torch.tensor([7, 7, -1, -1], dtype=torch.int32, device="cuda") - infer_state.max_kv_seq_len = 4 - state = TritonDecodeAttState( - backend=SimpleNamespace(model=model), - _infer_state_weak_ref=weakref.ref(infer_state), - ) - - state.init_state() - - assert state.b_mark_mtp_shared_group.cpu().tolist() == [0, 2, 1, 1] - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") -@pytest.mark.parametrize("state_class", [Fa3DecodeAttState, MlaFa3DecodeAttState]) -def test_fa3_dynamic_decode_state_builds_group_markers(state_class): - model = SimpleNamespace(req_manager=SimpleNamespace(HOLD_REQUEST_ID=-1)) - infer_state = InferStateInfo() - infer_state.b_req_idx = torch.tensor([7, 7, -1, -1], dtype=torch.int32, device="cuda") - infer_state.b_seq_len = torch.tensor([3, 4, 2, 2], dtype=torch.int32, device="cuda") - infer_state.batch_size = 4 - state = state_class( - backend=SimpleNamespace(model=model), - _infer_state_weak_ref=weakref.ref(infer_state), - ) - - b_att_req_idx = state._init_dynamic_spec_verify_state(draft_step=2) - torch.cuda.synchronize() - - assert b_att_req_idx.cpu().tolist() == [7, -1, -1, -1] - assert state.b_att_seq_len.cpu().tolist() == [4, 2, 2, 0] - assert state.cu_seqlens_q.cpu().tolist() == [0, 2, 3, 4, 4] - - -@pytest.mark.parametrize( - "model_type, spec_mode, expected_class_name", - [ - ("deepseek_v3", "vanilla_with_att", "Deepseek3MTPModel"), - ("deepseek_v3", "eagle_with_att", "Deepseek3MTPModel"), - ("glm4_moe_lite", "vanilla_with_att", "Glm4MoeLiteMTPModel"), - ("glm4_moe_lite", "eagle_with_att", "Glm4MoeLiteMTPModel"), - ("mistral", "vanilla_no_att", "MistralMTPModel"), - ("mistral", "eagle_no_att", "MistralMTPModel"), - ("qwen3_moe", "vanilla_no_att", "Qwen3MOEMTPModel"), - ("qwen3_moe", "eagle_no_att", "Qwen3MOEMTPModel"), - ("qwen3_5", "vanilla_with_att", "Qwen3_5MTPModel"), - ("qwen3_5_text", "eagle_with_att", "Qwen3_5MTPModel"), - ("qwen3_5_moe", "vanilla_with_att", "Qwen3_5MoeMTPModel"), - ("qwen3_5_moe_text", "eagle_with_att", "Qwen3_5MoeMTPModel"), - ("qwen3", "dflash", "Qwen3DFlashModel"), - ("qwen3", "dflash2", "Qwen3DFlash2Model"), - ("qwen3_5", "dflash", "Qwen3_5DFlashModel"), - ("qwen3_5_text", "dflash", "Qwen3_5DFlashModel"), - ("qwen3", "dspark", "Qwen3DSparkModel"), - ("qwen3_5", "dspark", "Qwen3_5DSparkModel"), - ("qwen3_5_text", "dspark", "Qwen3_5DSparkModel"), - ("qwen3", "eagle3", "Qwen3EagleModel"), - ], -) -def test_draft_model_registry(model_type, spec_mode, expected_class_name): - model_class = get_draft_model_class( - model_cfg={"model_type": model_type}, - spec_mode=spec_mode, - ) - - assert model_class.__name__ == expected_class_name - - -def test_draft_model_registry_rejects_unsupported_model_type(): - with pytest.raises(ValueError, match="Unsupported speculative draft model"): - get_draft_model_class( - model_cfg={"model_type": "gemma4"}, - spec_mode="dspark", - ) - - -@pytest.mark.parametrize( - "model_type, spec_mode", - [ - ("deepseek_v3", "eagle3"), - ("deepseek_v3", "dspark"), - ("deepseek_v3", "dflash"), - ("glm4_moe_lite", "eagle3"), - ("glm4_moe_lite", "dspark"), - ("glm4_moe_lite", "dflash"), - ("qwen3_5", "eagle3"), - ("qwen3_5_moe", "eagle3"), - ("qwen3_5_moe", "dspark"), - ("qwen3_5_moe", "dflash"), - ("qwen3_5", "dflash2"), - ("qwen3", "eagle_no_att"), - ], -) -def test_draft_model_registry_rejects_unsupported_mode(model_type, spec_mode): - with pytest.raises(ValueError, match="Unsupported speculative draft model"): - get_draft_model_class( - model_cfg={"model_type": model_type}, - spec_mode=spec_mode, - ) - - -def test_hidden_collector_reads_target_layer_ids(monkeypatch): - config_reads = [] - - def get_config_dict(path): - config_reads.append(path) - return {"target_layer_ids": [1, 20, 36]}, {} - - monkeypatch.setattr( - hidden_collector_module.PretrainedConfig, - "get_config_dict", - get_config_dict, - ) - monkeypatch.setattr( - hidden_collector_module, - "get_env_start_args", - lambda: SimpleNamespace(mtp_draft_model_dir=["/models/dspark"]), - ) - model = SimpleNamespace(is_mtp_draft_model=False, layers_num=40) - - collector = hidden_collector_module.LayerHiddenCollector(model=model) - new_collector = collector.new_instance() - - assert collector.layer_ids == frozenset((1, 20, 36)) - assert new_collector.layer_ids == collector.layer_ids - assert config_reads == ["/models/dspark"] - - -@pytest.mark.parametrize( - "mtp_mode, mtp_step, expected_layer_num", - [ - (None, 0, 0), - ("vanilla_no_att", 7, 0), - ("eagle_no_att", 7, 0), - ("vanilla_with_att", 7, 7), - ("eagle_with_att", 7, 1), - ], -) -def test_fixed_added_mtp_kv_layer_num_by_mode(monkeypatch, mtp_mode, mtp_step, expected_layer_num): - monkeypatch.setattr( - envs_utils, - "get_env_start_args", - lambda: SimpleNamespace(mtp_mode=mtp_mode, mtp_step=mtp_step), - ) - envs_utils.get_added_mtp_kv_layer_num.cache_clear() - - assert envs_utils.get_added_mtp_kv_layer_num() == expected_layer_num - - -@pytest.mark.parametrize("mtp_mode", ["dflash", "dflash2"]) -def test_dflash_added_kv_layers_come_from_draft_config(tmp_path, mtp_mode): - config_path = tmp_path / "config.json" - config_path.write_text(json.dumps({"num_hidden_layers": 5})) - - envs_utils.get_env_start_args.cache_clear() - envs_utils.get_added_mtp_kv_layer_num.cache_clear() - envs_utils.set_env_start_args( - { - "mtp_mode": mtp_mode, - "mtp_step": 7, - "mtp_dynamic_verify": False, - "mtp_draft_model_dir": [str(tmp_path)], - } - ) - - assert envs_utils.get_added_mtp_kv_layer_num() == 5 - - -def test_qwen35_dflash_added_kv_layers_come_from_nested_draft_config(tmp_path): - config_path = tmp_path / "config.json" - config_path.write_text( - json.dumps( - { - "num_hidden_layers": 48, - "dflash_config": {"num_hidden_layers": 5}, - } - ) - ) - - envs_utils.get_env_start_args.cache_clear() - envs_utils.get_added_mtp_kv_layer_num.cache_clear() - envs_utils.set_env_start_args( - { - "mtp_mode": "dflash", - "mtp_step": 7, - "mtp_dynamic_verify": False, - "mtp_draft_model_dir": [str(tmp_path)], - } - ) - - assert envs_utils.get_added_mtp_kv_layer_num() == 5 - - -@pytest.mark.parametrize( - "draft_config,expected_layer_num", - [ - ({"num_hidden_layers": 6}, 6), - ( - { - "model_type": "deepseek_v4", - "num_hidden_layers": 61, - "compress_ratios": [4] * 61 + [0] * 3, - "dspark_target_layer_ids": [15, 31, 47], - "dspark_block_size": 8, - }, - 3, - ), - ], -) -def test_dspark_added_kv_layers_come_from_draft_config(tmp_path, draft_config, expected_layer_num): - config_path = tmp_path / "config.json" - config_path.write_text(json.dumps(draft_config)) - - envs_utils.get_env_start_args.cache_clear() - envs_utils.get_added_mtp_kv_layer_num.cache_clear() - envs_utils.set_env_start_args( - { - "mtp_mode": "dspark", - "mtp_step": 7, - "mtp_dynamic_verify": True, - "mtp_draft_model_dir": [str(tmp_path)], - } - ) - - assert envs_utils.get_added_mtp_kv_layer_num() == expected_layer_num - - -def test_eagle3_added_kv_layers_come_from_draft_config(tmp_path): - config_path = tmp_path / "config.json" - config_path.write_text(json.dumps({"num_hidden_layers": 2})) - - envs_utils.get_env_start_args.cache_clear() - envs_utils.get_added_mtp_kv_layer_num.cache_clear() - envs_utils.set_env_start_args( - { - "mtp_mode": "eagle3", - "mtp_step": 7, - "mtp_dynamic_verify": False, - "mtp_draft_model_dir": [str(tmp_path)], - } - ) - - assert envs_utils.get_added_mtp_kv_layer_num() == 2 diff --git a/unit_tests/utils/test_start_utils.py b/unit_tests/utils/test_start_utils.py deleted file mode 100644 index 816e4000f7..0000000000 --- a/unit_tests/utils/test_start_utils.py +++ /dev/null @@ -1,319 +0,0 @@ -from types import SimpleNamespace - -import pytest - -from lightllm.utils import start_utils - - -class FakeHttpServerProcess: - def __init__(self, return_code=None): - self.return_code = return_code - self.pid = 4321 - self.sent_signals = [] - self.wait_timeouts = [] - - def poll(self): - return self.return_code - - def send_signal(self, sig): - self.sent_signals.append(sig) - - def wait(self, timeout=None): - self.wait_timeouts.append(timeout) - return self.return_code - - -class FakeProcess: - def __init__(self, pid, running=True, children=None, name="process", exitcode=None, wait_timeout=False): - self.pid = pid - self.running = running - self._children = children or [] - self._name = name - self.exitcode = exitcode - self.wait_timeout = wait_timeout - - def is_running(self): - return self.running - - def children(self, recursive=False): - return self._children - - def name(self): - return self._name - - def wait(self, timeout=None): - if self.wait_timeout: - raise start_utils.psutil.TimeoutExpired(timeout, pid=self.pid, name=self._name) - return self.exitcode - - -def test_start_submodule_processes_returns_and_manages_psutil_processes(monkeypatch): - class FakePipeReader: - def recv(self): - # 子进程应在等待初始化结果之前就进入 manager，保证此时 Ctrl-C 可以清理它们。 - assert len(process_manager.processes) == 2 - return "init ok" - - class FakeMpProcess: - next_pid = 1000 - - def __init__(self, target, args): - self.pid = self.next_pid - FakeMpProcess.next_pid += 1 - - def start(self): - pass - - def is_alive(self): - return True - - monkeypatch.setattr(start_utils.mp, "Pipe", lambda duplex: (FakePipeReader(), object())) - monkeypatch.setattr(start_utils.mp, "Process", FakeMpProcess) - monkeypatch.setattr( - start_utils.psutil, - "Process", - lambda pid: FakeProcess(pid=pid, name=f"process-{pid}"), - ) - process_manager = start_utils.SubmoduleManager() - - processes = process_manager.start_submodule_processes( - start_funcs=[lambda pipe_writer: None, lambda pipe_writer: None], - start_args=[(), ()], - ) - - assert processes == process_manager.processes - assert [process.pid for process in processes] == [1000, 1001] - assert process_manager.process_names == { - processes[0]: "process-1000", - processes[1]: "process-1001", - } - - -def test_register_process_tree_adds_recursive_descendants(monkeypatch): - descendants = [ - FakeProcess(pid=1001, name="lightllm::model_infer"), - FakeProcess(pid=1002, name="lightllm::pd_manager"), - FakeProcess(pid=1003, name="lightllm::pd_worker"), - ] - router_process = FakeProcess(pid=1000, children=descendants) - process_manager = start_utils.SubmoduleManager() - monkeypatch.setattr(start_utils, "is_process_active", lambda pid: True) - - process_manager.register_process_tree(router_process) - - assert process_manager.processes == descendants - assert process_manager.process_names == { - descendants[0]: "lightllm::model_infer", - descendants[1]: "lightllm::pd_manager", - descendants[2]: "lightllm::pd_worker", - } - - -def test_register_process_tree_filters_short_lived_helper_processes(monkeypatch): - model_process = FakeProcess(pid=1001, name="lightllm::model_infer") - compile_worker = FakeProcess(pid=1002, name="python") - pd_process = FakeProcess(pid=1003, name="lightllm::decode_trans") - router_process = FakeProcess(pid=1000, children=[model_process, compile_worker, pd_process]) - process_manager = start_utils.SubmoduleManager() - monkeypatch.setattr(start_utils, "is_process_active", lambda pid: True) - - process_manager.register_process_tree(router_process) - - assert process_manager.processes == [model_process, pd_process] - assert process_manager.process_names == { - model_process: "lightllm::model_infer", - pd_process: "lightllm::decode_trans", - } - - -def test_register_process_tree_ignores_processes_that_exit_during_scan(): - class ExitedProcess(FakeProcess): - def name(self): - raise start_utils.psutil.NoSuchProcess(self.pid) - - exited_process = ExitedProcess(pid=1001) - router_process = FakeProcess(pid=1000, children=[exited_process]) - process_manager = start_utils.SubmoduleManager() - - process_manager.register_process_tree(router_process) - - assert process_manager.processes == [] - assert process_manager.process_names == {} - - -def test_setup_signal_handlers_registers_and_handles_sigterm(monkeypatch): - http_server_process = FakeHttpServerProcess() - process_manager = start_utils.SubmoduleManager() - registered_handlers = {} - terminate_calls = [] - monkeypatch.setattr( - start_utils.signal, - "signal", - lambda sig, handler: registered_handlers.__setitem__(sig, handler), - ) - monkeypatch.setattr(process_manager, "terminate_all_processes", lambda: terminate_calls.append(True)) - - process_manager.setup_signal_handlers(http_server_process) - - assert set(registered_handlers) == { - start_utils.signal.SIGTERM, - start_utils.signal.SIGINT, - start_utils.signal.SIGHUP, - } - with pytest.raises(SystemExit) as exc_info: - registered_handlers[start_utils.signal.SIGTERM](start_utils.signal.SIGTERM, None) - - assert exc_info.value.code == 0 - assert http_server_process.sent_signals == [start_utils.signal.SIGTERM] - assert http_server_process.wait_timeouts == [60] - assert terminate_calls == [True] - - -def test_setup_signal_handlers_uses_configured_graceful_shutdown_signal(monkeypatch): - http_server_process = FakeHttpServerProcess() - process_manager = start_utils.SubmoduleManager() - registered_handlers = {} - monkeypatch.setattr( - start_utils.signal, - "signal", - lambda sig, handler: registered_handlers.__setitem__(sig, handler), - ) - monkeypatch.setattr(process_manager, "terminate_all_processes", lambda: None) - - process_manager.setup_signal_handlers( - http_server_process, - graceful_shutdown_signal=start_utils.signal.SIGUSR1, - ) - - with pytest.raises(SystemExit) as exc_info: - registered_handlers[start_utils.signal.SIGTERM](start_utils.signal.SIGTERM, None) - - assert exc_info.value.code == 0 - assert http_server_process.sent_signals == [start_utils.signal.SIGUSR1] - assert http_server_process.wait_timeouts == [60] - - -def test_setup_exit_controller_starts_cleanup_process_without_registering_signals(monkeypatch): - process_manager = start_utils.SubmoduleManager() - cleanup_process_calls = [] - registered_handlers = {} - monkeypatch.setattr(start_utils, "get_unique_server_name", lambda: "service_0") - monkeypatch.setattr( - start_utils, - "start_launcher_shm_cleanup_process", - lambda service_name: cleanup_process_calls.append(service_name), - ) - monkeypatch.setattr( - start_utils.signal, - "signal", - lambda sig, handler: registered_handlers.__setitem__(sig, handler), - ) - - process_manager.setup_exit_controller() - - assert cleanup_process_calls == ["service_0"] - assert registered_handlers == {} - - -def test_supervisor_fails_when_http_server_exits(monkeypatch): - http_server_process = FakeHttpServerProcess(return_code=0) - process_manager = start_utils.SubmoduleManager() - terminate_calls = [] - kill_calls = [] - monkeypatch.setattr(process_manager, "terminate_all_processes", lambda: terminate_calls.append(True)) - monkeypatch.setattr(start_utils, "kill_recursive", lambda process: kill_calls.append(process)) - - with pytest.raises(RuntimeError, match="HTTP server exited unexpectedly with return code 0"): - process_manager.supervise_processes(http_server_process) - - assert kill_calls == [] - assert terminate_calls == [True] - - -def test_supervisor_fails_and_cleans_up_when_submodule_exits(monkeypatch): - http_server_process = FakeHttpServerProcess() - dead_process = FakeProcess(pid=1234, running=False, name="router", exitcode=-9) - process_manager = start_utils.SubmoduleManager() - process_manager.processes = [dead_process] - process_manager.process_names = {dead_process: dead_process.name()} - terminate_calls = [] - kill_calls = [] - monkeypatch.setattr(process_manager, "terminate_all_processes", lambda: terminate_calls.append(True)) - monkeypatch.setattr(start_utils, "is_process_active", lambda pid: True) - monkeypatch.setattr(start_utils, "kill_recursive", lambda process: kill_calls.append(process)) - - with pytest.raises( - RuntimeError, - match="Critical LightLLM submodule exited unexpectedly: name=router pid=1234 exitcode=-9", - ): - process_manager.supervise_processes(http_server_process) - - assert kill_calls == [http_server_process] - assert terminate_calls == [True] - - -def test_supervisor_treats_zombie_submodule_as_dead(monkeypatch): - http_server_process = FakeHttpServerProcess() - zombie_process = FakeProcess(pid=1234, name="router", exitcode=-9) - process_manager = start_utils.SubmoduleManager() - process_manager.processes = [zombie_process] - process_manager.process_names = {zombie_process: zombie_process.name()} - terminate_calls = [] - kill_calls = [] - monkeypatch.setattr(process_manager, "terminate_all_processes", lambda: terminate_calls.append(True)) - monkeypatch.setattr(start_utils, "is_process_active", lambda pid: False) - monkeypatch.setattr(start_utils, "kill_recursive", lambda process: kill_calls.append(process)) - - with pytest.raises( - RuntimeError, - match="Critical LightLLM submodule exited unexpectedly: name=router pid=1234 exitcode=-9", - ): - process_manager.supervise_processes(http_server_process) - - assert kill_calls == [http_server_process] - assert terminate_calls == [True] - - -def test_supervisor_keeps_polling_while_all_processes_are_alive(monkeypatch): - http_server_process = FakeHttpServerProcess() - child_process = FakeProcess(pid=1234) - process_manager = start_utils.SubmoduleManager() - process_manager.processes = [child_process] - process_manager.process_names = {child_process: child_process.name()} - terminate_calls = [] - sleep_calls = [] - monkeypatch.setattr(process_manager, "terminate_all_processes", lambda: terminate_calls.append(True)) - - def stop_after_first_poll(interval): - sleep_calls.append(interval) - http_server_process.return_code = 2 - - monkeypatch.setattr(start_utils, "is_process_active", lambda pid: True) - monkeypatch.setattr(start_utils.time, "sleep", stop_after_first_poll) - - with pytest.raises(RuntimeError, match="HTTP server exited unexpectedly with return code 2"): - process_manager.supervise_processes(http_server_process) - - assert sleep_calls == [5.0] - assert terminate_calls == [True] - - -def test_supervisor_supports_submodule_only_processes(monkeypatch): - dead_process = FakeProcess(pid=1234, running=False, name="model_infer", wait_timeout=True) - process_manager = start_utils.SubmoduleManager() - process_manager.processes = [dead_process] - process_manager.process_names = {dead_process: dead_process.name()} - terminate_calls = [] - kill_calls = [] - monkeypatch.setattr(process_manager, "terminate_all_processes", lambda: terminate_calls.append(True)) - monkeypatch.setattr(start_utils, "is_process_active", lambda pid: False) - monkeypatch.setattr(start_utils, "kill_recursive", lambda process: kill_calls.append(process)) - - with pytest.raises( - RuntimeError, - match="Critical LightLLM submodule exited unexpectedly: name=model_infer pid=1234 exitcode=None", - ): - process_manager.supervise_processes() - - assert kill_calls == [] - assert terminate_calls == [True] diff --git a/unit_tests/utils/test_startup_process_lifecycle.py b/unit_tests/utils/test_startup_process_lifecycle.py deleted file mode 100644 index 9b9d38fb60..0000000000 --- a/unit_tests/utils/test_startup_process_lifecycle.py +++ /dev/null @@ -1,135 +0,0 @@ -import ctypes -import importlib -import multiprocessing as mp -import os -import signal -import sys -import time -import traceback -from types import SimpleNamespace - -import psutil -import pytest - -from lightllm.utils import process_check, start_utils - - -pytestmark = pytest.mark.skipif(sys.platform != "linux", reason="requires Linux subreaper and fork") - - -def _kill_and_reap_children(): - for child in psutil.Process().children(recursive=True): - try: - os.kill(child.pid, signal.SIGKILL) - except ProcessLookupError: - pass - while True: - try: - os.waitpid(-1, 0) - except ChildProcessError: - break - - -def _run_metrics_lifecycle(failure_mode, result_writer): - """Own all launcher descendants, including orphans, in an isolated process.""" - result = None - try: - libc = ctypes.CDLL(None, use_errno=True) - if libc.prctl(36, 1, 0, 0, 0) != 0: # PR_SET_CHILD_SUBREAPER - raise OSError(ctypes.get_errno(), "could not become a child subreaper") - - mp.set_start_method("fork", force=True) - ctx = mp.get_context("fork") - events, event_writer = ctx.Pipe(duplex=False) - metrics = importlib.import_module("lightllm.server.metrics.manager") - envs_utils = importlib.import_module("lightllm.utils.envs_utils") - - # Keep the real metrics entry point and RPyC server, without monitor/shared-memory setup. - metrics.Monitor = lambda args: SimpleNamespace() - metrics.get_unique_server_name = lambda: "metrics-lifecycle-regression" - metrics.get_shm_port_args = lambda: SimpleNamespace(metric_port=0) - envs_utils.get_env_start_args = lambda: SimpleNamespace(enable_mps=False) - # Speed up only the monitor's polling, without changing the shared time module. - process_check.time = SimpleNamespace(sleep=lambda _: time.sleep(0.01)) - - def failing_submodule(pipe_writer): - event_writer.send(("failing", os.getpid())) - pipe_writer.send("injected initialization failure") - while True: - time.sleep(100) - - def launcher_main(): - manager = start_utils.SubmoduleManager() - metric_processes = manager.start_submodule_processes( - start_funcs=[metrics.start_metric_manager], - start_args=[(SimpleNamespace(metric_gateway=None, push_interval=1),)], - ) - event_writer.send(("metrics", metric_processes[0].pid)) - if failure_mode == "parent-exit": - # Bypass multiprocessing joins to exercise the metrics parent monitor itself. - os._exit(2) - manager.start_submodule_processes(start_funcs=[failing_submodule], start_args=[()]) - - launcher = ctx.Process(target=launcher_main) - launcher.start() - reported = {} - expected = {"metrics", "failing"} if failure_mode == "startup-failure" else {"metrics"} - deadline = time.monotonic() + 8 - while set(reported) != expected: - assert time.monotonic() < deadline, "launcher did not start the expected submodules" - if events.poll(0.05): - name, pid = events.recv() - reported[name] = pid - - launcher.join(timeout=8) - assert not launcher.is_alive(), "launcher hung while joining an earlier metrics process" - assert launcher.exitcode == (1 if failure_mode == "startup-failure" else 2) - - if failure_mode == "parent-exit": - # The real monitor must kill metrics after its launcher's immediate exit. - deadline = time.monotonic() + 5 - while True: - reaped_pid, status = os.waitpid(reported["metrics"], os.WNOHANG) - if reaped_pid: - assert os.WIFSIGNALED(status) - assert os.WTERMSIG(status) == signal.SIGKILL - break - assert time.monotonic() < deadline, "metrics did not detect its dead launcher" - time.sleep(0.01) - - assert all(not psutil.pid_exists(pid) for pid in reported.values()), "submodule leaked after launcher exit" - result = "ok" - except BaseException: - result = traceback.format_exc() - finally: - _kill_and_reap_children() - result_writer.send(result) - - -@pytest.mark.parametrize("failure_mode", ["startup-failure", "parent-exit"]) -def test_metrics_cleanup_on_launcher_failure(failure_mode): - # Import before forking so framework imports happen only once per pytest run. - importlib.import_module("lightllm.server.metrics.manager") - ctx = mp.get_context("fork") - result_reader, result_writer = ctx.Pipe(duplex=False) - controller = ctx.Process(target=_run_metrics_lifecycle, args=(failure_mode, result_writer)) - controller.start() - try: - controller.join(timeout=20) - assert not controller.is_alive(), "metrics lifecycle regression controller timed out" - assert controller.exitcode == 0 - assert result_reader.poll(timeout=1), "metrics lifecycle controller returned no result" - result = result_reader.recv() - assert result == "ok", result - finally: - if controller.is_alive(): - for child in psutil.Process(controller.pid).children(recursive=True): - try: - os.kill(child.pid, signal.SIGKILL) - except ProcessLookupError: - pass - controller.kill() - controller.join(timeout=5) - controller.close() - result_reader.close() - result_writer.close() From df0b3fe9429d1bac2d2beba1a70e11ed076886d7 Mon Sep 17 00:00:00 2001 From: sufubao Date: Fri, 13 Nov 2026 20:37:41 +0800 Subject: [PATCH 2/6] test: retain critical unit coverage --- .../common/att_resources/manager_fixtures.py | 45 + .../common/att_resources/model_helpers.py | 342 +++++ .../att_resources/test_backend_startup.py | 374 ++++++ .../att_resources/test_bounded_state.py | 440 +++++++ .../common/att_resources/test_capacity.py | 1113 +++++++++++++++++ .../att_resources/test_compute_binding.py | 147 +++ .../att_resources/test_dsv4_integration.py | 294 +++++ .../common/att_resources/test_dsv4_state.py | 505 ++++++++ .../att_resources/test_fp8_compute_binding.py | 201 +++ .../test_fp8_e8m0_packed_mla_state.py | 235 ++++ .../att_resources/test_full_att_cpu_kernel.py | 412 ++++++ .../test_full_att_dp_transfer.py | 221 ++++ .../att_resources/test_full_att_transfer.py | 338 +++++ .../common/att_resources/test_full_state.py | 283 +++++ .../test_full_storage_initialization.py | 71 ++ .../att_resources/test_gqa_layout_identity.py | 114 ++ .../att_resources/test_gqa_pd_kernel.py | 283 +++++ .../common/att_resources/test_gqa_transfer.py | 192 +++ .../test_hybrid_layer_mapping.py | 130 ++ .../common/att_resources/test_layer_groups.py | 164 +++ .../test_layer_state_selection.py | 216 ++++ .../att_resources/test_manager_helpers.py | 316 +++++ .../att_resources/test_model_calibration.py | 261 ++++ .../test_model_config_normalization.py | 235 ++++ .../test_model_initialization.py | 185 +++ .../att_resources/test_model_resources.py | 214 ++++ .../test_model_storage_definition.py | 276 ++++ .../att_resources/test_mtp_state_commit.py | 114 ++ .../att_resources/test_network_transport.py | 320 +++++ .../att_resources/test_no_attention_draft.py | 84 ++ .../att_resources/test_page_pool_storage.py | 167 +++ .../att_resources/test_paged_storage.py | 22 + .../test_quantized_full_state.py | 388 ++++++ .../att_resources/test_runtime_transport.py | 574 +++++++++ .../att_resources/test_state_creation.py | 212 ++++ .../att_resources/test_state_identity.py | 356 ++++++ .../att_resources/test_state_queries.py | 189 +++ .../common/att_resources/test_tp_config.py | 223 ++++ .../att_resources/test_transfer_views.py | 233 ++++ .../test_validate_model_cache.py | 49 + .../test_window_compressor_state.py | 195 +++ unit_tests/common/basemodel/__init__.py | 0 .../basemodel/attention/fa3/test_neo.py | 119 ++ .../flashinfer/test_decode_graph_state.py | 100 ++ .../flashinfer/test_decode_wrapper_init.py | 132 ++ .../flashinfer/test_mla_cuda_graph_plan.py | 178 +++ .../attention/flashinfer/test_mla_nope.py | 106 ++ .../attention/flashinfer/test_workspace.py | 56 + .../basemodel/attention/linear/test_gdn.py | 438 +++++++ .../basemodel/attention/test_dsv4_flashmla.py | 377 ++++++ .../attention/test_fp8_logical_layer.py | 117 ++ .../basemodel/attention/test_fp8_page_size.py | 45 + .../basemodel/attention/test_lazy_import.py | 248 ++++ .../attention/test_state_lifetime.py | 174 +++ .../basemodel/test_cuda_graph_layout.py | 133 ++ .../common/basemodel/test_hidden_collector.py | 193 +++ .../common/basemodel/test_model_input.py | 212 ++++ .../common/basemodel/test_model_output.py | 344 +++++ .../basemodel/test_moe_route_info_manager.py | 263 ++++ .../common/basemodel/test_mtp_manager.py | 133 ++ .../common/basemodel/test_overlap_utils.py | 213 ++++ .../gqa/mtp_diverse/test_mtp_diverse.py | 196 +++ .../mtp_diverse/test_mtp_diverse_autotune.py | 301 +++++ .../gqa/mtp_diverse/test_stage2_autotune.py | 220 ++++ .../gqa/test_flash_decoding_autotune.py | 270 ++++ .../test_int8kv_flash_decoding_diverse.py | 139 ++ ...st_int8kv_flash_decoding_diverse_stage3.py | 35 + .../decode_att/int8kv/test_normal_autotune.py | 300 +++++ .../test_context_flashattention_nopad.py | 104 ++ .../test_context_flashattention_nopad1.py | 196 +++ .../att/prefill_att/test_fa3_neo.py | 561 +++++++++ .../linear_att/test_acceptance_fusion.py | 272 ++++ .../test_fused_recurrent_strided.py | 71 ++ .../test_mtp_fused_recurrent_equiv.py | 163 +++ .../linear_att/test_mtp_state_params.py | 65 + .../linear_att/test_replayssm_execution.py | 530 ++++++++ .../basemodel/triton_kernel/mhc/test_mhc.py | 275 ++++ .../fused_moe/test_activation_config.py | 132 ++ .../test_append_shared_expert_topk.py | 222 ++++ unit_tests/common/fused_moe/test_deepep.py | 311 +++++ .../fused_moe/test_deepgemm_activation.py | 96 ++ .../fused_moe/test_ep_reduce_decode_output.py | 36 + .../fused_moe/test_grouped_fused_moe.py | 275 ++++ .../common/fused_moe/test_grouped_topk.py | 100 ++ .../kv_trans_kernel/test_kv_trans_v2.py | 123 ++ .../kv_trans_kernel/test_nixl_kv_trans.py | 141 +++ .../common/quantization/test_no_quant.py | 43 + .../test_ue8m0_fp8_quantization.py | 513 ++++++++ unit_tests/common/test_autotuner.py | 486 +++++++ .../common/test_cache_copy_control_flow.py | 98 ++ .../test_hybrid_req_manager_pd_sharing.py | 148 +++ unit_tests/common/test_linear_state.py | 334 +++++ .../common/test_linear_state_contract.py | 241 ++++ unit_tests/common/test_mem_manager_ipc.py | 114 ++ .../common/test_req_manager_hold_page.py | 27 + .../triton_kernel/test_hadamard_transform.py | 83 ++ .../test_topk_index_to_mem_index.py | 24 + .../models/deepseek_v4/test_c128_workspace.py | 548 ++++++++ unit_tests/models/deepseek_v4/test_kernels.py | 342 +++++ unit_tests/models/deepseek_v4/test_model.py | 301 +++++ .../deepseek_v4/test_shared_expert_overlap.py | 247 ++++ .../models/deepseek_v4/test_tokenizer.py | 79 ++ unit_tests/models/deepseek_v4/test_vision.py | 350 ++++++ .../deepseek_v4/test_vision_integration.py | 321 +++++ .../models/deepseek_v4/test_workspace.py | 97 ++ unit_tests/models/glm5_next/test_cache.py | 270 ++++ .../glm5_next/test_compact_topk_modes.py | 65 + .../glm5_next/test_direct_nsa_indices.py | 257 ++++ .../glm5_next/test_fp8_compile_cache.py | 94 ++ unit_tests/models/glm5_next/test_fp8_kv.py | 1073 ++++++++++++++++ .../models/glm5_next/test_fused_mtp_norm.py | 126 ++ .../models/glm5_next/test_indexer_topk.py | 161 +++ .../glm5_next/test_indexer_workspace.py | 78 ++ .../models/glm5_next/test_kda_conv_qkv.py | 186 +++ unit_tests/models/glm5_next/test_kernels.py | 665 ++++++++++ .../models/glm5_next/test_kpool_gather.py | 141 +++ unit_tests/models/glm5_next/test_layers.py | 451 +++++++ unit_tests/models/glm5_next/test_mtp.py | 308 +++++ .../glm5_next/test_mtp_prefill_cache.py | 278 ++++ unit_tests/models/glm5_next/test_pd_cache.py | 447 +++++++ .../glm5_next/test_prefill_workspace.py | 116 ++ .../glm5_next/test_token_major_mla_bmm.py | 92 ++ unit_tests/models/glm5_next/test_visual.py | 210 ++++ unit_tests/models/qwen2_vl/test_mrope.py | 66 + .../models/qwen3-vl/test_deepstack_emb.py | 82 ++ .../models/sensenova_loop/test_gdn_state.py | 151 +++ unit_tests/models/sensenova_loop/test_loop.py | 376 ++++++ .../models/sensenova_loop/test_loop_mhc.py | 107 ++ .../models/test_qwen3_dspark_model_output.py | 517 ++++++++ .../models/test_qwen3_eagle_model_input.py | 100 ++ unit_tests/models/test_yarn_rotary.py | 142 +++ .../vit/test_flash_attention_forward.py | 139 ++ .../core/objs/test_out_token_circlequeue.py | 127 ++ unit_tests/server/core/objs/test_req.py | 127 ++ unit_tests/server/core/objs/test_shm_array.py | 74 ++ .../core/objs/test_token_ids_validation.py | 109 ++ unit_tests/server/grammar_helpers.py | 154 +++ .../server/httpserver/test_pd_continuation.py | 93 ++ .../httpserver/test_pd_generate_error.py | 506 ++++++++ .../httpserver/test_pd_kv_status_channel.py | 109 ++ .../test_pd_master_cached_tokens.py | 771 ++++++++++++ .../httpserver/test_pd_master_discovery.py | 371 ++++++ .../httpserver/test_pd_split_connections.py | 294 +++++ .../httpserver/test_pd_websocket_lifecycle.py | 81 ++ .../test_running_request_lifecycle.py | 565 +++++++++ .../test_cache_placement_controller.py | 367 ++++++ unit_tests/server/pd_master_fault_launcher.py | 94 ++ unit_tests/server/pd_master_test_utils.py | 188 +++ .../router/dynamic_prompt/test_radix_cache.py | 722 +++++++++++ .../mode_backend/test_dp_big_pages.py | 286 +++++ .../test_dp_overlap_spec_engine.py | 491 ++++++++ .../mode_backend/test_generic_pre_process.py | 176 +++ .../mode_backend/test_multi_level_kv_cache.py | 507 ++++++++ .../mode_backend/test_nccl_control_channel.py | 151 +++ .../test_paged_side_allocations.py | 253 ++++ .../mode_backend/test_pd_dynamic_split.py | 244 ++++ .../mode_backend/test_pd_manager_startup.py | 174 +++ .../mode_backend/test_pd_notify.py | 223 ++++ .../mode_backend/test_pd_page_io.py | 53 + .../mode_backend/test_pd_prefill_abort.py | 199 +++ .../mode_backend/test_pd_token_indices.py | 126 ++ .../mode_backend/test_pd_transfer_progress.py | 176 +++ .../test_pd_transfer_readiness.py | 38 + .../test_prefill_queue_strategy.py | 401 ++++++ .../mtp_speculative/test_dflash.py | 180 +++ .../mtp_speculative/test_dflash2.py | 117 ++ .../mtp_speculative/test_eagle_overlap.py | 363 ++++++ .../mtp_speculative/test_eagle_with_att.py | 213 ++++ .../mtp_speculative/test_planner.py | 964 ++++++++++++++ .../mtp_speculative/test_vanilla_no_att.py | 211 ++++ .../mtp_speculative/test_vanilla_overlap.py | 147 +++ .../mtp_speculative/test_vanilla_prefill.py | 237 ++++ .../structured_output/test_error_output.py | 66 + .../model_infer/structured_output/test_mtp.py | 444 +++++++ .../structured_output/test_overlap.py | 152 +++ .../structured_output/test_reasoning.py | 194 +++ .../structured_output/test_sampling.py | 320 +++++ .../structured_output/test_state.py | 407 ++++++ .../model_infer/test_cpu_resource_layout.py | 513 ++++++++ .../test_hybrid_cpu_state_cache.py | 362 ++++++ .../model_infer/test_multimodal_chunking.py | 404 ++++++ .../req_queue/test_pd_queue_selection.py | 108 ++ unit_tests/server/test_api_responses.py | 313 +++++ unit_tests/server/test_application_stop.py | 200 +++ unit_tests/server/test_detokenization.py | 114 ++ unit_tests/server/test_dp_request_capacity.py | 109 ++ .../server/test_generation_error_response.py | 91 ++ unit_tests/server/test_hybrid_start_args.py | 138 ++ unit_tests/server/test_mtp_start_args.py | 55 + unit_tests/server/test_pd_cache_aware.py | 82 ++ unit_tests/server/test_pd_local_selection.py | 919 ++++++++++++++ .../server/test_pd_master_haproxy_proxy.py | 434 +++++++ unit_tests/server/test_pd_master_mode.py | 612 +++++++++ .../server/test_pd_master_multiprocess_e2e.py | 496 ++++++++ .../test_pd_master_multiprocess_launcher.py | 399 ++++++ unit_tests/server/test_pd_master_req_ids.py | 149 +++ .../server/test_pd_master_topology_api.py | 96 ++ unit_tests/server/test_pd_node_stats.py | 321 +++++ unit_tests/server/test_pd_start_args.py | 311 +++++ .../server/test_reasoning_token_usage.py | 222 ++++ unit_tests/server/test_ssm_state_mode.py | 165 +++ unit_tests/utils/test_config_utils.py | 53 + unit_tests/utils/test_envs_utils.py | 68 + unit_tests/utils/test_service_shm_cleanup.py | 280 +++++ unit_tests/utils/test_sgl_utils.py | 288 +++++ unit_tests/utils/test_shm_utils.py | 88 ++ unit_tests/utils/test_speculative_utils.py | 413 ++++++ unit_tests/utils/test_start_utils.py | 319 +++++ 208 files changed, 50929 insertions(+) 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_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_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_compute_binding.py create mode 100644 unit_tests/common/att_resources/test_fp8_e8m0_packed_mla_state.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_dp_transfer.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_layout_identity.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_hybrid_layer_mapping.py create mode 100644 unit_tests/common/att_resources/test_layer_groups.py create mode 100644 unit_tests/common/att_resources/test_layer_state_selection.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_config_normalization.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_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_tp_config.py create mode 100644 unit_tests/common/att_resources/test_transfer_views.py create mode 100644 unit_tests/common/att_resources/test_validate_model_cache.py create mode 100644 unit_tests/common/att_resources/test_window_compressor_state.py create mode 100644 unit_tests/common/basemodel/__init__.py create mode 100644 unit_tests/common/basemodel/attention/fa3/test_neo.py create mode 100644 unit_tests/common/basemodel/attention/flashinfer/test_decode_graph_state.py create mode 100644 unit_tests/common/basemodel/attention/flashinfer/test_decode_wrapper_init.py create mode 100644 unit_tests/common/basemodel/attention/flashinfer/test_mla_cuda_graph_plan.py create mode 100644 unit_tests/common/basemodel/attention/flashinfer/test_mla_nope.py create mode 100644 unit_tests/common/basemodel/attention/flashinfer/test_workspace.py create mode 100644 unit_tests/common/basemodel/attention/linear/test_gdn.py create mode 100644 unit_tests/common/basemodel/attention/test_dsv4_flashmla.py create mode 100644 unit_tests/common/basemodel/attention/test_fp8_logical_layer.py create mode 100644 unit_tests/common/basemodel/attention/test_fp8_page_size.py create mode 100644 unit_tests/common/basemodel/attention/test_lazy_import.py create mode 100644 unit_tests/common/basemodel/attention/test_state_lifetime.py create mode 100644 unit_tests/common/basemodel/test_cuda_graph_layout.py create mode 100644 unit_tests/common/basemodel/test_hidden_collector.py create mode 100644 unit_tests/common/basemodel/test_model_input.py create mode 100644 unit_tests/common/basemodel/test_model_output.py create mode 100644 unit_tests/common/basemodel/test_moe_route_info_manager.py create mode 100644 unit_tests/common/basemodel/test_mtp_manager.py create mode 100644 unit_tests/common/basemodel/test_overlap_utils.py create mode 100644 unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_mtp_diverse.py create mode 100644 unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_mtp_diverse_autotune.py create mode 100644 unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_stage2_autotune.py create mode 100644 unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/test_flash_decoding_autotune.py create mode 100644 unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse.py create mode 100644 unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage3.py create mode 100644 unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_normal_autotune.py create mode 100644 unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_flashattention_nopad.py create mode 100644 unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_flashattention_nopad1.py create mode 100644 unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_fa3_neo.py create mode 100644 unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py create mode 100644 unit_tests/common/basemodel/triton_kernel/linear_att/test_fused_recurrent_strided.py create mode 100644 unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_equiv.py create mode 100644 unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_state_params.py create mode 100644 unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py create mode 100644 unit_tests/common/basemodel/triton_kernel/mhc/test_mhc.py create mode 100644 unit_tests/common/fused_moe/test_activation_config.py create mode 100644 unit_tests/common/fused_moe/test_append_shared_expert_topk.py create mode 100644 unit_tests/common/fused_moe/test_deepep.py create mode 100644 unit_tests/common/fused_moe/test_deepgemm_activation.py create mode 100644 unit_tests/common/fused_moe/test_ep_reduce_decode_output.py create mode 100644 unit_tests/common/fused_moe/test_grouped_fused_moe.py create mode 100755 unit_tests/common/fused_moe/test_grouped_topk.py create mode 100644 unit_tests/common/kv_trans_kernel/test_kv_trans_v2.py create mode 100644 unit_tests/common/kv_trans_kernel/test_nixl_kv_trans.py create mode 100644 unit_tests/common/quantization/test_no_quant.py create mode 100644 unit_tests/common/quantization/test_ue8m0_fp8_quantization.py create mode 100644 unit_tests/common/test_autotuner.py create mode 100644 unit_tests/common/test_cache_copy_control_flow.py create mode 100644 unit_tests/common/test_hybrid_req_manager_pd_sharing.py create mode 100644 unit_tests/common/test_linear_state.py create mode 100644 unit_tests/common/test_linear_state_contract.py create mode 100644 unit_tests/common/test_mem_manager_ipc.py create mode 100644 unit_tests/common/test_req_manager_hold_page.py create mode 100644 unit_tests/models/deepseek3_2/triton_kernel/test_hadamard_transform.py create mode 100644 unit_tests/models/deepseek3_2/triton_kernel/test_topk_index_to_mem_index.py create mode 100644 unit_tests/models/deepseek_v4/test_c128_workspace.py create mode 100644 unit_tests/models/deepseek_v4/test_kernels.py create mode 100644 unit_tests/models/deepseek_v4/test_model.py create mode 100644 unit_tests/models/deepseek_v4/test_shared_expert_overlap.py create mode 100644 unit_tests/models/deepseek_v4/test_tokenizer.py create mode 100644 unit_tests/models/deepseek_v4/test_vision.py create mode 100644 unit_tests/models/deepseek_v4/test_vision_integration.py create mode 100644 unit_tests/models/deepseek_v4/test_workspace.py create mode 100644 unit_tests/models/glm5_next/test_cache.py create mode 100644 unit_tests/models/glm5_next/test_compact_topk_modes.py create mode 100644 unit_tests/models/glm5_next/test_direct_nsa_indices.py create mode 100644 unit_tests/models/glm5_next/test_fp8_compile_cache.py create mode 100644 unit_tests/models/glm5_next/test_fp8_kv.py create mode 100644 unit_tests/models/glm5_next/test_fused_mtp_norm.py create mode 100644 unit_tests/models/glm5_next/test_indexer_topk.py create mode 100644 unit_tests/models/glm5_next/test_indexer_workspace.py create mode 100644 unit_tests/models/glm5_next/test_kda_conv_qkv.py create mode 100644 unit_tests/models/glm5_next/test_kernels.py create mode 100644 unit_tests/models/glm5_next/test_kpool_gather.py create mode 100644 unit_tests/models/glm5_next/test_layers.py create mode 100644 unit_tests/models/glm5_next/test_mtp.py create mode 100644 unit_tests/models/glm5_next/test_mtp_prefill_cache.py create mode 100644 unit_tests/models/glm5_next/test_pd_cache.py create mode 100644 unit_tests/models/glm5_next/test_prefill_workspace.py create mode 100644 unit_tests/models/glm5_next/test_token_major_mla_bmm.py create mode 100644 unit_tests/models/glm5_next/test_visual.py create mode 100644 unit_tests/models/qwen2_vl/test_mrope.py create mode 100644 unit_tests/models/qwen3-vl/test_deepstack_emb.py create mode 100644 unit_tests/models/sensenova_loop/test_gdn_state.py create mode 100644 unit_tests/models/sensenova_loop/test_loop.py create mode 100644 unit_tests/models/sensenova_loop/test_loop_mhc.py create mode 100644 unit_tests/models/test_qwen3_dspark_model_output.py create mode 100644 unit_tests/models/test_qwen3_eagle_model_input.py create mode 100644 unit_tests/models/test_yarn_rotary.py create mode 100644 unit_tests/models/vit/test_flash_attention_forward.py create mode 100644 unit_tests/server/core/objs/test_out_token_circlequeue.py create mode 100644 unit_tests/server/core/objs/test_req.py create mode 100644 unit_tests/server/core/objs/test_shm_array.py create mode 100644 unit_tests/server/core/objs/test_token_ids_validation.py create mode 100644 unit_tests/server/grammar_helpers.py create mode 100644 unit_tests/server/httpserver/test_pd_continuation.py create mode 100644 unit_tests/server/httpserver/test_pd_generate_error.py create mode 100644 unit_tests/server/httpserver/test_pd_kv_status_channel.py create mode 100644 unit_tests/server/httpserver/test_pd_master_cached_tokens.py create mode 100644 unit_tests/server/httpserver/test_pd_master_discovery.py create mode 100644 unit_tests/server/httpserver/test_pd_split_connections.py create mode 100644 unit_tests/server/httpserver/test_pd_websocket_lifecycle.py create mode 100644 unit_tests/server/httpserver/test_running_request_lifecycle.py create mode 100644 unit_tests/server/multi_level_kv_cache/test_cache_placement_controller.py create mode 100644 unit_tests/server/pd_master_fault_launcher.py create mode 100644 unit_tests/server/pd_master_test_utils.py create mode 100644 unit_tests/server/router/dynamic_prompt/test_radix_cache.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_dp_overlap_spec_engine.py create mode 100644 unit_tests/server/router/model_infer/mode_backend/test_generic_pre_process.py create mode 100644 unit_tests/server/router/model_infer/mode_backend/test_multi_level_kv_cache.py create mode 100644 unit_tests/server/router/model_infer/mode_backend/test_nccl_control_channel.py create mode 100644 unit_tests/server/router/model_infer/mode_backend/test_paged_side_allocations.py create mode 100644 unit_tests/server/router/model_infer/mode_backend/test_pd_dynamic_split.py create mode 100644 unit_tests/server/router/model_infer/mode_backend/test_pd_manager_startup.py create mode 100644 unit_tests/server/router/model_infer/mode_backend/test_pd_notify.py create mode 100644 unit_tests/server/router/model_infer/mode_backend/test_pd_page_io.py create mode 100644 unit_tests/server/router/model_infer/mode_backend/test_pd_prefill_abort.py create mode 100644 unit_tests/server/router/model_infer/mode_backend/test_pd_token_indices.py create mode 100644 unit_tests/server/router/model_infer/mode_backend/test_pd_transfer_progress.py create mode 100644 unit_tests/server/router/model_infer/mode_backend/test_pd_transfer_readiness.py create mode 100644 unit_tests/server/router/model_infer/mode_backend/test_prefill_queue_strategy.py create mode 100644 unit_tests/server/router/model_infer/mtp_speculative/test_dflash.py create mode 100644 unit_tests/server/router/model_infer/mtp_speculative/test_dflash2.py create mode 100644 unit_tests/server/router/model_infer/mtp_speculative/test_eagle_overlap.py create mode 100644 unit_tests/server/router/model_infer/mtp_speculative/test_eagle_with_att.py create mode 100644 unit_tests/server/router/model_infer/mtp_speculative/test_planner.py create mode 100644 unit_tests/server/router/model_infer/mtp_speculative/test_vanilla_no_att.py create mode 100644 unit_tests/server/router/model_infer/mtp_speculative/test_vanilla_overlap.py create mode 100644 unit_tests/server/router/model_infer/mtp_speculative/test_vanilla_prefill.py create mode 100644 unit_tests/server/router/model_infer/structured_output/test_error_output.py create mode 100644 unit_tests/server/router/model_infer/structured_output/test_mtp.py create mode 100644 unit_tests/server/router/model_infer/structured_output/test_overlap.py create mode 100644 unit_tests/server/router/model_infer/structured_output/test_reasoning.py create mode 100644 unit_tests/server/router/model_infer/structured_output/test_sampling.py create mode 100644 unit_tests/server/router/model_infer/structured_output/test_state.py create mode 100644 unit_tests/server/router/model_infer/test_cpu_resource_layout.py create mode 100644 unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py create mode 100644 unit_tests/server/router/model_infer/test_multimodal_chunking.py create mode 100644 unit_tests/server/router/req_queue/test_pd_queue_selection.py create mode 100644 unit_tests/server/test_api_responses.py create mode 100644 unit_tests/server/test_application_stop.py create mode 100644 unit_tests/server/test_detokenization.py create mode 100644 unit_tests/server/test_dp_request_capacity.py create mode 100644 unit_tests/server/test_generation_error_response.py create mode 100644 unit_tests/server/test_hybrid_start_args.py create mode 100644 unit_tests/server/test_mtp_start_args.py create mode 100644 unit_tests/server/test_pd_cache_aware.py create mode 100644 unit_tests/server/test_pd_local_selection.py create mode 100644 unit_tests/server/test_pd_master_haproxy_proxy.py create mode 100644 unit_tests/server/test_pd_master_mode.py create mode 100644 unit_tests/server/test_pd_master_multiprocess_e2e.py create mode 100644 unit_tests/server/test_pd_master_multiprocess_launcher.py create mode 100644 unit_tests/server/test_pd_master_req_ids.py create mode 100644 unit_tests/server/test_pd_master_topology_api.py create mode 100644 unit_tests/server/test_pd_node_stats.py create mode 100644 unit_tests/server/test_pd_start_args.py create mode 100644 unit_tests/server/test_reasoning_token_usage.py create mode 100644 unit_tests/server/test_ssm_state_mode.py create mode 100644 unit_tests/utils/test_config_utils.py create mode 100644 unit_tests/utils/test_envs_utils.py create mode 100644 unit_tests/utils/test_service_shm_cleanup.py create mode 100644 unit_tests/utils/test_sgl_utils.py create mode 100644 unit_tests/utils/test_shm_utils.py create mode 100644 unit_tests/utils/test_speculative_utils.py create mode 100644 unit_tests/utils/test_start_utils.py 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..cfcf1ad8bf --- /dev/null +++ b/unit_tests/common/att_resources/manager_fixtures.py @@ -0,0 +1,45 @@ +"""Concrete main/draft composition used to exercise manager transfer operations.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING +from dataclasses import replace +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 +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 + +if TYPE_CHECKING: + import torch + + +class MainDraftTestResourceManager(AttentionStateManager): + 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( + 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.register_state("draft0", draft_state) + self.register_state("main", full_state(device=device)) + prepared_full_state = MLAFullAttentionState( + layout_config=MLAConfig( + logical_layer_groups=self.main_gqa_fp32_h1_d4_state.config.logical_layer_groups, + latent_dim=self.main_gqa_fp32_h1_d4_state.config.head_dim, + latent_dtype=self.main_gqa_fp32_h1_d4_state.config.dtype, + ), + page_size=self.main_gqa_fp32_h1_d4_state.page_size, + device=device, + ) + prepared_full_state.allocate_storage(self.main_gqa_fp32_h1_d4_state.capacity) + self.register_state("draft0", prepared_full_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_registration() 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..86e62cb671 --- /dev/null +++ b/unit_tests/common/att_resources/model_helpers.py @@ -0,0 +1,342 @@ +"""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.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, +) + + +if TYPE_CHECKING: + from lightllm.models.qwen3next.config import Qwen3NextResourceConfig + from lightllm.common.attention_state.allocators import FullAttKVAllocator + from collections.abc import Mapping, Sequence + from lightllm.common.attention_state.manager import AttentionStateManager + 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, + config: ( + GQAAttentionConfig + | MLAConfig + | GroupQuantizedAttentionConfig + | StaticFP8AttentionConfig + | PackedMLAAttentionConfig + | CompressedAttentionConfig + | FP8E8M0PackedMLAConfig + | 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, + ) + 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, + ) + 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.allocate_storage(token_capacity) + 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( + sorted( + (state for state in attributes.values() if isinstance(state, state_class)), + key=lambda state: state.state_name, + ) + ) + + +def main_linear_state(resources: AttentionStateManager): + (state,) = (state for state in resources.get_bounded_attention_states("main") if hasattr(state, "conv_state")) + return state + + +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"))) + 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_attention_states() + + +def create_full_attention_roles( + 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" + states = [] + for start, end in ((0, boundary), (boundary, len(layers))): + if start == end: + states.append(None) + continue + 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, + 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.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.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") + ) + return fields + return { + 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, + ssm_state_mode: str = "native", + ssm_replay_config: dict | None = None, +) -> AttentionStateManager: + """测试按既定几何参数组装通用管理器，不依赖旧模型管理器的构造/绑定流程。""" + 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.bounded_att import WindowCompressorStateConfig, create_window_compressor_state + from lightllm.models.glm5_next.config import Glm5NextResourceConfig + + 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( + 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", + ssm_state_mode=ssm_state_mode, + ssm_replay_config=ssm_replay_config, + gate_lower_bound=config.gate_lower_bound if glm else -5.0, + ) + bounded_state.initialize_page_pools( + 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)): + if state is None: + continue + manager.register_state(role, state) + if glm: + window = create_window_compressor_state( + config=WindowCompressorStateConfig( + logical_layer_ids=tuple(layer for group in state.config.logical_layer_groups for layer in group), + 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, + ), + request_capacity=request_capacity, + mtp_step=mtp_step, + device="cuda", + ) + 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.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_backend_startup.py b/unit_tests/common/att_resources/test_backend_startup.py new file mode 100644 index 0000000000..457fdfa1c3 --- /dev/null +++ b/unit_tests/common/att_resources/test_backend_startup.py @@ -0,0 +1,374 @@ +"""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.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: + 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) + request_allocations = [] + manager_creations = [] + from lightllm.common.attention_state import AttentionStateManager + + def create_empty_manager(request_capacity: int): + 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.AttentionStateManager", create_empty_manager) + + 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 已创建完成，可通过缓存查询取得资源。 + for item in models: + 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_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 + + def warmup(model: TpPartBaseModel): + assert model.backends_bound + assert all(item.req_manager is mapping for item in models) + 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)]} + 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 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.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): + assert all(item.req_manager is mapping for item in models) + 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 + 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 + + 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_state_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: + 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), + ) + for item in models: + item.bind_req_manager(req_manager=req_manager) + for index, item in enumerate(models): + 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_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) + 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 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 []) + 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) + 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"), + 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_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( + torch.tensor([0], dtype=torch.int32, device="cuda"), + torch.tensor([i + 2], 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.logical_layer_groups) == 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_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] + 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_runtime_to_pages( + req_indices=[0], + state_token_lens=[4], + page_kind="big", + 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_runtime_from_pages( + req_indices=[0], + state_token_lens=[4], + page_kind="big", + 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_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_state.py b/unit_tests/common/att_resources/test_bounded_state.py new file mode 100644 index 0000000000..e2c08b1461 --- /dev/null +++ b/unit_tests/common/att_resources/test_bounded_state.py @@ -0,0 +1,440 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING +from dataclasses import replace + +import pytest +import torch + +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.bounded_att import ( + GDNPagesPool, + GDNState, + KDAPagesPool, + KDAState, + ConvSSMStateConfig, +) +from lightllm.common.attention_state.types import ( + AxisMapping, + StateTransferLayout, + TransferRegion, +) + +if TYPE_CHECKING: + from lightllm.common.attention_state.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", + 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, + ) + state = cls( + config=config, + request_capacity=3, + mtp_step=mtp_step, + device=device, + ) + 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 + + +def page_layout(state: KDAState | GDNState): + cfg = state.config + 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( + state_name=state.state_name, + 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_mappings=( + AxisMapping( + axis="channel", + local_start=0, + external_start=0, + length=cfg.conv_channels, + ), + AxisMapping( + axis="window", + local_start=0, + external_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_mappings=( + AxisMapping( + axis="head", + local_start=0, + external_start=0, + length=cfg.num_v_heads, + ), + AxisMapping( + axis="key_dim", + local_start=0, + external_start=0, + length=cfg.key_dim, + ), + AxisMapping( + axis="value_dim", + local_start=0, + external_start=0, + length=cfg.value_dim, + ), + ), + ), + ), + ) + + +@pytest.mark.parametrize("cls,pool_cls", [(KDAState, KDAPagesPool), (GDNState, GDNPagesPool)]) +@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], + 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( + 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, + ) + 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() + 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_runtime_to_pages( + req_indices=[0], + state_token_lens=[8], + page_kind=page_kind, + page_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_runtime_from_pages( + req_indices=[1], + state_token_lens=[8], + page_kind=page_kind, + 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]) + 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=config, + request_capacity=3, + mtp_step=2, + device="cuda", + ) + 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), + ) + state.conv_state[:, 0, :, :2].copy_(conv) + state.ssm_state[:, 0].copy_(ssm) + for page_kind in ("big", "small"): + state.save_bounded_att_runtime_to_pages( + req_indices=[0], + state_token_lens=[8], + page_kind=page_kind, + page_slot_ids=[0], + ) + state.reset_bounded_att_state([0]) + for req, page_kind in enumerate(("big", "small"), start=1): + state.restore_bounded_att_runtime_from_pages( + req_indices=[req], + state_token_lens=[8], + page_kind=page_kind, + page_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"): + 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, device="cuda") + state.conv_state[:, 0].fill_(7) + state.ssm_state[:, 0:3].fill_(11) + state.save_bounded_att_runtime_to_pages( + req_indices=[0], + state_token_lens=[6], + page_kind="small", + page_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), + 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, device="cuda") + source.conv_state[:, 0].fill_(13) + source.ssm_state[:, 0].fill_(29) + source.save_bounded_att_runtime_to_pages( + req_indices=[0], + state_token_lens=[8], + page_kind="big", + page_slot_ids=[0], + ) + layout = page_layout(source) + buffer = torch.full((3, layout.byte_end), 255, dtype=torch.uint8, pin_memory=True) + source.write_state_pages_to_cpu( + page_kind="big", + page_slot_ids=[0], + cpu_page_buffers=buffer, + cpu_page_indices=[2], + layout=layout, + ) + torch.cuda.synchronize() + assert (buffer[:2] == 255).all() and (buffer[2, :16] == 255).all() + target = make_state(cls, device="cuda") + 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_runtime_from_pages( + req_indices=[2], + state_token_lens=[8], + page_kind="big", + page_slot_ids=[1], + ) + assert (target.conv_state[:, 2, :, :2] == 13).all() + assert (target.ssm_state[:, 6] == 29).all() + # 页池只保存快照数据；前缀关联由调用方维护，读回的快照可以再次写出。 + target.read_state_pages_from_cpu( + cpu_page_buffers=buffer, + cpu_page_indices=[2], + big_page_slot_ids=[0], + layout=layout, + ) + target.write_state_pages_to_cpu( + page_kind="big", + page_slot_ids=[0], + cpu_page_buffers=buffer, + cpu_page_indices=[0], + layout=layout, + ) + torch.cuda.synchronize() + assert torch.equal(buffer[0], buffer[2]) + + +@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], + 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( + size=(3, layout.byte_end), + fill_value=254, + dtype=torch.uint8, + device=device, + ) + page = allocation[1] + source.write_to_pd_page( + 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) + target.read_from_pd_page( + 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() + 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", ["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_runtime_state_from_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, 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) + 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_pages_from_dp( + 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]) + 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..b6dff3dc0e --- /dev/null +++ b/unit_tests/common/att_resources/test_capacity.py @@ -0,0 +1,1113 @@ +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 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.bounded_att import WindowCompressorState, WindowCompressorStateConfig +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.common.attention_state.full_att import ( + FullAttentionState, + GQAFullAttentionState, + GQAAttentionConfig, + MLAConfig, + create_gqa_state, + create_mla_state, + GroupQuantizedAttentionConfig, + GroupQuantizedFullAttentionState, + MLAFullAttentionState, + PackedMLAFullAttentionState, + StaticFP8FullAttentionState, +) + +if TYPE_CHECKING: + from easydict import EasyDict + + +def definition(kind: str): + if kind == "mla": + return create_mla_state( + 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, + format="packed_group128_fp8" if kind == "packed_fp8" else "packed_bf16", + ), + 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 + 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 + + +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=template, + config=template.config, + token_capacity=20, + device="cpu", + ) + actual = allocated_bytes(state, token_only=True) + 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_operator().get_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 = AttentionStateManager(request_capacity=1) + manager.register_state("main", state) + 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_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_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) + monkeypatch.setattr( + "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 = 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_registration() + if expected is None: + with pytest.raises(MemoryError, match="HOLD"): + resources.initialize_storage( + 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_storage( + args=args, + token_capacity=None, + mem_fraction=1.0, + device="cpu", + ) + else: + resources.initialize_storage( + args=args, + token_capacity=None, + mem_fraction=1.0, + device="cpu", + ) + assert resources.full_att_allocator.size == expected + assert resources.get_pd_operator().get_page_bytes() == manager.get_pd_operator().get_page_bytes() + + +def test_shared_allocator_accounts_for_distinct_state_stores(): + main, draft = definition("int4"), definition("static_head") + 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( + template=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]) +@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], + 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, + ssm_dtype=torch.float32, + ) + tail = WindowCompressorState( + config=WindowCompressorStateConfig( + logical_layer_ids=(1,), + window_size=4, + compress_stride=4, + target="indexer", + candidate_dim=3, + dtype=torch.bfloat16, + ), + request_capacity=3, + mtp_step=3, + device="cuda", + ) + reference = state_cls( + config=config, + request_capacity=3, + mtp_step=3, + device="cuda", + ) + 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 + ) + 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_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_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.get_pd_operator().write_page( + 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.buffer.zero_() + manager.get_pd_operator().read_page( + 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.buffer[:, 0]) + for position in range(4, 8): + 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( + 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] + + +@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 = AttentionStateManager(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), + ): + if state_cls is MLAFullAttentionState: + config = MLAConfig( + 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, + ) + 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_operator().get_page_bytes() == 72 + monkeypatch.setattr( + "lightllm.common.attention_state.manager.FullAttKVAllocator", partial(FullAttKVAllocator, pin_memory=False) + ) + manager.allocate_storage( + token_capacity=4, + page_size=1, + big_page_tokens=4, + device="cuda", + ) + 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_operator().get_page_bytes()), dtype=torch.uint8, device="cuda") + page = pages[1] + 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.get_pd_operator().write_page( + 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.get_pd_operator().read_page( + 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() + + +@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( + logical_layer_groups=((0,),), + global_head_num=1, + head_dim=128, + scale_dtype=torch.float32, + bits=8, + group_size=128, + ) + manager = AttentionStateManager(request_capacity=1) + prepared_full_state = AlignedFieldsState(config, page_size=1, device="cuda") + 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.allocate_storage(8) + manager.register_state("draft0", prepared_full_state) + manager.full_att_allocator = FullAttKVAllocator(8, 1, pin_memory=False) + 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_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() + 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() + assert get_layouts() is layouts + pages = torch.full( + 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.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)] + 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.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.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.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.get_pd_operator().read_page( + 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]) + 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]) +@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, + 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 + for region in super().get_pd_transfer_regions() + ) + + manager = AttentionStateManager(request_capacity=1) + if full_page_tokens is not None: + pd_args.pd_kv_page_size = full_page_tokens + config = GQAAttentionConfig( + logical_layer_groups=((0,),), + global_head_num=1, + head_dim=1, + dtype=torch.float32, + ) + manager.register_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, + ), + request_capacity=1, + device="cuda", + ) + 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_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.pd_operator.logger.warning", warnings.append) + expected = max(80, (full_page_tokens or 0) * 8) + 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.get_pd_operator().write_page( + 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.get_pd_operator().read_page( + 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) + if warnings: + assert "64 bytes" in warnings[0] and "using 80 bytes per transfer page" in warnings[0] + # 相同 manager 的容量查询命中缓存，不重复打印 warning。 + assert manager.get_pd_operator().get_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 = AttentionStateManager(request_capacity=1) + config = GQAAttentionConfig( + logical_layer_groups=((0,),), + global_head_num=1, + head_dim=3, + dtype=torch.bfloat16, + ) + 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 + ) + assert get_page_bytes() == 48 + monkeypatch.setattr( + "lightllm.common.attention_state.manager.FullAttKVAllocator", partial(FullAttKVAllocator, pin_memory=False) + ) + 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_cache_operator().get_page_layouts() + if transfer == "cpu" + else manager.get_pd_operator().get_pd_full_att_page_layouts() + ) + assert layouts["main_gqa_bf16_h1_d3_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 = AttentionStateManager(request_capacity=1) + config = GQAAttentionConfig( + logical_layer_groups=((0,),), + global_head_num=1, + head_dim=1, + dtype=torch.float32, + ) + full_state = GQAFullAttentionState(config, page_size=1, device="meta") + manager.register_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, + ) + state = cls(config, request_capacity=1, device="meta") + manager.register_state(role, state) + if with_tail: + state = WindowCompressorState( + config=WindowCompressorStateConfig( + logical_layer_ids=(0,), + window_size=4, + compress_stride=4, + target="indexer", + candidate_dim=1, + dtype=torch.float32, + ), + request_capacity=1, + mtp_step=0, + device="meta", + ) + 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_registration() + 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_h1_d1_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) + 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) + for role in ("main", "draft0"): + 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( + 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( + config=bounded_config, + request_capacity=3, + mtp_step=2, + device="meta", + ) + 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_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 + } + 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_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): + 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) + if region.field == "ssm" + else region + for region in super().get_cpu_transfer_regions() + ) + + manager = AttentionStateManager(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, + ), + request_capacity=1, + device="meta", + ) + manager.register_state("main", state) + manager.finish_state_registration() + # Conv 占 6 字节，SSM 位于 [72, 76)，State 与物理页终点均对齐到 80。 + 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 + + +@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=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 = 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_registration() + 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) + 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( + size=(3, manager.get_cpu_cache_operator().get_page_bytes()), + fill_value=165, + dtype=torch.uint8, + pin_memory=True, + ) + 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, + ) + torch.cuda.synchronize() + assert (pages[1, data_end:gqa_start] == 165).all() + assert (pages[[0, 2]] == 165).all() + 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, + ) + 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]) +@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 = 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)): + 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, + ) + state = state_cls( + config=config, + request_capacity=1, + mtp_step=2, + device="cuda", + ) + state.initialize_page_pools(big_page_capacity=2, small_page_capacity=1, pin_memory=True) + manager.register_state(role, state) + manager.finish_state_registration() + states = manager.get_bounded_attention_states() + 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_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) + state.ssm_state.fill_(i + 7) + source_slots = manager.bounded_att_big_page_allocator.alloc(1) + manager.save_bounded_att_runtime_to_pages( + req_indices=[0], + state_token_lens=[8], + page_kind="big", + page_slot_ids=source_slots, + ) + 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], + 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]) + target_slots = manager.bounded_att_big_page_allocator.alloc(1) + manager.get_cpu_cache_operator().read_req( + 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, + ) + 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_token", "pd_runtime"]) +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 = AttentionStateManager(request_capacity=1) + config = GQAAttentionConfig( + logical_layer_groups=((0,),), + global_head_num=1, + head_dim=3, + 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, + ), + request_capacity=1, + device="cpu", + ), + ) + manager.finish_state_registration() + 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_h1_d3_state + before_layout = before[state.state_name] + monkeypatch.setattr( + "lightllm.common.attention_state.manager.FullAttKVAllocator", partial(FullAttKVAllocator, pin_memory=False) + ) + manager.allocate_storage( + token_capacity=8, + page_size=4, + big_page_tokens=8, + device="cpu", + ) + after = getter() + 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_h1_d3_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() + ) + 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()) + 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): + @deepcopy_result + @lru_cache(maxsize=None) + 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 = GQAAttentionConfig( + logical_layer_groups=((0,),), + global_head_num=1, + head_dim=2, + dtype=torch.float32, + ) + manager = AttentionStateManager(request_capacity=1) + prepared_full_state = ReversedCpuRegionsState(config, page_size=4, device="cuda") + 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.allocate_storage(8) + 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() + ) + 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.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 + + 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, + device="cpu" if transfer == "cpu" else "cuda", + pin_memory=transfer == "cpu", + ) + if transfer == "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.get_cpu_cache_operator().read_req( + 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.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.get_pd_operator().read_page( + 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 new file mode 100644 index 0000000000..85988d5edd --- /dev/null +++ b/unit_tests/common/att_resources/test_compute_binding.py @@ -0,0 +1,147 @@ +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 unit_tests.common.att_resources.model_helpers import replace_logical_layer_groups + +from dataclasses import replace +from types import SimpleNamespace + +import pytest +import torch + +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 ( + 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=template, + config=replace_logical_layer_groups(template.config, ((3,), (8,))), + token_capacity=8, + device="cuda", + ) + for storage_buffer in full_state_test_buffers(state).values(): + storage_buffer.zero_() + 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_registration() + state = getattr(manager, f"{role}_{template.get_state_type()}_state") + infer = InferStateInfo() + infer.att_state_manager = manager + infer.is_draft_model = role == "draft0" + 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.latent_dim if kind == "packed_glm" else state.config.head_dim + 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( + 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( + 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() + 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 state.get_storage_layer_index(logical_layer_id=8) == 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=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( + size=(2, 1, 576), + dtype=torch.bfloat16, + device="cuda", + ) + state.write_full_att_kv(0, history_ids, history) + 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, + 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( + 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_dsv4_integration.py b/unit_tests/common/att_resources/test_dsv4_integration.py new file mode 100644 index 0000000000..4a99ef89bc --- /dev/null +++ b/unit_tests/common/att_resources/test_dsv4_integration.py @@ -0,0 +1,294 @@ +"""通过模型工厂和生产 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.models.deepseek_v4.workspace import DeepseekV4Workspace +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) + model.dsv4_workspace = DeepseekV4Workspace(model) + model.dsv4_workspace.c128_compress_scratch = torch.empty( + (1, model.batch_max_tokens, model.config["head_dim"]), dtype=torch.float32, device="cuda" + ) + manager.dsv4_workspace = model.dsv4_workspace + 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.dsv4_workspace = manager.dsv4_workspace + 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", + 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() + 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.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 = [ + 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..ae2ba27179 --- /dev/null +++ b/unit_tests/common/att_resources/test_dsv4_state.py @@ -0,0 +1,505 @@ +"""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.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 create_compressed_att_state + + definition = make_full(ratio=4) + 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 + 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_cache_operator().get_page_layouts() + assert bounded[-1].state_name not in layouts + 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.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"), + 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.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 + + +@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.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.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 new file mode 100644 index 0000000000..ec5a247576 --- /dev/null +++ b/unit_tests/common/att_resources/test_fp8_compute_binding.py @@ -0,0 +1,201 @@ +"""FP8 backends obtain calibration from the role's matching storage scheme.""" + +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 unit_tests.common.att_resources.model_helpers import replace_logical_layer_groups + +from typing import TYPE_CHECKING +from dataclasses import replace +from types import SimpleNamespace +import weakref + +import pytest +import torch + +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 +from unit_tests.common.att_resources.test_capacity import definition + +if TYPE_CHECKING: + from collections.abc import Sequence + + +@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=template, + config=replace_logical_layer_groups(template.config, ((3,), (8,))), + token_capacity=8, + device="cpu", + scales=scales, + q_scales=q_scales, + ) + ) + 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") + plain = rebuild_full_state( + template=plain, + config=plain.config, + token_capacity=8, + device="cpu", + ) + manager.register_state(role, plain) + other = rebuild_full_state( + template=template, + config=replace(template.config, logical_layer_groups=((20,), (21,)), global_head_num=4), + token_capacity=8, + device="cpu", + scales=torch.full((2, 8 if kind == "static_head" else 2), 99.0), + q_scales=torch.full((2, 4), 99.0) if kind == "static_head" else None, + ) + manager.register_state(role, other) + 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_state_manager = manager + infer.is_draft_model = role == "draft0" + 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="cpu") + 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="cpu") + + def allocate(shape: Sequence[int], dtype: torch.dtype, device: str | torch.device | int): + return torch.empty(shape, dtype=dtype, device="cpu") + + if prefill: + att.prefill_att( + q=q, + k=k, + v=v, + att_control=AttControl(logical_layer_id=8), + alloc_func=allocate, + ) + else: + att.decode_att( + q=q, + k=k, + v=v, + att_control=AttControl(logical_layer_id=8), + 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.skipif(not torch.cuda.is_available(), reason="CUDA required") +@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.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 + 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, + scale_mode=scale_mode, + ), + 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.allocate_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(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) + 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_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..8e2ee61214 --- /dev/null +++ b/unit_tests/common/att_resources/test_fp8_e8m0_packed_mla_state.py @@ -0,0 +1,235 @@ +"""独立 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, + PackedMLALogicalConfig, + FullAttentionState, + PackedMLAFullAttentionState, + create_packed_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_packed_mla_state( + config=PackedMLALogicalConfig( + logical_layer_groups=((0, 4), (8,)), + latent_dim=512, + format="packed_group64_e8m0", + tp_rank=tp_rank, + tp_world_size=2, + ), + page_size=4, + device=device, + ) + 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 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 + 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="latent 576"): + create_packed_mla_state( + config=PackedMLALogicalConfig( + logical_layer_groups=((0,),), + latent_dim=512, + format="packed_group128_fp8", + ), + device="meta", + ) + with pytest.raises(AssertionError, match="512-dimensional"): + create_packed_mla_state( + config=PackedMLALogicalConfig( + logical_layer_groups=((0,),), + latent_dim=576, + format="packed_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.packed_mla.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_cpu_kernel.py b/unit_tests/common/att_resources/test_full_att_cpu_kernel.py new file mode 100644 index 0000000000..0197b17daa --- /dev/null +++ b/unit_tests/common/att_resources/test_full_att_cpu_kernel.py @@ -0,0 +1,412 @@ +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.attention_state.types import AxisMapping +from lightllm.common.attention_state.full_att import GQAAttentionConfig, 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 = GQAAttentionConfig( + 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") +@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( + 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] + 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( + 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)) + 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( + AxisMapping( + axis="head", + local_start=1, + external_start=0, + length=1, + ) + if item.axis == "head" + else item + for item in region.axis_mappings + ) + 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") + 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( + 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() + 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( + 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( + 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: + 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( + token_indices=ids, + cpu_page_buffers=pages, + cpu_page_indices=[0], + layout=layout, + ) + else: + 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") +@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( + 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() + 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.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) + with pytest.raises(AssertionError, match="token indices must reside on CUDA"): + 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, + 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_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_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_page( + 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_dp_transfer.py b/unit_tests/common/att_resources/test_full_att_dp_transfer.py new file mode 100644 index 0000000000..78bc924405 --- /dev/null +++ b/unit_tests/common/att_resources/test_full_att_dp_transfer.py @@ -0,0 +1,221 @@ +"""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.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 + + +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]) +@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_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_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() 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..8c11087762 --- /dev/null +++ b/unit_tests/common/att_resources/test_full_att_transfer.py @@ -0,0 +1,338 @@ +"""各 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.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 + +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": "packed", + "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( + state_name=source.state_name, + 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_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() + + # 比较整行原始字节：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( + 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) + 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_page( + token_indices=source_ids, + pd_page_buffers=pages, + pd_page_index=1, + layout=layout, + ) + state.read_from_pd_page( + 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( + state_name=state.state_name, + 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( + 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=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): + 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=token_indices, + cpu_page_buffers=pages[:, 16:-16], + cpu_page_indices=page_indices, + layout=layout, + cpu_page_ready_flags=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..1339ebb4b0 --- /dev/null +++ b/unit_tests/common/att_resources/test_full_state.py @@ -0,0 +1,283 @@ +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.attention_state.full_att import ( + FullAttentionState, + GQAFullAttentionState, + MLAFullAttentionState, + GQAAttentionConfig, + MLAConfig, +) +from lightllm.common.attention_state.types import ( + 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, + ) + 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 + + +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_mappings=( + AxisMapping( + axis="head", + local_start=0, + external_start=0, + length=1, + ), + AxisMapping( + axis="dim", + local_start=0, + external_start=0, + length=4, + ), + ), + ) + ) + offset += 128 + return StateTransferLayout( + state_name=state.state_name, + 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( + 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_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_page( + pd_page_buffers=pages, + token_indices=ids, + pd_page_index=1, + layout=layout, + ) + return + state.write_to_pd_page( + 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_page( + 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() + + +@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( + 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( + 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() + + +@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( + size=(layout.byte_end + 256,), + fill_value=53, + dtype=torch.uint8, + device="cuda", + ) + with pytest.raises(RuntimeError, match="shape"): + 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, + 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 = GQAAttentionConfig( + 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.allocate_storage(12) + target = GQAFullAttentionState(config, page_size=4, device=device) + 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_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_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, logical_layer_groups=config.logical_layer_groups[:1]), page_size=4, device=device + ) + 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_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..c943a78bee --- /dev/null +++ b/unit_tests/common/att_resources/test_full_storage_initialization.py @@ -0,0 +1,71 @@ +"""Real model initialization, compute writes, and profiled allocation.""" + +from __future__ import annotations + +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.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.usefixtures("runtime"), +] + + +@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( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + mode: str, + kind: str, + encoding: str, + model_cls: type[TpPartBaseModel], +): + _, (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 + 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)) + 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_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_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_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_layout_identity.py b/unit_tests/common/att_resources/test_gqa_layout_identity.py new file mode 100644 index 0000000000..f4c3174f44 --- /dev/null +++ b/unit_tests/common/att_resources/test_gqa_layout_identity.py @@ -0,0 +1,114 @@ +from dataclasses import replace +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.full_att import GQAAttentionConfig, create_gqa_state +from lightllm.common.basemodel.attention.base_att import AttControl + + +@pytest.mark.parametrize("kv_type", ["None", "int4kv", "int8kv", "fp8kv_spt", "fp8kv_sph"]) +def test_gqa_layouts_coexist_and_match_across_tp_and_registration_order(kv_type, pd_args, monkeypatch): + config = GQAAttentionConfig(logical_layer_groups=((0,),), global_head_num=8, head_dim=128, dtype=torch.bfloat16) + configs = [ + config, + replace(config, logical_layer_groups=((1,),), global_head_num=16), + replace(config, logical_layer_groups=((2,),), head_dim=256), + ] + source, target = (AttentionStateManager(request_capacity=0) for _ in range(2)) + for manager, items, tp in ((source, configs, 1), (target, reversed(configs), 2)): + for item in items: + state = create_gqa_state(replace(item, tp_rank=tp - 1, tp_world_size=tp), kv_type=kv_type, device="meta") + manager.register_state("main", state) + manager.finish_state_registration() + states = source.get_full_attention_states("main") + assert len({state.state_name for state in states}) == 3 + for layer in range(3): + state = source.get_state("main", "gqa", layer) + peer = target.get_state("main", "gqa", layer) + assert state.state_name == peer.state_name + assert state.config.global_head_num == configs[layer].global_head_num + assert state.config.head_dim == configs[layer].head_dim + for getter in ( + lambda m: m.get_cpu_cache_operator().get_page_layouts(), + lambda m: m.get_pd_operator().get_pd_full_att_page_layouts(), + ): + layouts, peer_layouts = getter(source), getter(target) + assert set(layouts) == set(peer_layouts) == {state.state_name for state in states} + ordered = sorted(layouts.values(), key=lambda layout: layout.byte_start) + assert all(a.byte_end <= b.byte_start for a, b in zip(ordered, ordered[1:])) + copied = [] + for state in states: + monkeypatch.setattr( + state, + "copy_tokens_from_dp", + lambda *, remote_dp_full_att_state, **_: copied.append(remote_dp_full_att_state), + ) + source.get_dp_shared_operator().copy_from_dp( + remote_att_manager=target, + src_token_indices=torch.tensor([0]), + dst_token_indices=torch.tensor([1]), + src_big_page_ids=[], + dst_big_page_ids=[], + src_req_idx=0, + dst_req_idx=1, + ) + assert copied == [getattr(target, state.state_name) for state in states] + + +@pytest.mark.parametrize("prefill", [False, True]) +def test_fa3_reads_descale_for_each_head_layout(prefill, monkeypatch): + from lightllm.common.basemodel.attention.fa3 import fp8 as fa3_fp8 + + manager = AttentionStateManager(request_capacity=2) + for layer, heads in enumerate((2, 4)): + config = GQAAttentionConfig( + logical_layer_groups=((layer,),), global_head_num=heads, head_dim=128, dtype=torch.bfloat16 + ) + state = create_gqa_state(config, kv_type="fp8kv_sph", device="cpu") + state.scales = torch.arange(1, 2 * heads + 1, dtype=torch.float32).view(1, -1) + state.q_scales = torch.arange(1, heads + 1, dtype=torch.float32).view(1, heads) + manager.register_state("main", state) + manager.finish_state_registration() + # prefill 必须按 b_seq_len 取 batch，不能读这个未设置的请求 batch。 + infer = SimpleNamespace( + att_state_manager=manager, + att_state_role="main", + batch_size=None, + b_seq_len=torch.ones(2, dtype=torch.int32), + max_q_seq_len=1, + ) + cls = fa3_fp8.Fp8Fa3PrefillAttState if prefill else fa3_fp8.Fp8Fa3DecodeAttState + parent = fa3_fp8.Fa3PrefillAttState if prefill else fa3_fp8.Fa3DecodeAttState + monkeypatch.setattr(parent, "init_state", lambda self: None) + monkeypatch.setattr(fa3_fp8, "q_per_head_static_fp8_quant", lambda q, scales: q) + observed = {} + + def attend(**kwargs): + observed.update(kwargs) + return kwargs["q"] + + monkeypatch.setattr(fa3_fp8, "flash_attn_with_kvcache", attend) + att = cls(backend=None, _infer_state_weak_ref=lambda: infer) + att.b_att_seq_len = torch.ones(3, dtype=torch.int32) + att.page_table = att.cu_seqlens_q = att.cu_seqlens_k = None + att.decode_max_q_seq_len = 1 + att.causal = True + att.init_state() + batch_size = infer.b_seq_len.shape[0] if prefill else att.b_att_seq_len.shape[0] + for state in manager.get_full_attention_states("main"): + heads = state.config.head_num + layer = state.config.logical_layer_groups[0][0] + q = k = v = torch.zeros((2, heads, state.config.head_dim)) + observed.clear() + run = att.prefill_att if prefill else att.decode_att + run(q=q, k=k, v=v, att_control=AttControl(logical_layer_id=layer)) + row = state.scales[0] + k_descale = observed["k_descale"] + v_descale = observed["v_descale"] + assert k_descale.data_ptr() == row[:heads].data_ptr() + assert v_descale.data_ptr() == row[heads:].data_ptr() + torch.testing.assert_close(k_descale, row[:heads].expand(batch_size, heads)) + torch.testing.assert_close(v_descale, row[heads:].expand(batch_size, heads)) 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..2870ca126e --- /dev/null +++ b/unit_tests/common/att_resources/test_gqa_pd_kernel.py @@ -0,0 +1,283 @@ +"""GQA PD 直接搬运：原始位模式、物理页选择、TP 切片及有效 token 范围。""" + +from __future__ import annotations + +from dataclasses import replace +from typing import TYPE_CHECKING + +import pytest +import torch + +from lightllm.common.attention_state.full_att import GQAAttentionConfig, GQAFullAttentionState +from lightllm.common.attention_state.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=GQAAttentionConfig( + logical_layer_groups=((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.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): + 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( + 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), + 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_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=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_page( + 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=GQAAttentionConfig( + 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", + ) + state.allocate_storage(token_capacity=8) + state.k_storage.fill_(19) + state.v_storage.fill_(19) + regions = state.get_pd_transfer_regions() + layout = 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, + ) + 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_page( + token_indices=ids, + pd_page_buffers=pages, + pd_page_index=1, + layout=layout, + ) + state.read_from_pd_page( + 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=GQAAttentionConfig( + logical_layer_groups=((0,), (2,)), + global_head_num=2, + head_dim=128, + dtype=torch.bfloat16, + ), + page_size=4, + device="cuda:1", + ) + state.allocate_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( + state_name=state.state_name, + 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_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_page( + 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..ac19dc39ab --- /dev/null +++ b/unit_tests/common/att_resources/test_gqa_transfer.py @@ -0,0 +1,192 @@ +"""GQA 共用搬运 kernel 在受限 grid 下完整覆盖多轮、跨页和尾块数据。""" + +from typing import TYPE_CHECKING + +import pytest +import torch + +from lightllm.common.attention_state.full_att import GQAAttentionConfig, GQAFullAttentionState +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: + 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 = GQAAttentionConfig( + 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.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( + state_name=source.state_name, + 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() + + +@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.attention_state.full_att import gqa_transfer + + pd_args.pd_kv_page_size = 2 + config = GQAAttentionConfig( + 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.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: + 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( + state_name=source.state_name, + 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() 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..fb9ce7ba5b --- /dev/null +++ b/unit_tests/common/att_resources/test_hybrid_layer_mapping.py @@ -0,0 +1,130 @@ +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("mode", ["native", "replay"]) +@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, + mode: str, +) -> 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", + mtp_step=2, + ssm_state_mode=mode, + ssm_replay_config={"cache_len": 8} if mode == "replay" else {}, + ) + 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", + ) + if mode == "native": + stride = bounded_state.mtp_step + 1 + assert bounded_state.ssm_state.shape[1] == bounded_state.request_capacity * stride + if mode == "replay": + assert bounded_state.capacity == 8 + 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 new file mode 100644 index 0000000000..0f55fe5986 --- /dev/null +++ b/unit_tests/common/att_resources/test_layer_groups.py @@ -0,0 +1,164 @@ +"""逻辑层共享物理 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.attention_state import AttentionStateManager +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, + replace_logical_layer_groups, +) +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_logical_layer_groups(template.config, ((0, 4), (8, 11))) + state = rebuild_full_state( + template=template, + config=config, + token_capacity=8, + device="cpu", + ) + 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) + + 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] + + 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(kind: str, groups: tuple) -> None: + with pytest.raises(AssertionError): + replace_logical_layer_groups(definition(kind).config, groups) + + +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,) = 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，保持四份物理层存储。 + assert state.config.logical_layer_groups == ((10,), (11,), (12, 14), (13, 15)) + 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") +@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=template, + config=replace_logical_layer_groups(template.config, ((0, 4), (8,))), + token_capacity=8, + device="cuda", + ) + target = rebuild_full_state( + template=template, + config=replace_logical_layer_groups(template.config, ((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 = 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") + if transfer == "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_cache_operator().get_page_layouts(), + ) + 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_cache_operator().get_page_layouts(), + ) + else: + 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_operator().get_pd_full_att_page_layouts(), + ) + 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_operator().get_pd_full_att_page_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_layer_state_selection.py b/unit_tests/common/att_resources/test_layer_state_selection.py new file mode 100644 index 0000000000..52337199a7 --- /dev/null +++ b/unit_tests/common/att_resources/test_layer_state_selection.py @@ -0,0 +1,216 @@ +"""按角色、层归属及方案查询资源，不在推理层保存绑定状态。""" + +from contextlib import nullcontext +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.attention_state import AttentionStateManager +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_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 + + +def test_creation_does_not_bind_layers(monkeypatch: pytest.MonkeyPatch) -> None: + 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.max_req_num = 0 + model.torch_memory_saver = SimpleNamespace(region=lambda **_: nullcontext()) + state = GQAFullAttentionState( + layout_config=GQAAttentionConfig( + logical_layer_groups=((0,),), + global_head_num=1, + head_dim=128, + dtype=torch.bfloat16, + ), + device="meta", + ) + monkeypatch.setattr( + LlamaTpPartModel, + "create_full_att_states", + classmethod(lambda cls, **kwargs: [state]), + ) + + model.prepare_attention_states() + + assert model.att_state_manager.main_gqa_bf16_h1_d128_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 = 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", + ), + ) + 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_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 = [] + for logical_layer_id in range(4): + layer = TransformerLayerInfer.__new__(TransformerLayerInfer) + layer.layer_num_ = logical_layer_id + model.layers_infer.append(layer) + + model.prepare_attention_states() + + assert all(not hasattr(layer, "full_att_state_name") for layer in model.layers_infer) + manager.finish_state_registration() + state = manager.get_state(role, "gqa", 0) + assert state.state_name == f"{role}_gqa_bf16_h1_d128_state" + assert manager.get_state(role, "gqa", 2) is state + assert manager.get_state(role, "gqa", 3).state_name == f"{role}_gqa_fp16_h1_d128_state" + + +@pytest.mark.parametrize("finalized", [False, True]) +def test_prefix_and_logical_layer_select_quantized_state(finalized: bool) -> None: + from dataclasses import replace + + manager = AttentionStateManager(request_capacity=1) + plain = definition("gqa") + quantized = definition("int8") + quantized = type(quantized)( + layout_config=replace_logical_layer_groups(quantized.config, ((10, 12),)), + device="meta", + ) + 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 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_state("main", "gqa", 99) + + +def test_selection_requires_finished_creation_and_reuses_cache() -> None: + manager = AttentionStateManager(request_capacity=1) + plain = definition("gqa") + manager.register_state("main", plain) + logical_layer_id = plain.config.logical_layer_groups[0][0] + 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 = AttentionStateManager(request_capacity=1) + plain = definition("gqa") + 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_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 = 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_state("main", state.get_state_type(), 3) is state + + +def test_mla_selection_excludes_same_layer_tail() -> None: + from dataclasses import replace + + manager = AttentionStateManager(request_capacity=3) + tail = make_window() + template = definition("mla") + state = type(template)( + layout_config=replace_logical_layer_groups(template.config, ((3, 7),)), + device="meta", + ) + 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_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 + ) + 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) + for model_name in ("draft1", "draft10"): + manager.register_state(model_name, definition("gqa")) + manager.finish_state_registration() + assert manager.get_state(role, "gqa", 0).state_name == f"{role}_gqa_bf16_h2_d16_state" + + +@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_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})) + state = type(template)(config=config, request_capacity=3, device="cpu") + manager.register_state("main", state) + states.append(state) + manager.finish_state_registration() + logical_layer_id = template.config.logical_layer_ids[0] + for state in states: + 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_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 new file mode 100644 index 0000000000..9405b46c89 --- /dev/null +++ b/unit_tests/common/att_resources/test_manager_helpers.py @@ -0,0 +1,316 @@ +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.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 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 + +if TYPE_CHECKING: + 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"): + return MainDraftTestResourceManager(device) + + +def layouts_for(manager: AttentionStateManager, kind: str = "cpu"): + layouts = [] + if kind != "att_state": + layouts += [layout_for(state) for state in (manager.main_gqa_fp32_h1_d4_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) + ] + + offset = 0 + result = {} + for layout in layouts: + state_end = offset + layout.byte_end + result[layout.state_name] = 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: AttentionStateManager): + for i, state in enumerate((manager.main_gqa_fp32_h1_d4_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_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) + state.ssm_state[:, 3:6].fill_(47 + i) + slots = manager.bounded_att_big_page_allocator.alloc(2) + manager.save_bounded_att_runtime_to_pages( + req_indices=[0, 1], + state_token_lens=[4, 8], + page_kind="big", + page_slot_ids=slots, + ) + if manager.main_kda_conv_fp32_ssm_fp32_state.conv_state.is_cuda: + torch.cuda.synchronize() + 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(device="cuda") + slots = init_source(manager) + manager.restore_bounded_att_runtime_from_pages( + req_indices=[2], + state_token_lens=[4], + page_kind="big", + 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() + assert (state.ssm_state[:, 6] == 31 + i).all() + + def fail(*args: object, **kwargs: object): + raise RuntimeError("simulated copy failure") + + 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_runtime_to_pages( + req_indices=[0, 1], + state_token_lens=[4, 8], + page_kind="big", + page_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 + + +@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() + calls = [] + + def record_lengths( + req_indices: list[int], + state_token_lens: list[int], + page_kind: PageKind, + page_slot_ids: list[int], + ) -> None: + 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) + getattr(manager, operation)( + req_indices=[0, 2], + state_token_lens=[17, 29], + page_kind="big", + page_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() + 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_state_pages_from_cpu", lambda **kwargs: None) + restored = [] + + def record_restore( + req_indices: list[int], + state_token_lens: list[int], + page_kind: PageKind, + page_slot_ids: list[int], + ) -> None: + restored.append((req_indices, state_token_lens, page_kind, page_slot_ids)) + + monkeypatch.setattr(manager, "restore_bounded_att_runtime_from_pages", record_restore) + # 只加载绝对位置 [64, 70)，本批长度 6；恢复位置必须是 70，且末尾碎页使用 LOAD。 + 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"), + 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_page", "read_from_pd_page"]) +def test_pd_bounded_only_forwards_prefix_length(operation: str, monkeypatch: pytest.MonkeyPatch): + manager = composition() + 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.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", + 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") +@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"), composition("cuda") + 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_runtime_to_pages( + req_indices=[page_count - 1], + state_token_lens=[token_count], + page_kind="small", + page_slot_ids=small_slots, + ) + torch.cuda.synchronize() + source.get_cpu_cache_operator().write_req( + 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() + + # 保留末槽给 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.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"), + state_token_len=token_count, + 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.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") + big_slots = init_source(source) + small_slots = source.bounded_att_small_page_allocator.alloc(1) + source.save_bounded_att_runtime_to_pages( + req_indices=[1], + state_token_lens=[6], + page_kind="small", + page_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( + 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.get_cpu_cache_operator().write_req( + 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() + 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.get_cpu_cache_operator().write_req( + token_indices=token_ids, + cpu_page_buffers=pages, + cpu_page_indices=page_ids, + layouts=layouts, + cpu_page_ready_flags=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..9544084bb6 --- /dev/null +++ b/unit_tests/common/att_resources/test_model_calibration.py @@ -0,0 +1,261 @@ +"""统一标定阶段不依赖 State 注册顺序，不在布局定义阶段加载参数。""" + +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.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 +from lightllm.server.core.objs.start_args_type import StartArgs + +if TYPE_CHECKING: + from pathlib import Path + from easydict import EasyDict + + +@pytest.mark.parametrize("encoding", ["fp8kv_sph", "fp8kv_spt"]) +def test_calibration_binds_all_model_states_after_creation(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 + + # 反向创建、注册，标定仍按 main 和 draft 的数值顺序绑定。 + manager = AttentionStateManager(request_capacity=1) + for index in reversed(range(len(models))): + model = models[index] + 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( + 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.state_name is None + assert state.scales is None and state.q_scales is None + 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_state_role) + 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.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": + 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.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) + 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 = AttentionStateManager(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=GQAAttentionConfig( + 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.register_state(role, state) + if encoding != "None": + expected.append((state, row)) + with pytest.raises(AssertionError, match="State registration"): + manager.load_full_att_calibration(args=args) + 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) + 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.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)) + 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.attention_state.full_att import GQAAttentionConfig, create_gqa_state + + state = create_gqa_state( + config=GQAAttentionConfig( + logical_layer_groups=((0, 1), (2,)), + global_head_num=2, + head_dim=16, + dtype=torch.bfloat16, + ), + kv_type=encoding, + device="cpu", + ) + manager = AttentionStateManager(request_capacity=1) + manager.register_state("main", state) + manager.finish_state_registration() + with pytest.raises(AssertionError, match="calibration must be bound"): + state.allocate_storage(4) + if encoding == "fp8kv_sph": + 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.allocate_storage(4) + assert state.k_storage.shape[0] == 2 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..6e3dfbc92c --- /dev/null +++ b/unit_tests/common/att_resources/test_model_config_normalization.py @@ -0,0 +1,235 @@ +"""模型启动和 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] + + +@pytest.mark.parametrize( + "model_type,text_model_type,wrapper", + [ + ("internvl_chat", kind, "llm_config") + for kind in ("phi3", "internlm2", "llama", "qwen2", "deepseek_v2", "qwen3", "qwen3_moe") + ] + + [ + ("neo_chat", "qwen3", "llm_config"), + ("neo_chat", "qwen3_moe", "llm_config"), + ("qwen3_omni_moe", "qwen3_moe", "thinker_config"), + ("qwen3_vl", "qwen3", "text_config"), + ("llama", "llama", None), + ], +) +def test_multimodal_cpu_layout_matches_runtime(model_type, text_model_type, wrapper, tmp_path, monkeypatch): + from dataclasses import asdict + from lightllm.common.attention_state import AttentionStateManager + from lightllm.models import get_model_class + from lightllm.utils.cpu_cache_layout import build_cpu_cache_layouts + from lightllm.utils.envs_utils import get_env_start_args + from lightllm.utils.llm_utils import get_llm_model_class + + text = dict( + model_type=text_model_type, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + hidden_size=256, + kv_lora_rank=64, + qk_rope_head_dim=32, + ) + raw = ( + text + if wrapper is None + else { + "model_type": model_type, + wrapper: {"text_config": text} if wrapper == "thinker_config" else text, + } + ) + original = deepcopy(raw) + (tmp_path / "config.json").write_text(json.dumps(raw)) + args = StartArgs(model_dir=str(tmp_path), data_type="bfloat16", enable_cpu_cache=True, cpu_cache_token_page_size=8) + monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(asdict(args))) + get_env_start_args.cache_clear() + get_llm_model_class.cache_clear() + try: + layouts = build_cpu_cache_layouts() + cls = get_model_class(raw) + model = cls.__new__(cls) + model.weight_dir_, model.finetune_config = str(tmp_path), None + model._init_config() + normalized = cls.normalize_model_config(raw) + assert model.config == normalized + assert cls.normalize_model_config(normalized) == normalized + assert raw == original + manager = AttentionStateManager(request_capacity=0) + for state in cls.create_full_att_states( + model_config=model.config, + args=args, + model_logical_layer_ids=(0, 1), + ): + manager.register_state("main", state) + manager.finish_state_registration() + assert layouts and layouts == manager.get_cpu_cache_operator().get_page_layouts() + finally: + get_llm_model_class.cache_clear() + get_env_start_args.cache_clear() 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..09e88b4a93 --- /dev/null +++ b/unit_tests/common/att_resources/test_model_initialization.py @@ -0,0 +1,185 @@ +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.attention_state.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_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 + 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_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) + 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.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_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( + 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_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_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") + 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 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 + assert resources.bounded_att_small_page_allocator.size == 2 + assert not hasattr(model, "mtp_state_cursors") + linear_state = main_linear_state(resources) + assert linear_state.mtp_state_index is not None + assert linear_state.mtp_enabled == bool(mtp) + if mode == "prefill": + 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 + 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 + assert draft_state.capacity == full_state.capacity + 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 new file mode 100644 index 0000000000..a571817fae --- /dev/null +++ b/unit_tests/common/att_resources/test_model_resources.py @@ -0,0 +1,214 @@ +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 + +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.attention_state.allocators import FullAttKVAllocator +from unit_tests.common.att_resources.model_helpers import create_full_attention_roles +from lightllm.common.attention_state.full_att import ( + GQAFullAttentionState, + GQAAttentionConfig, +) +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 +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) + 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_head_dim=512 if glm else 16, + global_linear_k_heads=gk, + global_linear_v_heads=gv, + head_linear_k_dim=8, + head_linear_v_dim=8, + conv_kernel_size=3, + conv_state_dtype=torch.bfloat16, + ssm_state_dtype=torch.float32, + full_attention_interval=4 if glm else 2, + all_layer_num=4, + ) + if glm: + full = PackedMLAFullAttentionState( + PackedMLAAttentionConfig( + 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, + ) + full.allocate_storage(16) + else: + full = GQAFullAttentionState( + GQAAttentionConfig( + 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.allocate_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, + full_state=main, + full_allocator=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=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 + 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_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_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_runtime_to_pages( + req_indices=[0], + state_token_lens=[8], + page_kind="big", + 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())) + 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.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_cache_operator().get_page_layouts(), + big_page_slot_ids=[0], + ) + torch.cuda.synchronize() + for rank, (resources, full) in enumerate(targets): + if glm: + resources.main_compressor_indexer_w4_s4_bf16_state.buffer.fill_(13) + slot = resources.bounded_att_big_page_allocator.alloc(1)[0] + resources.get_cpu_cache_operator().read_req( + cpu_page_buffers=buffer, + cpu_page_indices=[1], + token_indices=torch.arange( + start=8, + end=16, + dtype=torch.long, + device="cuda", + ), + state_token_len=8, + layouts=resources.get_cpu_cache_operator().get_page_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: + # 无重叠窗口在组边界不依赖旧历史，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() + + +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( + glm=True, + tp=1, + rank=0, + mtp=0, + ) + 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 new file mode 100644 index 0000000000..70970756df --- /dev/null +++ b/unit_tests/common/att_resources/test_model_storage_definition.py @@ -0,0 +1,276 @@ +"""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.cpu_cache_layout import build_cpu_cache_layouts +from lightllm.utils.envs_utils import get_env_start_args +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: + 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, + ), + ( + 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, + ), +] + + +@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_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_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_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_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") + + +@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 = 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_registration() + args = StartArgs( + page_size=4, + run_mode="prefill", + pd_kv_page_size=8, + pd_kv_page_num=3, + ) + 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 + 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.attention_state.manager.get_current_device_id", lambda: 3) + monkeypatch.setattr( + "lightllm.common.attention_state.manager.FullAttKVAllocator", partial(FullAttKVAllocator, pin_memory=False) + ) + template = rebuild_full_state( + template=state, + config=state.config, + token_capacity=0, + device="meta", + ) + 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, + device="cpu", + ) + assert resources.full_att_allocator.size == 28 + assert resources.get_pd_operator().get_page_bytes() == reference.get_pd_operator().get_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.attention_state.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_state_manager.finish_state_registration() + runtime_layouts = model.att_state_manager.get_cpu_cache_operator().get_page_layouts() + assert set(cpu_layouts) == set(runtime_layouts) == {"main_gqa_bf16_h8_d16_state", "main_gqa_fp16_h8_d16_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 new file mode 100644 index 0000000000..ea114d2ba9 --- /dev/null +++ b/unit_tests/common/att_resources/test_mtp_state_commit.py @@ -0,0 +1,114 @@ +"""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.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 = 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_registration() + 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"), + torch.tensor([3], dtype=torch.int32, device="cuda"), + role="main", + ) + 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_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"), + torch.tensor([1], dtype=torch.int32, device="cuda"), + role="draft0", + ) + 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") +def test_commit_without_recurrent_state_and_invalid_role(): + 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"), + 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"), + 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"), + 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_state_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=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] + assert accepted_index.tolist() == expected_mask + 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 new file mode 100644 index 0000000000..f66767b9f4 --- /dev/null +++ b/unit_tests/common/att_resources/test_network_transport.py @@ -0,0 +1,320 @@ +"""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.attention_state.manager import AttentionStateManager + 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 + + +@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: AttentionStateManager, + 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_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 + + +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( + 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( + 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 + + transporter = NixlKVTransporter(os.getpid(), device, buffer) + connection.send( + ( + "metadata", + PDAgentMetadata( + agent_name=transporter.agent_name, + agent_metadata=transporter.agent_metadata, + num_pages=transporter.num_pages, + page_reg_desc=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_compressor_indexer_w4_s4_bf16_state.buffer) + ): + tensor.copy_(torch.arange(tensor.numel(), device=device).reshape(tensor.shape) % (23 + i * 7)) + 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): + copy_page( + 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), + task_kind=kind, + req_idx=1, + state_token_len=task.end_kv_index, + write_to_pd_page=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=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) + 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=transporter, + stage="done", + pending=pending, + key=request.get_key(), + ) + assert done.get_key() == request.get_key() + if backend == "nccl": + 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( + 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), + task_kind=kind, + req_idx=2, + state_token_len=request.end_kv_index, + write_to_pd_page=False, + ) + copy_stream.synchronize() + 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() + 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() + 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..d7d285493b --- /dev/null +++ b/unit_tests/common/att_resources/test_no_attention_draft.py @@ -0,0 +1,84 @@ +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.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 +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 = 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") + 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_state_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_attention_states() + resources.finish_state_registration() + draft._init_att_backend() + draft._init_att_backend1() + 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)) + + 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_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 + 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"]) +def test_attention_draft_prepares_its_own_state_before_shared_allocation(mode: str): + draft, resources = draft_model(mode) + 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.state_name == f"draft0_{state.get_state_type()}_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 new file mode 100644 index 0000000000..032f7674eb --- /dev/null +++ b/unit_tests/common/att_resources/test_page_pool_storage.py @@ -0,0 +1,167 @@ +from __future__ import annotations + +import gc +from typing import TYPE_CHECKING + +import pytest +import torch +import torch.multiprocessing as mp + +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 + 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( + 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) + 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( + 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) + 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..e146261f11 --- /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.attention_state.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..de13eb28fc --- /dev/null +++ b/unit_tests/common/att_resources/test_quantized_full_state.py @@ -0,0 +1,388 @@ +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 typing import TYPE_CHECKING +from math import prod + +import pytest +import torch + +from lightllm.common.attention_state.full_att import ( + GroupQuantizedAttentionConfig, + GroupQuantizedFullAttentionState, + StaticFP8AttentionConfig, + StaticFP8FullAttentionState, + PackedMLAAttentionConfig, + PackedMLALogicalConfig, + PackedMLAFullAttentionState, +) +from lightllm.common.attention_state.types import ( + AxisMapping, + StateTransferLayout, + TransferRegion, +) + +if TYPE_CHECKING: + from collections.abc import Sequence + from lightllm.common.attention_state.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( + logical_layer_groups=((1,), (3,)), + global_head_num=2, + head_dim=16, + scale_dtype=torch.float16, + bits=int(kind[-1]), + ) + prepared_full_state = GroupQuantizedFullAttentionState(config, page_size=4, device=device) + prepared_full_state.allocate_storage(8) + return prepared_full_state + if kind.startswith("fp8_"): + per_head = kind == "fp8_head" + config = StaticFP8AttentionConfig( + logical_layer_groups=((1,), (3,)), + global_head_num=2, + 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 + ) + 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.allocate_storage(8) + return prepared_full_state + fp8 = kind == "packed_fp8" + config = PackedMLAAttentionConfig( + 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) + 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_mappings=tuple( + AxisMapping( + axis=axis, + local_start=0, + external_start=0, + length=dim, + ) + for axis, dim in zip(field.axes, shape) + if axis not in ("layer", "token") + ), + ) + ) + offset += prod(shape) * itemsize + return StateTransferLayout( + state_name=state.state_name, + 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=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( + 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( + 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_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_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_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_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_page( + 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) + + 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_dp( + src_token_indices=src_indices, + dst_token_indices=dst_indices, + remote_dp_full_att_state=source, + ) + return + target.copy_tokens_from_dp( + src_token_indices=src_indices, + dst_token_indices=dst_indices, + remote_dp_full_att_state=source, + ) + assert_token_fields_equal( + source=source, + source_ids=[0, 1], + target=target, + target_ids=[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.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") + state.write_full_att_kv(1, ids, kv) + if isinstance(state, GroupQuantizedFullAttentionState): + 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) + 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( + actual=dequant, + expected=kv.float(), + rtol=0.13, + atol=0.03, + ) + else: + view = state.get_layer_view(1) + assert (view.indexer_k == 29).all() + if state.config.dtype != torch.float8_e4m3fn: + 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() + + +@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", + ) + 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 torch.equal(state.scales, source.scales) + state.device = torch.device("cpu") + 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)) 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..c44cc68b5a --- /dev/null +++ b/unit_tests/common/att_resources/test_runtime_transport.py @@ -0,0 +1,574 @@ +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 + +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 +from contextlib import nullcontext +import pickle +import traceback +import uuid + +import pytest +import torch +import torch.multiprocessing as mp + +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 +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 = AttentionStateManager(request_capacity=2) + manager.register_state("main", state) + manager.full_att_allocator = FullAttKVAllocator(8, 4, pin_memory=False) + 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()): + 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( + 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_page=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_page=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_page=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_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( + token_indices=source_indices, + cpu_page_buffers=cpu_pages, + cpu_page_indices=[0], + layouts=managers[0].get_cpu_cache_operator().get_page_layouts(), + ) + else: + 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_cache_operator().get_page_layouts(), + ) + if device == "cuda": + torch.cuda.synchronize() + 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_cache_operator().get_page_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].get_dp_shared_operator().copy_from_dp( + remote_att_manager=managers[0], + 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].get_dp_shared_operator().copy_from_dp( + remote_att_manager=managers[0], + 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, + ) + assert_token_fields_equal( + source=src, + source_ids=[1, 3], + target=dst, + target_ids=[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=glm, + tp=1, + rank=0, + mtp=2, + ) + with torch.cuda.device(0): + target, target_full = make_resources( + glm=glm, + tp=1, + 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) + 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_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): + 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_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, + ) + 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_state_manager=target), + dp_att_state_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=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 + 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 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") +@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__) + # Some containers lack pidfd_getfd, which recent torch releases require for + # expandable-segment CUDA IPC handles. + torch.cuda.memory._set_allocator_settings("expandable_segments:False") + source_device = int(p2p) + with torch.cuda.device(source_device): + 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(): + storage_buffer[:, 4].zero_() + main_linear_state(resources).conv_state.fill_(11) + 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_() + pool.ssm_state.zero_() + slot = resources.bounded_att_big_page_allocator.alloc(1) + with torch.cuda.device(source_device): + resources.save_bounded_att_runtime_to_pages( + req_indices=[0], + state_token_lens=[8], + page_kind="big", + page_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=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)) + 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() + 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() + 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( + 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): + 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( + 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"), + task_kind=kind, + req_idx=1, + state_token_len=8, + write_to_pd_page=True, + ) + copy_page( + 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"), + task_kind=kind, + req_idx=2, + state_token_len=8, + 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()): + 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 + + +@pytest.mark.parametrize("kind", ["pd", "dp"]) +def test_ipc_readers_accept_different_pickle_lengths(kind: str, monkeypatch: pytest.MonkeyPatch) -> None: + """各读者独立导出；模拟 IPC 计数变化引起的长度差异，并验证补齐后仍可逐份读取。""" + resources = SimpleNamespace( + 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 = [] + 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) + for kind in ipc_kinds: + resources, mapping = shared_att_resources_utils.load_att_resources(0, kind=kind) + 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 + if kind == "dp": + 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.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": + 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") + ) + for window in window_states(resources): + window.reset_bounded_att_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: + 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 new file mode 100644 index 0000000000..5e838a4931 --- /dev/null +++ b/unit_tests/common/att_resources/test_state_creation.py @@ -0,0 +1,212 @@ +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 + +import pytest +import torch + +from lightllm.common.attention_state.full_att import ( + GQAAttentionConfig, + MLAConfig, + 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, +) +from lightllm.common.attention_state.bounded_att import WindowCompressorStateConfig, create_window_compressor_state + + +@pytest.mark.parametrize( + "kv_type, expected_type, storage_dtype, width", + [ + ("None", "gqa_bf16_h8_d128", torch.bfloat16, 128), + ("int4kv", "gqa_int4_group16_bf16_scale_h8_d128", torch.int8, 64), + ("int8kv", "gqa_int8_group16_bf16_scale_h8_d128", torch.int8, 128), + ("fp8kv_spt", "gqa_fp8_static_per_tensor_h8_d128", torch.uint8, 128), + ("fp8kv_sph", "gqa_fp8_static_per_head_h8_d128", 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 = GQAAttentionConfig( + 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, + ) + assert state.get_state_type() == expected_type + 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 + 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 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) + 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) + 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_group128"), + ], +) +def test_mla_creation_computes_packed_width_from_logical_dimensions( + packed: bool, + kv_type: str, + kv_dim: int, + row_width: int, + expected_type: str, +): + 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 + 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 + 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 torch.bfloat16) + if packed: + # 保持旧算子读取的物理布局：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 + + +@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, + ) + state = create_state( + config=config, + request_capacity=3, + mtp_step=2, + device="cpu", + ) + 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,) + assert state.big_pages_pool is None and state.small_pages_pool is None + + +def test_kpool_creation_preserves_pool_and_mtp_dimensions(): + config = WindowCompressorStateConfig( + logical_layer_ids=(3, 7), + window_size=4, + compress_stride=4, + target="indexer", + candidate_dim=128, + dtype=torch.bfloat16, + ) + 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.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(): + config = GQAAttentionConfig( + 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") + packed_config = PackedMLALogicalConfig( + 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") + 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", + ) 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..856b4ef4af --- /dev/null +++ b/unit_tests/common/att_resources/test_state_identity.py @@ -0,0 +1,356 @@ +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 unit_tests.common.att_resources.model_helpers import replace_logical_layer_groups + +from typing import TYPE_CHECKING +from dataclasses import replace + +import pytest +import torch + +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 + +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, "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 + + 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.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 + assert state.get_token_bytes() == token_bytes + 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) + + +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.register_state(role, definition("gqa")) + with pytest.raises(AssertionError, match="already bound"): + 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_h2_d16_state,) + assert manager.get_full_attention_states("draft10") == (manager.draft10_gqa_bf16_h2_d16_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.attention_state.full_att import GQAFullAttentionState + + class CustomGQAState(GQAFullAttentionState): + def get_state_type(self) -> str: + return "gqa_custom_fp32" + + state = full_state(CustomGQAState) + 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.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_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_window_compressor_state import make_window + + 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_window() + storage = state.buffer + config = state.config + 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.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": + assert state.conv_state is storage and state.big_pages_pool is pool + else: + 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) + 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(): + 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_state_role == "draft35" + draft.mtp_previous_draft_models.append(object()) + with pytest.raises(AssertionError, match="at most 36"): + _ = draft.att_state_role + + +@pytest.mark.parametrize("role", ["draft", "draft36", "draft-1", "draft01"]) +def test_register_state_rejects_invalid_draft_numbers(role: str): + manager = AttentionStateManager(request_capacity=1) + with pytest.raises(AssertionError, match="unknown model role"): + manager.register_state(role, definition("gqa")) + + +def allocated_state(kind: str, device: str = "cpu"): + template = definition(kind) + 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") +@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.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(): + 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 = 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_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 = 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_registration() + 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_cache_operator().get_page_layouts(), + ) + 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_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_cache_operator().get_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 = 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_logical_layer_groups(template.config, ((4,), (6,)))) + manager.register_state("main", quant) + if overlap: + overlapping = allocated_state("static_head") + manager.register_state("main", overlapping) + manager.finish_state_registration() + assert manager.main_gqa_bf16_h2_d16_state is main + assert manager.draft0_gqa_bf16_h2_d16_state is draft + assert manager.main_gqa_int8_group8_bf16_scale_h2_d16_state is quant + assert manager.get_full_attention_states("draft0") == (draft,) + expected = [main, quant] + if overlap: + assert manager.main_gqa_fp8_static_per_head_h2_d16_state is overlapping + expected.append(overlapping) + 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") +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() == "gqa_bf16_h2_d16" + assert quant.get_state_type() == "gqa_int8_group8_bf16_scale_h2_d16" + 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 = AttentionStateManager(request_capacity=2) + 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") + ) + 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_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.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=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.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=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( + 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) + 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..0d96541acf --- /dev/null +++ b/unit_tests/common/att_resources/test_state_queries.py @@ -0,0 +1,189 @@ +"""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.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 + + +def test_created_state_queries_keep_the_same_objects_through_storage_initialization(monkeypatch: pytest.MonkeyPatch): + args = StartArgs(page_size=4) + definitions = [] + 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 + "_") + ) + ) + for query in ( + manager.get_full_attention_states, + manager.get_bounded_attention_states, + ): + with pytest.raises(AssertionError, match="State registration"): + query() + assert manager._get_states.cache_info().currsize == 0 + + cached = [] + allocate_storage = AttentionStateManager.allocate_storage + + 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) + allocate_storage(manager, *args, **kwargs) + + monkeypatch.setattr(AttentionStateManager, "allocate_storage", check_created_before_allocation) + monkeypatch.setattr( + "lightllm.common.attention_state.manager.FullAttKVAllocator", partial(FullAttKVAllocator, pin_memory=False) + ) + manager.finish_state_registration() + manager.initialize_storage( + 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) + 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) + 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_h1_d4_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,) + + 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 + 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() + + +def test_finished_state_creation_allows_unallocated_storage(): + manager = AttentionStateManager(request_capacity=1) + manager.register_state("main", definition("gqa")) + manager.full_att_allocator = FullAttKVAllocator(0, 4, pin_memory=False) + manager.finish_state_registration() + 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.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"): + 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.state_name) + 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.state_name) for state in source_states + ) + for state in source.get_full_attention_states(): + 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 + source = MainDraftTestResourceManager() + 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.get_dp_shared_operator().create_ipc_copy() + return + 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 + for state in source.get_bounded_attention_states(): + 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 + + +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.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 + 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, + ): + 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 + 2 + + +def test_clearing_query_cache_releases_manager_and_states(): + manager = MainDraftTestResourceManager() + manager.get_full_attention_states() + manager.get_bounded_attention_states() + manager_ref = weakref.ref(manager) + 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 + 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_tp_config.py b/unit_tests/common/att_resources/test_tp_config.py new file mode 100644 index 0000000000..1fd660fa8d --- /dev/null +++ b/unit_tests/common/att_resources/test_tp_config.py @@ -0,0 +1,223 @@ +"""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 unit_tests.common.att_resources.model_helpers import replace_logical_layer_groups + +from typing import TYPE_CHECKING +import pickle +from dataclasses import replace + +import pytest +import torch + +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 ( + GQAAttentionConfig, + MLAConfig, + GQAFullAttentionState, + MLAFullAttentionState, +) + +if TYPE_CHECKING: + from collections.abc import Sequence + from easydict import EasyDict + 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 = AttentionStateManager(request_capacity=1) + manager.register_state("main", state) + manager.finish_state_registration() + return manager.get_cpu_cache_operator().get_page_layouts()[state.state_name] + + +@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 + if state_cls is MLAFullAttentionState: + config = MLAConfig( + logical_layer_groups=((0,),), + latent_dim=4, + latent_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(): + 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.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(): + field.tensor.copy_(fields[name][:, :, start:end]) + 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.allocate_storage(4) + layout = cpu_layout(target) + assert layout.byte_end == reference_layout.byte_end + 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]) + + +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, + ) + 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( + 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 = GQAAttentionConfig( + 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)( + config=other_config, + request_capacity=source.request_capacity, + mtp_step=source.mtp_step, + device="meta", + ) + else: + other = rebuild_full_state(source, config=other_config) + 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,) + assert query() == (source, other) + + 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_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 + ) + 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 + ) + 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 new file mode 100644 index 0000000000..e3fa612818 --- /dev/null +++ b/unit_tests/common/att_resources/test_transfer_views.py @@ -0,0 +1,233 @@ +"""传输视图整理必须保留存储别名、位模式、未参与区域和独立寻址轴。""" + +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( + "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( + "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 + + +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/att_resources/test_validate_model_cache.py b/unit_tests/common/att_resources/test_validate_model_cache.py new file mode 100644 index 0000000000..5e31d2a0d5 --- /dev/null +++ b/unit_tests/common/att_resources/test_validate_model_cache.py @@ -0,0 +1,49 @@ +from importlib.util import module_from_spec, spec_from_file_location +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.full_att import create_compressed_att_state +from lightllm.models.deepseek_v4.config import DSV4AttentionConfig + +spec = spec_from_file_location( + "validate_model_cache", Path(__file__).parents[3] / "test/att_resources/validate_model_cache.py" +) +assert spec is not None and spec.loader is not None +validate_model_cache = module_from_spec(spec) +spec.loader.exec_module(validate_model_cache) + + +@pytest.mark.parametrize("draft_count,bad_role", [(0, None), (2, None), (2, "main"), (2, "draft1")]) +def test_validator_checks_all_compressed_states(monkeypatch, draft_count, bad_role): + resources = AttentionStateManager(request_capacity=0) + resources.full_att_allocator = SimpleNamespace(size=512) + req_manager = object() + models = [] + for index, role in enumerate(["main"] + [f"draft{i}" for i in range(draft_count)]): + config = DSV4AttentionConfig( + logical_layer_ids=(index * 2, index * 2 + 1), + compress_ratios=(4, 128), + head_dim=512, + rope_head_dim=64, + indexer_head_dim=128, + sliding_window=128, + ) + for ratio in (4, 128): + state = create_compressed_att_state(config.create_compressed_config(ratio), device="meta") + state.allocate_storage(256 if role == bad_role and ratio == 4 else 512) + resources.register_state(role, state) + models.append(SimpleNamespace(att_state_manager=resources, att_state_role=role, req_manager=req_manager)) + resources.finish_state_registration() + monkeypatch.setattr(validate_model_cache, "StaticBenchmarkExecutor", lambda **kwargs: None) + monkeypatch.setattr(validate_model_cache, "TokenSource", lambda args: None) + # Stop at input validation after checking resources, before any GPU execution. + args = SimpleNamespace(input_len=255, cpu_cache_token_page_size=256) + if bad_role is not None: + with pytest.raises(AssertionError): + validate_model_cache.validate(models[0], args, models[1:]) + else: + with pytest.raises(ValueError, match="input_len must contain complete CPU pages"): + validate_model_cache.validate(models[0], args, models[1:]) 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..106743b126 --- /dev/null +++ b/unit_tests/common/att_resources/test_window_compressor_state.py @@ -0,0 +1,195 @@ +"""统一窗口状态：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_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_cache_operator().get_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/common/basemodel/__init__.py b/unit_tests/common/basemodel/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/unit_tests/common/basemodel/attention/fa3/test_neo.py b/unit_tests/common/basemodel/attention/fa3/test_neo.py new file mode 100644 index 0000000000..9ef88130cf --- /dev/null +++ b/unit_tests/common/basemodel/attention/fa3/test_neo.py @@ -0,0 +1,119 @@ +import importlib.util +import sys +from types import ModuleType, SimpleNamespace + +import pytest +import torch + +from lightllm.common.basemodel.infer_struct import InferStateInfo +from lightllm.common.basemodel.attention.fa3 import fp, neo + + +def _load_neo(monkeypatch, attention_func): + interface = None + if attention_func is not None: + interface = ModuleType("flash_attn_interface") + interface.flash_attn_with_kvcache = attention_func + monkeypatch.setitem(sys.modules, "flash_attn_interface", interface) + name = "lightllm.common.basemodel.attention.fa3._neo_import_test" + spec = importlib.util.spec_from_file_location(name, neo.__file__) + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, name, module) + spec.loader.exec_module(module) + return module + + +def test_recognizes_neo_interface_without_replacing_regular_fa3(monkeypatch): + def neo_attention(q, image_token_end=None): + return q + + regular_attention = fp.flash_attn_with_kvcache + module = _load_neo(monkeypatch, neo_attention) + + assert module.HAS_FLASH_ATTN_INTERFACE + assert module.flash_attn_with_kvcache_neo is neo_attention + assert fp.flash_attn_with_kvcache is regular_attention + + +@pytest.mark.parametrize("interface_kind", ["missing", "regular", "legacy_tag"]) +def test_unavailable_neo_interface_provides_source_install_instructions(monkeypatch, interface_kind): + def regular_attention(q): + return q + + def legacy_attention(q, image_token_tag=None): + return q + + attention_func = {"missing": None, "regular": regular_attention, "legacy_tag": legacy_attention}[interface_kind] + with pytest.warns(UserWarning, match="Neo FA3 is unavailable") as warnings: + module = _load_neo(monkeypatch, attention_func) + + assert not module.HAS_FLASH_ATTN_INTERFACE + assert module.flash_attn_with_kvcache_neo is None + message = str(warnings[0].message) + assert "https://github.com/WANDY666/flash-attention.git" in message + assert "cd flash-attention/hopper" in message + assert "FLASH_ATTENTION_FORCE_BUILD=TRUE" in message + assert "--no-deps --force-reinstall" in message + assert "--llm_prefill_att_backend triton" in message + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="no CUDA") +@pytest.mark.parametrize("page_size", [1, 16, 64]) +def test_prefill_initializes_paged_kv_metadata_and_forwards_image_ends(monkeypatch, page_size): + def indices(values): + return torch.tensor(values, dtype=torch.int32, device="cuda") + + image_ends = indices([2 * page_size, 2 * page_size, 0]) + # Each request uses two nonadjacent physical pages, in a different order. + req_to_token_indexs = indices( + [ + list(range(2 * page_size, 3 * page_size)) + list(range(page_size)), + list(range(3 * page_size, 4 * page_size)) + list(range(page_size, 2 * page_size)), + ] + ) + infer_state = InferStateInfo() + infer_state.batch_size = 2 + infer_state.max_kv_seq_len = 2 * page_size + infer_state.max_q_seq_len = 2 + infer_state.input_ids = torch.zeros(3, dtype=torch.int64, device="cuda") + infer_state.b_req_idx = indices([1, 0]) + infer_state.b_seq_len = indices([2 * page_size, 2 * page_size - 1]) + infer_state.b1_cu_q_seq_len = indices([0, 2, 3]).long() + infer_state.b1_cu_kv_seq_len = indices([0, 2 * page_size, 4 * page_size - 1]).long() + infer_state.b_image_token_end = image_ends + infer_state.req_manager = SimpleNamespace(req_to_token_indexs=req_to_token_indexs) + backend = SimpleNamespace(uses_causal_attention=lambda: True, infer_page_size=page_size) + state = neo.NeoFa3AttBackend.create_att_prefill_state(backend, infer_state) + state.init_state() + + q = torch.randn(3, 2, 16, device="cuda") + k = torch.randn(4 * page_size, 1, 16, device="cuda") + v = torch.randn_like(k) + output = torch.empty_like(q) + + def neo_attention(*, image_token_end, **kwargs): + assert image_token_end is image_ends + assert kwargs["q"] is q + assert kwargs["k_cache"].shape == kwargs["v_cache"].shape == (4, page_size, 1, 16) + torch.testing.assert_close(kwargs["page_table"], indices([[3, 1], [2, 0]])) + token_indices = req_to_token_indexs[infer_state.b_req_idx] + torch.testing.assert_close(kwargs["k_cache"][kwargs["page_table"]].flatten(1, 2), k[token_indices]) + torch.testing.assert_close(kwargs["v_cache"][kwargs["page_table"]].flatten(1, 2), v[token_indices]) + torch.testing.assert_close(kwargs["cu_seqlens_q"], indices([0, 2, 3])) + torch.testing.assert_close(kwargs["cu_seqlens_k_new"], indices([0, 2 * page_size, 4 * page_size - 1])) + assert kwargs["cache_seqlens"] is infer_state.b_seq_len + assert kwargs["max_seqlen_q"] == 2 + assert kwargs["causal"] is True + assert "image_token_tag" not in kwargs + return output + + monkeypatch.setattr(neo, "HAS_FLASH_ATTN_INTERFACE", True) + monkeypatch.setattr(neo, "flash_attn_with_kvcache_neo", neo_attention) + assert state.prefill_att(q, k, v) is output + + +def test_prefill_missing_extension_reports_install_instructions(monkeypatch): + monkeypatch.setattr(neo, "HAS_FLASH_ATTN_INTERFACE", False) + state = neo.NeoFa3PrefillAttState() + with pytest.raises(ImportError, match="WANDY666/flash-attention"): + state.prefill_att(None, None, None) diff --git a/unit_tests/common/basemodel/attention/flashinfer/test_decode_graph_state.py b/unit_tests/common/basemodel/attention/flashinfer/test_decode_graph_state.py new file mode 100644 index 0000000000..190d8f1162 --- /dev/null +++ b/unit_tests/common/basemodel/attention/flashinfer/test_decode_graph_state.py @@ -0,0 +1,100 @@ +import sys +import weakref +from types import ModuleType, SimpleNamespace + +import torch + +from lightllm.common.basemodel.attention.flashinfer.fp import FlashInferDecodeAttState +from lightllm.common.basemodel.attention.flashinfer import mla +from lightllm.common.basemodel.infer_struct import InferStateInfo + + +def test_fp_graph_copy_refreshes_plan_with_runtime_length_and_preserves_buffers(monkeypatch): + plan_calls = [] + flashinfer = ModuleType("flashinfer") + decode = ModuleType("flashinfer.decode") + decode.fast_decode_plan = lambda wrapper, **kwargs: plan_calls.append((wrapper, kwargs)) + flashinfer.decode = decode + monkeypatch.setitem(sys.modules, "flashinfer", flashinfer) + monkeypatch.setitem(sys.modules, "flashinfer.decode", decode) + backend = SimpleNamespace( + infer_page_size=16, + tp_q_head_num=4, + tp_kv_head_num=2, + head_dim=8, + q_data_type=torch.bfloat16, + kv_data_type=torch.bfloat16, + ) + captured_owner = InferStateInfo() + captured_owner.batch_size = 2 + runtime_owner = InferStateInfo() + runtime_owner.max_kv_seq_len = 17 + captured = FlashInferDecodeAttState( + backend=backend, + _infer_state_weak_ref=weakref.ref(captured_owner), + kv_starts=torch.tensor([0, 4, 8], dtype=torch.int32), + kv_indices=torch.arange(8, dtype=torch.int32), + kv_last_page_len_buffer=torch.tensor([16, 16], dtype=torch.int32), + decode_wrapper=object(), + ) + runtime = FlashInferDecodeAttState( + backend=backend, + _infer_state_weak_ref=weakref.ref(runtime_owner), + kv_starts=torch.tensor([0, 2, 3], dtype=torch.int32), + kv_indices=torch.arange(7, -1, -1, dtype=torch.int32), + kv_last_page_len_buffer=torch.tensor([1, 5], dtype=torch.int32), + ) + buffer_names = ("kv_starts", "kv_indices", "kv_last_page_len_buffer") + captured_ptrs = {name: getattr(captured, name).data_ptr() for name in buffer_names} + + captured.copy_for_decode_cuda_graph(runtime) + + assert len(plan_calls) == 1 + wrapper, plan = plan_calls[0] + assert wrapper is captured.decode_wrapper + assert plan["indptr"] is captured.kv_starts + assert plan["indices"] is captured.kv_indices + assert plan["last_page_len"] is captured.kv_last_page_len_buffer + torch.testing.assert_close(plan["global_override_indptr_cpu"], torch.tensor([0, 2, 4], dtype=torch.int32)) + for name in buffer_names: + assert getattr(captured, name).data_ptr() == captured_ptrs[name] + torch.testing.assert_close(getattr(captured, name), getattr(runtime, name)) + + +def test_mla_graph_copy_refreshes_plan_with_runtime_lengths_and_preserves_buffers(monkeypatch): + plan_calls = [] + monkeypatch.setattr(mla, "fill_mla_decode_plan_for_cuda_graph", lambda **kwargs: plan_calls.append(kwargs)) + backend = SimpleNamespace(tp_q_head_num=4) + captured_owner = InferStateInfo() + captured_owner.batch_size = 2 + runtime_owner = InferStateInfo() + runtime_owner.b_seq_len = torch.tensor([17, 5], dtype=torch.int32) + runtime_owner.max_kv_seq_len = 17 + captured = mla.MlaFlashInferDecodeAttState( + backend=backend, + _infer_state_weak_ref=weakref.ref(captured_owner), + kv_starts=torch.tensor([0, 4, 8], dtype=torch.int32), + kv_indices=torch.arange(8, dtype=torch.int32), + decode_wrapper=object(), + ) + runtime = mla.MlaFlashInferDecodeAttState( + backend=backend, + _infer_state_weak_ref=weakref.ref(runtime_owner), + kv_starts=torch.tensor([0, 2, 3], dtype=torch.int32), + kv_indices=torch.arange(7, -1, -1, dtype=torch.int32), + ) + buffer_names = ("kv_starts", "kv_indices") + captured_ptrs = {name: getattr(captured, name).data_ptr() for name in buffer_names} + + captured.copy_for_decode_cuda_graph(runtime) + + assert len(plan_calls) == 1 + plan = plan_calls[0] + assert plan["decode_wrapper"] is captured.decode_wrapper + assert plan["kv_page_indptr"] is captured.kv_starts + assert plan["batch_size"] == 2 + assert plan["max_kv_len"] == 17 + torch.testing.assert_close(plan["kv_lens"], torch.tensor([17, 5], dtype=torch.int32)) + for name in buffer_names: + assert getattr(captured, name).data_ptr() == captured_ptrs[name] + torch.testing.assert_close(getattr(captured, name), getattr(runtime, name)) diff --git a/unit_tests/common/basemodel/attention/flashinfer/test_decode_wrapper_init.py b/unit_tests/common/basemodel/attention/flashinfer/test_decode_wrapper_init.py new file mode 100644 index 0000000000..1392d5e688 --- /dev/null +++ b/unit_tests/common/basemodel/attention/flashinfer/test_decode_wrapper_init.py @@ -0,0 +1,132 @@ +import sys +import weakref +from types import ModuleType, SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch + +from lightllm.common.basemodel.attention.flashinfer import fp, fp8, mla +from lightllm.common.basemodel.attention.flashinfer.utils import should_init_decode_wrapper +from lightllm.common.basemodel.infer_struct import InferStateInfo + + +def test_decode_wrapper_rejects_conflicting_graph_flags(): + infer_state = InferStateInfo() + infer_state.is_cuda_graph_capture = True + infer_state.is_cuda_graph_replay = True + + with pytest.raises(ValueError, match="capture=True, replay=True"): + should_init_decode_wrapper(infer_state) + + +@pytest.fixture( + params=[fp.FlashInferDecodeAttState, fp8.Fp8FlashInferDecodeAttState, mla.MlaFlashInferDecodeAttState], + ids=["fp", "fp8", "mla"], +) +def decode_init(request, monkeypatch): + state_cls = request.param + module = mla if state_cls is mla.MlaFlashInferDecodeAttState else fp + page_size = 1 if state_cls is fp8.Fp8FlashInferDecodeAttState else 4 + max_page_num = 64 // page_size + wrapper = SimpleNamespace(plan=Mock()) + create_wrapper = Mock(return_value=wrapper) + flashinfer = ModuleType("flashinfer") + flashinfer.decode = SimpleNamespace(BatchDecodeWithPagedKVCacheWrapper=create_wrapper) + flashinfer.mla = SimpleNamespace(BatchMLAPagedAttentionWrapper=create_wrapper) + monkeypatch.setitem(sys.modules, "flashinfer", flashinfer) + + def cpu_arange(*args, **kwargs): + kwargs["device"] = "cpu" + return torch.arange(*args, **kwargs) + + monkeypatch.setattr(module, "torch", SimpleNamespace(int32=torch.int32, arange=cpu_arange, empty=torch.empty)) + + def cumsum_params(q_lengths, kv_lengths): + return tuple( + torch.cat((torch.zeros(1, dtype=torch.int32), lengths.cumsum(0, dtype=torch.int32))) + for lengths in (q_lengths, kv_lengths) + ) + + monkeypatch.setattr(module, "gen_cumsum_pad0_tensor", cumsum_params) + + def repack(**kwargs): + kwargs["out_page_indices"].fill_(7) + + repack_indices = Mock(side_effect=repack) + monkeypatch.setattr(module, "repack_kv_index", repack_indices) + can_run = Mock(return_value=True) + backend = SimpleNamespace( + model=SimpleNamespace( + graph=SimpleNamespace(can_run=can_run), graph_max_batch_size=4, graph_max_len_in_batch=64 + ), + infer_page_size=page_size, + max_page_num=max_page_num, + kv_indices_buffer=[torch.empty(4 * max_page_num, dtype=torch.int32)], + tp_q_head_num=4, + tp_kv_head_num=2, + head_dim=8, + kv_lora_rank=8, + qk_rope_head_dim=2, + softmax_scale=0.25, + q_data_type=torch.bfloat16, + kv_data_type=torch.float8_e4m3fn if page_size == 1 else torch.bfloat16, + workspace_buffer_key="decode_init_test", + workspace_buffer_size=1, + get_gpu_workspace_buffer=Mock(return_value=torch.empty(1, dtype=torch.uint8)), + ) + return SimpleNamespace( + module=module, + state=state_cls(backend=backend), + backend=backend, + wrapper=wrapper, + create_wrapper=create_wrapper, + repack_indices=repack_indices, + can_run=can_run, + ) + + +@pytest.mark.parametrize("mode", ["eager", "capture", "replay"]) +def test_decode_init_uses_explicit_graph_mode_for_wrapper_creation(decode_init, mode): + infer_state = InferStateInfo() + infer_state.input_ids = torch.tensor([1, 2], dtype=torch.int64) + infer_state.batch_size = 2 + infer_state.b_seq_len = torch.tensor([5, 2], dtype=torch.int32) + infer_state.max_kv_seq_len = 5 + infer_state.b_req_idx = torch.tensor([2, 3], dtype=torch.int32) + infer_state.microbatch_index = 0 + infer_state.req_manager = SimpleNamespace(req_to_token_indexs=torch.zeros((4, 64), dtype=torch.int32)) + if mode == "capture": + infer_state.is_cuda_graph_capture = True + elif mode == "replay": + infer_state.is_cuda_graph_replay = True + + decode_init.state._infer_state_weak_ref = weakref.ref(infer_state) + decode_init.state.init_state() + + decode_init.can_run.assert_not_called() + expected_starts = [0, 5, 7] if decode_init.backend.infer_page_size == 1 else [0, 2, 3] + torch.testing.assert_close(decode_init.state.kv_starts, torch.tensor(expected_starts, dtype=torch.int32)) + assert decode_init.state.kv_indices.data_ptr() == decode_init.backend.kv_indices_buffer[0].data_ptr() + assert torch.all(decode_init.state.kv_indices == 7) + decode_init.repack_indices.assert_called_once() + repack_args = decode_init.repack_indices.call_args.kwargs + assert repack_args["b_req_idx"] is infer_state.b_req_idx + assert repack_args["b_token_len"] is infer_state.b_seq_len + if decode_init.module is fp: + expected_last_page_lengths = [1, 1] if decode_init.backend.infer_page_size == 1 else [1, 2] + torch.testing.assert_close( + decode_init.state.kv_last_page_len_buffer, torch.tensor(expected_last_page_lengths, dtype=torch.int32) + ) + else: + torch.testing.assert_close(decode_init.state.q_indptr, torch.tensor([0, 1, 2], dtype=torch.int32)) + + if mode == "replay": + assert decode_init.state.decode_wrapper is None + decode_init.create_wrapper.assert_not_called() + decode_init.backend.get_gpu_workspace_buffer.assert_not_called() + else: + assert decode_init.state.decode_wrapper is decode_init.wrapper + decode_init.create_wrapper.assert_called_once() + decode_init.backend.get_gpu_workspace_buffer.assert_called_once() + decode_init.wrapper.plan.assert_called_once() diff --git a/unit_tests/common/basemodel/attention/flashinfer/test_mla_cuda_graph_plan.py b/unit_tests/common/basemodel/attention/flashinfer/test_mla_cuda_graph_plan.py new file mode 100644 index 0000000000..ccf5dfd454 --- /dev/null +++ b/unit_tests/common/basemodel/attention/flashinfer/test_mla_cuda_graph_plan.py @@ -0,0 +1,178 @@ +from types import SimpleNamespace +import weakref + +import pytest +import torch + +flashinfer = pytest.importorskip("flashinfer") + +from lightllm.common.basemodel.attention.flashinfer.mla import ( # noqa: E402 + MlaFlashInferDecodeAttState, +) +from lightllm.common.basemodel.infer_struct import InferStateInfo # noqa: E402 + + +def _make_wrapper( + q_indptr, + kv_indptr_buf, + kv_indptr, + kv_indices, + kv_lens_buf, + num_heads, + head_dim_ckv, + head_dim_kpe, + sm_scale, + dtype, + page_size, + init_short, +): + workspace = torch.empty(128 * 1024 * 1024, dtype=torch.int8, device="cuda") + wrapper = flashinfer.mla.BatchMLAPagedAttentionWrapper( + workspace, + use_cuda_graph=True, + qo_indptr=q_indptr, + kv_indices=kv_indices, + kv_indptr=kv_indptr_buf, + kv_len_arr=kv_lens_buf, + ) + if init_short: + batch_size = kv_lens_buf.numel() + init_lens = torch.full((batch_size,), 2, dtype=torch.int32, device="cuda") + init_indptr = torch.arange(batch_size + 1, dtype=torch.int32, device="cuda") * ( + (2 + page_size - 1) // page_size + ) + wrapper.plan( + q_indptr, + init_indptr, + kv_indices, + init_lens, + num_heads, + head_dim_ckv, + head_dim_kpe, + page_size, + False, + sm_scale, + dtype, + dtype, + ) + else: + wrapper.plan( + q_indptr, + kv_indptr, + kv_indices, + kv_lens_buf, + num_heads, + head_dim_ckv, + head_dim_kpe, + page_size, + False, + sm_scale, + dtype, + dtype, + ) + return wrapper + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@pytest.mark.parametrize("page_size", [1, 3, 4, 16, 64, 128]) +@pytest.mark.parametrize("lengths,num_heads", [([5, 64, 65], 32), ([1000, 32768], 32), ([1000, 131073], 128)]) +def test_mla_cuda_graph_triton_plan_matches_flashinfer_plan(lengths, num_heads, page_size): + torch.manual_seed(0) + dtype = torch.bfloat16 + head_dim_ckv = 512 + head_dim_kpe = 64 + batch_size = len(lengths) + total_pages = sum((length + page_size - 1) // page_size for length in lengths) + sm_scale = (head_dim_ckv + head_dim_kpe) ** -0.5 + + q_nope = torch.randn((batch_size, num_heads, head_dim_ckv), dtype=dtype, device="cuda") + q_pe = torch.randn((batch_size, num_heads, head_dim_kpe), dtype=dtype, device="cuda") + ckv = torch.randn((total_pages, page_size, head_dim_ckv), dtype=dtype, device="cuda") + kpe = torch.randn((total_pages, page_size, head_dim_kpe), dtype=dtype, device="cuda") + kv_lens = torch.tensor(lengths, dtype=torch.int32, device="cuda") + kv_indptr = torch.empty(batch_size + 1, dtype=torch.int32, device="cuda") + kv_indptr[0] = 0 + kv_indptr[1:] = torch.cumsum((kv_lens + page_size - 1) // page_size, dim=0) + kv_indices = torch.randperm(total_pages, dtype=torch.int32, device="cuda") + q_indptr = torch.arange(batch_size + 1, dtype=torch.int32, device="cuda") + + ref_wrapper = _make_wrapper( + q_indptr, + kv_indptr.clone(), + kv_indptr, + kv_indices, + kv_lens.clone(), + num_heads, + head_dim_ckv, + head_dim_kpe, + sm_scale, + dtype, + page_size, + init_short=False, + ) + graph_wrapper = _make_wrapper( + q_indptr, + kv_indptr.clone(), + kv_indptr, + kv_indices, + kv_lens.clone(), + num_heads, + head_dim_ckv, + head_dim_kpe, + sm_scale, + dtype, + page_size, + init_short=True, + ) + backend = SimpleNamespace(tp_q_head_num=num_heads, infer_page_size=page_size) + captured_owner = InferStateInfo() + captured_owner.batch_size = batch_size + captured_owner.b_seq_len = graph_wrapper._kv_len_arr_buf + graph_state = MlaFlashInferDecodeAttState( + backend=backend, + _infer_state_weak_ref=weakref.ref(captured_owner), + kv_starts=graph_wrapper._kv_indptr_buf, + kv_indices=graph_wrapper._kv_indices_buf, + decode_wrapper=graph_wrapper, + ) + ref_out = torch.empty((batch_size, num_heads, head_dim_ckv), dtype=dtype, device="cuda") + graph_out = torch.empty_like(ref_out) + graph_wrapper.run(q_nope, q_pe, ckv, kpe, out=graph_out, return_lse=False) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_wrapper.run(q_nope, q_pe, ckv, kpe, out=graph_out, return_lse=False) + + # 复用同一 graph，覆盖页表变化、末页空位及 split/exact 切换后的 plan 更新。 + for replay_lengths in (lengths, [1] * batch_size, lengths): + replay_lens = torch.tensor(replay_lengths, dtype=torch.int32, device="cuda") + replay_indptr = torch.zeros_like(kv_indptr) + replay_indptr[1:] = torch.cumsum((replay_lens + page_size - 1) // page_size, dim=0) + runtime_owner = InferStateInfo() + runtime_owner.batch_size = batch_size + runtime_owner.b_seq_len = replay_lens + runtime_owner.max_kv_seq_len = max(replay_lengths) + new_state = MlaFlashInferDecodeAttState( + backend=backend, + _infer_state_weak_ref=weakref.ref(runtime_owner), + kv_starts=replay_indptr, + kv_indices=kv_indices, + ) + graph_state.copy_for_decode_cuda_graph(new_state) + graph.replay() + ref_wrapper.plan( + q_indptr, + replay_indptr, + kv_indices, + replay_lens, + num_heads, + head_dim_ckv, + head_dim_kpe, + page_size, + False, + sm_scale, + dtype, + dtype, + ) + ref_wrapper.run(q_nope, q_pe, ckv, kpe, out=ref_out, return_lse=False) + torch.testing.assert_close(ref_out, graph_out, atol=1e-2, rtol=1e-2) diff --git a/unit_tests/common/basemodel/attention/flashinfer/test_mla_nope.py b/unit_tests/common/basemodel/attention/flashinfer/test_mla_nope.py new file mode 100644 index 0000000000..4c33d8cf8a --- /dev/null +++ b/unit_tests/common/basemodel/attention/flashinfer/test_mla_nope.py @@ -0,0 +1,106 @@ +"""Check the FlashInfer SM90 path used for unpadded GLM NoPE MLA.""" + +import inspect + +import pytest +import torch + +flashinfer = pytest.importorskip("flashinfer") + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 9, + reason="FlashInfer NoPE FA3 MLA requires Hopper", +) + + +@pytest.mark.parametrize( + "num_tokens,num_heads,kv_dtype", + [ + (1, 16, torch.bfloat16), + (8, 16, torch.bfloat16), + (128, 16, torch.bfloat16), + (8, 64, torch.bfloat16), + (8, 16, torch.float8_e4m3fn), + ], +) +def test_nope_sparse_mla_eager_and_graph_replay(num_tokens, num_heads, kv_dtype): + wrapper_cls = flashinfer.mla.BatchMLAPagedAttentionWrapper + assert "ckv_scale_arr" in inspect.signature(wrapper_cls.run).parameters, "FlashInfer >= 0.6.18 is required" + + torch.manual_seed(17) + dim, num_slots, width = 512, 4096, 2176 + sm_scale = 0.0625 + q = torch.randn(num_tokens, num_heads, dim, dtype=torch.bfloat16, device="cuda") + q_pe = q[..., :0] + # BF16 KV shares its token row with the 144-byte index region. The + # attention kernel must respect the resulting non-contiguous token stride. + tail = 72 if kv_dtype == torch.bfloat16 else 0 + packed = torch.full((num_slots, 1, dim + tail), 7.0, dtype=kv_dtype, device="cuda") + ckv = packed[..., :dim] + raw_kv = torch.randn(num_slots, 1, dim, dtype=torch.bfloat16, device="cuda") + scale = 0.125 if kv_dtype == torch.float8_e4m3fn else 1.0 + ckv.copy_((raw_kv.float() / scale).to(kv_dtype)) + kpe = packed[..., dim:dim] + reference_kv = (ckv.float() * scale).to(torch.bfloat16).float() + scale_kwargs = {"ckv_scale": scale, "kpe_scale": 1.0} if scale != 1.0 else {} + + qo_indptr = torch.arange(num_tokens + 1, dtype=torch.int32) + kv_indptr = qo_indptr * width + slots = torch.zeros(num_tokens, width, dtype=torch.int32, device="cuda") + wrapper = wrapper_cls( + torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device="cuda"), + use_cuda_graph=True, + qo_indptr=qo_indptr.to("cuda"), + kv_indptr=kv_indptr.to("cuda"), + kv_indices=slots.flatten(), + kv_len_arr=torch.empty(num_tokens, dtype=torch.int32, device="cuda"), + backend="fa3", + ) + + def plan(step): + choices = [1, 3, 127, 128, 511, 2048, 2049, 2051] + lengths = [choices[(row + step) % len(choices)] for row in range(num_tokens)] + slots.zero_() + for row, length in enumerate(lengths): + slots[row, :length] = (torch.arange(length, device="cuda") * 5 + row * 31 + step * 7) % num_slots + # Lengths are compiled into the host-side plan. Replan before graph + # replay when they change, while keeping all captured buffers stable. + wrapper.plan( + qo_indptr, + kv_indptr, + slots.flatten(), + torch.tensor(lengths, dtype=torch.int32), + num_heads, + dim, + 0, + 1, + False, + sm_scale, + torch.bfloat16, + kv_dtype, + ) + return lengths + + def check(output, lengths): + expected = [] + for row, length in enumerate(lengths): + keys = reference_kv[slots[row, :length].long(), 0] + scores = q[row].float() @ keys.T * sm_scale + expected.append(scores.softmax(-1) @ keys) + torch.testing.assert_close(output.float(), torch.stack(expected), atol=0.012, rtol=0.015) + + lengths = plan(0) + output = torch.empty_like(q) + wrapper.run(q, q_pe, ckv, kpe, out=output, **scale_kwargs) + check(output, lengths) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + wrapper.run(q, q_pe, ckv, kpe, out=output, **scale_kwargs) + graph.replay() + check(output, lengths) + + lengths = plan(1) + q.normal_() + graph.replay() + check(output, lengths) diff --git a/unit_tests/common/basemodel/attention/flashinfer/test_workspace.py b/unit_tests/common/basemodel/attention/flashinfer/test_workspace.py new file mode 100644 index 0000000000..924dead570 --- /dev/null +++ b/unit_tests/common/basemodel/attention/flashinfer/test_workspace.py @@ -0,0 +1,56 @@ +from lightllm.common.basemodel.attention import base_att +from lightllm.common.basemodel.attention.base_att import BaseAttBackend +from lightllm.common.basemodel.attention.flashinfer.fp import FlashInferAttBackend +from lightllm.common.basemodel.attention.flashinfer.fp8 import Fp8FlashInferAttBackend +from lightllm.common.basemodel.attention.flashinfer.mla import MlaFlashInferAttBackend + + +def test_workspace_is_shared_by_backend_family_and_device(monkeypatch): + allocations = [] + current_device_id = 0 + + def fake_empty(size, dtype, device): + workspace = object() + allocations.append((workspace, size, dtype, device)) + return workspace + + monkeypatch.setattr(base_att, "get_current_device_id", lambda: current_device_id) + monkeypatch.setattr(base_att.torch, "empty", fake_empty) + monkeypatch.setattr(BaseAttBackend, "_workspace_buffers", {}) + + def get_workspace(backend_cls): + return backend_cls.get_gpu_workspace_buffer( + key_name=backend_cls.workspace_buffer_key, + workspace_size=backend_cls.workspace_buffer_size, + ) + + fp_workspace = get_workspace(FlashInferAttBackend) + fp8_workspace = get_workspace(Fp8FlashInferAttBackend) + mla_workspace = get_workspace(MlaFlashInferAttBackend) + + assert fp8_workspace is fp_workspace + assert mla_workspace is not fp_workspace + assert [allocation[1] for allocation in allocations] == [ + 512 * 1024 * 1024, + 256 * 1024 * 1024, + ] + + current_device_id = 1 + other_device_workspace = get_workspace(FlashInferAttBackend) + + assert other_device_workspace is not fp_workspace + assert allocations[-1][3] == 1 + + +def test_workspace_key_includes_size_and_dtype(monkeypatch): + monkeypatch.setattr(base_att, "get_current_device_id", lambda: 0) + monkeypatch.setattr(base_att.torch, "empty", lambda *args, **kwargs: object()) + monkeypatch.setattr(BaseAttBackend, "_workspace_buffers", {}) + + small_workspace = BaseAttBackend.get_gpu_workspace_buffer("shared", 1024) + large_workspace = BaseAttBackend.get_gpu_workspace_buffer("shared", 2048) + fp16_workspace = BaseAttBackend.get_gpu_workspace_buffer("shared", 1024, dtype=base_att.torch.float16) + + assert BaseAttBackend.get_gpu_workspace_buffer("shared", 1024) is small_workspace + assert large_workspace is not small_workspace + assert fp16_workspace is not small_workspace diff --git a/unit_tests/common/basemodel/attention/linear/test_gdn.py b/unit_tests/common/basemodel/attention/linear/test_gdn.py new file mode 100644 index 0000000000..344bfea57f --- /dev/null +++ b/unit_tests/common/basemodel/attention/linear/test_gdn.py @@ -0,0 +1,438 @@ +import weakref +import inspect +import sys +from types import ModuleType, SimpleNamespace + +import pytest +import torch + +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState +from lightllm.common.basemodel.infer_struct import InferStateInfo +import lightllm.common.basemodel.attention.linear.gdn as gdn +import lightllm.common.basemodel.attention.create_linear_utils as linear_create_utils +import lightllm.common.basemodel.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.config import Qwen3NextResourceConfig +from lightllm.server.api_cli import make_argument_parser +import lightllm.utils.backend_validator as backend_validator + + +@pytest.mark.parametrize("cache_dtype", [torch.bfloat16, torch.float32]) +def test_prefill_casts_final_state_to_cache_dtype(monkeypatch, cache_dtype): + ssm_states = torch.zeros((1, 1, 2, 2), dtype=cache_dtype) + final_state = torch.ones((1, 1, 2, 2), dtype=torch.float32) + + monkeypatch.setattr(gdn, "fused_gdn_gating", lambda _log, a, b, _bias: (a, b)) + monkeypatch.setattr(gdn, "causal_conv1d_fn", lambda mixed, *args, **kwargs: mixed) + qkv = torch.zeros((1, 3), dtype=cache_dtype) + q = torch.zeros((1, 1, 1, 1), dtype=cache_dtype) + backend = SimpleNamespace( + mtp_step=0, + activation="silu", + ssm_state_dtype=cache_dtype, + _rearrange_mixed_qkv=lambda mixed: (q, q, q), + prefill_kernel=lambda *args, **kwargs: (None, final_state), + ) + infer_state = InferStateInfo() + infer_state.b1_cu_q_seq_len = torch.tensor([0, 1], dtype=torch.int32) + infer_state.b_ready_cache_len = 0 + state = gdn.LinearAttPrefillAttState( + backend=backend, + _infer_state_weak_ref=weakref.ref(infer_state), + ) + state.b_conv_buffer_idx = torch.tensor([0], dtype=torch.int64) + state.b_ssm_buffer_idx = torch.tensor([0], dtype=torch.int64) + layer_weight = SimpleNamespace( + linear_A_log=SimpleNamespace(weight=None), + linear_dt_bias=SimpleNamespace(weight=None), + linear_conv1d=SimpleNamespace(mm_param=SimpleNamespace(weight=None), bias=None), + ) + + state._gdn_prefill_kernel( + qkv, + torch.zeros((1, 3), dtype=cache_dtype), + ssm_states, + torch.zeros((1, 1), dtype=cache_dtype), + torch.zeros((1, 1), dtype=cache_dtype), + state.infer_state, + layer_weight, + ) + + assert ssm_states.dtype == cache_dtype + assert torch.equal(ssm_states, final_state.to(cache_dtype)) + + +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) + + state.init_state() + + assert state.b_conv_buffer_idx is b_req_idx + assert state.b_ssm_buffer_idx is b_req_idx + assert state.b1_mtp_cu_q_seq_len is None + assert state.b_num_accepted_tokens is None + + +@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( + draft_step=2, + dynamic_layout=False, + b_req_idx=b_req_idx, + req_to_mtp_state_index=req_to_mtp_state_index, + role=role, + ) + + state.init_state() + + torch.testing.assert_close(state.b1_mtp_cu_q_seq_len, torch.tensor([0, 3, 6], dtype=torch.int32)) + torch.testing.assert_close(state.b_conv_buffer_idx, torch.tensor([2, 4], dtype=torch.int32)) + torch.testing.assert_close(state.b_num_accepted_tokens, torch.tensor([2, 3], dtype=torch.int32)) + torch.testing.assert_close( + state.b_ssm_buffer_idx, + torch.tensor([[6, 7, 8], [12, 13, 14]], dtype=torch.int32), + ) + + +@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) + expected_conv_buffer_idx = torch.tensor([2, 4, 5, 5, 5], dtype=torch.int32) + expected_num_accepted_tokens = torch.tensor([2, 3, 1, 1, 1], dtype=torch.int32) + build_calls = [] + + def build_params(**kwargs): + build_calls.append(kwargs) + return expected_cu_q_seq_len, expected_conv_buffer_idx, expected_num_accepted_tokens + + monkeypatch.setattr(gdn, "build_dynamic_mtp_linear_att_state_params", build_params) + state, infer_state = _create_decode_state( + draft_step=2, + dynamic_layout=True, + b_req_idx=b_req_idx, + req_to_mtp_state_index=req_to_mtp_state_index, + role=role, + ) + + state.init_state() + + assert state.b1_mtp_cu_q_seq_len is expected_cu_q_seq_len + assert state.b_conv_buffer_idx is expected_conv_buffer_idx + assert state.b_num_accepted_tokens is expected_num_accepted_tokens + torch.testing.assert_close( + state.b_ssm_buffer_idx, + torch.tensor( + [[6, 7, 8], [12, 13, 14], [15, 16, 17], [15, 16, 17], [15, 16, 17]], + dtype=torch.int32, + ), + ) + assert len(build_calls) == 1 + assert build_calls[0]["b_req_idx"] is state.infer_state.b_req_idx + assert build_calls[0]["b_mtp_index"] is state.infer_state.b_mtp_index + assert build_calls[0]["req_to_mtp_state_index"] is req_to_mtp_state_index + assert build_calls[0]["hold_req_id"] == 5 + + +def test_linear_backend_is_abstract_and_triton_supplies_chunk_kernel(): + assert inspect.isabstract(gdn.LinearAttBackend) + backend = object.__new__(TritonLinearAttBackend) + assert backend.get_prefill_kernel() is fla_ops.chunk_gated_delta_rule + + +@pytest.fixture +def auto_linear_backend_args(monkeypatch): + monkeypatch.setattr( + linear_create_utils, + "get_env_start_args", + lambda: SimpleNamespace( + llm_prefill_att_backend=["fa3"], + llm_decode_att_backend=["flashinfer"], + ssm_state_mode="native", + ), + ) + + +def test_missing_linear_backend_arg_uses_auto_selection(monkeypatch, auto_linear_backend_args): + validate_calls = [] + flashqla = ModuleType("flash_qla") + flashqla.chunk_gated_delta_rule = lambda **kwargs: ("flashqla", kwargs) + monkeypatch.setitem(sys.modules, "flash_qla", flashqla) + monkeypatch.setattr( + linear_create_utils, + "validate", + lambda name: validate_calls.append(name) or True, + ) + + backend_class = linear_create_utils.get_gdn_linear_prefill_att_backend_class(index=1) + + assert backend_class is FlashQlaLinearAttBackend + backend = object.__new__(backend_class) + assert backend.get_prefill_kernel()(q="q")[0] == "flashqla" + assert validate_calls == ["flashqla"] + + +def test_replay_mode_selects_standalone_gdn_backends(monkeypatch): + monkeypatch.setattr( + linear_create_utils, + "get_env_start_args", + lambda: SimpleNamespace( + llm_prefill_att_backend=["fa3", "flashqla"], + llm_decode_att_backend=["flashinfer", "triton"], + ssm_state_mode="replay", + ), + ) + from lightllm.common.basemodel.attention.base_att import BaseAttBackend + from lightllm.common.basemodel.attention.linear.gdn import LinearAttBackend + from lightllm.common.basemodel.attention.linear.replay_gdn import ( + ReplayGDNFlashQlaAttBackend, + ReplayGDNTritonAttBackend, + ) + + prefill_backend = linear_create_utils.get_gdn_linear_prefill_att_backend_class(index=1) + decode_backend = linear_create_utils.get_gdn_linear_decode_att_backend_class(index=1) + + assert prefill_backend is ReplayGDNFlashQlaAttBackend + assert decode_backend is ReplayGDNTritonAttBackend + assert prefill_backend.__bases__ == (BaseAttBackend,) + assert decode_backend.__bases__ == (BaseAttBackend,) + assert not issubclass(prefill_backend, LinearAttBackend) + assert not issubclass(decode_backend, LinearAttBackend) + + +def test_kda_entry_selects_native_or_replay(monkeypatch): + from lightllm.common.basemodel.attention.linear.kda import KDALinearAttBackend + from lightllm.common.basemodel.attention.linear.replay_kda import ReplayKDALinearAttBackend + + monkeypatch.setattr( + linear_create_utils, + "get_env_start_args", + lambda: SimpleNamespace(ssm_state_mode="native"), + ) + assert linear_create_utils.get_kda_linear_att_backend_class() is KDALinearAttBackend + + monkeypatch.setattr( + linear_create_utils, + "get_env_start_args", + lambda: SimpleNamespace(ssm_state_mode="replay"), + ) + assert linear_create_utils.get_kda_linear_att_backend_class() is ReplayKDALinearAttBackend + + +def test_flashqla_validation_loads_linear_cache_config(monkeypatch): + linear_config = SimpleNamespace( + 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, + ssm_state_dtype=torch.float32, + ) + load_calls = [] + monkeypatch.setattr( + Qwen3NextResourceConfig, + "load_from_args", + staticmethod(lambda: load_calls.append(True) or linear_config), + ) + + kernel_calls = [] + + def kernel(**kwargs): + kernel_calls.append(kwargs) + return kwargs["q"], kwargs["initial_state"] + + flashqla = ModuleType("flash_qla") + flashqla.chunk_gated_delta_rule = kernel + monkeypatch.setitem(sys.modules, "flash_qla", flashqla) + monkeypatch.setattr(fla_ops, "chunk_gated_delta_rule", kernel) + + original_randn = torch.randn + original_rand = torch.rand + original_tensor = torch.tensor + + def cpu_randn(*args, **kwargs): + kwargs.pop("device", None) + return original_randn(*args, **kwargs) + + def cpu_rand(*args, **kwargs): + kwargs.pop("device", None) + return original_rand(*args, **kwargs) + + def cpu_tensor(*args, **kwargs): + kwargs.pop("device", None) + return original_tensor(*args, **kwargs) + + monkeypatch.setattr(backend_validator.torch, "randn", cpu_randn) + monkeypatch.setattr(backend_validator.torch, "rand", cpu_rand) + monkeypatch.setattr(backend_validator.torch, "tensor", cpu_tensor) + monkeypatch.setattr(backend_validator.torch.cuda, "synchronize", lambda: None) + + success, error = backend_validator._validate_flashqla() + + assert success is True + assert error is None + assert load_calls == [True] + assert len(kernel_calls) == 2 + assert kernel_calls[0]["q"].shape == (1, 64, 2, 8) + assert kernel_calls[0]["v"].shape == (1, 64, 4, 16) + assert kernel_calls[0]["q"].dtype == torch.bfloat16 + assert kernel_calls[0]["initial_state"].shape == (1, 4, 8, 16) + assert kernel_calls[0]["initial_state"].dtype == torch.float32 + + +def test_explicit_linear_backend_arg_skips_auto_selection(monkeypatch): + monkeypatch.setattr( + linear_create_utils, + "get_env_start_args", + lambda: SimpleNamespace( + llm_prefill_att_backend=["fa3", "triton"], + llm_decode_att_backend=["flashinfer", "triton"], + ssm_state_mode="native", + ), + ) + monkeypatch.setattr( + linear_create_utils, + "validate", + lambda *args: pytest.fail("an explicit linear backend must not be auto-validated"), + ) + + prefill_backend_class = linear_create_utils.get_gdn_linear_prefill_att_backend_class(index=1) + decode_backend_class = linear_create_utils.get_gdn_linear_decode_att_backend_class(index=1) + + assert prefill_backend_class is TritonLinearAttBackend + assert decode_backend_class is TritonLinearAttBackend + + +def test_cli_accepts_linear_backend_at_index_one(): + args = make_argument_parser().parse_args( + [ + "--llm_prefill_att_backend", + "fa3", + "flashqla", + "--llm_decode_att_backend", + "flashinfer", + "triton", + ] + ) + + assert args.llm_prefill_att_backend == ["fa3", "flashqla"] + assert args.llm_decode_att_backend == ["flashinfer", "triton"] + + +@pytest.mark.parametrize( + "get_backend_class", + [ + linear_create_utils.get_gdn_linear_prefill_att_backend_class, + linear_create_utils.get_gdn_linear_decode_att_backend_class, + ], +) +def test_linear_backend_rejects_non_linear_index(auto_linear_backend_args, get_backend_class): + with pytest.raises(AssertionError, match="index must be 1"): + get_backend_class(index=0) + + +def test_missing_linear_decode_backend_arg_defaults_to_triton(auto_linear_backend_args): + backend_class = linear_create_utils.get_gdn_linear_decode_att_backend_class(index=1) + + assert backend_class is TritonLinearAttBackend + + +def test_gdn_decode_backend_uses_priority_list(monkeypatch, auto_linear_backend_args): + validate_calls = [] + monkeypatch.setattr( + linear_create_utils, + "gdn_decode_backends", + {"optimized": FlashQlaLinearAttBackend, "triton": TritonLinearAttBackend}, + ) + monkeypatch.setattr( + linear_create_utils, + "validate", + lambda name: validate_calls.append(name) or True, + ) + + backend_class = linear_create_utils.get_gdn_linear_decode_att_backend_class( + index=1, priority_list=("optimized", "triton") + ) + + assert backend_class is FlashQlaLinearAttBackend + assert validate_calls == ["optimized"] + + +def test_gdn_prefill_backend_falls_back_when_flashqla_validation_fails(monkeypatch, auto_linear_backend_args): + monkeypatch.setattr(linear_create_utils, "validate", lambda *args: False) + + backend_class = linear_create_utils.get_gdn_linear_prefill_att_backend_class(index=1) + + assert backend_class is TritonLinearAttBackend + + +def test_gdn_prefill_backend_tries_candidates_in_order(monkeypatch, auto_linear_backend_args): + validate_calls = [] + monkeypatch.setattr( + linear_create_utils, + "gdn_prefill_backends", + {"flashqla2": TritonLinearAttBackend, "flashqla3": FlashQlaLinearAttBackend}, + ) + monkeypatch.setattr( + linear_create_utils, + "validate", + lambda name, *args: validate_calls.append(name) or name == "flashqla3", + ) + + backend_class = linear_create_utils.get_gdn_linear_prefill_att_backend_class( + index=1, priority_list=("flashqla2", "flashqla3") + ) + + 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, role="main"): + 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, + ) + infer_state.att_state_role = role + if req_to_mtp_state_index is not None: + 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, + ) + 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.register_state(model_role, linear) + resources.finish_state_registration() + infer_state.att_state_manager = resources + 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 diff --git a/unit_tests/common/basemodel/attention/test_dsv4_flashmla.py b/unit_tests/common/basemodel/attention/test_dsv4_flashmla.py new file mode 100644 index 0000000000..ae848f13de --- /dev/null +++ b/unit_tests/common/basemodel/attention/test_dsv4_flashmla.py @@ -0,0 +1,377 @@ +import gc +import os +from pathlib import Path +import subprocess +import sys +from types import SimpleNamespace + +import pytest +import torch + +pytest.importorskip("flash_mla") + +from lightllm.common.basemodel.attention.nsa import dsv4_fp8_flashmla_sparse as attention +from lightllm.common.basemodel.layer_weights.meta_weights import TpAttSinkWeight + + +@pytest.fixture(autouse=True) +def _tp_env(monkeypatch): + monkeypatch.setenv("LIGHTLLM_DP_WORLD_SIZE", "1") + monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_DP", "0") + monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") + yield + # Collect infer-state cycles here so old graphs are not destroyed during the next capture. + gc.collect() + + +def test_backend_import_without_vllm(): + script = """ +import importlib.abc +import sys + +class BlockVllm(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path=None, target=None): + if fullname == 'vllm' or fullname.startswith('vllm.'): + raise ModuleNotFoundError('vllm deliberately unavailable') + +sys.meta_path.insert(0, BlockVllm()) +from lightllm.common.basemodel.attention.nsa import dsv4_fp8_flashmla_sparse as attention +assert attention.flashmla.__name__ == 'flash_mla' +assert not any(name == 'vllm' or name.startswith('vllm.') for name in sys.modules) +""" + root = Path(__file__).resolve().parents[4] + env = dict(os.environ, PYTHONPATH=str(root)) + result = subprocess.run([sys.executable, "-c", script], cwd=root, env=env, capture_output=True, text=True) + 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") + rows = torch.randn((pages, page_size, 512), dtype=torch.bfloat16, device="cuda") * 0.1 + nope = rows[..., :448].to(torch.float8_e4m3fn) + data = buffer[:, : page_size * 576].view(pages, page_size, 576) + data[..., :448] = nope.view(torch.uint8) + data[..., 448:] = rows[..., 448:].contiguous().view(torch.uint8) + buffer[:, page_size * 576 : page_size * 584].view(pages, page_size, 8)[..., :7] = 127 + decoded = torch.cat((nope.float(), rows[..., 448:].float()), dim=-1).view(-1, 512) + return buffer, decoded + + +def _case(heads, ratio): + torch.manual_seed(6014) + tokens = 5 + model = type("TestModel", (), {})() + model.config = {"num_attention_heads": heads, "compress_ratios": [0, 4, 128]} + 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 + q[:, heads:] = 0 + pool, decoded = _pool(128, 2) + extra, extra_decoded = _pool(64 if ratio == 4 else 2, 3) if ratio else (None, None) + indices = torch.full((tokens, 1, 128), -1, dtype=torch.int32, device="cuda") + indices[:, 0, :4] = torch.tensor([0, 3, -1, 7], dtype=torch.int32, device="cuda") + lengths = torch.tensor([1, 2, 3, 4, 4], dtype=torch.int32, device="cuda") + extra_indices = indices.clone() if ratio else None + if ratio: + extra_indices[:, 0, :4] = torch.tensor([0, 1, -1, 5], dtype=torch.int32, device="cuda") + attn_sink = TpAttSinkWeight( + all_q_head_num=heads, weight_name="attn_sink", data_type=torch.float32, head_padding_multiple=64 + ) + attn_sink.load_hf_weights({"attn_sink": torch.linspace(-0.5, 0.5, heads, dtype=torch.float32)}) + params = dict( + compress_ratio=ratio, + head_dim_v=512, + softmax_scale=512 ** -0.5, + swa_indices=indices, + swa_lengths=lengths, + attn_sink=attn_sink.weight, + extra_indices=extra_indices, + extra_lengths=lengths.clone() if ratio else None, + ) + # The attention backend receives all cache inputs through k, without a memory manager. + state_info = SimpleNamespace() + return backend, state_info, q, (pool, extra), decoded, extra_decoded, params + + +def _reference(q, decoded, extra_decoded, params, heads): + outputs = [] + for row in range(q.shape[0]): + indexes = params["swa_indices"][row, 0, : params["swa_lengths"][row].item()].long() + selected = [decoded[indexes[indexes >= 0]]] + if extra_decoded is not None: + indexes = params["extra_indices"][row, 0, : params["extra_lengths"][row].item()].long() + selected.append(extra_decoded[indexes[indexes >= 0]]) + kv = torch.cat(selected) + logits = q[row, :heads].float() @ kv.T * params["softmax_scale"] + logits = torch.cat((logits, params["attn_sink"][:heads, None]), dim=-1) + weights = logits.softmax(-1)[:, :-1] + outputs.append((weights @ kv).to(q.dtype)) + return torch.stack(outputs) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("heads", [16, 64, 128]) +@pytest.mark.parametrize("ratio", [0, 4, 128]) +@pytest.mark.parametrize("provide_out", [False, True]) +def test_prefill_matches_reference_and_preserves_out(monkeypatch, heads, ratio, provide_out): + backend, state_info, q, pool, decoded, extra_decoded, params = _case(heads, ratio) + state = backend.create_att_prefill_state(state_info) + state.init_state() + allocations = [] + + def alloc_output(shape, **kwargs): + buffer = torch.empty(shape, **kwargs) + allocations.append(buffer) + return buffer + + kernel = attention.flashmla.flash_mla_with_kvcache + + def flashmla_with_out(**kwargs): + assert kwargs["out"] is not None + return kernel(**kwargs) + + monkeypatch.setattr(attention.flashmla, "flash_mla_with_kvcache", flashmla_with_out) + out = torch.full_like(q, float("nan")) if provide_out else None + actual = state.prefill_att( + q, + pool, + None, + att_control=attention.AttControl(nsa_prefill=True, nsa_prefill_dict=params), + alloc_func=alloc_output, + out=out, + ) + assert params["attn_sink"].shape == (backend.padded_q_head_num,) + assert torch.count_nonzero(params["attn_sink"][heads:]) == 0 + assert len(allocations) == (0 if provide_out else 1) + output_buffer = out if provide_out else allocations[0] + assert actual.data_ptr() == output_buffer.data_ptr() + assert actual.stride() == output_buffer.stride() + assert actual.shape == (q.shape[0], heads, 512) + torch.testing.assert_close(actual, _reference(q, decoded, extra_decoded, params, heads), rtol=0.02, atol=5e-4) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("heads", [16, 64]) +@pytest.mark.parametrize("ratio", [0, 4, 128]) +@pytest.mark.parametrize("overlap", [False, True]) +def test_decode_cuda_graph_reads_updated_query_and_lengths(monkeypatch, heads, ratio, overlap): + from lightllm.common.basemodel import cuda_graph as cuda_graph_module + + args = SimpleNamespace( + enable_decode_microbatch_overlap=overlap, enable_tpsp_mix_mode=False, enable_torch_memory_saver=False + ) + monkeypatch.setattr(cuda_graph_module, "get_env_start_args", lambda: args) + runner = cuda_graph_module.CudaGraph(1, 2, 1, max_batch_size=5, max_len_in_batch=16384) + cases = [] + for _ in range(2 if overlap else 1): + backend, state_info, q, packed_kv, decoded, extra_decoded, params = _case(heads, ratio) + indices = torch.arange(128, dtype=torch.int32, device=q.device) + params["swa_indices"].copy_(indices) + params["swa_lengths"].fill_(1) + if ratio: + page_size = 64 if ratio == 4 else 2 + extra, extra_decoded = _pool(page_size, 128 // page_size) + packed_kv = (packed_kv[0], extra) + params["extra_indices"].copy_(indices) + params["extra_lengths"].fill_(1) + state_info.input_ids = torch.zeros(q.shape[0], dtype=torch.int64, device=q.device) + state_info.q = q + state_info.packed_kv = packed_kv + state_info.att_control = attention.AttControl(nsa_decode=True, nsa_decode_dict=params) + state_info.decode_att_state = backend.create_att_decode_state(state_info) + state_info.decode_att_state1 = None + state_info.decode_att_state.init_state() + cases.append((state_info, decoded, extra_decoded, params)) + + def decode_func(*states): + outputs = tuple( + state.decode_att_state.decode_att(state.q, state.packed_kv, None, att_control=state.att_control) + for state in states + ) + return outputs if overlap else outputs[0] + + # Exercise the framework's warmup + capture directly, with no operator-specific reset. + outputs = runner.capture_decode(decode_func, *(case[0] for case in cases)) + outputs = outputs if overlap else (outputs,) + for state_info, _, _, _ in cases: + state = state_info.decode_att_state + eager_meta = state.sched_meta_by_ratio[ratio] + capture_meta = state.capture_sched_meta_by_ratio[ratio] + assert eager_meta.tile_scheduler_metadata.data_ptr() != capture_meta.tile_scheduler_metadata.data_ptr() + + # Reuse the same graph as lengths grow, shrink, and vary within the batch. + for lengths in ([128] * 5, [64] * 5, [1] * 5, [1, 64, 65, 127, 128]): + for index, (state_info, _, _, params) in enumerate(cases): + state_info.q.mul_(1.25 + index * 0.1) + new_lengths = torch.tensor(lengths, dtype=torch.int32, device=state_info.q.device) + if index: + new_lengths = new_lengths.flip(0) + params["swa_lengths"].copy_(new_lengths) + if ratio: + params["extra_lengths"].copy_(new_lengths) + runner.graph[5][0].replay() + + for actual, (state_info, decoded, extra_decoded, params) in zip(outputs, cases): + assert actual.shape == (state_info.q.shape[0], heads, 512) + assert actual.stride(0) == backend.padded_q_head_num * 512 + torch.testing.assert_close( + actual, _reference(state_info.q, decoded, extra_decoded, params, heads), rtol=0.02, atol=5e-4 + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("heads", [16, 64]) +def test_prefill_out_supports_strided_scratch_buffer(heads): + backend, state_info, q, pool, decoded, extra_decoded, params = _case(heads, 4) + storage = torch.full((q.shape[0], backend.padded_q_head_num, 1024), float("nan"), dtype=q.dtype, device=q.device) + out = storage[..., :512] + state = backend.create_att_prefill_state(state_info) + state.init_state() + actual = state.prefill_att( + q, pool, None, out=out, att_control=attention.AttControl(nsa_prefill=True, nsa_prefill_dict=params) + ) + assert actual.data_ptr() == out.data_ptr() + assert actual.stride() == out.stride() + assert storage[..., 512:].isnan().all() + torch.testing.assert_close(actual, _reference(q, decoded, extra_decoded, params, heads), rtol=0.02, atol=5e-4) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("heads", [16, 64]) +@pytest.mark.parametrize("ratio", [0, 4, 128]) +def test_prefill_cuda_graph_reads_new_request_state(monkeypatch, heads, ratio): + from lightllm.common.basemodel import prefill_cuda_graph as cuda_graph_module + from lightllm.common.basemodel.hidden_collector import NoopHiddenCollector + from lightllm.models.deepseek_v4.infer_struct import DeepseekV4InferStateInfo + from lightllm.models.deepseek_v4.layer_infer.transformer_layer_infer import DeepseekV4TransformerLayerInfer + + args = SimpleNamespace( + enable_prefill_microbatch_overlap=False, + prefill_cudagraph_max_handle_token=8, + batch_max_tokens=8, + enable_tpsp_mix_mode=False, + ) + monkeypatch.setattr(cuda_graph_module, "get_env_start_args", lambda: args) + runner = cuda_graph_module.PrefillCudaGraph(decode_cuda_graph=None, tp_world_size=1) + backend, _, q, packed_kv, decoded, extra_decoded, params = _case(heads, ratio) + indices = torch.arange(128, dtype=torch.int32, device=q.device) + params["swa_indices"].copy_(indices) + if ratio: + page_size = 64 if ratio == 4 else 2 + extra, extra_decoded = _pool(page_size, 128 // page_size) + 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.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) + infer_state.prefill_att_state = backend.create_att_prefill_state(infer_state) + infer_state.prefill_att_state.init_state() + return infer_state + + allocations = [] + callback_states = [] + layer = DeepseekV4TransformerLayerInfer.__new__(DeepseekV4TransformerLayerInfer) + layer.tp_q_head_num_ = heads + layer.head_dim_ = 512 + layer.v_head_dim = 512 + layer.compress_ratio = ratio + layer.softmax_scale = params["softmax_scale"] + layer.layer_num_ = 0 + layer_weight = SimpleNamespace(attn_sink_=SimpleNamespace(weight=params["attn_sink"])) + + def alloc_output(shape, **kwargs): + buffer = torch.empty(shape, **kwargs) + allocations.append(buffer) + return buffer + + def compress_and_index(q_lora, x, infer_state, layer_weight): + assert not torch.cuda.is_current_stream_capturing() + callback_states.append(infer_state) + return infer_state.attn_params + + layer.alloc_tensor = alloc_output + layer._compress_and_index = compress_and_index + q_lora = torch.empty((q.shape[0], 4), dtype=q.dtype, device=q.device) + x = torch.empty((q.shape[0], 8), dtype=q.dtype, device=q.device) + + def prefill_func(inputs, infer_state): + graph_q = inputs[0] * 1.0 + att_out = layer._context_attention_wrapper_run(graph_q, inputs[1], inputs[2], infer_state, layer_weight) + assert att_out.shape == (q.shape[0], heads, 512) + assert att_out.data_ptr() == allocations[-1].data_ptr() + assert att_out.stride(0) == backend.padded_q_head_num * 512 + # A graph segment after attention must read the output filled by the CPU callback. + return [att_out + 0.125] + + infer_state = new_request(1) + layer._context_attention_wrapper_run(q, q_lora, x, infer_state, layer_weight) + outputs = runner.capture_prefill(prefill_func, [q, q_lora, x], infer_state) + torch.testing.assert_close( + outputs[0], + _reference(q, decoded, extra_decoded, infer_state.attn_params, heads) + 0.125, + rtol=0.02, + atol=5e-4, + ) + assert len(infer_state.prefill_cuda_graph_exe_list) == 3 + previous_meta = infer_state.prefill_att_state.sched_meta_by_ratio[ratio] + for length in (128, 64, 1, 64): + q.mul_(1.25) + infer_state = new_request(length) + current_meta = infer_state.prefill_att_state.sched_meta_by_ratio[ratio] + assert current_meta.tile_scheduler_metadata is None + actual = runner.replay([q, q_lora, x], infer_state)[0] + assert callback_states[-1] is infer_state + assert current_meta.tile_scheduler_metadata.data_ptr() != previous_meta.tile_scheduler_metadata.data_ptr() + torch.testing.assert_close( + actual, + _reference(q, decoded, extra_decoded, infer_state.attn_params, heads) + 0.125, + rtol=0.02, + atol=5e-4, + ) + previous_meta = current_meta 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..9d02784676 --- /dev/null +++ b/unit_tests/common/basemodel/attention/test_fp8_logical_layer.py @@ -0,0 +1,117 @@ +"""FP8 后端按推理层传入的逻辑层号选择标定，覆盖稀疏层号、共享层和 draft。""" + +from dataclasses import replace +from types import SimpleNamespace +import weakref + +import pytest +import torch + +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 +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 + + +@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=template, + config=replace_logical_layer_groups(template.config, ((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 = AttentionStateManager(request_capacity=1) + manager.register_state(role, state) + manager.finish_state_registration() + infer = LlamaInferStateInfo() + 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 = [] + + # 仅替换外部计算内核，保留推理层、公开 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.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, + 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) + + 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/common/basemodel/attention/test_fp8_page_size.py b/unit_tests/common/basemodel/attention/test_fp8_page_size.py new file mode 100644 index 0000000000..107728c26d --- /dev/null +++ b/unit_tests/common/basemodel/attention/test_fp8_page_size.py @@ -0,0 +1,45 @@ +from types import SimpleNamespace +from unittest.mock import patch + +import pytest + +from lightllm.common.basemodel.attention.fa3 import fp8 as fa3_fp8 +from lightllm.common.basemodel.attention.fa3.fp import Fa3AttBackend +from lightllm.common.basemodel.attention.fa3.fp8 import Fp8Fa3AttBackend +from lightllm.common.basemodel.attention.fa3.mla import MlaFa3AttBackend +from lightllm.common.basemodel.attention.flashinfer import fp8 as flashinfer_fp8 +from lightllm.common.basemodel.attention.flashinfer.fp import FlashInferAttBackend +from lightllm.common.basemodel.attention.flashinfer.fp8 import Fp8FlashInferAttBackend +from lightllm.common.basemodel.attention.flashinfer.mla import MlaFlashInferAttBackend + + +@pytest.mark.parametrize( + "backend_class", + [Fa3AttBackend, MlaFa3AttBackend, FlashInferAttBackend, MlaFlashInferAttBackend], +) +def test_default_infer_page_size_matches_model_page_size(backend_class): + backend = object.__new__(backend_class) + backend.model = SimpleNamespace(args=SimpleNamespace(page_size=16)) + + backend._init_infer_page_size() + + assert backend.infer_page_size == 16 + + +@pytest.mark.parametrize( + ("backend_class", "backend_module"), + [ + (Fp8Fa3AttBackend, fa3_fp8), + (Fp8FlashInferAttBackend, flashinfer_fp8), + ], +) +def test_fp8_attention_backend_uses_single_token_infer_pages(backend_class, backend_module): + backend = object.__new__(backend_class) + backend.model = SimpleNamespace(args=SimpleNamespace(page_size=16)) + + with patch.object(backend_module.logger, "warning") as warning: + backend._init_infer_page_size() + + assert backend.infer_page_size == 1 + warning.assert_called_once() + assert "infer_page_size=1" in warning.call_args.args[0] diff --git a/unit_tests/common/basemodel/attention/test_lazy_import.py b/unit_tests/common/basemodel/attention/test_lazy_import.py new file mode 100644 index 0000000000..e7423d30f9 --- /dev/null +++ b/unit_tests/common/basemodel/attention/test_lazy_import.py @@ -0,0 +1,248 @@ +"""Check import isolation and selection without loading optional GPU dependencies.""" + +import importlib +import importlib.abc +import importlib.util +from pathlib import Path +import sys +from types import ModuleType, SimpleNamespace +import unittest +from unittest.mock import Mock, patch + + +ROOT = Path(__file__).resolve().parents[4] +ATTENTION = "lightllm.common.basemodel.attention" + + +class BackendImports(importlib.abc.MetaPathFinder, importlib.abc.Loader): + """Reject unselected implementations, and supply lightweight selected backends.""" + + def __init__(self, package=ATTENTION): + self.package = package + self.allowed = {} + self.imported = [] + + def allow(self, suffix, class_name, **attributes): + module = ModuleType(f"{self.package}.{suffix}") + setattr(module, class_name, type(class_name, (), {})) + module.__dict__.update(attributes) + self.allowed[module.__name__] = module + return getattr(module, class_name) + + def find_spec(self, fullname, path=None, target=None): + suffix = fullname.removeprefix(f"{self.package}.") + if not fullname.startswith(f"{self.package}.") or "." not in suffix: + return None + if suffix.split(".")[0] not in {"triton", "fa3", "flashinfer", "nsa", "linear", "sdpa", "xformers"}: + return None + self.imported.append(fullname) + if fullname not in self.allowed: + raise AssertionError(f"Unexpected attention backend import: {fullname}") + return importlib.util.spec_from_loader(fullname, self) + + def create_module(self, spec): + return self.allowed[spec.name] + + def exec_module(self, module): + pass + + +class TestLazyAttentionImports(unittest.TestCase): + def setUp(self): + # Isolate attention from basemodel's unrelated GPU initialization. + modules_patch = patch.dict(sys.modules) + modules_patch.start() + self.addCleanup(modules_patch.stop) + for name in list(sys.modules): + if name == ATTENTION or name.startswith(f"{ATTENTION}."): + del sys.modules[name] + for name in ("lightllm", "lightllm.common", "lightllm.common.basemodel", "lightllm.utils"): + module = ModuleType(name) + module.__path__ = [str(ROOT.joinpath(*name.split(".")))] + sys.modules[name] = module + + self.args = SimpleNamespace( + llm_kv_type="None", + llm_prefill_att_backend=["triton"], + llm_decode_att_backend=["triton"], + enable_ep_moe=False, + mtp_step=0, + ssm_state_mode="native", + ) + self.validate = Mock(return_value=True) + dependencies = { + f"{ATTENTION}.base_att": { + name: type(name, (), {}) + for name in ("BaseAttBackend", "BasePrefillAttState", "BaseDecodeAttState", "AttControl") + }, + "lightllm.utils.envs_utils": {"get_env_start_args": lambda: self.args}, + "lightllm.utils.backend_validator": {"validate": self.validate}, + "lightllm.utils.log_utils": {"init_logger": lambda name: Mock()}, + } + for name, attributes in dependencies.items(): + module = ModuleType(name) + module.__dict__.update(attributes) + sys.modules[name] = module + + self.imports = BackendImports() + sys.meta_path.insert(0, self.imports) + self.addCleanup(sys.meta_path.remove, self.imports) + self.attention = importlib.import_module(ATTENTION) + + def test_importing_selection_apis_loads_no_implementations(self): + importlib.import_module(f"{ATTENTION}.linear") + importlib.import_module(f"{ATTENTION}.nsa") + self.assertEqual(self.imports.imported, []) + + def test_explicit_regular_selection_skips_neo(self): + backend = self.imports.allow("fa3.fp", "Fa3AttBackend") + self.args.llm_prefill_att_backend = ["fa3"] + + self.assertIs(self.attention.get_prefill_att_backend_class(), backend) + self.assertEqual(self.imports.imported, [f"{ATTENTION}.fa3.fp"]) + self.validate.assert_not_called() + + def test_auto_selection_does_not_import_failed_candidates(self): + backend = self.imports.allow("flashinfer.fp", "FlashInferAttBackend") + self.args.llm_prefill_att_backend = ["auto"] + self.validate.side_effect = lambda name: name == "flashinfer" + + self.assertIs(self.attention.get_prefill_att_backend_class(), backend) + self.assertEqual([call.args[0] for call in self.validate.call_args_list], ["fa3", "flashinfer"]) + self.assertEqual(self.imports.imported, [f"{ATTENTION}.flashinfer.fp"]) + + def test_mtp_decode_preserves_fa3_priority(self): + backend = self.imports.allow("fa3.fp", "Fa3AttBackend") + self.args.llm_decode_att_backend = ["auto"] + self.args.mtp_step = 2 + + self.assertIs(self.attention.get_decode_att_backend_class(), backend) + self.validate.assert_called_once_with("fa3") + + def test_expert_parallelism_still_excludes_flashinfer(self): + backend = self.imports.allow("triton.fp", "TritonAttBackend") + self.args.llm_decode_att_backend = ["auto"] + self.args.enable_ep_moe = True + self.validate.side_effect = lambda name: name == "triton" + + self.assertIs(self.attention.get_decode_att_backend_class(), backend) + self.assertEqual([call.args[0] for call in self.validate.call_args_list], ["fa3", "triton"]) + + def test_auto_selection_falls_back_to_triton(self): + backend = self.imports.allow("triton.mla", "MlaTritonAttBackend") + self.args.llm_prefill_att_backend = ["auto"] + self.validate.return_value = False + + self.assertIs(self.attention.get_mla_prefill_att_backend_class(), backend) + self.assertEqual(self.imports.imported, [f"{ATTENTION}.triton.mla"]) + + def test_nsa_selection_does_not_import_other_kv_types(self): + backend = self.imports.allow("nsa.fp8_flashmla_sparse", "NsaFlashMlaFp8SparseAttBackend") + self.args.llm_kv_type = "fp8kv_dsa" + self.args.llm_decode_att_backend = ["flashmla_sparse"] + + self.assertIs(self.attention.get_nsa_decode_att_backend_class(), backend) + self.assertEqual(self.imports.imported, [f"{ATTENTION}.nsa.fp8_flashmla_sparse"]) + + def test_custom_backend_map_uses_import_paths(self): + backend = self.imports.allow("nsa.custom", "CustomBackend") + backend_map = {"None": {"custom": "nsa.custom:CustomBackend"}} + self.args.llm_prefill_att_backend = ["custom"] + self.args.llm_decode_att_backend = ["auto"] + + self.assertIs(self.attention.get_nsa_prefill_att_backend_class(backend_map=backend_map), backend) + self.assertIs( + self.attention.get_nsa_decode_att_backend_class(priority_list=["custom"], backend_map=backend_map), + backend, + ) + self.assertEqual(self.imports.imported, [f"{ATTENTION}.nsa.custom"]) + + def test_neo_triton_selection_does_not_probe_fa3(self): + backend = self.imports.allow("triton.neo", "NeoTritonAttBackend") + + self.assertIs(self.attention.get_neo_prefill_att_backend_class(), backend) + self.args.llm_prefill_att_backend = ["auto"] + self.assertIs(self.attention.get_neo_prefill_att_backend_class(priority_list=["triton"]), backend) + self.assertEqual(self.imports.imported, [f"{ATTENTION}.triton.neo"]) + + def test_neo_auto_selection_still_falls_back_without_image_token_interface(self): + self.imports.allow("fa3.neo", "NeoFa3AttBackend", HAS_FLASH_ATTN_INTERFACE=False) + backend = self.imports.allow("triton.neo", "NeoTritonAttBackend") + self.args.llm_prefill_att_backend = ["auto"] + + self.assertIs(self.attention.get_neo_prefill_att_backend_class(), backend) + self.validate.assert_called_once_with("triton") + self.args.llm_prefill_att_backend = ["fa3"] + with self.assertRaises(KeyError): + self.attention.get_neo_prefill_att_backend_class() + + def test_neo_auto_selection_uses_supported_fa3(self): + backend = self.imports.allow("fa3.neo", "NeoFa3AttBackend", HAS_FLASH_ATTN_INTERFACE=True) + self.args.llm_prefill_att_backend = ["auto"] + + self.assertIs(self.attention.get_neo_prefill_att_backend_class(), backend) + self.assertEqual(self.imports.imported, [f"{ATTENTION}.fa3.neo"]) + self.validate.assert_called_once_with("fa3") + + def _allow_linear_backend_classes(self): + flashqla = self.imports.allow("linear.flashqla", "FlashQlaLinearAttBackend") + triton = self.imports.allow("linear.triton", "TritonLinearAttBackend") + self.imports.allow("linear.replay_gdn", "ReplayGDNTritonAttBackend") + replay_gdn = self.imports.allowed[f"{ATTENTION}.linear.replay_gdn"] + replay_gdn.ReplayGDNFlashQlaAttBackend = type("ReplayGDNFlashQlaAttBackend", (), {}) + self.imports.allow("linear.kda", "KDALinearAttBackend") + self.imports.allow("linear.replay_kda", "ReplayKDALinearAttBackend") + return flashqla, triton + + def test_linear_prefill_auto_selects_validated_flashqla(self): + flashqla, _triton = self._allow_linear_backend_classes() + selection = importlib.import_module(f"{ATTENTION}.create_linear_utils") + + self.assertIs(selection.get_gdn_linear_prefill_att_backend_class(), flashqla) + self.validate.assert_called_once_with("flashqla") + + def test_linear_selection_falls_back_to_triton(self): + _flashqla, triton = self._allow_linear_backend_classes() + selection = importlib.import_module(f"{ATTENTION}.create_linear_utils") + self.validate.return_value = False + + self.assertIs(selection.get_gdn_linear_prefill_att_backend_class(), triton) + self.assertIs(selection.get_gdn_linear_decode_att_backend_class(), triton) + self.args.llm_prefill_att_backend.append("triton") + self.assertIs(selection.get_gdn_linear_prefill_att_backend_class(), triton) + + def test_vit_selection_loads_only_requested_implementation(self): + package = "lightllm.common.basemodel.attention_vit" + for name in list(sys.modules): + if name == package or name.startswith(f"{package}."): + del sys.modules[name] + base = ModuleType(f"{package}.base_att") + base.BaseVitAttBackend = type("BaseVitAttBackend", (), {}) + sys.modules[base.__name__] = base + sys.modules["lightllm.utils.backend_validator"]._validate = self.validate + imports = BackendImports(package) + sys.meta_path.insert(0, imports) + self.addCleanup(sys.meta_path.remove, imports) + selection = importlib.import_module(f"{package}.create_utils") + self.assertEqual(imports.imported, []) + + for name, class_name in ( + ("fa3", "Fa3VitAttBackend"), + ("triton", "TritonVitAttBackend"), + ("sdpa", "SdpaVitAttBackend"), + ("xformers", "XformersVitAttBackend"), + ): + backend = imports.allow(f"{name}.fp", class_name) + self.args.vit_att_backend = [name] + self.assertEqual(selection.init_vit_att_backend(), name) + self.assertIs(selection.get_vit_att_backend_class(name), backend) + self.validate.assert_not_called() + + self.args.vit_att_backend = ["auto"] + self.validate.side_effect = lambda name: name == "sdpa" + self.assertEqual(selection.init_vit_att_backend(), "sdpa") + self.assertEqual([call.args[0] for call in self.validate.call_args_list], ["fa3", "xformers", "sdpa"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/unit_tests/common/basemodel/attention/test_state_lifetime.py b/unit_tests/common/basemodel/attention/test_state_lifetime.py new file mode 100644 index 0000000000..9e049d54b4 --- /dev/null +++ b/unit_tests/common/basemodel/attention/test_state_lifetime.py @@ -0,0 +1,174 @@ +import gc +import weakref + +import pytest +import torch + +from lightllm.common.basemodel.attention.base_att import BaseDecodeAttState, BasePrefillAttState +from lightllm.common.basemodel.cuda_graph import CudaGraph +from lightllm.common.basemodel.infer_struct import InferStateInfo + + +class _PrefillState(BasePrefillAttState): + def init_state(self): + self.positions = self.infer_state.position_ids + + def prefill_att(self, q, k, v, **kwargs): + return q + + +class _DecodeState(BaseDecodeAttState): + def init_state(self): + self.positions = self.infer_state.position_ids + + def decode_att(self, q, k, v, **kwargs): + return q + + +@pytest.fixture +def no_cycle_gc(): + enabled = gc.isenabled() + gc.collect() + gc.disable() + try: + yield + finally: + gc.collect() + if enabled: + gc.enable() + + +def _create_state(is_prefill, secondary=False): + owner = InferStateInfo() + owner.is_prefill = is_prefill + owner.position_ids = torch.arange(4) + state_cls = _PrefillState if is_prefill else _DecodeState + slot = "prefill_att_state" if is_prefill else "decode_att_state" + setattr(owner, slot, state_cls(_infer_state_weak_ref=weakref.ref(owner))) + if secondary: + setattr(owner, slot + "1", state_cls(_infer_state_weak_ref=weakref.ref(owner))) + owner.init_att_state() + return owner + + +@pytest.mark.parametrize("is_prefill", [False, True]) +@pytest.mark.parametrize("secondary", [False, True]) +def test_forward_states_release_without_cycle_gc(no_cycle_gc, is_prefill, secondary): + owner = _create_state(is_prefill, secondary) + states = [state for state in vars(owner).values() if isinstance(state, (BasePrefillAttState, BaseDecodeAttState))] + owner_ref = weakref.ref(owner) + tensor_ref = weakref.ref(owner.position_ids) + state_refs = [weakref.ref(state) for state in states] + + # Even an external reference to attention metadata must not retain its owner. + del owner + assert owner_ref() is None + with pytest.raises(ReferenceError): + states[0].infer_state + assert tensor_ref() is not None + del states + assert all(ref() is None for ref in state_refs) + assert tensor_ref() is None + + +@pytest.mark.parametrize("state_cls", [_PrefillState, _DecodeState]) +def test_attention_getter_uses_a_weak_reference_and_is_read_only(state_cls): + state = state_cls() + assert state.infer_state is None + + owner = InferStateInfo() + state = state_cls(_infer_state_weak_ref=weakref.ref(owner)) + assert isinstance(state._infer_state_weak_ref, weakref.ReferenceType) + assert state.infer_state is owner + state.check_no_infer_state() + with pytest.raises(AttributeError): + state.infer_state = owner + + +@pytest.mark.parametrize("is_prefill", [False, True]) +@pytest.mark.parametrize("slot_suffix", ["", "1"]) +def test_check_rejects_infer_state_under_any_attribute_name(is_prefill, slot_suffix): + owner = _create_state(is_prefill, secondary=True) + slot = ("prefill_att_state" if is_prefill else "decode_att_state") + slot_suffix + attention = getattr(owner, slot) + + class _OtherInferState(InferStateInfo): + pass + + # Check the type, rather than just the identity of the expected owner. + attention.saved_context = _OtherInferState() + with pytest.raises(RuntimeError, match="saved_context"): + owner.check_att_states() + + +@pytest.mark.parametrize("is_prefill", [False, True]) +@pytest.mark.parametrize("slot_suffix", ["", "1"]) +def test_init_checks_for_accidental_back_reference(is_prefill, slot_suffix): + owner = _create_state(is_prefill, secondary=True) + slot = ("prefill_att_state" if is_prefill else "decode_att_state") + slot_suffix + base = _PrefillState if is_prefill else _DecodeState + + class _RetainingState(base): + def init_state(self): + self.saved_context = self.infer_state + + setattr(owner, slot, _RetainingState(_infer_state_weak_ref=weakref.ref(owner))) + try: + with pytest.raises(RuntimeError, match="saved_context"): + owner.init_att_state() + finally: + del getattr(owner, slot).saved_context + + +def test_decode_graph_replay_preserves_capture_buffers_and_owner_lifetime(no_cycle_gc): + captured = _create_state(is_prefill=False, secondary=True) + captured.input_ids = torch.zeros(4, dtype=torch.int64) + captured_ref = weakref.ref(captured) + capture_positions_ptr = captured.position_ids.data_ptr() + replay_calls = [] + + class _Graph: + def replay(self): + replay_calls.append(True) + + graph = CudaGraph.__new__(CudaGraph) + output = object() + graph.graph = {4: (_Graph(), captured, output)} + del captured + assert captured_ref() is not None + + runtime = _create_state(is_prefill=False, secondary=True) + runtime.input_ids = torch.ones(4, dtype=torch.int64) + runtime.position_ids.add_(7) + runtime_ref = weakref.ref(runtime) + assert graph._replay(runtime) is output + assert replay_calls == [True] + assert captured_ref().position_ids.data_ptr() == capture_positions_ptr + assert captured_ref().decode_att_state.infer_state is captured_ref() + assert captured_ref().decode_att_state1.infer_state is captured_ref() + torch.testing.assert_close(captured_ref().position_ids, runtime.position_ids) + + del runtime + 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) diff --git a/unit_tests/common/basemodel/test_cuda_graph_layout.py b/unit_tests/common/basemodel/test_cuda_graph_layout.py new file mode 100644 index 0000000000..30d4c54368 --- /dev/null +++ b/unit_tests/common/basemodel/test_cuda_graph_layout.py @@ -0,0 +1,133 @@ +import math +from types import SimpleNamespace + +import pytest + +import lightllm.common.basemodel.basemodel as basemodel_module +import lightllm.common.basemodel.cuda_graph as cuda_graph_module +import lightllm.common.basemodel.mtp_manager as mtp_manager_module +from lightllm.common.basemodel.basemodel import TpPartBaseModel +from lightllm.common.basemodel.cuda_graph import CudaGraph +from lightllm.common.basemodel.mtp_manager import MtpManager + + +@pytest.fixture(autouse=True) +def _graph_args(monkeypatch): + args = SimpleNamespace( + enable_decode_microbatch_overlap=False, + enable_tpsp_mix_mode=False, + enable_torch_memory_saver=False, + ) + monkeypatch.setattr(cuda_graph_module, "get_env_start_args", lambda: args) + return args + + +def _batch_sizes(max_batch_size, batch_stride=1): + physical_max_batch_size = max_batch_size * batch_stride + graph = CudaGraph( + batch_step_size_before_split=batch_stride, + split_batch_size=4 * batch_stride, + batch_step_size_after_split=2 * batch_stride, + max_batch_size=physical_max_batch_size, + ) + return graph.cuda_graph_batch_sizes + + +def test_dynamic_schedule_uses_compacted_physical_rows(_graph_args): + assert _batch_sizes(max_batch_size=128) == [1, 2, 3, 4, *range(6, 129, 2)] + + +def test_public_static_schedule_preserves_original_static_mtp_default(_graph_args): + assert CudaGraph.gen_cuda_graph_batch_sizes( + batch_step_size_before_split=8, + split_batch_size=32, + batch_step_size_after_split=16, + max_batch_size=32, + ) == [ + 8, + 16, + 24, + 32, + ] + + +def test_instance_and_public_static_schedule_match(_graph_args): + graph = CudaGraph( + batch_step_size_before_split=8, + split_batch_size=32, + batch_step_size_after_split=16, + max_batch_size=128, + ) + + assert graph.cuda_graph_batch_sizes == CudaGraph.gen_cuda_graph_batch_sizes( + batch_step_size_before_split=8, + split_batch_size=32, + batch_step_size_after_split=16, + max_batch_size=graph.max_batch_size, + tp_world_size=graph.tp_world_size, + ) + + +def test_batch_step_size_before_split_controls_capture_range(_graph_args): + assert _batch_sizes(max_batch_size=4, batch_stride=8) == [8, 16, 24, 32] + + +def test_batch_step_size_after_split_controls_capture_range(_graph_args): + assert _batch_sizes(max_batch_size=8, batch_stride=7) == [ + 7, + 14, + 21, + 28, + 42, + 56, + ] + + +@pytest.mark.parametrize("tp_size", [1, 2, 4, 8]) +@pytest.mark.parametrize("mtp_step", [1, 2]) +@pytest.mark.parametrize("dynamic,is_draft", [(False, False), (True, False), (False, True)]) +@pytest.mark.parametrize("overlap", [False, True]) +def test_mtp_tpsp_layout(monkeypatch, _graph_args, tp_size, mtp_step, dynamic, is_draft, overlap): + args = _graph_args + args.enable_tpsp_mix_mode = True + args.enable_decode_microbatch_overlap = overlap + args.mtp_mode = "eagle_with_att" + args.mtp_step = mtp_step + args.mtp_dynamic_verify = dynamic + monkeypatch.setattr(mtp_manager_module, "get_env_start_args", lambda: args) + model = TpPartBaseModel.__new__(TpPartBaseModel) + model.is_mtp_draft_model = is_draft + monkeypatch.setattr(basemodel_module, "get_env_start_args", lambda: args) + monkeypatch.setattr(basemodel_module, "get_llm_data_type", lambda: None) + monkeypatch.setattr(basemodel_module, "get_dp_world_size", lambda: tp_size) + monkeypatch.setattr(MtpManager, "_instance", MtpManager()) + + class StopBeforeWeights(Exception): + pass + + def stop_before_weights(): + raise StopBeforeWeights + + # 执行真实构造函数的容量计算，在读取模型配置、分配 GPU 权重前停止。 + monkeypatch.setattr(model, "_init_config", stop_before_weights) + with pytest.raises(StopBeforeWeights): + model.__init__(dict(run_mode="normal", weight_dir="", max_total_token_num=1024, graph_max_batch_size=7)) + + width = 1 if dynamic or is_draft else mtp_step + 1 + logical_max = 7 // 2 if overlap else 7 + physical_max = logical_max * (1 if is_draft else mtp_step + 1) + alignment = math.lcm(width, tp_size) + assert model.graph_max_batch_size % alignment == 0 + assert physical_max <= model.graph_max_batch_size < physical_max + alignment + + sizes = CudaGraph.gen_cuda_graph_batch_sizes( + batch_step_size_before_split=width, + split_batch_size=4 * width, + batch_step_size_after_split=2 * width, + max_batch_size=model.graph_max_batch_size, + tp_world_size=tp_size, + ) + assert sizes[-1] == model.graph_max_batch_size + assert all(size % width == 0 and size % tp_size == 0 for size in sizes) + if tp_size == 8 and width == 3: + assert sizes == [24] diff --git a/unit_tests/common/basemodel/test_hidden_collector.py b/unit_tests/common/basemodel/test_hidden_collector.py new file mode 100644 index 0000000000..22748b7294 --- /dev/null +++ b/unit_tests/common/basemodel/test_hidden_collector.py @@ -0,0 +1,193 @@ +from inspect import isabstract +from types import SimpleNamespace + +import pytest +import torch + +import lightllm.common.basemodel.hidden_collector as hidden_collector_module +from lightllm.common.basemodel.hidden_collector import ( + FinalHiddenCollector, + HiddenCollector, + LayerHiddenCollector, + MtpHeadOutputCollector, + NoopHiddenCollector, +) + + +class _IdentityPreInfer: + @staticmethod + def _tpsp_allgather(input, infer_state): + del infer_state + return input + + +def _mock_target_layer_ids(monkeypatch, layer_ids): + monkeypatch.setattr( + hidden_collector_module, + "get_env_start_args", + lambda: SimpleNamespace(mtp_draft_model_dir=["/models/draft"]), + ) + monkeypatch.setattr( + hidden_collector_module.PretrainedConfig, + "get_config_dict", + lambda _: ({"target_layer_ids": layer_ids}, {}), + ) + + +def test_hidden_collector_is_base_class_for_implementations(): + assert isabstract(HiddenCollector) + assert issubclass(NoopHiddenCollector, HiddenCollector) + assert issubclass(FinalHiddenCollector, HiddenCollector) + assert issubclass(LayerHiddenCollector, HiddenCollector) + + +def test_final_hidden_collectors_are_independent_instances(): + prototype = FinalHiddenCollector() + collector0 = prototype.new_instance() + collector1 = prototype.new_instance() + hidden0 = torch.randn(2, 3) + hidden1 = torch.randn(2, 3) + infer_state = SimpleNamespace(need_dp_prefill_balance=False) + + collector0.add_final_hidden(hidden0) + collected = collector0.finish_output(infer_state=infer_state).spec_hidden + assert collected.data_ptr() == hidden0.data_ptr() + assert collector0.final_hidden is None + + collector0.add_final_hidden(hidden0) + collector1.add_final_hidden(hidden1) + collected0 = collector0.finish_output(infer_state=infer_state).spec_hidden + collected1 = collector1.finish_output(infer_state=infer_state).spec_hidden + assert collected0.data_ptr() == hidden0.data_ptr() + assert collected1.data_ptr() == hidden1.data_ptr() + + +def test_layer_hidden_collector_keeps_microbatch_state_separate(monkeypatch): + _mock_target_layer_ids(monkeypatch, [0]) + model = SimpleNamespace(is_mtp_draft_model=False, layers_num=2, pre_infer=_IdentityPreInfer()) + prototype = LayerHiddenCollector(model=model) + collector0 = prototype.new_instance() + collector1 = prototype.new_instance() + hidden0 = torch.full((2, 3), 1.0) + hidden1 = torch.full((2, 3), 2.0) + infer_state = SimpleNamespace(need_dp_prefill_balance=False) + + collector0.add(layer_index=0, hidden=hidden0) + collector1.add(layer_index=0, hidden=hidden1) + + collected0 = collector0.finish_output(infer_state=infer_state).spec_hidden + collected1 = collector1.finish_output(infer_state=infer_state).spec_hidden + + assert torch.equal(collected0, hidden0) + assert torch.equal(collected1, hidden1) + + +def test_layer_hidden_collector_requires_target_layer_ids_from_config(monkeypatch): + _mock_target_layer_ids(monkeypatch, None) + model = SimpleNamespace(layers_num=2, pre_infer=_IdentityPreInfer()) + + with pytest.raises(AssertionError, match="target_layer_ids is required in draft config"): + LayerHiddenCollector(model=model) + + +def test_noop_collector_keeps_normal_forward_output_minimal(): + final_hidden = torch.randn(2, 3) + collector = NoopHiddenCollector() + + collector.add(layer_index=0, hidden=final_hidden) + collector.add_final_hidden(final_hidden) + + assert collector.finish_output(infer_state=None).spec_hidden is None + + +def test_mtp_head_output_collector_returns_and_clears_outputs(): + collector = MtpHeadOutputCollector() + draft_token_ids = torch.arange(6) + confidence_logits = torch.arange(6).view(2, 3) + + collector.add_mtp_outputs( + draft_token_ids=draft_token_ids, + confidence_logits=confidence_logits, + ) + output = collector.finish_output(infer_state=None) + + assert output.spec_hidden is None + assert output.draft_token_ids is draft_token_ids + assert output.confidence_logits is confidence_logits + assert collector.draft_token_ids is None + assert collector.confidence_logits is None + + +def test_final_collector_returns_final_hidden_without_layer_bookkeeping(): + final_hidden = torch.randn(2, 3) + collector = FinalHiddenCollector() + collector.add_final_hidden(final_hidden) + collected = collector.finish_output(infer_state=None).spec_hidden + + assert collected.data_ptr() == final_hidden.data_ptr() + + +def test_layer_collector_preserves_selected_layers_in_model_order(monkeypatch): + _mock_target_layer_ids(monkeypatch, [0, 2]) + layer0 = torch.full((2, 2), 1.0) + layer1 = torch.full((2, 2), 2.0) + layer2 = torch.full((2, 2), 3.0) + model = SimpleNamespace(layers_num=3, pre_infer=_IdentityPreInfer()) + collector = LayerHiddenCollector(model=model) + + collector.add(layer_index=0, hidden=layer0) + collector.add(layer_index=1, hidden=layer1) + collector.add(layer_index=2, hidden=layer2) + layer0.fill_(9.0) + + collected = collector.finish_output(infer_state=SimpleNamespace(need_dp_prefill_balance=False)).spec_hidden + + assert torch.equal(collected, torch.cat([torch.full((2, 2), 1.0), layer2], dim=-1)) + assert not collector.layer_hiddens + + collector.add(layer_index=0, hidden=layer0) + collector.add(layer_index=2, hidden=layer2) + collected = collector.finish_output(infer_state=SimpleNamespace(need_dp_prefill_balance=False)).spec_hidden + + assert torch.equal(collected, torch.cat([layer0, layer2], dim=-1)) + assert not collector.layer_hiddens + + +def test_layer_collector_restores_graph_state_without_sharing_runtime_container(monkeypatch): + _mock_target_layer_ids(monkeypatch, [0]) + model = SimpleNamespace(layers_num=2, pre_infer=_IdentityPreInfer()) + graph_collector = LayerHiddenCollector(model=model) + collector = graph_collector.new_instance() + infer_state = SimpleNamespace(need_dp_prefill_balance=False) + graph_collector.add(layer_index=0, hidden=torch.full((2, 3), 1.0)) + collector.restore_graph_state(graph_collector) + + collected = collector.finish_output(infer_state=infer_state).spec_hidden + + assert torch.equal(collected, torch.full((2, 3), 1.0)) + assert not collector.layer_hiddens + assert len(graph_collector.layer_hiddens) == 1 + + +def test_layer_collector_releases_graph_tensor_ownership_with_statistics(monkeypatch): + _mock_target_layer_ids(monkeypatch, [0, 2]) + model = SimpleNamespace(layers_num=3, pre_infer=_IdentityPreInfer()) + collector = LayerHiddenCollector(model=model) + hidden0 = torch.randn(2, 3) + hidden2 = torch.randn(2, 3) + converted = [] + + def to_no_ref(hidden): + converted.append(hidden) + return hidden + + monkeypatch.setattr(hidden_collector_module, "tensor_to_no_ref_tensor", to_no_ref) + collector.add(layer_index=0, hidden=hidden0) + collector.add(layer_index=2, hidden=hidden2) + + tensor_count, total_nbytes = collector.release_graph_tensor_ownership() + + assert tensor_count == 2 + assert total_nbytes == hidden0.numel() * hidden0.element_size() + hidden2.numel() * hidden2.element_size() + assert all(actual is expected for actual, expected in zip(converted, collector.layer_hiddens)) + assert NoopHiddenCollector().release_graph_tensor_ownership() == (0, 0) diff --git a/unit_tests/common/basemodel/test_model_input.py b/unit_tests/common/basemodel/test_model_input.py new file mode 100644 index 0000000000..6672cccc31 --- /dev/null +++ b/unit_tests/common/basemodel/test_model_input.py @@ -0,0 +1,212 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.basemodel.basemodel import TpPartBaseModel +from lightllm.common.basemodel.batch_objs import ModelInput + + +def _create_model_input(*, is_prefill=False): + batch_size = 2 + kwargs = dict( + batch_size=batch_size, + total_token_num=batch_size, + max_q_seq_len=1, + max_kv_seq_len=1, + b_req_idx=torch.arange(batch_size, dtype=torch.int32), + b_mtp_index=torch.zeros(batch_size, dtype=torch.int32), + b_seq_len=torch.ones(batch_size, dtype=torch.int32), + is_prefill=is_prefill, + multimodal_params=[{"images": [], "audios": []} for _ in range(batch_size)], + ) + if is_prefill: + kwargs["max_cache_len"] = 0 + kwargs["input_ids"] = torch.ones(batch_size, dtype=torch.int64) + kwargs["b_ready_cache_len"] = torch.zeros(batch_size, dtype=torch.int32) + kwargs["b_prefill_start_loc"] = torch.arange(batch_size, dtype=torch.int32) + kwargs["b_is_decode_req"] = torch.zeros(batch_size, dtype=torch.bool) + kwargs["b_prefill_has_output_cpu"] = [False] * batch_size + else: + kwargs["b_position_delta"] = torch.zeros(batch_size, dtype=torch.int32) + kwargs["b_shared_seq_len"] = torch.tensor([4, 4], dtype=torch.int32) + kwargs["b_shared_radix_node_id"] = torch.tensor([10, 10], dtype=torch.int64) + return ModelInput(**kwargs) + + +def test_decode_requires_shared_radix_metadata(): + with pytest.raises(AssertionError): + ModelInput( + batch_size=1, + total_token_num=1, + max_q_seq_len=1, + max_kv_seq_len=1, + b_req_idx=torch.zeros(1, dtype=torch.int32), + b_mtp_index=torch.zeros(1, dtype=torch.int32), + b_seq_len=torch.ones(1, dtype=torch.int32), + b_position_delta=torch.zeros(1, dtype=torch.int32), + is_prefill=False, + multimodal_params=[{"images": [], "audios": []}], + ) + + +def test_decode_carries_raw_shared_radix_metadata(): + model_input = _create_model_input() + + assert torch.equal(model_input.b_shared_seq_len, torch.tensor([4, 4], dtype=torch.int32)) + assert torch.equal(model_input.b_shared_radix_node_id, torch.tensor([10, 10], dtype=torch.int64)) + + +def test_prefill_does_not_require_shared_radix_metadata(): + model_input = _create_model_input(is_prefill=True) + + assert model_input.b_shared_seq_len is None + assert model_input.b_shared_radix_node_id is None + + +def test_decode_requires_position_delta(): + model_input = _create_model_input() + model_input.b_position_delta = None + + with pytest.raises(AssertionError): + model_input.check_input() + + +def test_prefill_requires_prefill_metadata(): + model_input = _create_model_input(is_prefill=True) + model_input.b_ready_cache_len = None + + with pytest.raises(AssertionError): + model_input.check_input() + + +def test_prefill_requires_prefill_output_markers(): + model_input = _create_model_input(is_prefill=True) + model_input.b_prefill_has_output_cpu = None + + with pytest.raises(AssertionError, match="prefill must provide b_prefill_has_output_cpu"): + model_input.check_input() + + +def test_prefill_requires_decode_request_markers(): + model_input = _create_model_input(is_prefill=True) + model_input.b_is_decode_req = None + + with pytest.raises(AssertionError): + model_input.check_input() + + +def test_prefill_rejects_position_delta(): + model_input = _create_model_input(is_prefill=True) + model_input.b_position_delta = torch.zeros(model_input.batch_size, dtype=torch.int32) + + with pytest.raises(AssertionError, match="prefill must not provide b_position_delta"): + model_input.check_input() + + +def test_padded_prefill_adds_non_decode_request_marker(): + model_input = ModelInput( + batch_size=1, + total_token_num=2, + max_q_seq_len=2, + max_kv_seq_len=2, + max_cache_len=0, + input_ids=torch.ones(2, dtype=torch.int64), + b_req_idx=torch.zeros(1, dtype=torch.int32), + b_mtp_index=torch.zeros(1, dtype=torch.int32), + b_seq_len=torch.full((1,), 2, dtype=torch.int32), + b_is_decode_req=torch.ones(1, dtype=torch.bool), + 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=[False], + multimodal_params=[{"images": [], "audios": []}], + ) + model = SimpleNamespace( + mem_manager=SimpleNamespace(HOLD_TOKEN_MEMINDEXES=(-1,), page_size=1), + req_manager=SimpleNamespace(HOLD_REQUEST_ID=-1), + ) + + padded_input = TpPartBaseModel._create_padded_prefill_model_input( + model, + model_input=model_input, + new_handle_token_num=4, + ) + + assert padded_input.b_is_decode_req.dtype == torch.bool + assert padded_input.b_is_decode_req.tolist() == [True, False] + + +def test_padded_prefill_builds_internal_request_for_empty_input(): + model_input = ModelInput( + batch_size=0, + total_token_num=0, + max_q_seq_len=0, + max_kv_seq_len=0, + max_cache_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_is_decode_req=torch.empty((0,), dtype=torch.bool), + b_ready_cache_len=torch.empty((0,), dtype=torch.int32), + b_prefill_start_loc=torch.empty((0,), dtype=torch.int32), + is_prefill=True, + b_prefill_has_output_cpu=[], + multimodal_params=[], + mtp_draft_input_hiddens=torch.empty((0, 4), dtype=torch.float32), + ) + model = SimpleNamespace( + mem_manager=SimpleNamespace(HOLD_TOKEN_MEMINDEXES=(99,), page_size=1), + req_manager=SimpleNamespace(HOLD_REQUEST_ID=88), + ) + + padded_input = TpPartBaseModel._create_padded_prefill_model_input( + model, + model_input=model_input, + new_handle_token_num=1, + ) + + assert model_input.batch_size == 0 + assert padded_input.batch_size == 1 + assert padded_input.input_ids.tolist() == [1] + assert padded_input.b_req_idx.tolist() == [88] + assert padded_input.b_seq_len.tolist() == [1] + assert padded_input.b_prefill_has_output_cpu == [False] + assert torch.equal(padded_input.mtp_draft_input_hiddens, torch.zeros((1, 4))) + + +def test_padded_decode_builds_internal_request_from_empty_token_tensor(): + model_input = 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=[], + ) + model = SimpleNamespace( + mem_manager=SimpleNamespace(HOLD_TOKEN_MEMINDEXES=(99,), page_size=1), + req_manager=SimpleNamespace(HOLD_REQUEST_ID=88), + ) + + padded_input = TpPartBaseModel._create_padded_decode_model_input( + model, + model_input=model_input, + new_batch_size=1, + ) + + assert model_input.batch_size == 0 + assert padded_input.batch_size == 1 + assert padded_input.input_ids.tolist() == [1] + assert padded_input.b_req_idx.tolist() == [88] + assert padded_input.b_seq_len.tolist() == [2] + assert padded_input.b_shared_seq_len.tolist() == [0] + assert padded_input.b_shared_radix_node_id.tolist() == [-1] diff --git a/unit_tests/common/basemodel/test_model_output.py b/unit_tests/common/basemodel/test_model_output.py new file mode 100644 index 0000000000..123f603ff2 --- /dev/null +++ b/unit_tests/common/basemodel/test_model_output.py @@ -0,0 +1,344 @@ +from types import SimpleNamespace +import weakref + +import torch +import pytest + +from lightllm.common.basemodel import basemodel +from lightllm.common.basemodel.attention.flashinfer.fp import FlashInferDecodeAttState +from lightllm.common.basemodel.attention.flashinfer.utils import should_init_decode_wrapper +from lightllm.common.basemodel.basemodel import TpPartBaseModel +from lightllm.common.basemodel.batch_objs import ModelInput, ModelMtpOutputCollector, ModelOutput +from lightllm.common.basemodel.infer_struct import InferStateInfo + + +def test_decode_unpad_slices_spec_output_with_logits(): + model = TpPartBaseModel.__new__(TpPartBaseModel) + output = ModelOutput( + logits=torch.arange(24).view(6, 4), + mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.arange(18).view(6, 3)), + ) + + unpadded = model._create_unpad_decode_model_output(output, origin_batch_size=4) + + assert unpadded.logits.shape == (4, 4) + assert unpadded.mtp_collector.spec_hidden.shape == (4, 3) + # Unpadding returns a shallow output copy and leaves the graph-owned + # tensors on the original ModelOutput intact. + assert output.logits.shape == (6, 4) + assert output.mtp_collector.spec_hidden.shape == (6, 3) + + +def test_prefill_unpad_uses_token_rows_for_spec_hidden(): + model = TpPartBaseModel.__new__(TpPartBaseModel) + output = ModelOutput( + logits=torch.arange(20).view(5, 4), + mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.arange(24).view(8, 3)), + prompt_logics=torch.arange(32).view(8, 4), + ) + + unpadded = model._create_unpad_prefill_model_output( + output, + origin_handle_token_num=6, + origin_batch_size=3, + ) + + assert unpadded.logits.shape == (3, 4) + assert unpadded.mtp_collector.spec_hidden.shape == (6, 3) + assert unpadded.prompt_logics.shape == (6, 4) + + +def test_decode_unpad_restores_empty_output(): + model = TpPartBaseModel.__new__(TpPartBaseModel) + output = ModelOutput( + logits=torch.arange(4).view(1, 4), + mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.arange(3).view(1, 3)), + ) + + unpadded = model._create_unpad_decode_model_output(output, origin_batch_size=0) + + assert unpadded.logits.shape == (0, 4) + assert unpadded.mtp_collector.spec_hidden.shape == (0, 3) + + +def test_prefill_unpad_restores_empty_output(): + model = TpPartBaseModel.__new__(TpPartBaseModel) + output = ModelOutput( + logits=torch.arange(4).view(1, 4), + mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.arange(6).view(2, 3)), + prompt_logics=torch.arange(8).view(2, 4), + ) + + unpadded = model._create_unpad_prefill_model_output( + output, + origin_handle_token_num=0, + origin_batch_size=0, + ) + + assert unpadded.logits.shape == (0, 4) + assert unpadded.mtp_collector.spec_hidden.shape == (0, 3) + assert unpadded.prompt_logics.shape == (0, 4) + + +@torch.no_grad() +def test_decode_pads_only_once_after_selecting_execution_path(monkeypatch): + execution_configs = ( + # eager 普通模式：空 batch 只补一个 dummy request。 + (None, False, 1, False, 1), + # eager TPSP 模式：dummy request 仍需对齐到 TP world size。 + (None, True, 2, False, 2), + # CUDA Graph replay：使用最终选中的 graph batch size。 + (4, False, 1, False, 4), + # CUDA Graph capture：必须在初始化 attention state 前设置 graph 标记。 + (4, False, 1, True, 4), + ) + for ( + graph_batch_size, + enable_tpsp_mix_mode, + tp_world_size, + need_capture, + expected_batch_size, + ) in execution_configs: + model = TpPartBaseModel.__new__(TpPartBaseModel) + model.args = SimpleNamespace(enable_tpsp_mix_mode=enable_tpsp_mix_mode, page_size=1) + model.tp_world_size_ = tp_world_size + model.is_mtp_draft_model = False + model.mtp_manager = SimpleNamespace(get_decode_batch_alignment=lambda _: 1) + model.req_manager = SimpleNamespace(HOLD_REQUEST_ID=88, req_to_token_indexs=object()) + + pad_batch_sizes = [] + + def pad_once(model_input, new_batch_size): + pad_batch_sizes.append(new_batch_size) + return TpPartBaseModel._create_padded_decode_model_input(model, model_input, new_batch_size) + + model._create_padded_decode_model_input = pad_once + + graph_flags_at_att_init = [] + executions = [] + + def create_infer_state(model_input): + infer_state = SimpleNamespace( + b_req_idx=model_input.b_req_idx, + b_seq_len=model_input.b_seq_len, + mem_index=torch.empty((model_input.batch_size,), dtype=torch.int32), + init_some_extra_state=lambda _: None, + is_cuda_graph_capture=False, + is_cuda_graph_replay=False, + ) + infer_state.init_att_state = lambda: graph_flags_at_att_init.append( + (infer_state.is_cuda_graph_capture, infer_state.is_cuda_graph_replay) + ) + return infer_state + + def execute(mode, infer_state): + executions.append((mode, infer_state.is_cuda_graph_capture, infer_state.is_cuda_graph_replay)) + return ModelOutput(logits=torch.ones((infer_state.b_req_idx.shape[0], 4))) + + model._create_inferstate = create_infer_state + model._token_forward = lambda infer_state: execute("eager", infer_state) + + graph = None + if graph_batch_size is not None: + graph = SimpleNamespace() + graph.can_run = lambda **kwargs: True + graph.find_closest_graph_batch_size = lambda batch_size: graph_batch_size + graph.need_capture = lambda batch_size: need_capture + graph.capture_decode = lambda decode_func, infer_state: execute("capture", infer_state) + graph.replay = lambda infer_state: execute("replay", infer_state) + model.graph = graph + + output = model._decode(_create_empty_decode_input()) + + expected_replay = graph_batch_size is not None and not need_capture + expected_mode = "eager" + if graph_batch_size is not None: + expected_mode = "capture" if need_capture else "replay" + assert pad_batch_sizes == [expected_batch_size] + assert graph_flags_at_att_init == [(need_capture, expected_replay)] + assert executions == [(expected_mode, need_capture, expected_replay)] + assert output.logits.shape == (0, 4) + + +@pytest.mark.parametrize( + "lengths,graph_mode,expected_mode", + [ + pytest.param((1024, 2048), None, "eager", id="graph-disabled"), + pytest.param((1024, 2048), "capture", "capture", id="capture"), + pytest.param((1024, 2048), "replay", "replay", id="replay"), + pytest.param((1024, 16384), "replay", "eager", id="long-second-microbatch"), + pytest.param((16384, 1024), "capture", "eager", id="long-first-microbatch"), + ], +) +def test_overlap_decode_initializes_attention_for_the_selected_execution_mode(lengths, graph_mode, expected_mode): + model = TpPartBaseModel.__new__(TpPartBaseModel) + model.args = SimpleNamespace(enable_tpsp_mix_mode=True, page_size=1) + model.tp_world_size_ = 2 + model.is_mtp_draft_model = False + model.mtp_manager = SimpleNamespace(get_decode_batch_alignment=lambda _: 1) + model.req_manager = SimpleNamespace(HOLD_REQUEST_ID=88) + initializations = [] + executions = [] + + def graph_flags(state): + return state.is_cuda_graph_capture, state.is_cuda_graph_replay + + class DecodeState(FlashInferDecodeAttState): + def init_state(self): + infer_state = self.infer_state + # Replace the external GPU wrapper allocation while preserving the + # production decision and InferStateInfo's initialization entry point. + if should_init_decode_wrapper(infer_state): + self.decode_wrapper = object() + initializations.append( + (infer_state.microbatch_index, graph_flags(infer_state), self.decode_wrapper is not None) + ) + + def create_state(model_input, microbatch_index=0): + state = InferStateInfo() + state.is_prefill = False + state.input_ids = model_input.input_ids + state.batch_size = model_input.batch_size + state.max_kv_seq_len = model_input.max_kv_seq_len + state.b_req_idx = model_input.b_req_idx + state.b_seq_len = model_input.b_seq_len + state.microbatch_index = microbatch_index + state.decode_att_state = DecodeState(_infer_state_weak_ref=weakref.ref(state)) + state.init_some_extra_state = lambda _: None + return state + + def outputs(state, state1): + return tuple(ModelOutput(logits=torch.stack((s.b_req_idx, s.b_seq_len), dim=1)) for s in (state, state1)) + + def forward(state, infer_state1): + # Eager and capture execute attention using these newly prepared states. + assert state.decode_att_state.decode_wrapper is not None + assert infer_state1.decode_att_state.decode_wrapper is not None + executions.append(("forward", graph_flags(state), graph_flags(infer_state1))) + return outputs(state, infer_state1) + + def capture(func, state, infer_state1): + executions.append(("capture", graph_flags(state), graph_flags(infer_state1))) + return func(state, infer_state1) + + def replay(state, infer_state1): + # Replay uses graph-owned wrappers, so temporary states need none. + assert state.decode_att_state.decode_wrapper is None + assert infer_state1.decode_att_state.decode_wrapper is None + executions.append(("replay", graph_flags(state), graph_flags(infer_state1))) + return outputs(state, infer_state1) + + model._create_inferstate = create_state + model._overlap_tpsp_token_forward = forward + model.graph = None + if graph_mode is not None: + model.graph = SimpleNamespace( + can_run=lambda batch_size, max_len_in_batch: batch_size <= 4 and max_len_in_batch <= 8192, + find_closest_graph_batch_size=lambda batch_size: 4, + need_capture=lambda batch_size: graph_mode == "capture", + capture_decode=capture, + replay=replay, + ) + + inputs = [] + for rows, length in zip((1, 3), lengths): + model_input = model._create_padded_decode_model_input(_create_empty_decode_input(), rows) + model_input.b_req_idx = torch.arange(rows, dtype=torch.int32) + model_input.b_seq_len.fill_(length) + model_input.max_kv_seq_len = length + model_input.total_token_num = rows * length + inputs.append(model_input) + + result = model._microbatch_overlap_decode_cuda(*inputs) + + expected_flags = {"eager": (False, False), "capture": (True, False), "replay": (False, True)}[expected_mode] + assert initializations == [(index, expected_flags, expected_mode != "replay") for index in (0, 1)] + expected_execution = {"eager": ["forward"], "capture": ["capture", "forward"], "replay": ["replay"]}[expected_mode] + assert executions == [(mode, expected_flags, expected_flags) for mode in expected_execution] + for output, model_input in zip(result, inputs): + torch.testing.assert_close(output.logits, torch.stack((model_input.b_req_idx, model_input.b_seq_len), dim=1)) + + +@pytest.mark.parametrize("overlap", [False, True]) +@pytest.mark.parametrize("dynamic", [False, True]) +@pytest.mark.parametrize("graph_mode", ["eager", "capture", "replay"]) +@pytest.mark.parametrize("rows", [0, 3, 9, 24, 27]) +def test_mtp_decode_preserves_groups_and_real_rows(overlap, dynamic, graph_mode, rows): + model = TpPartBaseModel.__new__(TpPartBaseModel) + model.args = SimpleNamespace(enable_tpsp_mix_mode=True) + model.tp_world_size_ = 8 + model.is_mtp_draft_model = False + model.mtp_manager = SimpleNamespace(get_decode_batch_alignment=lambda _: 1 if dynamic else 3) + model.req_manager = SimpleNamespace(HOLD_REQUEST_ID=88) + seen_sizes = [] + + def create_state(model_input, microbatch_index=0): + def check_layout(): + if not dynamic: + assert model_input.batch_size % 3 == 0 + assert model_input.batch_size % 8 == 0 + seen_sizes.append(model_input.batch_size) + + return SimpleNamespace( + b_req_idx=model_input.b_req_idx, + b_seq_len=model_input.b_seq_len, + init_some_extra_state=lambda _: None, + init_att_state=check_layout, + ) + + def forward(state): + return ModelOutput(logits=torch.stack((state.b_req_idx, state.b_seq_len), dim=1)) + + model._create_inferstate = create_state + model._token_forward = forward + model._overlap_tpsp_token_forward = lambda state, infer_state1: (forward(state), forward(infer_state1)) + model.graph = None + if graph_mode != "eager": + model.graph = SimpleNamespace( + can_run=lambda batch_size, max_len_in_batch: batch_size <= 24, + find_closest_graph_batch_size=lambda batch_size: ((batch_size + 7) // 8 * 8 if dynamic else 24), + need_capture=lambda batch_size: graph_mode == "capture", + capture_decode=lambda func, state, **kwargs: func(state, **kwargs), + replay=lambda state, infer_state1=None: ( + forward(state) if infer_state1 is None else (forward(state), forward(infer_state1)) + ), + ) + + model_input = model._create_padded_decode_model_input(_create_empty_decode_input(), rows) + model_input.b_req_idx = torch.arange(rows, dtype=torch.int32) // 3 + model_input.b_seq_len = 10 + torch.arange(rows, dtype=torch.int32) % 3 + if dynamic: + # 压缩后每请求分别保留 1、2、3 个 token，保持各请求内部的 token 顺序。 + request_rows = [(req, step) for req in range(rows) for step in range(req % 3 + 1)][:rows] + model_input.b_req_idx = torch.tensor([req for req, _ in request_rows], dtype=torch.int32) + model_input.b_mtp_index = torch.tensor([step for _, step in request_rows], dtype=torch.int32) + model_input.b_seq_len = 10 + model_input.b_mtp_index + expected = torch.stack((model_input.b_req_idx, model_input.b_seq_len), dim=1) + if overlap: + output, empty_output = model._microbatch_overlap_decode_cuda(model_input, _create_empty_decode_input()) + assert empty_output.logits.shape == (0, 2) + else: + output = model._decode(model_input) + + expected_size = (max(1, rows) + 7) // 8 * 8 if dynamic else (48 if rows > 24 else 24) + 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_moe_route_info_manager.py b/unit_tests/common/basemodel/test_moe_route_info_manager.py new file mode 100644 index 0000000000..de2bc7a857 --- /dev/null +++ b/unit_tests/common/basemodel/test_moe_route_info_manager.py @@ -0,0 +1,263 @@ +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +from lightllm.common.basemodel.infer_struct import InferStateInfo +from lightllm.common.basemodel.moe_route_info_manager import MoeRouteInfoManager, get_moe_capture_callback + + +def _skip_without_cuda(): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for moe route info capture.") + + +def test_get_moe_capture_callback_uses_global_manager(monkeypatch): + calls = [] + + class _Manager: + def is_buffer_initialized(self): + return True + + def get_moe_capture_callback(self, layer_index, mem_indexes): + calls.append((layer_index, mem_indexes)) + return "callback" + + mem_indexes = object() + infer_state = SimpleNamespace(mem_index=mem_indexes) + monkeypatch.setattr(MoeRouteInfoManager, "_instance", _Manager()) + + assert get_moe_capture_callback(infer_state, 7) == "callback" + assert calls == [(7, mem_indexes)] + + +class TestMoeRouteInfoManager: + def test_capture_and_extract_basic(self): + _skip_without_cuda() + + manager = MoeRouteInfoManager( + num_moe_layers=4, + topk=8, + dtype_id=1, + ) + manager.init_capture_buffer(kv_cache_size=1024) + mem_indexes = torch.arange(100, 110, device="cuda") + expected = np.zeros((10, 4, 8), dtype=np.uint8) + + for layer_idx in range(4): + topk_ids = torch.randint(0, 64, (10, 8), device="cuda") + manager.get_moe_capture_callback(layer_idx, mem_indexes)(topk_ids) + expected[:, layer_idx, :] = topk_ids.cpu().numpy().astype(np.uint8) + + result = manager.extract(mem_indexes) + assert result.shape == (10, 4, 8) + assert result.dtype == np.uint8 + np.testing.assert_array_equal(result, expected) + + def test_capture_writes_to_correct_kv_positions(self): + _skip_without_cuda() + + manager = MoeRouteInfoManager( + num_moe_layers=2, + topk=4, + dtype_id=1, + ) + manager.init_capture_buffer(kv_cache_size=256) + mem_indexes = torch.tensor([10, 50, 200], device="cuda") + topk_ids = torch.tensor([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]], device="cuda") + topk_ids_layer1 = topk_ids + 20 + + manager.get_moe_capture_callback(0, mem_indexes)(topk_ids) + manager.get_moe_capture_callback(1, mem_indexes)(topk_ids_layer1) + + result = manager.extract(mem_indexes) + np.testing.assert_array_equal(result[:, 0, :], topk_ids.cpu().numpy().astype(np.uint8)) + np.testing.assert_array_equal(result[:, 1, :], topk_ids_layer1.cpu().numpy().astype(np.uint8)) + + def test_capture_maps_transformer_layer_num_to_routing_slot(self): + _skip_without_cuda() + + manager = MoeRouteInfoManager( + num_moe_layers=2, + topk=2, + dtype_id=1, + layer_index_to_moe_index={3: 0, 7: 1}, + ) + manager.init_capture_buffer(kv_cache_size=256) + mem_indexes = torch.tensor([10, 11], device="cuda") + ids_layer3 = torch.tensor([[1, 2], [3, 4]], device="cuda") + ids_layer7 = torch.tensor([[5, 6], [7, 8]], device="cuda") + + manager.get_moe_capture_callback(3, mem_indexes)(ids_layer3) + manager.get_moe_capture_callback(7, mem_indexes)(ids_layer7) + assert manager.get_moe_capture_callback(4, mem_indexes) is None + + result = manager.extract(mem_indexes) + np.testing.assert_array_equal(result[:, 0, :], ids_layer3.cpu().numpy().astype(np.uint8)) + np.testing.assert_array_equal(result[:, 1, :], ids_layer7.cpu().numpy().astype(np.uint8)) + + def test_capture_rejects_unexpected_topk_width(self): + _skip_without_cuda() + + manager = MoeRouteInfoManager( + num_moe_layers=1, + topk=2, + dtype_id=1, + ) + manager.init_capture_buffer(kv_cache_size=256) + mem_indexes = torch.tensor([10, 11], device="cuda") + topk_ids_with_shared_expert = torch.tensor([[1, 2, 32], [3, 4, 32]], device="cuda") + + with pytest.raises(AssertionError): + manager.get_moe_capture_callback(0, mem_indexes)(topk_ids_with_shared_expert) + + def test_cuda_graph_replay_uses_copied_mem_indexes(self): + _skip_without_cuda() + + class _NoopDecodeState: + def copy_for_decode_cuda_graph(self, other): + return + + manager = MoeRouteInfoManager( + num_moe_layers=1, + topk=2, + dtype_id=1, + ) + manager.init_capture_buffer(kv_cache_size=256) + graph_infer_state = InferStateInfo() + graph_infer_state.decode_att_state = _NoopDecodeState() + graph_infer_state.mem_index = torch.tensor([10, 11], device="cuda") + moe_capture_callback = manager.get_moe_capture_callback(0, graph_infer_state.mem_index) + topk_ids = torch.tensor([[1, 2], [3, 4]], device="cuda") + + graph = torch.cuda.CUDAGraph() + torch.cuda.synchronize() + with torch.cuda.graph(graph): + moe_capture_callback(topk_ids) + + graph.replay() + result = manager.extract(graph_infer_state.mem_index) + np.testing.assert_array_equal(result[:, 0, :], topk_ids.cpu().numpy().astype(np.uint8)) + + new_infer_state = InferStateInfo() + new_infer_state.decode_att_state = _NoopDecodeState() + new_infer_state.mem_index = torch.tensor([20, 21], device="cuda") + new_topk_ids = torch.tensor([[5, 6], [7, 8]], device="cuda") + graph_infer_state.copy_for_cuda_graph(new_infer_state) + topk_ids.copy_(new_topk_ids) + + graph.replay() + + result = manager.extract(new_infer_state.mem_index) + np.testing.assert_array_equal(result[:, 0, :], new_topk_ids.cpu().numpy().astype(np.uint8)) + + def test_microbatch_isolation(self): + _skip_without_cuda() + + manager = MoeRouteInfoManager( + num_moe_layers=1, + topk=4, + dtype_id=1, + ) + manager.init_capture_buffer(kv_cache_size=256) + mem0 = torch.tensor([10, 11], device="cuda") + mem1 = torch.tensor([20, 21], device="cuda") + ids_0 = torch.ones((2, 4), dtype=torch.int64, device="cuda") + ids_1 = torch.ones((2, 4), dtype=torch.int64, device="cuda") * 2 + + capture0 = manager.get_moe_capture_callback(0, mem0) + capture1 = manager.get_moe_capture_callback(0, mem1) + capture0(ids_0) + capture1(ids_1) + + result0 = manager.extract(mem0) + result1 = manager.extract(mem1) + assert result0[0, 0, 0] == 1 + assert result1[0, 0, 0] == 2 + + def test_dtype_selection_uint8(self): + manager = MoeRouteInfoManager( + num_moe_layers=1, + topk=2, + dtype_id=1, + ) + assert manager.get_torch_dtype() == torch.uint8 + assert manager.get_np_dtype() == np.uint8 + assert manager.dtype_id == 1 + assert not manager.is_buffer_initialized() + + def test_dtype_selection_int16(self): + manager = MoeRouteInfoManager( + num_moe_layers=1, + topk=2, + dtype_id=2, + ) + assert manager.get_torch_dtype() == torch.int16 + assert manager.get_np_dtype() == np.int16 + assert manager.dtype_id == 2 + assert not manager.is_buffer_initialized() + + def test_extract_preserves_uint8_values(self): + _skip_without_cuda() + + manager = MoeRouteInfoManager( + num_moe_layers=1, + topk=4, + dtype_id=1, + ) + manager.init_capture_buffer(kv_cache_size=64) + mem_indexes = torch.tensor([0, 1, 2], device="cuda") + topk_ids = torch.tensor([[10, 20, 30, 40], [50, 60, 63, 1], [0, 5, 255, 3]], device="cuda") + + manager.get_moe_capture_callback(0, mem_indexes)(topk_ids) + + result = manager.extract(mem_indexes) + np.testing.assert_array_equal(result[:, 0, :], topk_ids.cpu().numpy().astype(np.uint8)) + + def test_routing_buffer_and_pointer_shape(self): + _skip_without_cuda() + + manager = MoeRouteInfoManager( + num_moe_layers=48, + topk=8, + dtype_id=1, + ) + assert not manager.is_buffer_initialized() + manager.init_capture_buffer(kv_cache_size=2048) + assert manager.routing_buffer.shape == (2048, 48, 8) + assert manager.routing_buffer.dtype == torch.uint8 + assert manager.routing_buffer.device.type == "cpu" + assert manager.routing_buffer.is_pinned() + assert manager.routing_buffer_ptr.shape == (1,) + assert manager.routing_buffer_ptr.dtype == torch.uint64 + assert manager.routing_buffer_ptr.device.type == "cuda" + + def test_partial_token_capture(self): + _skip_without_cuda() + + manager = MoeRouteInfoManager( + num_moe_layers=1, + topk=2, + dtype_id=1, + ) + manager.init_capture_buffer(kv_cache_size=128) + mem_indexes = torch.tensor([10, 11, 12, 13, 14], device="cuda") + topk_ids = torch.tensor([[1, 2], [3, 4], [5, 6]], device="cuda") + + manager.get_moe_capture_callback(0, mem_indexes[:3])(topk_ids) + + result_written = manager.extract(mem_indexes[:3]) + np.testing.assert_array_equal(result_written[:, 0, :], topk_ids.cpu().numpy().astype(np.uint8)) + + result_unwritten = manager.extract(mem_indexes[3:]) + np.testing.assert_array_equal(result_unwritten[:, 0, :], np.zeros((2, 2), dtype=np.uint8)) + + def test_phase1_does_not_allocate_capture_buffer(self): + manager = MoeRouteInfoManager( + num_moe_layers=4, + topk=8, + dtype_id=1, + ) + assert not manager.is_buffer_initialized() + assert manager.routing_buffer_ptr is None diff --git a/unit_tests/common/basemodel/test_mtp_manager.py b/unit_tests/common/basemodel/test_mtp_manager.py new file mode 100644 index 0000000000..f7e91f8c79 --- /dev/null +++ b/unit_tests/common/basemodel/test_mtp_manager.py @@ -0,0 +1,133 @@ +from types import SimpleNamespace + +import pytest + +import lightllm.common.basemodel.hidden_collector as hidden_collector_module +import lightllm.common.basemodel.mtp_manager as mtp_manager_module +from lightllm.common.basemodel.hidden_collector import ( + FinalHiddenCollector, + LayerHiddenCollector, + MtpHeadOutputCollector, + NoopHiddenCollector, +) +from lightllm.common.basemodel.mtp_manager import MtpManager + + +@pytest.fixture(autouse=True) +def _reset_mtp_manager(): + MtpManager._instance = None + yield + MtpManager._instance = None + + +def _decode_tokens_per_request(monkeypatch, spec_mode, *, is_draft_model, mtp_step=7): + args = SimpleNamespace( + mtp_mode=spec_mode, + mtp_step=mtp_step, + mtp_dynamic_verify=False, + ) + monkeypatch.setattr(mtp_manager_module, "get_env_start_args", lambda: args) + return MtpManager.get_instance().get_decode_tokens_per_request(is_draft_model) + + +@pytest.mark.parametrize( + "spec_mode,is_draft_model,expected", + [ + (None, False, 1), + ("eagle3", False, 8), + ("eagle3", True, 1), + ("vanilla_with_att", True, 1), + ("vanilla_no_att", True, 1), + ("eagle_with_att", True, 1), + ("eagle_no_att", True, 1), + ("dspark", True, 7), + ("dflash", True, 7), + ("dflash2", True, 8), + ], +) +def test_decode_tokens_per_request(monkeypatch, spec_mode, is_draft_model, expected): + assert _decode_tokens_per_request(monkeypatch, spec_mode, is_draft_model=is_draft_model) == expected + + +@pytest.mark.parametrize( + "spec_mode,dynamic_verify,is_draft_model,expected", + [ + ("vanilla_with_att", False, False, 8), + ("vanilla_with_att", True, False, 1), + ("vanilla_with_att", False, True, 1), + ("vanilla_with_att", True, True, 1), + ("dspark", True, True, 7), + ("dflash", True, True, 7), + ], +) +def test_decode_batch_alignment(monkeypatch, spec_mode, dynamic_verify, is_draft_model, expected): + args = SimpleNamespace( + mtp_mode=spec_mode, + mtp_step=7, + mtp_dynamic_verify=dynamic_verify, + ) + monkeypatch.setattr(mtp_manager_module, "get_env_start_args", lambda: args) + + assert MtpManager.get_instance().get_decode_batch_alignment(is_draft_model) == expected + + +@pytest.mark.parametrize( + "spec_mode,is_draft_model,expected", + [ + (None, False, 0), + ("eagle3", False, 7), + ("eagle3", True, 0), + ("vanilla_with_att", True, 0), + ("dspark", True, 6), + ("dflash", True, 6), + ("dflash2", True, 7), + ], +) +def test_decode_draft_step(monkeypatch, spec_mode, is_draft_model, expected): + args = SimpleNamespace(mtp_mode=spec_mode, mtp_step=7, mtp_dynamic_verify=False) + monkeypatch.setattr(mtp_manager_module, "get_env_start_args", lambda: args) + + assert MtpManager.get_instance().get_decode_draft_step(is_draft_model) == expected + + +def test_get_instance_returns_singleton(monkeypatch): + args = SimpleNamespace(mtp_mode="eagle3", mtp_step=7, mtp_dynamic_verify=False) + monkeypatch.setattr(mtp_manager_module, "get_env_start_args", lambda: args) + + assert MtpManager.get_instance() is MtpManager.get_instance() + + +@pytest.mark.parametrize( + "spec_mode,is_draft_model,expected_type", + [ + (None, False, NoopHiddenCollector), + ("vanilla_with_att", False, FinalHiddenCollector), + ("eagle3", False, LayerHiddenCollector), + ("dspark", False, LayerHiddenCollector), + ("eagle3", True, FinalHiddenCollector), + ("dspark", True, MtpHeadOutputCollector), + ("dflash2", True, MtpHeadOutputCollector), + ], +) +def test_create_hidden_collector_selects_implementation(monkeypatch, spec_mode, is_draft_model, expected_type): + args = SimpleNamespace( + mtp_mode=spec_mode, + mtp_step=7, + mtp_dynamic_verify=False, + mtp_draft_model_dir=["/models/draft"], + ) + monkeypatch.setattr(mtp_manager_module, "get_env_start_args", lambda: args) + monkeypatch.setattr(hidden_collector_module, "get_env_start_args", lambda: args) + monkeypatch.setattr( + hidden_collector_module.PretrainedConfig, + "get_config_dict", + lambda _: ({"target_layer_ids": [0]}, {}), + ) + model = SimpleNamespace(is_mtp_draft_model=is_draft_model, layers_num=2) + + prototype = MtpManager.get_instance().create_hidden_collector(model=model) + collector = prototype.new_instance() + + assert isinstance(prototype, expected_type) + assert isinstance(collector, expected_type) + assert collector is not prototype diff --git a/unit_tests/common/basemodel/test_overlap_utils.py b/unit_tests/common/basemodel/test_overlap_utils.py new file mode 100644 index 0000000000..ef7c18a061 --- /dev/null +++ b/unit_tests/common/basemodel/test_overlap_utils.py @@ -0,0 +1,213 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.basemodel import basemodel +from lightllm.common.basemodel.basemodel import TpPartBaseModel +from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput + + +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) + events = [] + original_to_cuda = ModelInput.to_cuda + + def record_to_cuda(self): + events.append("to_cuda") + original_to_cuda(self) + + monkeypatch.setattr(ModelInput, "to_cuda", record_to_cuda) + + def fake_gather(**kwargs): + events.append("gather") + kwargs["input_ids"][-1] = 90 + events.count("gather") + + monkeypatch.setattr(basemodel, "gather_token_prefill_decode_mixed", fake_gather) + + model = TpPartBaseModel.__new__(TpPartBaseModel) + model.args = SimpleNamespace(enable_prefill_decode_mixed=True, page_size=1) + model.req_manager = SimpleNamespace( + req_sampling_params_manager=SimpleNamespace(req_to_next_token_ids=object()), + ) + captured_inputs = [] + + def fake_overlap_forward(input0, input1): + events.append("forward") + captured_inputs.extend((input0, input1)) + return ( + ModelOutput(logits=torch.zeros((input0.batch_size, 1))), + ModelOutput(logits=torch.zeros((input1.batch_size, 1))), + ) + + model._microbatch_overlap_prefill_cuda = fake_overlap_forward + + outputs = model.microbatch_overlap_prefill(model_input0, model_input1) + + assert events == ["to_cuda", "gather", "to_cuda", "gather", "forward"] + assert captured_inputs == [model_input0, model_input1] + assert captured_inputs[0].input_ids.tolist() == [10, 91] + assert captured_inputs[1].input_ids.tolist() == [92] + assert len(outputs) == 2 + + +def test_base_model_decode_accepts_two_inputs_and_skips_empty_gather(monkeypatch): + model_input0 = _make_decode_input() + model_input0.input_ids = None + model_input1 = _make_empty_decode_input() + model_input1.input_ids = None + events = [] + original_to_cuda = ModelInput.to_cuda + + def record_to_cuda(self): + events.append("to_cuda") + original_to_cuda(self) + + monkeypatch.setattr(ModelInput, "to_cuda", record_to_cuda) + + def fake_gather(**kwargs): + events.append("gather") + return torch.arange( + start=40, + end=46, + dtype=torch.int64, + device="cuda", + ) + + monkeypatch.setattr(basemodel, "gather_token", fake_gather) + + model = TpPartBaseModel.__new__(TpPartBaseModel) + model.args = SimpleNamespace(page_size=1) + model.req_manager = SimpleNamespace( + req_sampling_params_manager=SimpleNamespace(req_to_next_token_ids=object()), + ) + captured_inputs = [] + + def fake_overlap_forward(input0, input1): + events.append("forward") + captured_inputs.extend((input0, input1)) + return ( + ModelOutput(logits=torch.zeros((input0.batch_size, 1))), + ModelOutput(logits=torch.zeros((input1.batch_size, 1))), + ) + + model._microbatch_overlap_decode_cuda = fake_overlap_forward + + model.microbatch_overlap_decode(model_input0, model_input1) + + assert events == ["to_cuda", "gather", "to_cuda", "forward"] + assert captured_inputs == [model_input0, model_input1] + assert captured_inputs[0].input_ids.tolist() == [40, 41, 42, 43, 44, 45] + assert captured_inputs[1].input_ids.numel() == 0 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_overlap_decode_cuda_pads_empty_side_and_unpads_outputs(monkeypatch): + model_input0 = _make_decode_input() + model_input0.batch_size = 1 + model_input0.total_token_num = 4 + model_input0.max_kv_seq_len = 4 + model_input0.input_ids = model_input0.input_ids[:1] + model_input0.b_req_idx = model_input0.b_req_idx[:1] + model_input0.b_mtp_index = model_input0.b_mtp_index[:1] + model_input0.b_seq_len = model_input0.b_seq_len[:1] + model_input0.b_position_delta = model_input0.b_position_delta[:1] + model_input0.b_shared_seq_len = model_input0.b_shared_seq_len[:1] + model_input0.b_shared_radix_node_id = model_input0.b_shared_radix_node_id[:1] + model_input0.multimodal_params = model_input0.multimodal_params[:1] + model_input0.check_input() + model_input1 = _make_empty_decode_input() + model_input0.to_cuda() + model_input1.to_cuda() + + model = TpPartBaseModel.__new__(TpPartBaseModel) + model.args = SimpleNamespace(enable_tpsp_mix_mode=True, page_size=1) + model.tp_world_size_ = 2 + 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()) + infer_batch_sizes = [] + + def fake_create_inferstate(model_input, microbatch_index): + infer_batch_sizes.append(model_input.batch_size) + return SimpleNamespace( + b_req_idx=model_input.b_req_idx, + b_seq_len=model_input.b_seq_len, + mem_index=torch.empty((model_input.batch_size,), dtype=torch.int32, device="cuda"), + init_some_extra_state=lambda _: None, + init_att_state=lambda: None, + ) + + model._create_inferstate = fake_create_inferstate + model._overlap_tpsp_token_forward = lambda infer_state0, infer_state1: ( + ModelOutput(logits=torch.zeros((infer_state0.b_req_idx.shape[0], 1), device="cuda")), + ModelOutput(logits=torch.zeros((infer_state1.b_req_idx.shape[0], 1), device="cuda")), + ) + output0, output1 = model._microbatch_overlap_decode_cuda(model_input0, model_input1) + + 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/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_mtp_diverse.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_mtp_diverse.py new file mode 100644 index 0000000000..7c1af6bbc2 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_mtp_diverse.py @@ -0,0 +1,196 @@ +""" +MTP Diverse Attention Unit Test + +测试 MTP diverse attention 算子的正确性和性能。 + +Single Token Mode: +- 每个请求只有 1 个 Q token +- 组内第 i 个请求只能看到前 i 个 KV +- b_mark_shared_group 标记： + - 0: 组内非最后一个请求（跳过） + - N>=1: 一个 N 人组的最后一个请求（需要计算） +""" +import pytest +import torch + + +def gqa_attention_reference( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + req_to_tokens: torch.Tensor, + b_req_idx: torch.Tensor, + b_seq_len: torch.Tensor, +) -> torch.Tensor: + """ + Reference GQA attention implementation for single token mode + + q: [batch_size, num_heads, head_dim] - 每个请求 1 个 Q token + k, v: [kv_pool_size, kv_head_num, head_dim] - KV 池 + req_to_tokens: [batch, max_kv_len] - 每个请求的 KV 索引 + b_seq_len: [batch] - 每个请求可见的 KV 数量 + """ + batch_size = b_seq_len.shape[0] + num_heads = q.shape[1] + kv_head_num = k.shape[1] + head_dim = q.shape[2] + gqa_group_size = num_heads // kv_head_num + + # 输出：[batch, num_heads, head_dim] + output = torch.zeros((batch_size, num_heads, head_dim), dtype=q.dtype, device=q.device) + + for b in range(batch_size): + seq_len = b_seq_len[b].item() # 这个请求可见的 KV 数量 + + # 获取这个请求的 KV 索引：[seq_len] + kv_indices = req_to_tokens[b, :seq_len] + + # 获取 K 和 V：[seq_len, kv_head_num, head_dim] + k_batch = k[kv_indices] + v_batch = v[kv_indices] + + # Process each Q head + for h in range(num_heads): + kv_h = h // gqa_group_size + k_h = k_batch[:, kv_h, :] # [seq_len, head_dim] + v_h = v_batch[:, kv_h, :] # [seq_len, head_dim] + + q_h = q[b, h] # [head_dim] + + # Compute attention scores + att = torch.matmul(q_h, k_h.transpose(0, 1)) # [seq_len] + att = att / (head_dim ** 0.5) + att = torch.softmax(att, dim=-1) + out_h = torch.matmul(att, v_h) # [head_dim] + + output[b, h] = out_h + + return output + + +def setup_mtp_test_data(kv_len, group_size, batch_groups, test_dtype=torch.bfloat16, device="cuda", seed=42): + """ + 设置 MTP 测试数据 - Single Token Mode + + Single Token Mode 的特点： + - 每个请求有 1 个 Q token + - 组内第 i 个请求只能看到前 i+1 个 KV + - 组内请求共享相同的 KV 前缀 + + b_mark_shared_group 标记： + - 0: 组内非最后一个请求（跳过计算） + - N>=1: 一个 N 人组的最后一个请求（需要计算） + """ + torch.manual_seed(seed) + + num_heads = 32 + kv_head_num = 4 # gqa_group_size = 8 + head_dim = 128 + + batch_size = batch_groups * group_size + + # KV 池：[batch_groups * kv_len, kv_head_num, head_dim] + # 每个组使用不同的 KV 范围 + kv_pool_size = batch_groups * kv_len + k = torch.randn(size=(kv_pool_size, kv_head_num, head_dim), dtype=test_dtype, device=device) + v = torch.randn(size=(kv_pool_size, kv_head_num, head_dim), dtype=test_dtype, device=device) + + # req_to_tokens: [batch, max_kv_len] - 每个请求的 KV 索引 + max_kv_len = group_size # 每个请求最多有 group_size 个 KV（可见范围） + req_to_tokens = torch.zeros((batch_size, max_kv_len), dtype=torch.int32, device=device) + + b_req_idx = torch.arange(batch_size, dtype=torch.int32, device=device) + b_seq_len = torch.zeros(batch_size, dtype=torch.int32, device=device) # 每个请求可见的 KV 数量 + b_mark_shared_group = torch.zeros(batch_size, dtype=torch.int32, device=device) + + # Q: [batch_size, num_heads, head_dim] - 每个请求 1 个 Q token + q = torch.randn(size=(batch_size, num_heads, head_dim), dtype=test_dtype, device=device) + + for group_idx in range(batch_groups): + group_start = group_idx * group_size + kv_base = group_idx * kv_len + + for member_idx in range(group_size): + batch_idx = group_start + member_idx + # 第 member_idx 个请求可以看到前 member_idx+1 个 KV + b_seq_len[batch_idx] = member_idx + 1 + + # 设置 KV 索引 - 组内请求共享相同的 KV 前缀 + for kv_pos in range(member_idx + 1): + req_to_tokens[batch_idx, kv_pos] = kv_base + kv_pos + + # b_mark_shared_group: + # - 0: 组内非最后一个请求 + # - group_size: 组内最后一个请求 + if member_idx == group_size - 1: + b_mark_shared_group[batch_idx] = group_size + else: + b_mark_shared_group[batch_idx] = 0 + + return q, k, v, req_to_tokens, b_req_idx, b_seq_len, b_mark_shared_group + + +def mtp_diverse_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + req_to_tokens: torch.Tensor, + b_req_idx: torch.Tensor, + b_seq_len: torch.Tensor, + b_mark_shared_group: torch.Tensor, +) -> torch.Tensor: + """ + MTP Diverse Attention 调用入口 + """ + from lightllm.common.basemodel.triton_kernel.att.decode_att.gqa.mtp_diverse import ( + token_decode_attention_mtp_diverse_single_token, + ) + + return token_decode_attention_mtp_diverse_single_token( + q=q, + k=k, + v=v, + Req_to_tokens=req_to_tokens, + B_req_idx=b_req_idx, + b_seq_len=b_seq_len, + b_mark_shared_group=b_mark_shared_group, + max_kv_len=int(b_seq_len.max().item()), + ) + + +@pytest.mark.parametrize("kv_len", [100, 256]) +@pytest.mark.parametrize("group_size", [2, 3]) +@pytest.mark.parametrize("batch_groups", [1, 2, 4]) +def test_mtp_diverse_vs_reference(kv_len, group_size, batch_groups): + """ + 测试 MTP diverse attention 与 reference 实现的正确性对比 + """ + test_dtype = torch.bfloat16 + device = "cuda" + + q, k, v, req_to_tokens, b_req_idx, b_seq_len, b_mark_shared_group = setup_mtp_test_data( + kv_len, group_size, batch_groups, test_dtype, device + ) + + # 运行 reference 实现（计算所有请求） + reference_out = gqa_attention_reference( + q=q, + k=k, + v=v, + req_to_tokens=req_to_tokens, + b_req_idx=b_req_idx, + b_seq_len=b_seq_len, + ) + + # 运行 MTP diverse（只计算组内最后一个请求） + diverse_out = mtp_diverse_attention( + q=q, + k=k, + v=v, + req_to_tokens=req_to_tokens, + b_req_idx=b_req_idx, + b_seq_len=b_seq_len, + b_mark_shared_group=b_mark_shared_group, + ) + + torch.testing.assert_close(diverse_out, reference_out, atol=1e-2, rtol=1e-2) diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_mtp_diverse_autotune.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_mtp_diverse_autotune.py new file mode 100644 index 0000000000..0989fe6d26 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_mtp_diverse_autotune.py @@ -0,0 +1,301 @@ +import weakref +import collections +import inspect +import json +import math +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F +from torch.nn.attention import SDPBackend, sdpa_kernel + +from lightllm.common.basemodel.infer_struct import InferStateInfo +from lightllm.common.basemodel import cuda_graph as graph_module +from lightllm.common.basemodel.attention.triton import fp as state_module +from lightllm.common.basemodel.triton_kernel.att.decode_att.gqa.mtp_diverse import ( + mtp_diverse_attn as decode_module, + stage1_single_token as stage1_module, + stage2_single_token as stage2_module, +) +from lightllm.common.triton_utils import autotuner as autotuner_module +from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType, AutotuneLevel +from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len + + +@pytest.fixture(autouse=True) +def autotune_environment(monkeypatch): + torch.manual_seed(42) + monkeypatch.delenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", raising=False) + monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", None) + monkeypatch.setattr(stage1_module, "get_triton_autotune_level", lambda: AutotuneLevel.ADAPTIVE_AUTOTUNE) + monkeypatch.setattr(stage2_module, "get_triton_autotune_level", lambda: AutotuneLevel.ADAPTIVE_AUTOTUNE) + monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: AutotuneLevel.ADAPTIVE_AUTOTUNE) + get_decode_attn_autotune_seq_len.cache_clear() + yield + get_decode_attn_autotune_seq_len.cache_clear() + + +@pytest.mark.parametrize("level", [AutotuneLevel.ADAPTIVE_AUTOTUNE, AutotuneLevel.FORCE_AUTOTUNE]) +@pytest.mark.parametrize("configured, expected", [(None, 32768), ("8193", 8704)]) +def test_tuning_key_uses_configured_length(monkeypatch, level, configured, expected): + if configured is not None: + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", configured) + monkeypatch.setattr(stage1_module, "get_triton_autotune_level", lambda: level) + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + assert stage1_module.get_run_key(torch.empty(7, 8, 64), 2) == 7_000_000_000 + expected + + +@pytest.mark.parametrize( + "phase,level", + [ + (None, AutotuneLevel.ADAPTIVE_AUTOTUNE), + (None, AutotuneLevel.FORCE_AUTOTUNE), + (AutotuneKernelType.GENERAL, AutotuneLevel.ADAPTIVE_AUTOTUNE), + (AutotuneKernelType.DECODE_ATTENTION, AutotuneLevel.USE_AUTOTUNE_HIS_CONFIG), + ], +) +def test_lookup_uses_actual_length(monkeypatch, phase, level): + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", "invalid") + monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", phase) + monkeypatch.setattr(stage1_module, "get_triton_autotune_level", lambda: level) + for actual, expected in [(2, 512), (511, 512), (512, 512), (513, 1024), (8193, 8704)]: + assert stage1_module.get_run_key(torch.empty(7, 8, 64), actual) == 7_000_000_000 + expected + + +def make_inputs(device="cpu", batch=7, block_batch=3, head_dim=128, token_count=17, table_width=16, block_num=2): + kv = torch.randn(token_count, 4, head_dim, dtype=torch.bfloat16, device=device) + return dict( + q=torch.randn(batch, 8, head_dim, dtype=kv.dtype, device=device), + k=kv[:, :2], + v=kv[:, 2:], + Req_to_tokens=torch.arange(6 * table_width, dtype=torch.int32, device=device).view(6, table_width) + % max(1, token_count), + B_req_idx=torch.full((batch,), 4, dtype=torch.int32, device=device), + b_seq_len=torch.full((batch,), 2, dtype=torch.int32, device=device), + b_mark_shared_group=torch.ones(batch, dtype=torch.int32, device=device), + max_kv_len=2, + mid_out=torch.full((batch, 8, block_num, head_dim), -100, dtype=kv.dtype, device=device), + mid_out_logsumexp=torch.full((batch, 8, block_num), -100, dtype=torch.float32, device=device), + block_batch=block_batch, + ) + + +def rebuild(inputs): + args, kwargs = stage1_module.rebuild_inputs(**inputs) + return inspect.signature(stage1_module.mtp_diverse_stage1_single_token.fn).bind(*args, **kwargs).arguments + + +@pytest.mark.parametrize("token_count", [3, 128]) +@pytest.mark.parametrize( + "batch, block_batch, length, req_ids, lengths, marks", + [ + (7, 3, 17, [0, 0, 0, 1, 1, 1, 2], [15, 16, 17, 15, 16, 17, 17], [0, 0, 3, 0, 0, 3, 1]), + (5, 4, 4, [0, 0, 0, 0, 1], [1, 2, 3, 4, 4], [0, 0, 0, 4, 1]), + (1, 4, 1, [0], [1], [1]), + ], +) +def test_rebuild_preserves_inputs_and_shared_prefixes( + monkeypatch, token_count, batch, block_batch, length, req_ids, lengths, marks +): + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", str(length)) + inputs = make_inputs(batch=batch, block_batch=block_batch, token_count=token_count) + snapshots = {name: value.clone() for name, value in inputs.items() if isinstance(value, torch.Tensor)} + rebuilt = rebuild(inputs) + assert rebuilt["B_req_idx"].tolist() == req_ids + assert rebuilt["b_seq_len"].tolist() == lengths + assert rebuilt["b_mark_shared_group"].tolist() == marks + assert rebuilt["max_kv_len"] == length + expected = torch.arange((max(req_ids) + 1) * length).view(-1, length) % token_count + torch.testing.assert_close(rebuilt["Req_to_tokens"].long(), expected) + for name in ["q", "k", "v", "mid_out", "mid_out_logsumexp"]: + assert rebuilt[name] is inputs[name] + assert rebuilt["block_batch"] == block_batch + for name, value in snapshots.items(): + torch.testing.assert_close(inputs[name], value) + + +@pytest.mark.parametrize("batch, block_batch, length", [(7, 3, 2), (2, 4, 1)]) +def test_rebuild_rejects_length_shorter_than_group(monkeypatch, batch, block_batch, length): + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", str(length)) + with pytest.raises(ValueError, match="at least the largest MTP group size"): + rebuild(make_inputs(batch=batch, block_batch=block_batch)) + + +def test_rebuild_rejects_empty_or_mismatched_cache(): + with pytest.raises(ValueError, match="non-empty KV cache"): + rebuild(make_inputs(token_count=0)) + inputs = make_inputs() + inputs["v"] = inputs["v"][:-1] + with pytest.raises(AssertionError, match="same number of tokens"): + rebuild(inputs) + + +def test_decode_passes_saved_length_and_selected_block_n(monkeypatch): + inputs = make_inputs() + infer_state = InferStateInfo() + infer_state.max_kv_seq_len = 8193 + infer_state.b_req_idx = inputs["B_req_idx"] + infer_state.b_seq_len = inputs["b_seq_len"] + infer_state.req_manager = SimpleNamespace(req_to_token_indexs=inputs["Req_to_tokens"]) + model = SimpleNamespace( + is_mtp_draft_model=False, + mtp_manager=SimpleNamespace(get_decode_draft_step=lambda _: 2), + req_manager=SimpleNamespace(HOLD_REQUEST_ID=-1), + ) + monkeypatch.setattr(state_module, "build_mtp_shared_group_markers", lambda *a, **kw: inputs["b_mark_shared_group"]) + state = state_module.TritonDecodeAttState( + backend=SimpleNamespace(model=model), _infer_state_weak_ref=weakref.ref(infer_state) + ) + state.init_state() + infer_state.max_kv_seq_len = 32768 + calls = [] + + def stage1(**kwargs): + calls.append(kwargs) + return 32 + + monkeypatch.setattr(decode_module, "mtp_diverse_stage1_single_token", stage1) + monkeypatch.setattr(decode_module, "mtp_diverse_stage2_single_token", lambda **kwargs: calls.append(kwargs)) + state.decode_att(q=inputs["q"], k=inputs["k"], v=inputs["v"]) + assert calls[0]["max_kv_len"] == 8193 + assert calls[0]["b_mark_shared_group"] is inputs["b_mark_shared_group"] + assert calls[1]["block_n"] == 32 + assert calls[1]["max_kv_len"] == 8193 + + +def reference_attention(inputs): + results = [] + for row, length in enumerate(inputs["b_seq_len"].tolist()): + indices = inputs["Req_to_tokens"][inputs["B_req_idx"][row], :length].long() + with sdpa_kernel(SDPBackend.MATH): + results.append( + F.scaled_dot_product_attention( + inputs["q"][row].float().unsqueeze(1), + inputs["k"][indices].float().transpose(0, 1), + inputs["v"][indices].float().transpose(0, 1), + enable_gqa=True, + ).squeeze(1) + ) + return torch.stack(results) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("kv_len", [17, 8192, 16384]) +@pytest.mark.parametrize("block_batch", [1, 3, 4]) +@pytest.mark.parametrize("head_dim", [64, 128, 256]) +def test_shared_groups_match_fp32_reference(monkeypatch, kv_len, block_batch, head_dim): + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", str(kv_len)) + inputs = rebuild(make_inputs("cuda", block_batch=block_batch, head_dim=head_dim, token_count=7 * kv_len)) + reference = reference_attention(inputs) + out = torch.empty_like(inputs["q"]) + # 仅两个 grid block，长 KV 必须循环多次；17 token 同时覆盖末块 mask 和无效分块的归约边界。 + for block_n in [16, 32, 64]: + selected_block_n = stage1_module.mtp_diverse_stage1_single_token.fn( + **inputs, + run_config={"BLOCK_N": block_n, "num_warps": 4, "num_stages": 2, "warp_specialize": False}, + ) + stage2_module.mtp_diverse_stage2_single_token.fn( + inputs["mid_out"], inputs["mid_out_logsumexp"], inputs["b_seq_len"], out, selected_block_n, kv_len + ) + torch.testing.assert_close(out.float(), reference, atol=2e-3, rtol=2e-2) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("kv_len", [8192, 16384]) +def test_full_autotune_before_graph_and_cached_reuse(tmp_path, monkeypatch, kv_len): + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", str(kv_len)) + monkeypatch.setattr(autotuner_module.dist, "is_initialized", lambda: False) + monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: AutotuneLevel.FORCE_AUTOTUNE) + monkeypatch.setattr(stage1_module, "get_triton_autotune_level", lambda: AutotuneLevel.FORCE_AUTOTUNE) + kernel = stage1_module.mtp_diverse_stage1_single_token + monkeypatch.setattr(kernel, "_cache_dir", str(tmp_path), raising=False) + monkeypatch.setattr(kernel, "cached_configs", {}) + monkeypatch.setattr(kernel, "fast_match_configs", collections.defaultdict(dict)) + monkeypatch.setattr(kernel, "warmuped_configs_set", set()) + assert kernel.mutates_args == [] + # 使用生产路径的小 batch 布局，让 stage2 在长请求下归约完整的 128 个分块。 + inputs = make_inputs("cuda", table_width=kv_len, token_count=3 * kv_len, block_num=128) + stage2 = stage2_module.mtp_diverse_stage2_single_token + stage2_cache = tmp_path / "stage2" + stage2_cache.mkdir() + monkeypatch.setattr(stage2, "_cache_dir", str(stage2_cache), raising=False) + monkeypatch.setattr(stage2, "cached_configs", {}) + monkeypatch.setattr(stage2, "fast_match_configs", collections.defaultdict(dict)) + monkeypatch.setattr(stage2, "warmuped_configs_set", set()) + snapshots = {name: value.clone() for name, value in inputs.items() if isinstance(value, torch.Tensor)} + original_rebuild, original_bench = kernel.rebuild_input_func, kernel._bench + rebuild_count, benchmark_count, valid_benchmark_count = 0, 0, 0 + + def checked_rebuild(*args, **kwargs): + nonlocal rebuild_count + assert not torch.cuda.is_current_stream_capturing() + rebuild_count += 1 + return original_rebuild(*args, **kwargs) + + def checked_bench(*args, **kwargs): + nonlocal benchmark_count, valid_benchmark_count + rebuilt = inspect.signature(kernel.fn).bind(*args, **kwargs).arguments + assert rebuilt["max_kv_len"] == kv_len + assert rebuilt["b_seq_len"].tolist() == [kv_len - 2, kv_len - 1, kv_len] * 2 + [kv_len] + assert rebuilt["b_mark_shared_group"].tolist() == [0, 0, 3, 0, 0, 3, 1] + assert rebuilt["Req_to_tokens"].shape == (3, kv_len) + elapsed = original_bench(*args, **kwargs) + # 不可编译或资源不足的候选按 Autotuner 原有规则返回 inf；完整搜索必须选出有效配置。 + if math.isfinite(elapsed): + valid_benchmark_count += 1 + benchmark_count += 1 + return elapsed + + monkeypatch.setattr(kernel, "rebuild_input_func", checked_rebuild) + monkeypatch.setattr(kernel, "_bench", checked_bench) + monkeypatch.setattr( + graph_module, + "get_env_start_args", + lambda: SimpleNamespace( + enable_decode_microbatch_overlap=False, enable_tpsp_mix_mode=False, enable_torch_memory_saver=False + ), + ) + out = torch.empty_like(inputs["q"]) + + def decode(state): + block_n = kernel(**inputs) + stage2(inputs["mid_out"], inputs["mid_out_logsumexp"], inputs["b_seq_len"], out, block_n, inputs["max_kv_len"]) + return out + + graph = graph_module.CudaGraph(1, 7, 1, max_batch_size=7, max_len_in_batch=32768) + graph.capture_decode(decode, SimpleNamespace(input_ids=torch.ones(7, device="cuda"))) + torch.cuda.synchronize() + assert rebuild_count == 1 and benchmark_count == len(kernel.configs_gen_func()) + assert valid_benchmark_count > 0 + # 两个 KV 的短请求沿用原 MTP 测试的 BF16 容差；下面的长请求仍使用更严格的 FP32 对照。 + torch.testing.assert_close(out.float(), reference_attention(inputs), atol=1e-2, rtol=1e-2) + for name, value in snapshots.items(): + if name not in ["mid_out", "mid_out_logsumexp"]: + torch.testing.assert_close(inputs[name], value) + cache_file = next(tmp_path.glob("*.json")) + saved = cache_file.read_bytes() + assert list(json.loads(saved)) == [str(7_000_000_000 + kv_len)] + assert json.loads(saved)[str(7_000_000_000 + kv_len)] is not None + stage2_configs = json.loads(next(stage2_cache.glob("*.json")).read_bytes()) + assert list(stage2_configs) == [str(7 * 8 * 1_000_000_000 + 128)] + assert next(iter(stage2_configs.values())) is not None + # 同一层配置以及从文件重新加载的配置，在 FORCE 模式下都不能再次搜索。 + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + kernel(**inputs) + kernel.cached_configs.clear() + kernel.fast_match_configs.clear() + kernel.warmuped_configs_set.clear() + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + kernel(**inputs) + assert rebuild_count == 1 and benchmark_count == len(kernel.configs_gen_func()) + assert cache_file.read_bytes() == saved + # 原缓冲区改为真实长请求和共享组后回放，验证 Graph 没有捕获 benchmark 的重建输入。 + inputs["B_req_idx"].copy_(torch.tensor([1, 1, 1, 3, 3, 3, 2], device="cuda", dtype=torch.int32)) + inputs["b_mark_shared_group"].copy_(torch.tensor([0, 0, 3, 0, 0, 3, 1], device="cuda", dtype=torch.int32)) + inputs["b_seq_len"].copy_( + torch.tensor([kv_len - 2, kv_len - 1, kv_len] * 2 + [kv_len], device="cuda", dtype=torch.int32) + ) + graph.graph[7][0].replay() + torch.cuda.synchronize() + torch.testing.assert_close(out.float(), reference_attention(inputs), atol=2e-3, rtol=2e-2) diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_stage2_autotune.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_stage2_autotune.py new file mode 100644 index 0000000000..48956c4f36 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/mtp_diverse/test_stage2_autotune.py @@ -0,0 +1,220 @@ +import collections +import inspect +import json +import math +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.basemodel import cuda_graph as graph_module +from lightllm.common.basemodel.triton_kernel.att.decode_att.gqa.mtp_diverse import stage2_single_token as stage2_module +from lightllm.common.triton_utils import autotuner as autotuner_module +from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType, AutotuneLevel +from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len + + +@pytest.fixture(autouse=True) +def autotune_environment(monkeypatch): + torch.manual_seed(42) + monkeypatch.delenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", raising=False) + monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", None) + monkeypatch.setattr(stage2_module, "get_triton_autotune_level", lambda: AutotuneLevel.ADAPTIVE_AUTOTUNE) + monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: AutotuneLevel.ADAPTIVE_AUTOTUNE) + get_decode_attn_autotune_seq_len.cache_clear() + yield + get_decode_attn_autotune_seq_len.cache_clear() + + +@pytest.mark.parametrize("level", [AutotuneLevel.ADAPTIVE_AUTOTUNE, AutotuneLevel.FORCE_AUTOTUNE]) +@pytest.mark.parametrize("configured, capacity, expected", [(None, 128, 128), ("8193", 256, 129), ("65", 128, 2)]) +def test_tuning_key_uses_effective_reduction_blocks(monkeypatch, level, configured, capacity, expected): + if configured is not None: + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", configured) + monkeypatch.setattr(stage2_module, "get_triton_autotune_level", lambda: level) + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + assert stage2_module.get_run_key(torch.empty(2, 4, capacity, 64), 64, 2) == 8_000_000_000 + expected + + +@pytest.mark.parametrize( + "phase,level", + [ + (None, AutotuneLevel.ADAPTIVE_AUTOTUNE), + (None, AutotuneLevel.FORCE_AUTOTUNE), + (AutotuneKernelType.GENERAL, AutotuneLevel.ADAPTIVE_AUTOTUNE), + (AutotuneKernelType.DECODE_ATTENTION, AutotuneLevel.USE_AUTOTUNE_HIS_CONFIG), + (AutotuneKernelType.DECODE_ATTENTION, AutotuneLevel.CLOSE_AUTOTUNE), + ], +) +def test_lookup_counts_real_blocks_without_reading_gpu_lengths(monkeypatch, phase, level): + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", "invalid") + monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", phase) + monkeypatch.setattr(stage2_module, "get_triton_autotune_level", lambda: level) + mid = torch.empty(2, 4, 128, 64) + for length, blocks in [(1, 1), (64, 1), (65, 2), (511, 8), (512, 8), (513, 9), (8192, 128), (16384, 128)]: + assert stage2_module.get_run_key(mid, 64, length) == 8_000_000_000 + blocks + + +def test_key_reuses_equal_work_and_distinguishes_capacity_limits(): + key = stage2_module.get_run_key + small, large = torch.empty(2, 4, 64, 64), torch.empty(2, 4, 128, 64) + assert key(small, 64, 513) == key(large, 64, 513) + assert key(small, 64, 8192) != key(large, 64, 8192) + assert key(large, 64, 8192) == key(large, 64, 16384) + + +def make_inputs(device="cpu", capacity=128, block_n=64, dtype=torch.bfloat16, head_dim=128): + # 模拟短请求：stage1 仅写入第一个分块，其余位置用 NaN 标记为未初始化。 + mid = torch.full((2, 4, capacity, head_dim), float("nan"), dtype=dtype, device=device) + lse = torch.full((2, 4, capacity), float("nan"), dtype=torch.float32, device=device) + mid[:, :, 0] = torch.randn(2, 4, head_dim, dtype=dtype, device=device) + lse[:, :, 0] = 0 + return dict( + mid_out=mid, + mid_out_logsumexp=lse, + B_Seqlen=torch.full((2,), 2, dtype=torch.int32, device=device), + out=torch.full((2, 4, head_dim), -100, dtype=dtype, device=device), + block_n=block_n, + max_kv_len=2, + ) + + +@pytest.mark.parametrize("length", [65, 8193, 16384]) +def test_rebuild_reuses_readonly_intermediate_inputs(monkeypatch, length): + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", str(length)) + inputs = make_inputs() + snapshots = {name: value.clone() for name, value in inputs.items() if isinstance(value, torch.Tensor)} + args, kwargs = stage2_module.rebuild_inputs(**inputs) + rebuilt = inspect.signature(stage2_module.mtp_diverse_stage2_single_token.fn).bind(*args, **kwargs).arguments + assert rebuilt["max_kv_len"] == length + assert rebuilt["B_Seqlen"].tolist() == [length, length] + assert rebuilt["block_n"] == inputs["block_n"] + assert rebuilt["out"] is inputs["out"] + for name in ["mid_out", "mid_out_logsumexp"]: + assert rebuilt[name] is inputs[name] + assert rebuilt[name].shape == inputs[name].shape + assert rebuilt[name].stride() == inputs[name].stride() + assert rebuilt[name].dtype == inputs[name].dtype + assert rebuilt[name].device == inputs[name].device + assert rebuilt["B_Seqlen"] is not inputs["B_Seqlen"] + assert torch.isfinite(rebuilt["B_Seqlen"]).all() + for name, snapshot in snapshots.items(): + torch.testing.assert_close(inputs[name], snapshot, equal_nan=True) + + +def reduction_reference(inputs): + results = [] + for row, length in enumerate(inputs["B_Seqlen"].tolist()): + blocks = min((length + inputs["block_n"] - 1) // inputs["block_n"], inputs["mid_out"].shape[2]) + weights = torch.softmax(inputs["mid_out_logsumexp"][row, :, :blocks].float(), dim=-1) + results.append((weights.unsqueeze(-1) * inputs["mid_out"][row, :, :blocks].float()).sum(dim=1)) + return torch.stack(results) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("block_n", [16, 64]) +@pytest.mark.parametrize("capacity", [2, 128]) +@pytest.mark.parametrize("dtype,head_dim", [(torch.bfloat16, 128), (torch.float16, 64), (torch.bfloat16, 256)]) +def test_all_configs_reduce_only_valid_blocks(block_n, capacity, dtype, head_dim): + lengths = [1, block_n, block_n + 1, capacity * block_n, capacity * block_n + 8192] + mid = torch.randn(len(lengths), 4, capacity, head_dim, dtype=dtype, device="cuda") + lse = torch.randn(len(lengths), 4, capacity, dtype=torch.float32, device="cuda") * 20 + for row, length in enumerate(lengths): + blocks = min((length + block_n - 1) // block_n, capacity) + mid[row, :, blocks:] = float("nan") + lse[row, :, blocks:] = float("nan") + inputs = dict( + mid_out=mid, + mid_out_logsumexp=lse, + B_Seqlen=torch.tensor(lengths, dtype=torch.int32, device="cuda"), + out=torch.empty(len(lengths), 4, head_dim, dtype=dtype, device="cuda"), + block_n=block_n, + max_kv_len=max(lengths), + ) + reference = reduction_reference(inputs) + for config in stage2_module.get_test_configs(): + stage2_module.mtp_diverse_stage2_single_token.fn(**inputs, run_config=config) + torch.testing.assert_close(inputs["out"].float(), reference, atol=2e-3, rtol=1e-2) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("length", [65, 8192, 16384]) +def test_full_tuning_graph_and_cache_reuse(tmp_path, monkeypatch, length): + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", str(length)) + monkeypatch.setattr(autotuner_module.dist, "is_initialized", lambda: False) + monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: AutotuneLevel.FORCE_AUTOTUNE) + monkeypatch.setattr(stage2_module, "get_triton_autotune_level", lambda: AutotuneLevel.FORCE_AUTOTUNE) + kernel = stage2_module.mtp_diverse_stage2_single_token + monkeypatch.setattr(kernel, "_cache_dir", str(tmp_path), raising=False) + monkeypatch.setattr(kernel, "cached_configs", {}) + monkeypatch.setattr(kernel, "fast_match_configs", collections.defaultdict(dict)) + monkeypatch.setattr(kernel, "warmuped_configs_set", set()) + inputs = make_inputs("cuda") + snapshots = {name: value.clone() for name, value in inputs.items() if isinstance(value, torch.Tensor)} + original_rebuild, original_bench = kernel.rebuild_input_func, kernel._bench + rebuild_count, benchmark_count = 0, 0 + + def checked_rebuild(*args, **kwargs): + nonlocal rebuild_count + assert not torch.cuda.is_current_stream_capturing() + rebuild_count += 1 + return original_rebuild(*args, **kwargs) + + def checked_bench(*args, **kwargs): + nonlocal benchmark_count + rebuilt = inspect.signature(kernel.fn).bind(*args, **kwargs).arguments + assert rebuilt["B_Seqlen"].tolist() == [length, length] + assert rebuilt["max_kv_len"] == length and rebuilt["block_n"] == 64 + for name in ["mid_out", "mid_out_logsumexp"]: + assert rebuilt[name] is inputs[name] + assert rebuilt[name].shape == inputs[name].shape + elapsed = original_bench(*args, **kwargs) + assert math.isfinite(elapsed) + torch.testing.assert_close(inputs["out"], snapshots["out"]) + benchmark_count += 1 + return elapsed + + monkeypatch.setattr(kernel, "rebuild_input_func", checked_rebuild) + monkeypatch.setattr(kernel, "_bench", checked_bench) + monkeypatch.setattr( + graph_module, + "get_env_start_args", + lambda: SimpleNamespace( + enable_decode_microbatch_overlap=False, enable_tpsp_mix_mode=False, enable_torch_memory_saver=False + ), + ) + + def decode(state): + kernel(**inputs) + return inputs["out"] + + graph = graph_module.CudaGraph(1, 2, 1, max_batch_size=2, max_len_in_batch=32768) + graph.capture_decode(decode, SimpleNamespace(input_ids=torch.ones(2, device="cuda"))) + torch.cuda.synchronize() + assert rebuild_count == 1 and benchmark_count == 20 + torch.testing.assert_close(inputs["out"].float(), reduction_reference(inputs), atol=2e-3, rtol=1e-2) + for name in ["mid_out", "mid_out_logsumexp", "B_Seqlen"]: + torch.testing.assert_close(inputs[name], snapshots[name], equal_nan=True) + cache_file = next(tmp_path.glob("*.json")) + saved = cache_file.read_bytes() + expected_key = str(8_000_000_000 + min((length + 63) // 64, 128)) + assert list(json.loads(saved)) == [expected_key] + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + kernel(**inputs) + kernel.cached_configs.clear() + kernel.fast_match_configs.clear() + kernel.warmuped_configs_set.clear() + # 8K/16K 都归约 128 个块，切换代表性长度后也应复用同一个 key。 + if length >= 8192: + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", str(24576 - length)) + get_decode_attn_autotune_seq_len.cache_clear() + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + kernel(**inputs) + assert rebuild_count == 1 and benchmark_count == 20 + assert cache_file.read_bytes() == saved + inputs["mid_out"].normal_() + inputs["mid_out_logsumexp"].normal_() + inputs["B_Seqlen"].fill_(length) + graph.graph[2][0].replay() + torch.cuda.synchronize() + torch.testing.assert_close(inputs["out"].float(), reduction_reference(inputs), atol=2e-3, rtol=1e-2) diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/test_flash_decoding_autotune.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/test_flash_decoding_autotune.py new file mode 100644 index 0000000000..b3068db5fb --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/gqa/test_flash_decoding_autotune.py @@ -0,0 +1,270 @@ +import weakref +import collections +import inspect +import math +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F +from frozendict import frozendict +from torch.nn.attention import SDPBackend, sdpa_kernel + +from lightllm.common.basemodel.infer_struct import InferStateInfo +from lightllm.common.basemodel.attention.triton.fp import TritonDecodeAttState +from lightllm.common.basemodel.triton_kernel.att.decode_att.gqa.flash_decoding import ( + gqa_flash_decoding as decode_module, + gqa_flash_decoding_stage1 as stage1_module, +) +from lightllm.common.basemodel.triton_kernel.att.decode_att.gqa.flash_decoding.gqa_flash_decoding_stage2 import ( + flash_decode_stage2, +) +from lightllm.common.triton_utils import autotuner as autotuner_module +from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType, AutotuneLevel +from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len + + +@pytest.fixture(autouse=True) +def autotune_environment(monkeypatch): + torch.manual_seed(42) + monkeypatch.delenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", raising=False) + monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", None) + monkeypatch.setattr(stage1_module, "get_triton_autotune_level", lambda: AutotuneLevel.ADAPTIVE_AUTOTUNE) + monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: AutotuneLevel.ADAPTIVE_AUTOTUNE) + get_decode_attn_autotune_seq_len.cache_clear() + yield + get_decode_attn_autotune_seq_len.cache_clear() + + +@pytest.mark.parametrize("level", [AutotuneLevel.ADAPTIVE_AUTOTUNE, AutotuneLevel.FORCE_AUTOTUNE]) +@pytest.mark.parametrize("configured, expected", [(None, 32768), ("8193", 8704)]) +def test_tuning_key_uses_bucketed_configured_length(monkeypatch, level, configured, expected): + if configured is not None: + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", configured) + monkeypatch.setattr(stage1_module, "get_triton_autotune_level", lambda: level) + q = torch.empty(2, 8, 64) + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + assert stage1_module.get_run_key(q, 2) == 2_000_000_000 + expected + + +@pytest.mark.parametrize( + "phase,level", + [ + (None, AutotuneLevel.ADAPTIVE_AUTOTUNE), + (None, AutotuneLevel.FORCE_AUTOTUNE), + (AutotuneKernelType.GENERAL, AutotuneLevel.ADAPTIVE_AUTOTUNE), + (AutotuneKernelType.DECODE_ATTENTION, AutotuneLevel.USE_AUTOTUNE_HIS_CONFIG), + ], +) +def test_non_tuning_key_uses_actual_length(monkeypatch, phase, level): + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", "invalid") + monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", phase) + monkeypatch.setattr(stage1_module, "get_triton_autotune_level", lambda: level) + q = torch.empty(2, 8, 64) + for actual, bucket in [(2, 512), (511, 512), (512, 512), (513, 1024), (8193, 8704)]: + assert stage1_module.get_run_key(q, actual) == 2_000_000_000 + bucket + + +def make_inputs(device="cpu", head_dim=64, token_count=17, table_width=16): + # K/V 是共享存储中的不连续视图；原请求索引也不是从 0 开始连续排列。 + kv = torch.randn(token_count, 4, head_dim, dtype=torch.bfloat16, device=device) + return dict( + q=torch.randn(2, 8, head_dim, dtype=kv.dtype, device=device), + k=kv[:, :2], + v=kv[:, 2:], + Req_to_tokens=torch.arange(5 * table_width, dtype=torch.int32, device=device).view(5, table_width) + % token_count, + B_req_idx=torch.tensor([4, 2], dtype=torch.int32, device=device), + B_Seqlen=torch.tensor([3, 2], dtype=torch.int32, device=device), + max_len_in_batch=3, + mid_out=torch.full((2, 8, 2, head_dim), -100, dtype=kv.dtype, device=device), + mid_out_logsumexp=torch.full((2, 8, 2), -100, dtype=torch.float32, device=device), + block_seq=256, + sliding_window=(-1, -1), + ) + + +def test_static_key_shares_heads_and_blocks_but_separates_windows(): + kernel = stage1_module.flash_decode_stage1 + inputs = make_inputs() + full_key = frozendict(kernel._static_key(**inputs)) + inputs.pop("sliding_window") + assert frozendict(kernel._static_key(**inputs)) == full_key + + inputs["q"] = torch.empty(2, 4, 64, dtype=torch.bfloat16) + inputs["k"] = torch.empty(17, 1, 64, dtype=torch.bfloat16) + inputs["mid_out"] = torch.empty(2, 4, 128, 64, dtype=torch.bfloat16) + assert frozendict(kernel._static_key(**inputs)) == full_key + + window_key = frozendict(kernel._static_key(**inputs, sliding_window=(511, 0))) + assert window_key != full_key + assert frozendict(kernel._static_key(**inputs, sliding_window=[511, 0])) == window_key + assert frozendict(kernel._static_key(**inputs, sliding_window=(1023, 0))) != window_key + + +@pytest.mark.parametrize("num_tokens", [3, 128]) +def test_rebuild_preserves_layout_and_bounds_indices(monkeypatch, num_tokens): + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", "17") + inputs = make_inputs(token_count=num_tokens) + inputs["sliding_window"] = (7, 0) + snapshots = {name: value.clone() for name, value in inputs.items() if isinstance(value, torch.Tensor)} + + args, kwargs = stage1_module.rebuild_inputs(**inputs) + rebuilt = inspect.signature(stage1_module.flash_decode_stage1.fn).bind(*args, **kwargs).arguments + assert rebuilt["Req_to_tokens"].shape == (2, 17) + assert rebuilt["Req_to_tokens"].min() >= 0 and rebuilt["Req_to_tokens"].max() < num_tokens + assert rebuilt["B_req_idx"].tolist() == [0, 1] + assert rebuilt["B_Seqlen"].tolist() == [17, 17] + assert rebuilt["max_len_in_batch"] == 17 + for name in ["q", "k", "v", "mid_out", "mid_out_logsumexp"]: + assert rebuilt[name] is inputs[name] + assert rebuilt["block_seq"] == 256 + assert rebuilt["sliding_window"] == (7, 0) + for name, value in snapshots.items(): + torch.testing.assert_close(inputs[name], value) + + +def test_state_retains_actual_length_before_graph_capture(monkeypatch): + infer_state = InferStateInfo() + infer_state.max_kv_seq_len = 8193 + state = TritonDecodeAttState( + backend=SimpleNamespace( + model=SimpleNamespace( + is_mtp_draft_model=False, mtp_manager=SimpleNamespace(get_decode_draft_step=lambda _: 0) + ) + ), + _infer_state_weak_ref=weakref.ref(infer_state), + ) + state.init_state() + state.infer_state.max_kv_seq_len = 32768 + calls = [] + monkeypatch.setattr( + decode_module, "gqa_token_decode_attention_flash_decoding", lambda **kwargs: calls.append(kwargs) + ) + state.decode_att( + q=torch.empty(2, 8, 64), + k=torch.empty(16, 2, 64), + v=torch.empty(16, 2, 64), + alloc_func=lambda shape, dtype: torch.empty(shape, dtype=dtype), + ) + assert calls[0]["max_len_in_batch"] == 8193 + + +def reference_attention(inputs): + outputs = [] + for i, length in enumerate(inputs["B_Seqlen"].tolist()): + start = max(0, length - 1 - inputs["sliding_window"][0]) if inputs["sliding_window"][0] >= 0 else 0 + indices = inputs["Req_to_tokens"][inputs["B_req_idx"][i], start:length].long() + with sdpa_kernel(SDPBackend.MATH): + outputs.append( + F.scaled_dot_product_attention( + inputs["q"][i].float().unsqueeze(1), + inputs["k"][indices].float().transpose(0, 1), + inputs["v"][indices].float().transpose(0, 1), + enable_gqa=True, + ).squeeze(1) + ) + return torch.stack(outputs) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("kv_len", [8192, 16384]) +@pytest.mark.parametrize("head_dim", [64, 256]) +@pytest.mark.parametrize("sliding_window", [(-1, -1), (511, 0)]) +def test_stage1_stage2_long_kv_matches_reference(kv_len, head_dim, sliding_window): + inputs = make_inputs("cuda", head_dim=head_dim, token_count=kv_len * 2, table_width=kv_len) + inputs["B_Seqlen"] = torch.tensor([kv_len, kv_len - 1], dtype=torch.int32, device="cuda") + inputs["max_len_in_batch"] = kv_len + inputs["sliding_window"] = sliding_window + reference = reference_attention(inputs) + # 只有两个 grid block，长 KV 需要每个 program 循环多次；stage2 仍归约同一布局。 + output = torch.empty_like(inputs["q"]) + for block_n in [16, 64, 128]: + stage1_module.flash_decode_stage1.fn(**inputs, run_config={"BLOCK_N": block_n, "num_warps": 4, "num_stages": 2}) + flash_decode_stage2( + inputs["mid_out"], + inputs["mid_out_logsumexp"], + inputs["B_Seqlen"], + output, + inputs["block_seq"], + sliding_window=sliding_window, + ) + assert torch.isfinite(output).all() + torch.testing.assert_close(output.float(), reference, atol=2e-3, rtol=2e-2) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("kv_len", [8192, 16384]) +def test_autotune_rebuilds_once_and_graph_uses_original_inputs(tmp_path, monkeypatch, kv_len): + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", str(kv_len)) + monkeypatch.setattr(autotuner_module.dist, "is_initialized", lambda: False) + kernel = stage1_module.flash_decode_stage1 + monkeypatch.setattr(kernel, "_cache_dir", str(tmp_path), raising=False) + monkeypatch.setattr(kernel, "cached_configs", {}) + monkeypatch.setattr(kernel, "fast_match_configs", collections.defaultdict(dict)) + monkeypatch.setattr(kernel, "warmuped_configs_set", set()) + assert kernel.mutates_args == [] + monkeypatch.setattr( + kernel, + "configs_gen_func", + lambda: [{"BLOCK_N": n, "num_warps": 4, "num_stages": 2} for n in [16, 64, 128]], + ) + inputs = make_inputs("cuda") + snapshots = {name: value.clone() for name, value in inputs.items() if isinstance(value, torch.Tensor)} + reference = reference_attention(inputs) + rebuild = kernel.rebuild_input_func + benchmark = kernel._bench + rebuild_count = 0 + benchmark_count = 0 + + def checked_rebuild(*args, **kwargs): + nonlocal rebuild_count + rebuild_count += 1 + return rebuild(*args, **kwargs) + + def checked_bench(*args, **kwargs): + nonlocal benchmark_count + rebuilt = inspect.signature(kernel.fn).bind(*args, **kwargs).arguments + assert rebuilt["B_Seqlen"].tolist() == [kv_len, kv_len] + assert rebuilt["Req_to_tokens"].shape == (2, kv_len) + assert rebuilt["B_req_idx"].tolist() == [0, 1] + assert rebuilt["max_len_in_batch"] == kv_len + assert rebuilt["mid_out"] is inputs["mid_out"] + elapsed = benchmark(*args, **kwargs) + assert math.isfinite(elapsed) + benchmark_count += 1 + return elapsed + + monkeypatch.setattr(kernel, "rebuild_input_func", checked_rebuild) + monkeypatch.setattr(kernel, "_bench", checked_bench) + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + kernel(**inputs) + assert rebuild_count == 1 and benchmark_count == 3 + assert all(list(configs) == [str(2_000_000_000 + kv_len)] for configs in kernel.cached_configs.values()) + # FORCE 在 decode 阶段仍复用缓存，不因模型层重复调用而再次调优。 + monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: AutotuneLevel.FORCE_AUTOTUNE) + monkeypatch.setattr(stage1_module, "get_triton_autotune_level", lambda: AutotuneLevel.FORCE_AUTOTUNE) + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + kernel(**inputs) + assert rebuild_count == 1 and benchmark_count == 3 + + output = torch.empty_like(inputs["q"]) + + def decode(): + kernel(**inputs) + flash_decode_stage2( + inputs["mid_out"], inputs["mid_out_logsumexp"], inputs["B_Seqlen"], output, inputs["block_seq"] + ) + + decode() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + decode() + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(output.float(), reference, atol=2e-3, rtol=2e-2) + assert rebuild_count == 1 and benchmark_count == 3 + for name, value in snapshots.items(): + if name not in ["mid_out", "mid_out_logsumexp"]: + torch.testing.assert_close(inputs[name], value) diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse.py new file mode 100644 index 0000000000..e1b6bc7758 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse.py @@ -0,0 +1,139 @@ +import pytest +from types import SimpleNamespace + +import torch + + +def alloc_tensor_func(shape, dtype, device): + """兼容的 tensor 分配函数""" + return torch.empty(shape, dtype=dtype, device=device) + + +class MockReqManager: + """Mock request manager for testing""" + + def __init__(self, req_to_token_indexs): + self.req_to_token_indexs = req_to_token_indexs + + +class MockInferState: + """Mock infer state for testing""" + + def __init__( + self, + batch_size, + max_kv_seq_len, + req_to_tokens, + b_req_idx, + b_seq_len, + ): + self.batch_size = batch_size + self.max_kv_seq_len = max_kv_seq_len + self.req_manager = MockReqManager(req_to_tokens) + self.b_req_idx = b_req_idx + self.b_seq_len = b_seq_len + + +# @pytest.mark.parametrize("shared_seq_len", [512]) +@pytest.mark.parametrize("shared_seq_len", [0, 77, 256, 311, 512, 550]) +@pytest.mark.parametrize("batch_size", list(range(6, 121, 6))) +def test_token_decode_attention_flash_decoding_diverse_matches_normal_decode(shared_seq_len, batch_size): + """ + diverse 与 normal 均为仓库内 Triton 实现，应数值一致（无外部 CUDA extension）。 + diverse：int8kv_flash_decoding_diverse；对照：int8kv/normal token_decode_attention_flash_decoding。 + """ + + from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.int8kv_flash_decoding_diverse import ( + token_decode_attention_flash_decoding as diverse_attention, + ) + from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.normal import ( + token_decode_attention_flash_decoding as normal_decode, + ) + + num_heads = 32 + kv_head_num = 8 + mark_shared_group_size = 3 + seq_len = 3547 + head_dim = 128 + quant_group_size = 8 + max_len_in_batch = 8192 + test_dtype = torch.bfloat16 + + # 创建测试数据 + kv_shape = (batch_size * max_len_in_batch, kv_head_num, head_dim) + kv_scale_shape = (batch_size * max_len_in_batch, kv_head_num, head_dim // quant_group_size) + + q = torch.randn(size=(batch_size, num_heads, head_dim), dtype=test_dtype, device="cuda") + + # 生成 cache_k 和 cache_v，使得每 mark_shared_group_size 个 batch 共享相同的 cache + + cache_k = torch.randint(low=-100, high=100, size=kv_shape, dtype=torch.int8, device="cuda") + cache_k_scale = torch.ones(size=kv_scale_shape, dtype=test_dtype, device="cuda") / 100.0 + cache_v = torch.randint(low=-100, high=100, size=kv_shape, dtype=torch.int8, device="cuda") + cache_v_scale = torch.ones(size=kv_scale_shape, dtype=test_dtype, device="cuda") / 100.0 + + req_to_tokens = torch.arange(0, max_len_in_batch * batch_size, dtype=torch.int32, device="cuda").view( + batch_size, max_len_in_batch + ) + for i in range(batch_size): + if i % mark_shared_group_size != 0: + req_to_tokens[i, :shared_seq_len] = req_to_tokens[i - 1, :shared_seq_len] + + b_req_idx = torch.arange(batch_size, dtype=torch.int32, device="cuda") + b_seq_len = torch.full((batch_size,), seq_len, dtype=torch.int32, device="cuda") + b_shared_seq_len = torch.full((batch_size,), shared_seq_len, dtype=torch.int32, device="cuda") + b_mark_shared_group = torch.zeros((batch_size,), dtype=torch.int32, device="cuda") + b_mark_shared_group[mark_shared_group_size - 1 :: mark_shared_group_size] = mark_shared_group_size + + # 标准 int8 decode（单路径 Triton） + baseline_infer_state = MockInferState( + batch_size=batch_size, + max_kv_seq_len=max_len_in_batch, + req_to_tokens=req_to_tokens, + b_req_idx=b_req_idx, + b_seq_len=b_seq_len, + ) + + # diverse：多流 + 共享前缀（Triton） + diverse_infer_state = MockInferState( + batch_size=batch_size, + max_kv_seq_len=max_len_in_batch, + req_to_tokens=req_to_tokens, + b_req_idx=b_req_idx, + b_seq_len=b_seq_len, + ) + diverse_infer_state.decode_att_state = SimpleNamespace( + b_shared_seq_len=b_shared_seq_len, + b_mark_shared_group=b_mark_shared_group, + ) + + # 运行 baseline + normal_out = normal_decode( + q=q.clone(), + infer_state=baseline_infer_state, + cache_k=cache_k, + cache_k_scale=cache_k_scale, + cache_v=cache_v, + cache_v_scale=cache_v_scale, + max_len_in_batch=seq_len, + alloc_tensor_func=alloc_tensor_func, + ) + # 运行 diverse 版本 + diverse_out = diverse_attention( + q=q.clone(), + infer_state=diverse_infer_state, + cache_k=cache_k, + cache_k_scale=cache_k_scale, + cache_v=cache_v, + cache_v_scale=cache_v_scale, + alloc_tensor_func=alloc_tensor_func, + ) + + print(f"\nshared_seq_len={shared_seq_len}\nbatch_size={batch_size}") + print(f"normal_out: {normal_out[0, 0, :4]}") + print(f"diverse_out: {diverse_out[0, 0, :4]}") + print(f"max diff: {(normal_out - diverse_out).abs().max()}") + + assert torch.allclose( + normal_out, diverse_out, atol=1e-2, rtol=1e-2 + ), f"diverse vs normal decode mismatch for shared_seq_len={shared_seq_len}" diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage3.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage3.py new file mode 100644 index 0000000000..c1a0ca1e58 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_int8kv_flash_decoding_diverse_stage3.py @@ -0,0 +1,35 @@ +import pytest +import torch +from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.int8kv_flash_decoding_diverse_stage3 import ( + flash_diverse_decode_stage3, +) + + +@pytest.mark.parametrize( + "batch, head_num, seq_len, shared_seq_len, block_seq, head_dim", + [ + (2, 4, 256, 256, 256, 128), + (1, 8, 256 * 2, 256, 256, 128), + (3, 2, 256 * 4, 256 * 2, 256, 128), + ], +) +def test_flash_diverse_decode_stage3(batch, head_num, seq_len, shared_seq_len, block_seq, head_dim): + # Initialize inputs + mid_out = torch.randn(batch, head_num, seq_len // block_seq + 2, head_dim, dtype=torch.bfloat16, device="cuda") + mid_out_logexpsum = torch.randn(batch, head_num, seq_len // block_seq + 2, dtype=torch.bfloat16, device="cuda") + B_Seqlen = torch.tensor([seq_len] * batch, dtype=torch.int32, device="cuda") + b_shared_seq_len = torch.tensor([shared_seq_len] * batch, dtype=torch.int32, device="cuda") + out = torch.zeros(batch, head_num, head_dim, dtype=torch.float32, device="cuda") + + # Call the function + flash_diverse_decode_stage3(mid_out, mid_out_logexpsum, B_Seqlen, b_shared_seq_len, out, block_seq) + + true_out = torch.zeros_like(out) + + from lightllm.common.basemodel.triton_kernel.att.decode_att.mha.flash_decoding.flash_decoding_stage2 import ( + flash_decode_stage2, + ) + + flash_decode_stage2(mid_out, mid_out_logexpsum, B_Seqlen, true_out, block_seq) + + assert torch.allclose(out, true_out, atol=1e-2) diff --git a/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_normal_autotune.py b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_normal_autotune.py new file mode 100644 index 0000000000..7033489f66 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/att/decode_att/int8kv/test_normal_autotune.py @@ -0,0 +1,300 @@ +import weakref +import collections +import inspect +import json +import math +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F +from torch.nn.attention import SDPBackend, sdpa_kernel + +from lightllm.common.basemodel.infer_struct import InferStateInfo +from lightllm.common.basemodel.attention.triton import int8kv as state_module +from lightllm.common.basemodel.attention.triton.int8kv import Int8kvTritonDecodeAttState +from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.normal import ( + int8kv_flash_decoding_stage1 as stage1_module, + int8kv_flash_decoding as decode_module, +) +from lightllm.common.basemodel.triton_kernel.att.decode_att.int8kv.normal.int8kv_flash_decoding_stage2 import ( + flash_decode_stage2, +) +from lightllm.common.triton_utils import autotuner as autotuner_module +from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType, AutotuneLevel +from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len + + +@pytest.fixture(autouse=True) +def autotune_environment(monkeypatch): + torch.manual_seed(42) + monkeypatch.setattr(state_module, "enable_diverse_mode_gqa_decode_fast_kernel", lambda: False) + monkeypatch.delenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", raising=False) + monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", None) + monkeypatch.setattr(stage1_module, "get_triton_autotune_level", lambda: AutotuneLevel.ADAPTIVE_AUTOTUNE) + monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: AutotuneLevel.ADAPTIVE_AUTOTUNE) + get_decode_attn_autotune_seq_len.cache_clear() + yield + get_decode_attn_autotune_seq_len.cache_clear() + + +def make_inputs( + device="cpu", head_dim=64, group_size=8, dtype=torch.bfloat16, token_count=17, block_num=2, scale_dtype=None +): + # 与内存管理器一致：K/V 和 scale 分别来自共享存储中的非连续视图。 + kv = torch.randint(-127, 128, (token_count, 4, head_dim), dtype=torch.int8, device=device) + scale = (torch.rand(token_count, 4, head_dim // group_size, device=device) * 0.006 + 0.002).to(scale_dtype or dtype) + return dict( + q=torch.randn(2, 8, head_dim, dtype=dtype, device=device), + k=kv[:, :2], + k_scale=scale[:, :2], + v=kv[:, 2:], + v_scale=scale[:, 2:], + Req_to_tokens=torch.arange(5 * 16, dtype=torch.int32, device=device).view(5, 16) % max(1, token_count), + B_req_idx=torch.tensor([4, 2], dtype=torch.int32, device=device), + B_seq_len=torch.tensor([3, 2], dtype=torch.int32, device=device), + max_len_in_batch=3, + mid_out=torch.full((2, 8, block_num, head_dim), -100, dtype=dtype, device=device), + mid_out_logsumexp=torch.full((2, 8, block_num), -100, dtype=torch.float32, device=device), + block_seq=256, + ) + + +def rebuild(inputs): + args, kwargs = stage1_module.rebuild_inputs(**inputs) + return inspect.signature(stage1_module.flash_decode_stage1.fn).bind(*args, **kwargs).arguments + + +@pytest.mark.parametrize("level", [AutotuneLevel.ADAPTIVE_AUTOTUNE, AutotuneLevel.FORCE_AUTOTUNE]) +@pytest.mark.parametrize("configured, expected", [(None, 32768), ("8193", 8704)]) +def test_tuning_key_uses_configured_length(monkeypatch, level, configured, expected): + if configured is not None: + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", configured) + monkeypatch.setattr(stage1_module, "get_triton_autotune_level", lambda: level) + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + assert stage1_module.get_run_key(torch.empty(2, 8, 64), 32768) == 2_000_000_000 + expected + + +@pytest.mark.parametrize( + "phase,level", + [ + (None, AutotuneLevel.ADAPTIVE_AUTOTUNE), + (None, AutotuneLevel.FORCE_AUTOTUNE), + (AutotuneKernelType.GENERAL, AutotuneLevel.ADAPTIVE_AUTOTUNE), + (AutotuneKernelType.DECODE_ATTENTION, AutotuneLevel.USE_AUTOTUNE_HIS_CONFIG), + ], +) +def test_lookup_uses_actual_length(monkeypatch, phase, level): + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", "invalid") + monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", phase) + monkeypatch.setattr(stage1_module, "get_triton_autotune_level", lambda: level) + for actual, expected in [(2, 512), (511, 512), (512, 512), (513, 1024), (8193, 8704)]: + assert stage1_module.get_run_key(torch.empty(2, 8, 64), actual) == 2_000_000_000 + expected + + +@pytest.mark.parametrize("token_count", [3, 128]) +def test_rebuild_preserves_int8_cache_and_scales(monkeypatch, token_count): + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", "17") + inputs = make_inputs(token_count=token_count) + snapshots = {name: value.clone() for name, value in inputs.items() if isinstance(value, torch.Tensor)} + rebuilt = rebuild(inputs) + torch.testing.assert_close(rebuilt["Req_to_tokens"].long(), torch.arange(34).view(2, 17) % token_count) + assert rebuilt["B_req_idx"].tolist() == [0, 1] + assert rebuilt["B_seq_len"].tolist() == [17, 17] + assert rebuilt["max_len_in_batch"] == 17 + assert rebuilt["block_seq"] == inputs["block_seq"] + for name in ["q", "k", "v", "k_scale", "v_scale", "mid_out", "mid_out_logsumexp"]: + assert rebuilt[name] is inputs[name] + for name, value in snapshots.items(): + torch.testing.assert_close(inputs[name], value) + + +@pytest.mark.parametrize("name", ["v", "k_scale", "v_scale"]) +def test_rebuild_rejects_inconsistent_token_capacity(name): + inputs = make_inputs() + inputs[name] = inputs[name][:-1] + with pytest.raises(AssertionError, match="same number of tokens"): + rebuild(inputs) + + +def test_rebuild_rejects_empty_cache(): + with pytest.raises(ValueError, match="non-empty KV cache"): + rebuild(make_inputs(token_count=0)) + + +def test_static_key_separates_quant_groups_without_head_or_block_counts(): + kernel = stage1_module.flash_decode_stage1 + inputs = make_inputs() + key = kernel._static_key(**inputs) + inputs["q"] = inputs["q"][:, :4] + for name in ["k", "k_scale", "v", "v_scale"]: + inputs[name] = inputs[name][:, :1] + inputs["mid_out"] = torch.empty(2, 4, 128, 64, dtype=torch.bfloat16) + assert kernel._static_key(**inputs) == key + inputs["k_scale"] = inputs["k_scale"][:, :, :2] + assert kernel._static_key(**inputs) != key + + +def test_decode_passes_saved_length_to_stage1(monkeypatch): + inputs = make_inputs() + infer_state = InferStateInfo() + infer_state.batch_size = 2 + infer_state.max_kv_seq_len = 8193 + infer_state.b_req_idx = inputs["B_req_idx"] + infer_state.b_seq_len = inputs["B_seq_len"] + infer_state.req_manager = SimpleNamespace(req_to_token_indexs=inputs["Req_to_tokens"]) + state = Int8kvTritonDecodeAttState( + backend=SimpleNamespace(), + _infer_state_weak_ref=weakref.ref(infer_state), + ) + state.init_state() + state.infer_state.max_kv_seq_len = 32768 + calls = [] + monkeypatch.setattr(decode_module, "flash_decode_stage1", lambda **kwargs: calls.append(kwargs)) + monkeypatch.setattr(decode_module, "flash_decode_stage2", lambda **kwargs: None) + state.decode_att( + inputs["q"], + (inputs["k"], inputs["k_scale"]), + (inputs["v"], inputs["v_scale"]), + alloc_func=lambda shape, dtype, device: torch.empty(shape, dtype=dtype), + ) + assert calls[0]["max_len_in_batch"] == 8193 + assert calls[0]["block_seq"] == 256 + + +def reference_attention(inputs): + def dequant(quantized, scale): + group_size = quantized.shape[-1] // scale.shape[-1] + # int8 直接乘分组 scale，转换为 Q 的 dtype 后再用 FP32 SDPA 验证 attention。 + return (quantized.float() * scale.float().repeat_interleave(group_size, dim=-1)).to(inputs["q"].dtype).float() + + k, v = dequant(inputs["k"], inputs["k_scale"]), dequant(inputs["v"], inputs["v_scale"]) + outputs = [] + for row, length in enumerate(inputs["B_seq_len"].tolist()): + indices = inputs["Req_to_tokens"][inputs["B_req_idx"][row], :length].long() + with sdpa_kernel(SDPBackend.MATH): + outputs.append( + F.scaled_dot_product_attention( + inputs["q"][row].float().unsqueeze(1), + k[indices].transpose(0, 1), + v[indices].transpose(0, 1), + enable_gqa=True, + ).squeeze(1) + ) + return torch.stack(outputs) + + +def reduce_output(inputs): + output = torch.empty_like(inputs["q"]) + flash_decode_stage2( + inputs["mid_out"], inputs["mid_out_logsumexp"], inputs["B_seq_len"], output, inputs["block_seq"] + ) + return output + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("kv_len", [257, 8192, 16384]) +@pytest.mark.parametrize( + "head_dim,group_size,dtype,scale_dtype", + [ + (64, 8, torch.bfloat16, torch.bfloat16), + (128, 32, torch.float16, torch.float16), + (256, 8, torch.bfloat16, torch.bfloat16), + (128, 8, torch.bfloat16, torch.float32), + ], +) +def test_long_kv_matches_dequantized_fp32_reference(monkeypatch, kv_len, head_dim, group_size, dtype, scale_dtype): + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", str(kv_len)) + inputs = rebuild(make_inputs("cuda", head_dim, group_size, dtype, token_count=2 * kv_len, scale_dtype=scale_dtype)) + inputs["B_seq_len"][1] -= 1 + reference = reference_attention(inputs) + # 两个中间分块迫使长请求循环处理多个 BLOCK_SEQ，同时验证末尾 mask。 + for block_n in [16, 32, 64, 128]: + stage1_module.flash_decode_stage1.fn(**inputs, run_config={"BLOCK_N": block_n, "num_warps": 4, "num_stages": 2}) + output = reduce_output(inputs) + assert torch.isfinite(output).all() + torch.testing.assert_close(output.float(), reference, atol=2e-3, rtol=2e-2) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("kv_len", [8193, 16384]) +def test_full_autotune_and_graph_reuse(tmp_path, monkeypatch, kv_len): + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", str(kv_len)) + monkeypatch.setattr(autotuner_module.dist, "is_initialized", lambda: False) + kernel = stage1_module.flash_decode_stage1 + monkeypatch.setattr(kernel, "_cache_dir", str(tmp_path), raising=False) + monkeypatch.setattr(kernel, "cached_configs", {}) + monkeypatch.setattr(kernel, "fast_match_configs", collections.defaultdict(dict)) + monkeypatch.setattr(kernel, "warmuped_configs_set", set()) + assert kernel.mutates_args == [] + inputs = make_inputs("cuda", token_count=2 * kv_len, block_num=128) + inputs["Req_to_tokens"] = torch.arange(5 * kv_len, dtype=torch.int32, device="cuda").view(5, kv_len) % (2 * kv_len) + snapshots = {name: value.clone() for name, value in inputs.items() if isinstance(value, torch.Tensor)} + reference = reference_attention(inputs) + rebuild_func, benchmark = kernel.rebuild_input_func, kernel._bench + rebuild_count, timings = 0, [] + + def checked_rebuild(*args, **kwargs): + nonlocal rebuild_count + rebuild_count += 1 + return rebuild_func(*args, **kwargs) + + def checked_bench(*args, **kwargs): + rebuilt = inspect.signature(kernel.fn).bind(*args, **kwargs).arguments + assert rebuilt["B_seq_len"].tolist() == [kv_len, kv_len] + assert rebuilt["Req_to_tokens"].shape == (2, kv_len) + assert rebuilt["max_len_in_batch"] == kv_len + for name in ["q", "k", "v", "k_scale", "v_scale", "mid_out", "mid_out_logsumexp"]: + assert rebuilt[name] is inputs[name] + elapsed = benchmark(*args, **kwargs) + timings.append(elapsed) + return elapsed + + monkeypatch.setattr(kernel, "rebuild_input_func", checked_rebuild) + monkeypatch.setattr(kernel, "_bench", checked_bench) + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + kernel(**inputs) + assert rebuild_count == 1 and len(timings) == 48 and any(math.isfinite(t) for t in timings) + run_key = str(2_000_000_000 + (kv_len + 511) // 512 * 512) + assert all(list(configs) == [run_key] for configs in kernel.cached_configs.values()) + assert len(list(tmp_path.glob("*.json"))) == 1 + assert run_key in json.loads(next(tmp_path.glob("*.json")).read_text()) + torch.testing.assert_close(reduce_output(inputs).float(), reference, atol=2e-3, rtol=2e-2) + + # 清空内存配置后从文件重载，FORCE 下仍不能跨层重复搜索。 + monkeypatch.setattr(kernel, "cached_configs", {}) + monkeypatch.setattr(kernel, "fast_match_configs", collections.defaultdict(dict)) + monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: AutotuneLevel.FORCE_AUTOTUNE) + monkeypatch.setattr(stage1_module, "get_triton_autotune_level", lambda: AutotuneLevel.FORCE_AUTOTUNE) + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + kernel(**inputs) + assert rebuild_count == 1 and len(timings) == 48 + + output = torch.empty_like(inputs["q"]) + + def decode(): + kernel(**inputs) + flash_decode_stage2( + inputs["mid_out"], inputs["mid_out_logsumexp"], inputs["B_seq_len"], output, inputs["block_seq"] + ) + + decode() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + decode() + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(output.float(), reference, atol=2e-3, rtol=2e-2) + assert rebuild_count == 1 and len(timings) == 48 + for name, snapshot in snapshots.items(): + if name not in ["mid_out", "mid_out_logsumexp"]: + torch.testing.assert_close(inputs[name], snapshot) + + # 捕获后更新为真正的长请求，验证 Graph 读取原始输入缓冲区，stage2 仍按 BLOCK_SEQ 归约。 + inputs["B_seq_len"].fill_(kv_len) + inputs["max_len_in_batch"] = kv_len + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(output.float(), reference_attention(inputs), atol=2e-3, rtol=2e-2) + assert rebuild_count == 1 and len(timings) == 48 diff --git a/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_flashattention_nopad.py b/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_flashattention_nopad.py new file mode 100644 index 0000000000..d1a53f873f --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_flashattention_nopad.py @@ -0,0 +1,104 @@ +import torch +import time +import pytest +from lightllm.common.basemodel.triton_kernel.att.prefill_att.context_flashattention_nopad import ( + context_attention_fwd_contiguous_kv, +) +from lightllm.utils.log_utils import init_logger + +logger = init_logger(__name__) + + +def torch_context_attention_fwd2(q, k, v, o, b_start_loc, b_kv_start_loc, b_seq_len, b_prompt_cache_len): + batch = b_start_loc.shape[0] + + for i in range(batch): + start_loc = b_start_loc[i] + kv_start_loc = b_kv_start_loc[i] + seq_len = b_seq_len[i] + prompt_cache_len = b_prompt_cache_len[i] + cur_q = q[start_loc : start_loc + seq_len - prompt_cache_len, :, :] + cur_q = cur_q.clone().to(torch.float32) + cur_k = k[kv_start_loc : (kv_start_loc + seq_len), :, :] + cur_k = cur_k.clone().to(torch.float32) + + cur_v = v[kv_start_loc : (kv_start_loc + seq_len), :, :] + cur_v = cur_v.clone().to(torch.float32) + + dk = cur_q.shape[-1] + cur_q = cur_q.permute(1, 0, 2) + cur_k = cur_k.permute(1, 2, 0) + cur_v = cur_v.permute(1, 0, 2) + dk = cur_q.shape[-1] + + p = torch.matmul(cur_q, cur_k) / torch.sqrt(torch.tensor(dk, dtype=torch.float32)) + + q_index = (torch.arange(cur_q.shape[1]).to(p.device) + prompt_cache_len).view(-1, 1) + k_index = torch.arange(seq_len).to(p.device).view(1, -1) + + p[:, (q_index < k_index)] = float("-inf") + + s = torch.nn.functional.softmax(p, dim=-1) + + o[start_loc : start_loc + seq_len - prompt_cache_len, :, :] = torch.matmul(s, cur_v).transpose(0, 1) + + +@pytest.mark.parametrize( + "B, H, N_CTX, D_HEAD, prompt_cache_len", + [ + (b, H, N_CTX, D_HEAD, prompt_cache_len) + for b in [1, 2, 4] + for H in [1, 8] + for N_CTX in [3, 10, 1024] + for D_HEAD in [64, 128] + for prompt_cache_len in [0, 56, 200] + ], +) +def test_context_attention_fwd_contiguous_kv(B, H, N_CTX, D_HEAD, prompt_cache_len): + dtype = torch.float16 + prompt_cache_len = 0 + if prompt_cache_len >= N_CTX - 1: + return + + q = torch.empty((B * (N_CTX - prompt_cache_len), H, D_HEAD), dtype=dtype, device="cuda").normal_(mean=0.1, std=0.2) + kv = torch.empty((B * N_CTX, 2 * H, D_HEAD), dtype=dtype, device="cuda").normal_(mean=0.4, std=0.2) + k = kv[:, :H, :] + v = kv[:, H:, :] + + o = torch.empty((B * (N_CTX - prompt_cache_len), H, D_HEAD), dtype=dtype, device="cuda") + torch_o = torch.empty((B * (N_CTX - prompt_cache_len), H, D_HEAD), dtype=dtype, device="cuda") + + max_q_input_len = N_CTX - prompt_cache_len + + b_seq_len = torch.ones((B,), dtype=torch.int32, device="cuda") + b_start_loc = torch.zeros((B,), dtype=torch.int32, device="cuda") + b_seq_len = torch.ones((B,), dtype=torch.int32, device="cuda") + b_prompt_cache_len = torch.zeros(B, dtype=torch.int32, device="cuda") + + for i in range(B): + b_seq_len[i] = N_CTX + if i != 0: + b_start_loc[i] = b_start_loc[i - 1] + N_CTX - prompt_cache_len + b_prompt_cache_len[i] = prompt_cache_len + + b_kv_start_loc = torch.cumsum(b_seq_len, dim=0, dtype=torch.int32) - b_seq_len + torch_context_attention_fwd2(q, k, v, torch_o, b_start_loc, b_kv_start_loc, b_seq_len, b_prompt_cache_len) + context_attention_fwd_contiguous_kv( + q=q, + k=k, + v=v, + o=o, + b_start_loc=b_start_loc, + b_kv_start_loc=b_kv_start_loc, + b_seq_len=b_seq_len, + max_q_input_len=max_q_input_len, + b_prompt_cache_len=b_prompt_cache_len, + ) + + assert torch.allclose(torch_o, o, atol=1e-2, rtol=0) + cos = torch.nn.CosineSimilarity(0) + assert cos(o.flatten().float(), torch_o.flatten().float()) > 0.99 + + +if __name__ == "__main__": + pytest.main() diff --git a/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_flashattention_nopad1.py b/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_flashattention_nopad1.py new file mode 100644 index 0000000000..541594306d --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_context_flashattention_nopad1.py @@ -0,0 +1,196 @@ +import torch +import time +import pytest +import numpy as np +import torch.nn.functional as F +import flashinfer +from lightllm.utils.log_utils import init_logger +from lightllm.common.basemodel.triton_kernel.att.prefill_att.context_flashattention_nopad import ( + context_attention_fwd, + context_attention_fwd_no_prompt_cache, +) +from lightllm.models.llama.infer_struct import LlamaInferStateInfo + +logger = init_logger(__name__) + +seed = 42 +torch.manual_seed(seed) + +if torch.cuda.is_available(): + torch.cuda.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + + +@pytest.mark.parametrize( + "batch, seqlen, q_heads, kv_heads, head_dim", + [ + (a, b, c, d, e) + for a in [1, 16, 32, 128, 512] + for b in [16, 32, 512, 1024] + for c in [28] + for d in [4] + for e in [128] + ], +) +def test_context_attention_fwd(batch, seqlen, q_heads, kv_heads, head_dim): + Z, N_CTX, Q_HEADS, KV_HEADS, HEAD_DIM = batch, seqlen, q_heads, kv_heads, head_dim + dtype = torch.bfloat16 + kv = torch.randn((Z * N_CTX, 2 * KV_HEADS, HEAD_DIM), dtype=dtype, device="cuda") + + max_input_len = Z * N_CTX + req_to_token_indexs = torch.randperm(max_input_len, dtype=torch.int32).cuda().view(Z, N_CTX) + b_seq_len = torch.ones((Z,), dtype=torch.int32, device="cuda") * N_CTX + b_ready_cache_len = torch.zeros_like(b_seq_len, dtype=torch.int32, device="cuda") + b_ready_cache_len = torch.randint_like(b_seq_len, high=N_CTX - 1, dtype=torch.int32, device="cuda") + b_req_idx = torch.randperm(Z, dtype=torch.int32).cuda() + q_lens = b_seq_len - b_ready_cache_len + q_start_loc = q_lens.cumsum(0) - q_lens + kv_start_loc = b_seq_len.cumsum(0) - b_seq_len + + q = torch.randn((q_lens.sum(), Q_HEADS, HEAD_DIM), dtype=dtype, device="cuda") + o = torch.zeros((q_lens.sum(), Q_HEADS, HEAD_DIM), dtype=dtype, device="cuda") + o1 = torch.zeros((q_lens.sum(), Q_HEADS, HEAD_DIM), dtype=dtype, device="cuda") + + infer_state = LlamaInferStateInfo() + infer_state.batch_size = Z + infer_state.max_q_seq_len = N_CTX + infer_state.total_token_num = Z * N_CTX + infer_state.req_manager = type("Object", (), {})() + 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 + infer_state.b_ready_cache_len = b_ready_cache_len + infer_state.b_q_start_loc = q_start_loc + + context_attention_fwd( + q, + kv[:, :KV_HEADS, :], + kv[:, KV_HEADS:, :], + o, + infer_state.b_req_idx, + infer_state.b_q_start_loc, + infer_state.b_seq_len, + infer_state.b_ready_cache_len, + infer_state.max_q_seq_len, + infer_state.req_manager.req_to_token_indexs, + ) + + batch_size = Z + head_dim = HEAD_DIM + q_heads = Q_HEADS + kv_heads = KV_HEADS + page_size = 1 + workspace_buffer = torch.empty(256 * 1024 * 1024, dtype=torch.int8).to(0) + q_starts = torch.zeros((Z + 1,)).int().cuda() + q_starts[1:] = torch.cumsum(b_seq_len - b_ready_cache_len, dim=0) + kv_starts = torch.zeros_like(q_starts) + kv_starts[1:] = torch.cumsum(b_seq_len, dim=0) + q_indptr = q_starts.int() + kv_indptr = kv_starts.int() + kv_indices = torch.arange(Z * N_CTX).cuda().int() + for b, sl, start in zip(b_req_idx, b_seq_len, kv_start_loc): + kv_indices[start : start + sl] = req_to_token_indexs[b][:sl] + kv_last_page_len_buffer = torch.empty(batch_size, device="cuda:0", dtype=torch.int32) + wrapper = flashinfer.prefill.BatchPrefillWithPagedKVCacheWrapper( + workspace_buffer, + qo_indptr_buf=q_indptr, + paged_kv_indptr_buf=kv_indptr, + paged_kv_indices_buf=kv_indices, + paged_kv_last_page_len_buf=kv_last_page_len_buffer, + ) + kv_last_page_len = torch.full((batch_size,), page_size, dtype=torch.int32) + wrapper.plan( + q_indptr, + kv_indptr, + kv_indices, + kv_last_page_len, + q_heads, + kv_heads, + head_dim, + page_size, + causal=True, + pos_encoding_mode="NONE", + logits_soft_cap=0.0, + q_data_type=q.dtype, + kv_data_type=kv.dtype, + ) + kv = kv.unsqueeze(1) + wrapper.run(q, (kv[:, :, :KV_HEADS, :], kv[:, :, KV_HEADS:, :]), out=o1, return_lse=False) + + # assert torch.allclose(o, o1, atol=1e-2, rtol=0) + cos_sim1 = F.cosine_similarity(o, o1).mean() + assert cos_sim1 == 1.0 + + +@pytest.mark.parametrize( + "batch, seqlen, q_heads, kv_heads, head_dim", + [ + (a, b, c, d, e) + for a in [ + 1, + 16, + 32, + ] + for b in [16, 32, 512, 1024] + for c in [28] + for d in [4] + for e in [128] + ], +) +def test_context_attention_fwd_no_prompt_cache(batch, seqlen, q_heads, kv_heads, head_dim): + Z, N_CTX, Q_HEADS, KV_HEADS, HEAD_DIM = batch, seqlen, q_heads, kv_heads, head_dim + dtype = torch.bfloat16 + q = torch.randn((Z * N_CTX, Q_HEADS, HEAD_DIM), dtype=dtype, device="cuda") + k = torch.randn((Z * N_CTX, KV_HEADS, HEAD_DIM), dtype=dtype, device="cuda") + v = torch.randn((Z * N_CTX, KV_HEADS, HEAD_DIM), dtype=dtype, device="cuda") + + b_seq_len = torch.ones((Z,), dtype=torch.int32, device="cuda") * N_CTX + b_start_loc = b_seq_len.cumsum(0) - b_seq_len + + o = torch.zeros((Z * N_CTX, Q_HEADS, HEAD_DIM), dtype=dtype, device="cuda") + o1 = torch.zeros((Z * N_CTX, Q_HEADS, HEAD_DIM), dtype=dtype, device="cuda") + + infer_state = LlamaInferStateInfo() + infer_state.batch_size = Z + infer_state.max_q_seq_len = N_CTX + infer_state.b_seq_len = b_seq_len + infer_state.b_q_start_loc = b_start_loc + + context_attention_fwd_no_prompt_cache( + q, + k, + v, + o, + infer_state.b_q_start_loc, + infer_state.b_seq_len, + infer_state.max_q_seq_len, + ) + + head_dim = HEAD_DIM + q_heads = Q_HEADS + kv_heads = KV_HEADS + workspace_buffer = torch.empty(256 * 1024 * 1024, dtype=torch.int8).to(0) + q_starts = torch.zeros((Z + 1,)).int().cuda() + q_starts[1:] = torch.cumsum(b_seq_len, dim=0) + kv_starts = torch.zeros_like(q_starts) + kv_starts[1:] = torch.cumsum(b_seq_len, dim=0) + q_indptr = q_starts.int() + kv_indptr = kv_starts.int() + wrapper = flashinfer.prefill.BatchPrefillWithRaggedKVCacheWrapper( + workspace_buffer, + ) + wrapper.plan( + qo_indptr=q_indptr, + kv_indptr=kv_indptr, + num_qo_heads=q_heads, + num_kv_heads=kv_heads, + head_dim_qk=head_dim, + head_dim_vo=head_dim, + q_data_type=dtype, + causal=True, + ) + wrapper.run(q, k, v, out=o1, return_lse=False) + + # assert torch.allclose(o, o1, atol=1e-2, rtol=0) + cos_sim1 = F.cosine_similarity(o, o1).mean() + assert cos_sim1 == 1.0 diff --git a/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_fa3_neo.py b/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_fa3_neo.py new file mode 100644 index 0000000000..eaa19ac711 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/att/prefill_att/test_fa3_neo.py @@ -0,0 +1,561 @@ +"""Unit test for the FA3-based prefill path with image-token support. + +Neo FA3's ``flash_attn_with_kvcache`` accepts an ``image_token_end`` tensor +of exclusive KV end positions. Image queries can see their entire image span +and preceding tokens, while text queries retain the causal mask. + +Torch reference expresses the *semantics* of the attention, not FA3's internal +tiling — it has no notion of BLOCK_N / BLOCK_M. For each batch element we +gather K/V for the whole request (prompt + new tokens) and apply:: + + allow[m, k] = (k <= q_pos[m]) OR (k < image_end[m]) + +i.e. image queries cannot see subsequent text or later images. + +Run directly for quick debugging: + + python unit_tests/common/basemodel/triton_kernel/att/prefill_att/\ + test_fa3_neo.py + +or via pytest: + + pytest unit_tests/common/basemodel/triton_kernel/att/prefill_att/\ + test_fa3_neo.py -x -s +""" + +import math +import pytest +import torch + +from lightllm.common.basemodel.attention.fa3.neo import flash_attn_with_kvcache_neo as flash_attn_with_kvcache + +try: + import triton + import triton.testing as triton_testing +except ImportError: + triton = None + triton_testing = None + +try: + from lightllm.common.basemodel.triton_kernel.att.prefill_att.context_attention_fwd_neo import ( + context_attention_fwd_neo, + ) +except ImportError: + context_attention_fwd_neo = None + + +def torch_reference_context_attention_neo( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + b_req_idx: torch.Tensor, + b_q_start_loc: torch.Tensor, + b_seq_len: torch.Tensor, + b_prompt_cache_len: torch.Tensor, + req_to_token_indexs: torch.Tensor, + b_image_token_end: torch.Tensor, +) -> torch.Tensor: + device = q.device + dtype = q.dtype + _, Hq, D = q.shape + Hk = k.shape[1] + kv_group = Hq // Hk + scale = 1.0 / math.sqrt(D) + + out = torch.empty_like(q) + + for b in range(b_seq_len.shape[0]): + req_idx = int(b_req_idx[b].item()) + seq_len = int(b_seq_len[b].item()) + prompt_cache_len = int(b_prompt_cache_len[b].item()) + q_seq_len = seq_len - prompt_cache_len + if q_seq_len <= 0: + continue + + q_start = int(b_q_start_loc[b].item()) + q_blk = q[q_start : q_start + q_seq_len] # [M, Hq, D] + image_end = b_image_token_end[q_start : q_start + q_seq_len].to(torch.int64) + + token_locs = req_to_token_indexs[req_idx, :seq_len].to(torch.int64) + k_blk = k[token_locs] # [seq_len, Hk, D] + v_blk = v[token_locs] + + q_pos = torch.arange(prompt_cache_len, seq_len, device=device, dtype=torch.int64) # [M] + k_pos = torch.arange(0, seq_len, device=device, dtype=torch.int64) # [seq_len] + causal = k_pos[None, :] <= q_pos[:, None] + allow = causal | (k_pos[None, :] < image_end[:, None]) + + out_blk = torch.empty_like(q_blk) + for h in range(Hq): + h_k = h // kv_group + q_h = q_blk[:, h, :].to(torch.float32) + k_h = k_blk[:, h_k, :].to(torch.float32) + v_h = v_blk[:, h_k, :].to(torch.float32) + + scores = (q_h @ k_h.transpose(0, 1)) * scale + scores = torch.where(allow, scores, torch.full_like(scores, -1.0e8)) + probs = torch.softmax(scores, dim=-1) + out_h = (probs @ v_h).to(dtype) + out_blk[:, h, :] = out_h + + out[q_start : q_start + q_seq_len] = out_blk + + return out + + +def _build_inputs( + batch: int, + Hq: int, + Hk: int, + D: int, + dtype: torch.dtype, + device: str, + max_q_seq_len: int = 256, + max_prompt_cache_len: int = 512, + image_prob: float = 0.7, + num_image_spans_max: int = 3, + image_span_len_max: int = 24, + kv_pool_slack: int = 4096, + seed: int = 0, +): + """Build one realistic prefill batch. + + Naming matches lightllm's infer_state: + - ``q_seq_len`` = number of new Q tokens in this prefill call + - ``prompt_cache_len`` = length of the already-cached prefix for this req + - ``seq_len`` = prompt_cache_len + q_seq_len (total KV length) + """ + g = torch.Generator(device="cpu").manual_seed(seed) + + q_seq_lens = torch.randint(low=1, high=max_q_seq_len + 1, size=(batch,), generator=g) + prompt_cache_lens = torch.randint(low=0, high=max_prompt_cache_len + 1, size=(batch,), generator=g) + seq_lens = q_seq_lens + prompt_cache_lens + + sum_q = int(q_seq_lens.sum().item()) + sum_total = int(seq_lens.sum().item()) + max_seq_len_in_batch = int(seq_lens.max().item()) + max_q_seq_len_in_batch = int(q_seq_lens.max().item()) + + b_q_start_loc = torch.zeros(batch, dtype=torch.int32) + cur = 0 + for i in range(batch): + b_q_start_loc[i] = cur + cur += int(q_seq_lens[i].item()) + + # Permute so batch idx != request idx: exercises the page_table indexing. + b_req_idx = torch.randperm(batch, generator=g).to(torch.int32) + + # Global KV pool with scattered, non-contiguous slot assignment per request. + base = 1024 + kv_pool_size = base + sum_total + kv_pool_slack + pool = torch.randperm(kv_pool_size - base, generator=g)[:sum_total] + base + + req_to_token_indexs = torch.zeros((batch, max_seq_len_in_batch), dtype=torch.int32) + p = 0 + for r_logical, req_id in enumerate(b_req_idx.tolist()): + L = int(seq_lens[r_logical].item()) + req_to_token_indexs[req_id, :L] = pool[p : p + L].to(torch.int32) + p += L + + # Randomly place contiguous image-token spans inside each batch's new-Q region. + b_image_token_end = torch.zeros(sum_q, dtype=torch.int32) + for i in range(batch): + M = int(q_seq_lens[i].item()) + if M < 2: + continue + if torch.rand((), generator=g).item() > image_prob: + continue + n_spans = int(torch.randint(1, num_image_spans_max + 1, (1,), generator=g).item()) + start_pack = int(b_q_start_loc[i].item()) + for _ in range(n_spans): + span_len = int(torch.randint(1, max(2, image_span_len_max) + 1, (1,), generator=g).item()) + span_len = min(span_len, M) + s_rel = int(torch.randint(0, M - span_len + 1, (1,), generator=g).item()) + image_span = b_image_token_end[start_pack + s_rel : start_pack + s_rel + span_len] + if (image_span > 0).any(): + continue + image_span.fill_(int(prompt_cache_lens[i].item()) + s_rel + span_len) + + b_seq_len = seq_lens.to(torch.int32) + b_prompt_cache_len = prompt_cache_lens.to(torch.int32) + + q = torch.randn((sum_q, Hq, D), dtype=dtype, device=device) + k = torch.randn((kv_pool_size, Hk, D), dtype=dtype, device=device) + v = torch.randn((kv_pool_size, Hk, D), dtype=dtype, device=device) + + return dict( + q=q, + k=k, + v=v, + b_req_idx=b_req_idx.to(device), + b_q_start_loc=b_q_start_loc.to(device), + b_seq_len=b_seq_len.to(device), + b_prompt_cache_len=b_prompt_cache_len.to(device), + max_seq_len_in_batch=max_seq_len_in_batch, + max_q_seq_len_in_batch=max_q_seq_len_in_batch, + req_to_token_indexs=req_to_token_indexs.to(device), + b_image_token_end=b_image_token_end.to(device), + q_seq_lens=q_seq_lens, + prompt_cache_lens=prompt_cache_lens, + ) + + +def _fa3_prefill_with_image_end(inputs: dict) -> torch.Tensor: + """Drive ``flash_attn_with_kvcache`` with the same prefill semantics as + ``Fa3PrefillAttState._nomarl_prefill_att`` plus an optional + ``image_token_end`` kwarg for image-token bidirectional attention. + """ + q = inputs["q"] + k = inputs["k"] + v = inputs["v"] + device = q.device + + # Build page_table[b, p] = req_to_token_indexs[b_req_idx[b], p]. + page_table = inputs["req_to_token_indexs"][inputs["b_req_idx"].long(), : inputs["max_seq_len_in_batch"]].to( + torch.int32 + ) + + q_seq_lens_t = inputs["b_seq_len"].to(torch.int32) - inputs["b_prompt_cache_len"].to(torch.int32) + + cu_seqlens_q = torch.zeros(q_seq_lens_t.shape[0] + 1, dtype=torch.int32, device=device) + cu_seqlens_q[1:] = q_seq_lens_t.cumsum(0).to(torch.int32) + + cu_seqlens_k = torch.zeros(inputs["b_seq_len"].shape[0] + 1, dtype=torch.int32, device=device) + cu_seqlens_k[1:] = inputs["b_seq_len"].cumsum(0).to(torch.int32) + + sm_scale = 1.0 / math.sqrt(q.shape[-1]) + + # page_size = 1 paged KV cache view. + k_cache = k.view(k.shape[0], 1, k.shape[1], k.shape[2]) + v_cache = v.view(v.shape[0], 1, v.shape[1], v.shape[2]) + + o = flash_attn_with_kvcache( + q=q, + k_cache=k_cache, + v_cache=v_cache, + page_table=page_table, + cache_seqlens=inputs["b_seq_len"], + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k_new=cu_seqlens_k, + max_seqlen_q=inputs["max_q_seq_len_in_batch"], + softmax_scale=sm_scale, + causal=True, + window_size=(-1, -1), + softcap=0.0, + k_descale=None, + v_descale=None, + return_softmax_lse=False, + # Packed like q: int32 exclusive KV end for each image query; + # zero means ordinary causal attention. + image_token_end=inputs["b_image_token_end"], + ) + return o + + +def _report_per_batch_error(out_fa3, out_ref, q_seq_lens, b_q_start_loc, image_end, tag=""): + print(f"\n[{tag}] per-batch error breakdown (abs / rel / cos):") + for i in range(q_seq_lens.shape[0]): + s = int(b_q_start_loc[i].item()) + m = int(q_seq_lens[i].item()) + if m == 0: + continue + a = out_fa3[s : s + m].float() + b = out_ref[s : s + m].float() + abs_err = (a - b).abs().max().item() + denom = b.abs().max().item() + 1e-6 + rel_err = abs_err / denom + cos = torch.nn.functional.cosine_similarity(a.flatten(), b.flatten(), dim=0).item() + n_img = int((image_end[s : s + m] > 0).sum().item()) + print( + f" batch {i:02d} | M={m:4d} | image_tokens={n_img:4d} | " + f"max_abs={abs_err:.4e} | max_rel={rel_err:.4e} | cos={cos:.6f}" + ) + + +def _run_case( + batch: int, + Hq: int, + Hk: int, + D: int, + dtype: torch.dtype, + seed: int, + max_q_seq_len: int, + max_prompt_cache_len: int, + atol: float = 5e-2, + rtol: float = 5e-2, + cos_threshold: float = 0.99, + verbose: bool = True, +): + assert Hq % Hk == 0 + device = "cuda" + + inputs = _build_inputs( + batch=batch, + Hq=Hq, + Hk=Hk, + D=D, + dtype=dtype, + device=device, + max_q_seq_len=max_q_seq_len, + max_prompt_cache_len=max_prompt_cache_len, + seed=seed, + ) + + out_fa3 = _fa3_prefill_with_image_end(inputs) + + out_ref = torch_reference_context_attention_neo( + inputs["q"], + inputs["k"], + inputs["v"], + inputs["b_req_idx"], + inputs["b_q_start_loc"], + inputs["b_seq_len"], + inputs["b_prompt_cache_len"], + inputs["req_to_token_indexs"], + inputs["b_image_token_end"], + ) + + a = out_fa3.float().reshape_as(out_ref.float()) + b = out_ref.float() + abs_err = (a - b).abs().max().item() + denom = b.abs().max().item() + 1e-6 + rel_err = abs_err / denom + cos = torch.nn.functional.cosine_similarity(a.flatten(), b.flatten(), dim=0).item() + + n_image = int((inputs["b_image_token_end"] > 0).sum().item()) + n_tokens = int(inputs["b_image_token_end"].numel()) + if verbose: + print( + f"\ncase: batch={batch} Hq={Hq} Hk={Hk} D={D} dtype={dtype} " + f"seed={seed} image_tokens={n_image}/{n_tokens}" + ) + print( + f" global: max_abs={abs_err:.4e} max_rel={rel_err:.4e} cos={cos:.6f} " + f"(allclose atol={atol}, rtol={rtol}? " + f"{torch.allclose(a, b, atol=atol, rtol=rtol)})" + ) + _report_per_batch_error( + out_fa3, + out_ref, + inputs["q_seq_lens"], + inputs["b_q_start_loc"], + inputs["b_image_token_end"], + tag=f"seed={seed}", + ) + + return abs_err, rel_err, cos + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="no CUDA") +@pytest.mark.skipif(flash_attn_with_kvcache is None, reason="fa3 not available") +@pytest.mark.parametrize( + "batch,Hq,Hk,D,dtype,seed,max_q_seq_len,max_prompt_cache_len", + [ + (4, 8, 2, 128, torch.bfloat16, 0, 128, 256), + (4, 8, 2, 128, torch.bfloat16, 1, 256, 512), + (8, 16, 4, 128, torch.bfloat16, 2, 256, 512), + (16, 28, 4, 128, torch.bfloat16, 3, 128, 256), + (4, 8, 2, 128, torch.float16, 4, 256, 512), + (4, 8, 8, 64, torch.bfloat16, 5, 128, 256), + (3, 8, 2, 128, torch.bfloat16, 6, 8, 1024), + ], +) +def test_fa3_neo_prefill_with_image_end(batch, Hq, Hk, D, dtype, seed, max_q_seq_len, max_prompt_cache_len): + abs_err, rel_err, cos = _run_case( + batch=batch, + Hq=Hq, + Hk=Hk, + D=D, + dtype=dtype, + seed=seed, + max_q_seq_len=max_q_seq_len, + max_prompt_cache_len=max_prompt_cache_len, + verbose=True, + ) + assert cos > 0.99, f"cosine similarity too low: {cos}" + assert rel_err < 5e-2, f"max relative error too large: {rel_err}" + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="no CUDA") +@pytest.mark.parametrize("backend", ["triton", "fa3"]) +def test_image_queries_stop_at_their_own_image_end(backend): + if backend == "fa3" and flash_attn_with_kvcache is None: + pytest.skip("Neo FA3 with image_token_end support is not available") + if backend == "triton" and context_attention_fwd_neo is None: + pytest.skip("Triton attention is not available") + + def indices(values): + return torch.tensor(values, dtype=torch.int32, device="cuda") + + # Two cached text tokens, then text / image / text / image / text. + # Zero Q/K make each output the mean of its visible V values. + inputs = dict( + q=torch.zeros((8, 2, 64), dtype=torch.float16, device="cuda"), + k=torch.zeros((10, 1, 64), dtype=torch.float16, device="cuda"), + v=torch.arange(10, dtype=torch.float16, device="cuda").view(10, 1, 1).expand(10, 1, 64).contiguous(), + b_req_idx=indices([0]), + b_q_start_loc=indices([0]), + b_seq_len=indices([10]), + b_prompt_cache_len=indices([2]), + req_to_token_indexs=indices([list(range(10))]), + b_image_token_end=indices([0, 5, 5, 0, 8, 8, 0, 0]), + max_seq_len_in_batch=10, + max_q_seq_len_in_batch=8, + ) + if backend == "fa3": + output = _fa3_prefill_with_image_end(inputs) + else: + output = torch.empty_like(inputs["q"]) + context_attention_fwd_neo( + inputs["q"], + inputs["k"], + inputs["v"], + output, + inputs["b_req_idx"], + inputs["b_q_start_loc"], + inputs["b_seq_len"], + inputs["b_prompt_cache_len"], + inputs["max_q_seq_len_in_batch"], + inputs["req_to_token_indexs"], + inputs["b_image_token_end"], + ) + expected = torch.tensor([1, 2, 2, 2.5, 3.5, 3.5, 4, 4.5], dtype=output.dtype, device="cuda") + torch.testing.assert_close(output, expected[:, None, None].expand_as(output), rtol=1e-3, atol=1e-3) + + +def _bench_case( + batch: int, + Hq: int, + Hk: int, + D: int, + dtype: torch.dtype, + seed: int, + max_q_seq_len: int, + max_prompt_cache_len: int, + rep_ms: int = 100, + warmup_iters: int = 3, +): + """Compare FA3 (with image_token_end) vs the original Triton + ``context_attention_fwd_neo`` using ``triton.testing.do_bench_cudagraph``. + + Both kernels are captured into a CUDA graph so scheduling/launch overhead + is minimized and the measurement reflects the kernel cost. + """ + assert Hq % Hk == 0 + device = "cuda" + + inputs = _build_inputs( + batch=batch, + Hq=Hq, + Hk=Hk, + D=D, + dtype=dtype, + device=device, + max_q_seq_len=max_q_seq_len, + max_prompt_cache_len=max_prompt_cache_len, + seed=seed, + ) + + # --- fa3 runner: output tensor is allocated inside flash_attn_with_kvcache. + def fa3_run(): + return _fa3_prefill_with_image_end(inputs) + + # --- triton runner: pre-allocate o so the graph captures only the kernel launch. + o_triton = torch.empty_like(inputs["q"]) + + def triton_run(): + context_attention_fwd_neo( + inputs["q"], + inputs["k"], + inputs["v"], + o_triton, + inputs["b_req_idx"], + inputs["b_q_start_loc"], + inputs["b_seq_len"], + inputs["b_prompt_cache_len"], + inputs["max_q_seq_len_in_batch"], + inputs["req_to_token_indexs"], + inputs["b_image_token_end"], + ) + + # Warm up outside the graph capture so lazy allocations / autotune happen. + for _ in range(warmup_iters): + fa3_run() + triton_run() + torch.cuda.synchronize() + + fa3_ms = triton_testing.do_bench_cudagraph(fa3_run, rep=rep_ms) + triton_ms = triton_testing.do_bench_cudagraph(triton_run, rep=rep_ms) + + n_image = int((inputs["b_image_token_end"] > 0).sum().item()) + n_tokens = int(inputs["b_image_token_end"].numel()) + sum_kv = int(inputs["b_seq_len"].sum().item()) + speedup = triton_ms / fa3_ms if fa3_ms > 0 else float("inf") + + print( + f"bench: batch={batch} Hq={Hq} Hk={Hk} D={D} dtype={str(dtype).split('.')[-1]:<8s} " + f"max_q_seq_len={max_q_seq_len:4d} max_prompt_cache_len={max_prompt_cache_len:4d} " + f"image_tokens={n_image:4d}/{n_tokens:5d} sum_kv={sum_kv:6d} | " + f"fa3 {fa3_ms*1000:8.1f} us | triton {triton_ms*1000:8.1f} us | " + f"speedup {speedup:5.2f}x" + ) + + return fa3_ms, triton_ms + + +if __name__ == "__main__": + if not torch.cuda.is_available(): + print("No CUDA available.") + raise SystemExit(0) + if flash_attn_with_kvcache is None: + print("Neo FA3 flash_attn_with_kvcache with image_token_end support is not available.") + raise SystemExit(0) + + torch.manual_seed(0) + + cases = [ + dict(batch=4, Hq=8, Hk=2, D=128, dtype=torch.bfloat16, seed=0, max_q_seq_len=128, max_prompt_cache_len=256), + dict(batch=8, Hq=16, Hk=4, D=128, dtype=torch.bfloat16, seed=1, max_q_seq_len=256, max_prompt_cache_len=512), + dict(batch=16, Hq=28, Hk=4, D=128, dtype=torch.bfloat16, seed=2, max_q_seq_len=128, max_prompt_cache_len=256), + # FP16 case disabled: current build compiled without FP16 support + # dict(batch=4, Hq=8, Hk=2, D=128, dtype=torch.float16, seed=3, max_q_seq_len=256, max_prompt_cache_len=512), + ] + + print("=" * 100) + print("Correctness") + print("=" * 100) + for cfg in cases: + _run_case(**cfg, verbose=True) + + if triton_testing is None or context_attention_fwd_neo is None: + print("\nSkipping benchmark: triton or context_attention_fwd_neo not available.") + raise SystemExit(0) + + print("\n" + "=" * 100) + print("Benchmark (triton.testing.do_bench_cudagraph)") + print("=" * 100) + + # Cold-prefill sweep: max_prompt_cache_len=0 so seq_len == q_seq_len. + # Head shape matches neo_chat_moe / Qwen3 llm_config: + # num_attention_heads = 32, num_key_value_heads = 8, head_dim = 128 + # (GQA ratio 4:1) + bench_batches = [8, 16, 32, 64, 128] + bench_q_seq_lens = [1024, 4096, 8192] + + bench_cases = [ + dict( + batch=b, + Hq=32, + Hk=8, + D=128, + dtype=torch.bfloat16, + seed=0, + max_q_seq_len=s, + max_prompt_cache_len=0, + ) + for b in bench_batches + for s in bench_q_seq_lens + ] + + for cfg in bench_cases: + _bench_case(**cfg) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py new file mode 100644 index 0000000000..7fbb63df5d --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py @@ -0,0 +1,272 @@ +import pytest +import torch + +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import advance_cursor +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_fold import fold_prefix +from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix +from lightllm.common.basemodel.triton_kernel.mtp_utils import linear_att_mtp_state_index_update as update +from unit_tests.common.basemodel.triton_kernel.linear_att.replay_history import ( + accept_replay_window, + attach_replay_history as ReplaySSMExecutor, + replay_layer, +) + + +def ssm_snapshot(backend, req): + """折进一份 SSM 副本再读出这个请求。不改 canonical SSM，也不动 cursor。""" + + saved = backend.ssm_state + clone = saved.clone() + backend.ssm_state = clone + try: + fold_prefix(backend, torch.tensor([req], dtype=torch.int32, device=saved.device)) + finally: + backend.ssm_state = saved + return clone[:, req] + + +from types import SimpleNamespace +from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState + + +def make(n, width): + hold = n + 3 + ids = list(reversed(range(n))) + [hold] + lengths = [1 + i % width for i in range(n)] + [1] + starts = [] + reqs = [] + mtps = [] + for r, l in zip(ids, lengths): + starts.append(len(reqs)) + reqs.extend([r] * l) + mtps.extend(range(l)) + tensor = lambda x: torch.tensor(x, device="cuda", dtype=torch.int32) + args = [tensor(starts), tensor(reqs), tensor(mtps), tensor([1] * len(reqs)), width] + ssm_backend = ReplayGDNState.__new__(ReplayGDNState) + ssm_backend.cursors = tensor([7] * (hold + 1)) + ssm_backend.hold_req_id = hold + # 这个用例不触发折回，容量要盖住多步累加后的 cursor。 + ssm_backend.capacity = 4096 + ssm_backend.verify_width = width + manager = SimpleNamespace() + manager.mtp_state_index = tensor([-9] * (hold + 1)) + manager._ssm_executor = ssm_backend + return manager, args, (ids, lengths, starts) + + +def _replay_cu(args): + """args[0] 是每个请求的起点，最后一个请求接到本批 token 末尾。""" + + starts = args[0] + return torch.cat((starts, starts.new_full((1,), args[3].numel()))) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_request_counts_share_compiled_kernel(monkeypatch): + from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import _commit_prefix, commit_ssm_prefix + + reqs = torch.arange(512, dtype=torch.int32, device="cuda") + lengths = reqs % 4 + 1 + window = lengths + 4 + indices = torch.zeros(513, dtype=torch.int32, device="cuda") + cursors = torch.zeros_like(indices) + hashes = set() + run = _commit_prefix.run + + def record(*args, **kwargs): + kernel = run(*args, **kwargs) + hashes.add(kernel.hash) + return kernel + + monkeypatch.setattr(_commit_prefix, "run", record) + for count in (1, 2, 7, 16, 23, 24, 25, 31, 32, 33, 255, 256, 257, 511): + indices.fill_(-1) + cursors.zero_() + cursors[:count] = window[:count] + commit_ssm_prefix(indices, reqs[:count], lengths[:count], window[:count], cursors) + assert torch.equal(indices[:count], lengths[:count] - 1) + assert (indices[count:] == -1).all() + assert torch.equal(cursors[:count], lengths[:count]) + assert not cursors[count:].any() + assert len(hashes) == 1 + + +def accept_unfused(m, args): + update(m.mtp_state_index, *args) + reqs = args[1][args[0].long()] + if m._ssm_executor.verify_width > 1: + accept_replay_window(m._ssm_executor, reqs, m.mtp_state_index, _replay_cu(args)) + else: + accept_replay_window(m._ssm_executor, reqs) + + +def accept_fused(m, args): + reqs = args[1][args[0].long()] + ends = torch.cat((args[0][1:], args[0].new_full((1,), args[3].numel()))) + cumulative = torch.cat((args[3].new_zeros(1), args[3].cumsum(0))) + lengths = cumulative[ends.long()] - cumulative[args[0].long()] + executor = m._ssm_executor + if executor.verify_width > 1: + cu = _replay_cu(args) + advance_cursor(executor, reqs, cu) + commit_ssm_prefix(m.mtp_state_index, reqs, lengths, cu[1:] - cu[:-1], executor.cursors) + else: + # 宽度为 1 时正式路径不提交。对照只写下标，并把每个请求的 cursor 加一。 + m.mtp_state_index[reqs] = (lengths - 1).to(m.mtp_state_index.dtype) + advance_cursor(executor, reqs, torch.arange(reqs.numel() + 1, dtype=torch.int32, device=reqs.device)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_acceptance_fusion_matches_original_and_cpu_under_graph(): + checks = 0 + for width in [1, 2, 4, 8, 16]: + for n in [1, 7, 33, 128]: + m, args, (ids, lengths, starts) = make(n, width) + ref, _, _ = make(n, width) + # Compile both paths before graph capture, then reset before replay. + accept_unfused(ref, args) + accept_fused(m, args) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + accept_fused(m, args) + expected_index = [-9] * len(m.mtp_state_index) + expected_cursor = [7] * len(expected_index) + m.mtp_state_index.fill_(-9) + ref.mtp_state_index.fill_(-9) + m._ssm_executor.cursors.fill_(7) + ref._ssm_executor.cursors.fill_(7) + for step in range(24): + mask = [] + for i, (r, l) in enumerate(zip(ids, lengths)): + count = (step + i) % (l + 1) + mask.extend([1] * count + [0] * (l - count)) + expected_index[r] = count - 1 + if r != m._ssm_executor.hold_req_id: + expected_cursor[r] += count if width > 1 else 1 + args[3].copy_(torch.tensor(mask, device="cuda", dtype=torch.int32)) + accept_unfused(ref, args) + if step % 2: + graph.replay() + else: + accept_fused(m, args) + assert torch.equal(m.mtp_state_index, ref.mtp_state_index) + assert torch.equal(m._ssm_executor.cursors, ref._ssm_executor.cursors) + assert m.mtp_state_index.tolist() == expected_index + assert m._ssm_executor.cursors.tolist() == expected_cursor + checks += 1 + + print("ACCEPT_EQUIVALENCE_CHECKS=" + str(checks), flush=True) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("compute_mode", ["inline", "precompute"]) +def test_fused_acceptance_preserves_folds_snapshots_and_slot_reuse(dtype, compute_mode): + torch.manual_seed(93) + initial = torch.randn(2, 8, 4, 32, 32, device="cuda").to(dtype) * 0.01 + ssm_backends = [ + ReplaySSMExecutor( + initial.clone(), + 8, + 4, + torch.bfloat16, + num_key_heads=2, + compute_mode=compute_mode, + ) + for _ in range(2) + ] + managers = [] + for ssm_backend in ssm_backends: + manager = SimpleNamespace() + manager._ssm_executor = ssm_backend + manager.mtp_state_index = torch.zeros(8, dtype=torch.int32, device="cuda") + managers.append(manager) + reqs = torch.full((4,), 7, dtype=torch.int32, device="cuda") + cu = torch.tensor([0, 3, 6, 9, 12], dtype=torch.int32, device="cuda") + token_reqs = torch.full((12,), 7, dtype=torch.int32, device="cuda") + mtp_indexes = torch.arange(12, dtype=torch.int32, device="cuda") % 3 + accepted_mask = torch.ones(12, dtype=torch.int32, device="cuda") + q = torch.randn(12, 2, 32, dtype=torch.bfloat16, device="cuda") + k = torch.randn_like(q) + v = torch.randn(12, 4, 32, dtype=torch.bfloat16, device="cuda") + a = torch.full((12, 4), -3.0, dtype=torch.bfloat16, device="cuda") + b = torch.randn_like(a) + alog = torch.zeros(4, device="cuda") + args = [cu[:-1], token_reqs, mtp_indexes, accepted_mask, 4] + + def cycle(i): + ssm_backend = ssm_backends[i] + outputs = [ + replay_layer( + ssm_backend, + layer, + q, + k, + v, + a, + b, + alog, + alog, + reqs, + cu, + run_config={"BV": 32, "num_warps": 1, "num_stages": 1}, + ) + for layer in range(2) + ] + (accept_unfused if i == 0 else accept_fused)(managers[i], args) + return outputs + + graphs, outputs = [], [] + for i in range(2): + cycle(i) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = cycle(i) + graphs.append(graph) + outputs.append(output) + for iteration in range(40): + ids = [1, 3, 6] if iteration % 2 else [6, 1, 3] + lengths = [4, 4, 3] if iteration % 3 == 0 else [3, 3, 3] + lengths.append(12 - sum(lengths)) + ids.append(7) + starts, tokens, indexes, masks = [0], [], [], [] + for i, (req, length) in enumerate(zip(ids, lengths)): + count = (iteration + i) % (length + 1) if req != 7 else 0 + tokens.extend([req] * length) + indexes.extend(range(length)) + masks.extend([1] * count + [0] * (length - count)) + starts.append(len(tokens)) + for tensor, values in [ + (reqs, ids), + (cu, starts), + (token_reqs, tokens), + (mtp_indexes, indexes), + (accepted_mask, masks), + ]: + tensor.copy_(tensor.new_tensor(values)) + for tensor in (q, k, v, b): + tensor.normal_() + for graph in graphs: + graph.replay() + for left, right in zip(outputs[0], outputs[1]): + assert torch.equal(left, right) + assert torch.equal(ssm_backends[0].ssm_state, ssm_backends[1].ssm_state) + assert torch.equal(ssm_backends[0].cursors, ssm_backends[1].cursors) + assert torch.equal(managers[0].mtp_state_index, managers[1].mtp_state_index) + if iteration % 7 == 0: + assert torch.equal(ssm_snapshot(ssm_backends[0], 3), ssm_snapshot(ssm_backends[1], 3)) + if iteration % 11 == 10: + for ssm_backend in ssm_backends: + ssm_backend.fold_accepted_prefix(reqs[:3]) + assert torch.equal(ssm_backends[0].ssm_state, ssm_backends[1].ssm_state) + if iteration == 19: + for manager in managers: + manager._ssm_executor.cursors[1].zero_() + manager._ssm_executor.ssm_state[:, 1].zero_() + manager.mtp_state_index[1] = 0 + for ssm_backend in ssm_backends: + ssm_backend.fold_accepted_prefix(reqs[:3]) + assert torch.equal(ssm_backends[0].ssm_state, ssm_backends[1].ssm_state) + assert torch.equal(ssm_backends[0].ssm_state[:, [0, 2, 4, 5, 7]], initial[:, [0, 2, 4, 5, 7]]) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_fused_recurrent_strided.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_fused_recurrent_strided.py new file mode 100644 index 0000000000..969a3f354d --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_fused_recurrent_strided.py @@ -0,0 +1,71 @@ +import pytest +import torch + +from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.fused_recurrent import ( + fused_recurrent_gated_delta_rule, +) + +if not torch.cuda.is_available(): + pytest.skip("CUDA required", allow_module_level=True) + + +@pytest.mark.parametrize("batch", [1, 2, 16]) +def test_decode_strided_views_match_contiguous(batch): + """q/k/v/a/b passed as column views of one projection output (the decode + path layout) must produce the same result as contiguous copies.""" + torch.manual_seed(0) + H, HV, K, V = 2, 8, 128, 128 + key_dim, value_dim = H * K, HV * V + qkv_dim = 2 * key_dim + value_dim + total_dim = qkv_dim + value_dim + 2 * HV # qkv + z + b + a + cache_slots = 64 + + mixed = torch.randn(batch, total_dim, device="cuda", dtype=torch.bfloat16) + mixed_qkv = mixed[:, :qkv_dim] + b_raw = mixed[:, qkv_dim + value_dim : qkv_dim + value_dim + HV] + a_raw = mixed[:, qkv_dim + value_dim + HV :] + + query, key, value = torch.split(mixed_qkv, [key_dim, key_dim, value_dim], dim=-1) + q = query.view(batch, 1, H, K) + k = key.view(batch, 1, H, K) + v = value.view(batch, 1, HV, V) + + A_log = torch.randn(HV, device="cuda", dtype=torch.float32) * 0.1 + dt_bias = torch.randn(HV, device="cuda", dtype=torch.float32) * 0.1 + ssm_state = torch.randn(cache_slots, HV, K, V, device="cuda", dtype=torch.bfloat16) + idx = torch.randperm(cache_slots, device="cuda")[:batch].to(torch.int32) + + def run(q_, k_, v_, a_, b_, state): + out, _ = fused_recurrent_gated_delta_rule( + q=q_, + k=k_, + v=v_, + initial_state=state, + inplace_final_state=True, + ssm_state_indices=idx, + use_qk_l2norm_in_kernel=True, + A_log=A_log, + dt_bias=dt_bias, + a_raw=a_, + b_raw=b_, + ) + return out + + state_ref = ssm_state.clone() + out_ref = run(q.contiguous(), k.contiguous(), v.contiguous(), a_raw.contiguous(), b_raw.contiguous(), state_ref) + state_strided = ssm_state.clone() + out_strided = run(q, k, v, a_raw, b_raw, state_strided) + + assert torch.equal(out_ref, out_strided) + assert torch.equal(state_ref, state_strided) + + +# NOTE: the decode-only `cu_seqlens is None` contract from upstream #1349 was +# intentionally lifted on this branch so the Qwen3Next MTP verify path can drive +# the kernel with variable-length verify chunks (cu_seqlens + 2D SSM index +# rows). That varlen path is exercised end-to-end by the MTP GSM8K accuracy +# check rather than a hand-rolled unit test. + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_equiv.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_equiv.py new file mode 100644 index 0000000000..5e03665d43 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_equiv.py @@ -0,0 +1,163 @@ +import pytest +import torch + +from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.fused_recurrent import ( + fused_recurrent_gated_delta_rule, +) +from lightllm.common.basemodel.triton_kernel.linear_att.mtp_fused_recurrent import ( + mtp_fused_recurrent_gated_delta_rule, +) + +if not torch.cuda.is_available(): + pytest.skip("CUDA required", allow_module_level=True) + + +def _run_both( + q, + k, + v, + initial_state, + cu_seqlens, + ssm_state_indices, + ssm_state_write_indices, + num_accepted_tokens, + A_log, + dt_bias, + a_raw, + b_raw, +): + """Run old (via autograd.Function) and new (direct kernel) side-by-side.""" + state_old = initial_state.clone() + state_new = initial_state.clone() + + o_old, fs_old = fused_recurrent_gated_delta_rule( + q=q, + k=k, + v=v, + initial_state=state_old, + inplace_final_state=True, + cu_seqlens=cu_seqlens, + ssm_state_indices=ssm_state_indices, + ssm_state_write_indices=ssm_state_write_indices, + num_accepted_tokens=num_accepted_tokens, + use_qk_l2norm_in_kernel=True, + A_log=A_log, + dt_bias=dt_bias, + a_raw=a_raw, + b_raw=b_raw, + ) + + o_new, fs_new = mtp_fused_recurrent_gated_delta_rule( + q=q, + k=k, + v=v, + initial_state=state_new, + cu_seqlens=cu_seqlens, + ssm_state_indices=ssm_state_indices, + ssm_state_write_indices=ssm_state_write_indices, + num_accepted_tokens=num_accepted_tokens, + A_log=A_log, + dt_bias=dt_bias, + a_raw=a_raw, + b_raw=b_raw, + ) + + return o_old, o_new, fs_old, fs_new + + +@pytest.mark.parametrize("mtp_step", [1, 2, 3]) +def test_mtp_verify_path(mtp_step): + """MTP verify path with cu_seqlens and 2D SSM indices.""" + torch.manual_seed(123) + batch, H, HV, K, V = 2, 2, 8, 64, 64 + seqlen = mtp_step + 1 + num_tokens = batch * seqlen + cache_slots = 64 + + q = torch.randn(1, num_tokens, H, K, device="cuda", dtype=torch.bfloat16) + k = torch.randn(1, num_tokens, H, K, device="cuda", dtype=torch.bfloat16) + v = torch.randn(1, num_tokens, HV, V, device="cuda", dtype=torch.bfloat16) + A_log = torch.randn(HV, device="cuda", dtype=torch.float32) * 0.1 + dt_bias = torch.randn(HV, device="cuda", dtype=torch.float32) * 0.1 + a_raw = torch.randn(num_tokens, HV, device="cuda", dtype=torch.bfloat16) + b_raw = torch.randn(num_tokens, HV, device="cuda", dtype=torch.bfloat16) + ssm_state = torch.randn(cache_slots, HV, K, V, device="cuda", dtype=torch.bfloat16) + ssm_idx = torch.randint(0, cache_slots, (batch, seqlen), device="cuda", dtype=torch.int32) + cu_seqlens = torch.arange(batch + 1, device="cuda", dtype=torch.int32) * seqlen + num_accepted = torch.full((batch,), seqlen, device="cuda", dtype=torch.int32) + + o_old, o_new, fs_old, fs_new = _run_both( + q, + k, + v, + ssm_state, + cu_seqlens.to(torch.long), + ssm_idx, + ssm_idx, + num_accepted, + A_log, + dt_bias, + a_raw, + b_raw, + ) + + torch.testing.assert_close(o_old, o_new, rtol=torch.finfo(o_old.dtype).eps, atol=2e-3) + if not torch.equal(fs_old, fs_new): + assert torch.allclose(fs_old.float(), fs_new.float(), rtol=1e-2, atol=5.0), ( + f"final_state mismatch at mtp_step={mtp_step}, " + f"max diff={torch.abs(fs_old.float() - fs_new.float()).max().item():.6f}" + ) + + +@pytest.mark.parametrize("seed", [1, 2, 3, 4]) +def test_variable_length_cu_seqlens(seed): + """Variable-length sequences via cu_seqlens: individual lengths vary, + each ≤ mtp_step+1, with some padded to shorter lengths.""" + torch.manual_seed(seed) + batch, mtp_step, H, HV, K, V = 4, 3, 2, 8, 64, 64 + max_len = mtp_step + 1 # 4 + cache_slots = 64 + + # Random per-sequence lengths in [1, max_len] + lengths = torch.randint(1, max_len + 1, (batch,), device="cpu") + cu_seqlens = torch.zeros(batch + 1, dtype=torch.int32) + torch.cumsum(lengths, dim=0, out=cu_seqlens[1:]) + + num_tokens = int(cu_seqlens[-1].item()) + q = torch.randn(1, num_tokens, H, K, device="cuda", dtype=torch.bfloat16) + k = torch.randn(1, num_tokens, H, K, device="cuda", dtype=torch.bfloat16) + v = torch.randn(1, num_tokens, HV, V, device="cuda", dtype=torch.bfloat16) + A_log = torch.randn(HV, device="cuda", dtype=torch.float32) * 0.1 + dt_bias = torch.randn(HV, device="cuda", dtype=torch.float32) * 0.1 + a_raw = torch.randn(num_tokens, HV, device="cuda", dtype=torch.bfloat16) + b_raw = torch.randn(num_tokens, HV, device="cuda", dtype=torch.bfloat16) + ssm_state = torch.randn(cache_slots, HV, K, V, device="cuda", dtype=torch.bfloat16) + # 2D indices: S+1 = max_len columns, unused columns for short seqs are irrelevant + ssm_idx = torch.randint(0, cache_slots, (batch, max_len), device="cuda", dtype=torch.int32) + num_accepted = lengths.to(device="cuda", dtype=torch.int32) + + o_old, o_new, fs_old, fs_new = _run_both( + q, + k, + v, + ssm_state, + cu_seqlens.to(torch.long).cuda(), + ssm_idx, + ssm_idx, + num_accepted, + A_log, + dt_bias, + a_raw, + b_raw, + ) + + torch.testing.assert_close(o_old, o_new, rtol=torch.finfo(o_old.dtype).eps, atol=2e-3) + if not torch.equal(fs_old, fs_new): + assert torch.allclose(fs_old.float(), fs_new.float(), rtol=1e-2, atol=5.0), ( + f"seed={seed}: final_state mismatch, " + f"max diff={torch.abs(fs_old.float() - fs_new.float()).max().item():.6f}" + ) + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_state_params.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_state_params.py new file mode 100644 index 0000000000..8252c445a1 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_state_params.py @@ -0,0 +1,65 @@ +import pytest +import torch + +from lightllm.common.attention_state.bounded_att import GDNState, ConvSSMStateConfig + +from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( + build_dynamic_mtp_linear_att_state_params, +) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("request_capacity,token_bucket", [(4, 3), (4, 16), (128, 512)]) +def test_dynamic_sequences_stay_within_request_capacity_under_graph(request_capacity, token_bucket): + hold = request_capacity + reqs = torch.full((token_bucket,), hold, dtype=torch.int32, device="cuda") + mtp_index = torch.zeros_like(reqs) + config = ConvSSMStateConfig( + logical_layer_ids=(0,), + global_k_heads=1, + global_v_heads=1, + key_dim=2, + value_dim=2, + conv_width=3, + conv_dtype=torch.float32, + ssm_dtype=torch.float32, + ) + manager = GDNState(config, request_capacity + 1, mtp_step=3) + manager.mtp_state_index.copy_(torch.arange(request_capacity + 1, dtype=torch.int32, device="cuda") % 4) + sequences = min(request_capacity, token_bucket) + + def build(): + return build_dynamic_mtp_linear_att_state_params(reqs, mtp_index, manager.mtp_state_index, hold) + + build() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + cu, ids, counts = build() + assert cu.numel() == token_bucket + 1 + assert ids.numel() == counts.numel() == token_bucket + + # Reuse one graph for full request pools, mixed verify lengths, and HOLD-only padding. + for lengths in ([1] * sequences, [min(4, token_bucket // sequences)] * sequences, [2, 1], []): + lengths = lengths[:request_capacity] + if sum(lengths) > token_bucket: + continue + request_ids = list(reversed(range(len(lengths)))) + host_reqs = [req for req, length in zip(request_ids, lengths) for _ in range(length)] + host_mtp = [j for length in lengths for j in range(length)] + tokens = len(host_reqs) + host_reqs += [hold] * (token_bucket - tokens) + host_mtp += [0] * (token_bucket - tokens) + reqs.copy_(torch.tensor(host_reqs, dtype=torch.int32, device="cuda")) + mtp_index.copy_(torch.tensor(host_mtp, dtype=torch.int32, device="cuda")) + graph.replay() + + expected_cu = [0] + for length in lengths: + expected_cu.append(expected_cu[-1] + length) + expected_cu += [tokens] * (token_bucket + 1 - len(expected_cu)) + expected_ids = request_ids + [hold] * (token_bucket - len(request_ids)) + assert cu.tolist() == expected_cu + assert ids.tolist() == expected_ids + expected_counts = [req % 4 + 1 for req in request_ids] + [1] * (token_bucket - len(request_ids)) + assert counts.tolist() == expected_counts diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py new file mode 100644 index 0000000000..02797b61b7 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py @@ -0,0 +1,530 @@ +import pytest +import torch + +from unit_tests.common.basemodel.triton_kernel.linear_att.replay_history import ( + accept_replay_window, + attach_replay_history as ReplaySSMExecutor, + replay_layer, +) +from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( + build_dynamic_mtp_linear_att_state_params, +) +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import advance_cursor +from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix +from types import SimpleNamespace + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) +def test_decode_compilation_uses_shape_buckets(monkeypatch, kda): + from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import _advance_cursor + from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_decode import _decode_replay + + kernels = {"_decode_replay": _decode_replay, "_advance_cursor": _advance_cursor} + hashes = {name: set() for name in kernels} + launches = [] + for name, kernel in kernels.items(): + + def record(*args, run=kernel.run, name=name, **kwargs): + compiled = run(*args, **kwargs) + hashes[name].add(compiled.hash) + launches.append(name) + return compiled + + monkeypatch.setattr(kernel, "run", record) + + torch.manual_seed(137) + initial = torch.randn(1, 513, 1, 32, 32, device="cuda") * 0.01 + eager, captured = [ReplaySSMExecutor(initial.clone(), 8, 1, torch.float32, kda=kda) for _ in range(2)] + reqs = torch.arange(512, device="cuda", dtype=torch.int32) + q = torch.randn(512, 1, 32, device="cuda") * 0.2 + k, v = torch.randn_like(q), torch.randn_like(q) + a = torch.full((512, 32 if kda else 1), -3.0, device="cuda") + beta = torch.randn(512, 1, device="cuda") + alog, bias = torch.zeros(1, device="cuda"), torch.zeros(a.shape[-1], device="cuda") + + def cycle(backend, ids): + count = ids.numel() + cu = torch.arange(count + 1, dtype=torch.int32, device=ids.device) + out = replay_layer( + backend, + 0, + q[:count], + k[:count], + v[:count], + a[:count], + beta[:count], + alog, + bias, + ids, + cu, + run_config={"BV": 32, "num_warps": 1, "num_stages": 3}, + ) + accept_replay_window(backend, ids, cu_seqlens=cu) + return out + + graphs = {} + for bucket in (32, 512): + ids = reqs[:bucket].clone() + cycle(captured, ids) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = cycle(captured, ids) + graphs[bucket] = graph, ids, out + captured.ssm_state.copy_(initial) + captured.cursors.zero_() + counts = (1, 2, 7, 16, 23, 24, 25, 31, 32, 33, 255, 256, 257, 511) + for count in counts: + graph, ids, output = graphs[32 if count <= 32 else 512] + ids.fill_(captured.hold_req_id) + ids[:count].copy_(reqs[:count]) + for _ in range(8): + live = reqs[:count] + if (eager.cursors[live.long()] + 1 > eager.capacity).any(): + eager.fold_accepted_prefix(live) + captured.fold_accepted_prefix(ids[:count]) + expected = cycle(eager, reqs[:count]) + calls = len(launches) + graph.replay() + assert len(launches) == calls + torch.testing.assert_close(output[:count], expected, atol=3e-5, rtol=3e-5) + assert torch.equal(captured.cursors, eager.cursors) + torch.testing.assert_close(captured.ssm_state, eager.ssm_state, atol=3e-5, rtol=3e-5) + assert len(hashes["_decode_replay"]) == 1 + assert len(hashes["_advance_cursor"]) == 1 + assert captured.cursors[-1] == 0 + assert torch.equal(captured.ssm_state[:, -1], initial[:, -1]) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("capture_metadata", [False, True]) +@pytest.mark.parametrize("varlen", [False, True]) +@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) +def test_request_manager_accept_with_replay_graph(dtype, capture_metadata, varlen, kda): + torch.manual_seed(827) + config = {"BV": 32, "num_warps": 1, "num_stages": 1} + initial = torch.randn(3, 5, 4, 128, 128, device="cuda", dtype=dtype) * 0.01 + ssm_backends = [ + cls( + initial.clone(), + 8, + 4, + torch.bfloat16, + num_key_heads=2, + kda=kda, + compute_mode="precompute", + ) + for cls in [ReplaySSMExecutor, ReplaySSMExecutor] + ] + manager = SimpleNamespace() + manager._ssm_executor = ssm_backends[1] + manager.mtp_state_index = torch.zeros(5, device="cuda", dtype=torch.int32) + rows = torch.full((12,), 4, device="cuda", dtype=torch.int32) + mtp = torch.zeros_like(rows) + flags = torch.zeros_like(rows) + starts = torch.zeros(3, device="cuda", dtype=torch.int32) + accepted = torch.full((5,), -1, device="cuda", dtype=torch.int32) + q = torch.randn(12, 2, 128, device="cuda", dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn(12, 4, 128, device="cuda", dtype=torch.bfloat16) + a = torch.full((12, 4 * 128 if kda else 4), -3.0, device="cuda", dtype=torch.bfloat16) + beta = torch.randn(a.shape[0], 4, device="cuda", dtype=a.dtype) + alog = torch.zeros(4, device="cuda") + bias = torch.zeros(4 * 128 if kda else 4, device="cuda") + + def inputs(iteration): + ids = [0, 1, 3] if iteration % 2 == 0 else [3, 0, 1] + lengths = [1 + (iteration + j) % 4 for j in range(3)] if varlen else [4] * 3 + counts = [1 + (iteration // 2 + j) % length for j, length in enumerate(lengths)] + offsets, req_rows, indexes, mask = [], [], [], [] + table = [-1] * 5 + for req, length, count in zip(ids, lengths, counts): + offsets.append(len(req_rows)) + req_rows.extend([req] * length) + indexes.extend(range(length)) + mask.extend([int(j < count) for j in range(length)]) + table[req] = count - 1 + logical = len(req_rows) + req_rows.extend([4] * (12 - logical)) + indexes.extend([0] * (12 - logical)) + mask.extend([0] * (12 - logical)) + rows.copy_(rows.new_tensor(req_rows)) + mtp.copy_(mtp.new_tensor(indexes)) + flags.copy_(flags.new_tensor(mask)) + starts.copy_(starts.new_tensor(offsets)) + accepted.copy_(accepted.new_tensor(table)) + return ids, lengths, counts, logical + + def metadata(ssm_backend): + cu, reqs, _counts = build_dynamic_mtp_linear_att_state_params(rows, mtp, manager.mtp_state_index, 4) + return cu, reqs, ssm_backend.cursors.index_select(0, reqs.long()) + + def forward_accept(index, cu, reqs): + ssm_backend = ssm_backends[index] + out = [ + replay_layer(ssm_backend, layer, q, k, v, a, beta, alog, bias, reqs, cu, run_config=config) + for layer in range(3) + ] + if index == 0: + accept_replay_window(ssm_backend, reqs, accepted, cu) + else: + count = starts.numel() + commit_reqs = rows[starts.long()] + window = (cu[1 : count + 1] - cu[:count]).contiguous() + advance_cursor(ssm_backend, commit_reqs, cu[: count + 1]) + commit_ssm_prefix( + manager.mtp_state_index, + commit_reqs, + accepted[commit_reqs.long()] + 1, + window, + ssm_backend.cursors, + ) + return out + + inputs(0) + graphs, captured, outputs = [], [], [] + for index, ssm_backend in enumerate(ssm_backends): + cu, reqs, pos = metadata(ssm_backend) + forward_accept(index, cu, reqs) + ssm_backend.ssm_state.copy_(initial) + ssm_backend.cursors.zero_() + manager.mtp_state_index.zero_() + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + if capture_metadata: + cu, reqs, pos = metadata(ssm_backend) + out = forward_accept(index, cu, reqs) + graphs.append(g) + captured.append((cu, reqs, pos)) + outputs.append(out) + for ssm_backend in ssm_backends: + ssm_backend.ssm_state.copy_(initial) + ssm_backend.cursors.zero_() + manager.mtp_state_index.zero_() + folds = 0 + for iteration in range(32): + ids, lengths, counts, logical = inputs(iteration) + for tensor in [q, k, v, beta]: + tensor.normal_() + overflow = [ + req for req, length in zip(ids, lengths) if ssm_backends[0].cursors[req] + length > ssm_backends[0].capacity + ] + if overflow: + overflow_reqs = torch.tensor(overflow, dtype=torch.int32, device="cuda") + for ssm_backend in ssm_backends: + ssm_backend.fold_accepted_prefix(overflow_reqs) + folds += len(overflow) + for index, ssm_backend in enumerate(ssm_backends): + if not capture_metadata: + for dst, src in zip(captured[index], metadata(ssm_backend)): + dst.copy_(src) + graphs[index].replay() + torch.cuda.synchronize() + for expected, actual in zip(captured[0], captured[1]): + assert torch.equal(actual, expected) + assert torch.equal(ssm_backends[0].cursors, ssm_backends[1].cursors) + for expected, actual in zip(outputs[0], outputs[1]): + torch.testing.assert_close( + actual[:, :logical].float(), expected[:, :logical].float(), atol=0.006, rtol=0.02 + ) + torch.testing.assert_close( + ssm_backends[1].ssm_state.float(), + ssm_backends[0].ssm_state.float(), + atol=0.002, + rtol=torch.finfo(dtype).eps, + ) + table = manager.mtp_state_index.cpu().tolist() + for req, count in zip(ids, counts): + assert table[req] == count - 1 + assert table[2] == table[4] == 0 + assert ssm_backends[1].cursors[4].item() == 0 + assert folds > 0 + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("capture_metadata", [False, True]) +@pytest.mark.parametrize("stagger", [False, True]) +@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) +def test_decode_accept_order_preserves_forward_snapshot(dtype, capture_metadata, stagger, kda): + torch.manual_seed(731) + initial = torch.randn(3, 3, 4, 128, 128, device="cuda", dtype=dtype) * 0.01 + ssm_backends = [ + ReplaySSMExecutor( + initial.clone(), + 8, + 1, + dtype, + num_key_heads=2, + kda=kda, + ) + for _ in range(2) + ] + for ssm_backend in ssm_backends: + ssm_backend.fold_programs = 3 + reqs = torch.tensor([0, 1, 2], device="cuda", dtype=torch.int32) + q = torch.randn(3, 2, 128, device="cuda", dtype=dtype) + k = torch.randn_like(q) + v = torch.randn(3, 4, 128, device="cuda", dtype=dtype) + a = torch.full((3, 4 * 128 if kda else 4), -3.0, device="cuda", dtype=dtype) + beta = torch.randn(a.shape[0], 4, device="cuda", dtype=a.dtype) + alog = torch.zeros(4, device="cuda") + bias = torch.zeros(4 * 128 if kda else 4, device="cuda") + + cu = torch.arange(reqs.numel() + 1, dtype=torch.int32, device=reqs.device) + + def forward(ssm_backend): + return [ + replay_layer( + ssm_backend, + layer, + q, + k, + v, + a, + beta, + alog, + bias, + reqs, + cu, + run_config={"BV": 32, "num_warps": 1, "num_stages": 1}, + ) + for layer in range(3) + ] + + def warm(ssm_backend): + out = forward(ssm_backend) + accept_replay_window(ssm_backend, reqs) + return out + + graphs, positions, outputs = [], [], [] + for ssm_backend in ssm_backends: + warm(ssm_backend) + ssm_backend.ssm_state.copy_(initial) + ssm_backend.cursors.zero_() + if stagger: + reqs.copy_(reqs.new_tensor([0, 2, 2])) + for _ in range(3): + warm(ssm_backend) + reqs.copy_(reqs.new_tensor([0, 1, 2])) + pos = torch.zeros_like(reqs) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + if capture_metadata: + pos = ssm_backend.cursors.index_select(0, reqs.long()) + out = forward(ssm_backend) + if capture_metadata: + accept_replay_window(ssm_backend, reqs) + graphs.append(graph) + positions.append(pos) + outputs.append(out) + for iteration in range(24): + reqs.copy_(reqs.new_tensor([0, 1, 2] if iteration % 2 == 0 else [1, 0, 2])) + for tensor in [q, k, v, beta]: + tensor.normal_() + overflow = [int(req) for req in reqs.tolist() if ssm_backends[0].cursors[req] + 1 > ssm_backends[0].capacity] + if overflow: + overflow_reqs = torch.tensor(sorted(set(overflow)), dtype=torch.int32, device="cuda") + for ssm_backend in ssm_backends: + ssm_backend.fold_accepted_prefix(overflow_reqs) + old = ssm_backends[0].cursors[reqs.long()].clone() + for ssm_backend, graph, pos in zip(ssm_backends, graphs, positions): + if not capture_metadata: + pos.copy_(ssm_backend.cursors.index_select(0, reqs.long())) + graph.replay() + if not capture_metadata: + accept_replay_window(ssm_backend, reqs) + torch.cuda.synchronize() + assert torch.equal(positions[0], old) + assert torch.equal(positions[0], positions[1]) + assert torch.equal(ssm_backends[0].cursors, ssm_backends[1].cursors) + for expected, actual in zip(outputs[0], outputs[1]): + torch.testing.assert_close(actual.float(), expected.float(), atol=0.006, rtol=0.02) + torch.testing.assert_close( + ssm_backends[1].ssm_state.float(), + ssm_backends[0].ssm_state.float(), + atol=0.002, + rtol=torch.finfo(dtype).eps, + ) + + +@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) +def test_cross_layer_fold_with_state_offset_above_int32(kda): + if torch.cuda.mem_get_info()[0] < 6 * 1024 ** 3: + pytest.skip("large-stride regression requires 6 GiB of free device memory") + torch.manual_seed(791) + shape = (48, 3, 12, 128, 128) + strides = (50_528_256, 196608, 16384, 128, 1) + assert strides[0] < 2 ** 31 < 47 * strides[0] + dense = torch.randn(shape, device="cuda", dtype=torch.bfloat16) * 0.01 + sparse = torch.empty_strided(shape, strides, device="cuda", dtype=dense.dtype) + sparse.copy_(dense) + config = {"BV": 32, "num_warps": 1, "num_stages": 1} + ssm_backends = [ + ReplaySSMExecutor(state, 8, 4, torch.bfloat16, num_key_heads=4, kda=kda) for state in [dense, sparse] + ] + q = torch.randn(8, 4, 128, device="cuda", dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn(8, 12, 128, device="cuda", dtype=torch.bfloat16) + a = torch.full((8, 12 * 128 if kda else 12), -3.0, device="cuda", dtype=torch.bfloat16) + beta = torch.randn(8, 12, device="cuda", dtype=torch.bfloat16) + alog = torch.zeros(12, device="cuda") + bias = torch.zeros(12 * 128 if kda else 12, device="cuda") + reqs = torch.tensor([0, 2], device="cuda", dtype=torch.int32) + cu = reqs.new_tensor([0, 4, 8]) + + def step(ssm_backend, accepted): + overflow = [int(req) for req in reqs.tolist() if ssm_backend.cursors[req] + 4 > ssm_backend.capacity] + if overflow: + ssm_backend.fold_accepted_prefix(reqs.new_tensor(overflow)) + pos = ssm_backend.cursors.index_select(0, reqs.long()) + out = [ + replay_layer(ssm_backend, layer, q, k, v, a, beta, alog, bias, reqs, cu, run_config=config) + for layer in range(48) + ] + accept_replay_window(ssm_backend, reqs, accepted, cu) + return pos, out + + for count in [2, 3]: + accepted = reqs.new_tensor([count - 1, -1, -1]) + for ssm_backend in ssm_backends: + step(ssm_backend, accepted) + reqs.copy_(reqs.new_tensor([0, 1])) + expected_pos, expected_out = step(ssm_backends[0], reqs.new_tensor([0, 0, -1])) + actual_pos, actual_out = step(ssm_backends[1], reqs.new_tensor([0, 0, -1])) + torch.cuda.synchronize() + assert torch.equal(actual_pos, expected_pos) + assert torch.equal(ssm_backends[0].cursors, ssm_backends[1].cursors) + for layer in range(48): + torch.testing.assert_close(actual_out[layer].float(), expected_out[layer].float(), atol=0.006, rtol=0.02) + torch.testing.assert_close( + sparse[layer].float(), dense[layer].float(), atol=0.002, rtol=torch.finfo(torch.bfloat16).eps + ) + + +@pytest.mark.parametrize("batch", [1, 8, 12]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) +def test_small_batch_matches_larger_graph_padding(batch, dtype, kda): + torch.manual_seed(1087 + batch) + initial = torch.randn(2, 14, 4, 128, 128, device="cuda", dtype=dtype) * 0.01 + config = {"BV": 32, "num_warps": 1, "num_stages": 1} + ssm_backends = [ + ReplaySSMExecutor(initial.clone(), 8, 4, num_key_heads=2, kda=kda, compute_mode="precompute") for _ in range(2) + ] + reqs = torch.arange(batch, device="cuda", dtype=torch.int32) + padded = torch.cat((reqs, reqs.new_full((13 - batch,), 13))) + cu = torch.arange(batch + 1, device="cuda", dtype=torch.int32) * 4 + padded_cu = torch.cat((cu, cu.new_full((13 - batch,), batch * 4))) + q = torch.randn(batch * 4, 2, 128, device="cuda", dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn(batch * 4, 4, 128, device="cuda", dtype=torch.bfloat16) + a = torch.full((batch * 4, 4 * 128 if kda else 4), -3.0, device="cuda", dtype=torch.bfloat16) + beta = torch.randn(a.shape[0], 4, device="cuda", dtype=a.dtype) + alog = torch.zeros(4, device="cuda") + bias = torch.zeros(4 * 128 if kda else 4, device="cuda") + accepted = torch.zeros(14, device="cuda", dtype=torch.int32) + graphs, outputs = [], [] + for ssm_backend, ids, offsets in zip(ssm_backends, [reqs, padded], [cu, padded_cu]): + + def step(): + out = [ + replay_layer(ssm_backend, layer, q, k, v, a, beta, alog, bias, ids, offsets, run_config=config) + for layer in range(2) + ] + accept_replay_window(ssm_backend, ids, accepted, offsets) + return out + + step() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = step() + graphs.append(graph) + outputs.append(out) + ssm_backend.ssm_state.copy_(initial) + ssm_backend.cursors.zero_() + for iteration in range(32): + accepted[:batch] = (torch.arange(batch, device="cuda") + iteration) % 4 + for value in [q, k, v, beta]: + value.normal_() + overflow = [idx for idx in range(batch) if ssm_backends[0].cursors[idx] + 4 > ssm_backends[0].capacity] + if overflow: + overflow_reqs = torch.tensor(overflow, dtype=torch.int32, device="cuda") + for ssm_backend in ssm_backends: + ssm_backend.fold_accepted_prefix(overflow_reqs) + for graph in graphs: + graph.replay() + for small, padded_out in zip(outputs[0], outputs[1]): + torch.testing.assert_close(small.float(), padded_out.float(), atol=0.006, rtol=0.02) + assert torch.equal(ssm_backends[0].cursors, ssm_backends[1].cursors) + torch.testing.assert_close( + ssm_backends[0].ssm_state.float(), + ssm_backends[1].ssm_state.float(), + atol=0.002, + rtol=torch.finfo(dtype).eps, + ) + + +@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) +def test_full_prefix_does_not_address_past_capacity(kda): + torch.manual_seed(4) + capacity = 4 + backend = ReplaySSMExecutor( + torch.randn(1, 3, 2, 32, 32, device="cuda") * 0.01, + capacity, + 1, + torch.bfloat16, + kda=kda, + ) + reqs = torch.tensor([0, 1], device="cuda", dtype=torch.int32) + cu = torch.arange(reqs.numel() + 1, dtype=torch.int32, device=reqs.device) + q = torch.randn(2, 1, 32, device="cuda") * 0.2 + k, v = torch.randn_like(q), torch.randn_like(q) + a = torch.full((2, 32 if kda else 1), -3.0, device="cuda") + beta = torch.randn(2, 1, device="cuda") + alog, bias = torch.zeros(1, device="cuda"), torch.zeros(a.shape[-1], device="cuda") + backend.ssm_state.zero_() + backend.cursors.zero_() + for name in ("keys", "deltas", "gates", "raw_keys", "raw_values", "betas"): + getattr(backend, name).zero_() + backend.cursors[reqs.long()] = capacity + 2 + replay_layer( + backend, + 0, + q, + k, + v, + a, + beta, + alog, + bias, + reqs, + cu, + run_config={"BV": 32, "num_warps": 1, "num_stages": 1}, + ) + backend.fold_accepted_prefix(reqs) + torch.cuda.synchronize() + assert torch.equal(backend.cursors[reqs.long()], torch.zeros(2, device="cuda", dtype=torch.int32)) + assert torch.isfinite(backend.ssm_state).all() + for _ in range(12): + replay_layer( + backend, + 0, + q, + k, + v, + a, + beta, + alog, + bias, + reqs, + cu, + run_config={"BV": 32, "num_warps": 1, "num_stages": 1}, + ) + accept_replay_window(backend, reqs) + overflow = reqs[backend.cursors[reqs.long()] + backend.verify_width > backend.capacity] + if overflow.numel(): + backend.fold_accepted_prefix(overflow) + torch.cuda.synchronize() + assert int(backend.cursors[reqs.long()].max()) <= capacity + assert torch.isfinite(backend.ssm_state).all() diff --git a/unit_tests/common/basemodel/triton_kernel/mhc/test_mhc.py b/unit_tests/common/basemodel/triton_kernel/mhc/test_mhc.py new file mode 100644 index 0000000000..dd5cece4af --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/mhc/test_mhc.py @@ -0,0 +1,275 @@ +# SPDX-License-Identifier: Apache-2.0 + +import sys +from types import ModuleType + +import pytest +import torch +import torch.nn.functional as F +import triton + +from lightllm.common.basemodel.triton_kernel.mhc import ( + hc_contract, + hc_expand, + hc_head, + hc_post, + hc_post_pre_norm, + hc_pre_norm, +) +from lightllm.common.basemodel.triton_kernel.mhc import pre_norm +from lightllm.common.basemodel.triton_kernel.norm.rmsnorm import rmsnorm_forward + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +def _hc_pre_reference(x, fn, scale, base, streams, rms_eps, hc_eps, sinkhorn_iters, post_multiplier=2.0): + """Unfused FP32 reference; returns the merged input before sublayer RMSNorm.""" + assert x.ndim == 2 and x.shape[-1] % streams == 0 + tokens, flattened_hidden = x.shape + hidden = flattened_hidden // streams + residual = x.view(tokens, streams, hidden) + + x_fp32 = x.float() + inv_rms = torch.rsqrt(x_fp32.square().mean(dim=-1, keepdim=True) + rms_eps) + mixes = F.linear(x_fp32, fn) * inv_rms + + pre_raw = mixes[:, :streams] + post_raw = mixes[:, streams : 2 * streams] + residual_raw = mixes[:, 2 * streams :].view(tokens, streams, streams) + + pre = torch.sigmoid(pre_raw * scale[0] + base[:streams]) + hc_eps + post = post_multiplier * torch.sigmoid(post_raw * scale[1] + base[streams : 2 * streams]) + residual_mix = (residual_raw * scale[2] + base[2 * streams :].view(streams, streams)).softmax(dim=-1) + residual_mix = residual_mix + hc_eps + residual_mix = residual_mix / (residual_mix.sum(dim=-2, keepdim=True) + hc_eps) + for _ in range(sinkhorn_iters - 1): + residual_mix = residual_mix / (residual_mix.sum(dim=-1, keepdim=True) + hc_eps) + residual_mix = residual_mix / (residual_mix.sum(dim=-2, keepdim=True) + hc_eps) + + layer_input = (pre.unsqueeze(-1) * residual.float()).sum(dim=1).to(x.dtype) + return layer_input, residual_mix, post + + +def _hc_post_reference(layer_output, residual, residual_mix, post_mix, streams): + """Unfused FP32 reference; residual_mix maps input streams to output streams.""" + tokens, hidden = layer_output.shape + residual_3d = residual.view(tokens, streams, hidden) + mixed_residual = (residual_mix.unsqueeze(-1) * residual_3d.float().unsqueeze(2)).sum(dim=1) + out = post_mix.unsqueeze(-1) * layer_output.float().unsqueeze(1) + mixed_residual + return out.to(layer_output.dtype).reshape(tokens, streams * hidden) + + +@pytest.fixture(autouse=True) +def setup(): + torch.manual_seed(1525) + triton.set_allocator(lambda size, alignment, stream: torch.empty(size, device="cuda", dtype=torch.int8)) + + +@pytest.mark.parametrize("tokens", [3, 19, 3073]) +@pytest.mark.parametrize("disable_deepgemm", [False, True], ids=["deepgemm", "torch"]) +def test_mhc_matches_reference(tokens, disable_deepgemm, monkeypatch): + monkeypatch.setattr(pre_norm, "LIGHTLLM_DISABLE_DEEPGEMM_MHC", disable_deepgemm) + if disable_deepgemm: + monkeypatch.setitem(sys.modules, "deep_gemm", None) + streams, hidden = 4, 4096 + x = torch.randn(tokens, streams * hidden, device="cuda", dtype=torch.bfloat16) + fn = torch.randn(24, streams * hidden, device="cuda") * 0.005 + scale = torch.randn(3, device="cuda") + base = torch.randn(24, device="cuda") + norm = torch.randn(hidden, device="cuda", dtype=torch.bfloat16) + expected = _hc_pre_reference(x, fn, scale, base, streams, 1e-5, 1e-6, 20) + actual = hc_pre_norm(x, fn, scale, base, norm, streams, 1e-5, 1e-5, 1e-6, 20) + torch.testing.assert_close(actual[0], rmsnorm_forward(expected[0], norm, 1e-5), atol=0.04, rtol=0.03) + for a, b in zip(actual[1:], expected[1:]): + # DeepGEMM's mHC projection uses TF32 inputs and FP32 accumulation. + torch.testing.assert_close(a, b, atol=1e-4, rtol=1e-3) + layer_out = torch.randn(tokens, hidden, device="cuda", dtype=torch.bfloat16) + torch.testing.assert_close( + hc_post(layer_out, x, *actual[1:], streams), + _hc_post_reference(layer_out, x, *actual[1:], streams), + atol=0.04, + rtol=0.02, + ) + + +def test_empty_stream_expansion_and_contraction(): + embeddings = torch.empty(0, 4096, device="cuda", dtype=torch.bfloat16) + expanded = hc_expand(embeddings, 4) + assert expanded.shape == (0, 16384) + torch.testing.assert_close(hc_contract(expanded, 4), embeddings) + + +@pytest.mark.parametrize("missing", ["module", "kernel"]) +def test_mhc_reports_missing_deepgemm(monkeypatch, missing): + monkeypatch.setattr(pre_norm, "LIGHTLLM_DISABLE_DEEPGEMM_MHC", False) + monkeypatch.setitem(sys.modules, "deep_gemm", None if missing == "module" else ModuleType("deep_gemm")) + x = torch.empty(1, 16, dtype=torch.bfloat16) + fn = torch.empty(24, 16) + scale = torch.empty(3) + base = torch.empty(24) + norm = torch.empty(4, dtype=torch.bfloat16) + + with pytest.raises(ImportError, match="deep_gemm"): + hc_pre_norm(x, fn, scale, base, norm, 4, 1e-5, 1e-5, 1e-6, 20) + + +@pytest.mark.parametrize("tokens,hidden", [(1, 4096), (8, 4096), (16, 6144), (17, 4096), (129, 4096), (3073, 4096)]) +@pytest.mark.parametrize("disable_deepgemm", [False, True], ids=["fused", "torch"]) +def test_hc_post_pre_norm_matches_reference(tokens, hidden, disable_deepgemm, monkeypatch): + monkeypatch.setattr(pre_norm, "LIGHTLLM_DISABLE_DEEPGEMM_MHC", disable_deepgemm) + streams = 4 + residual = torch.randn(tokens, streams * hidden, device="cuda", dtype=torch.bfloat16) + layer_output = torch.randn(tokens, hidden, device="cuda", dtype=torch.bfloat16) + # Deliberately asymmetric mixing checks input/output stream orientation. + residual_mix = torch.randn(tokens, streams, streams, device="cuda").softmax(-1) + post_mix = torch.rand(tokens, streams, device="cuda") * 2 + fn = torch.randn(24, streams * hidden, device="cuda") * 0.005 + scale = torch.randn(3, device="cuda") + base = torch.randn(24, device="cuda") + norm = torch.randn(hidden, device="cuda", dtype=torch.bfloat16) + residual_before = residual.clone() + expected_residual = _hc_post_reference(layer_output, residual, residual_mix, post_mix, streams) + expected_pre = _hc_pre_reference(expected_residual, fn, scale, base, streams, 1e-6, 1e-6, 20) + actual = hc_post_pre_norm( + layer_output, residual, residual_mix, post_mix, fn, scale, base, norm, streams, 1e-6, 1e-5, 1e-6, 20 + ) + torch.testing.assert_close(residual, residual_before, atol=0, rtol=0) + torch.testing.assert_close(actual[1], expected_residual, atol=0.04, rtol=0.02) + torch.testing.assert_close(actual[0], rmsnorm_forward(expected_pre[0], norm, 1e-5), atol=0.04, rtol=0.03) + for a, b in zip(actual[2:], expected_pre[1:]): + torch.testing.assert_close(a, b, atol=2e-4, rtol=2e-3) + + +@pytest.mark.parametrize("tokens,hidden", [(1, 4096), (9, 6144), (129, 4096), (3073, 4096)]) +@pytest.mark.parametrize("zeros", [False, True]) +def test_hc_head_learned_gates_and_preallocated_output(tokens, hidden, zeros): + streams = 4 + residual = torch.randn(tokens, streams * hidden, device="cuda", dtype=torch.bfloat16) + if zeros: + residual.zero_() + fn = torch.randn(streams, streams * hidden, device="cuda") * 0.005 + scale = torch.randn(1, device="cuda") + base = torch.randn(streams, device="cuda") + inv_rms = torch.rsqrt(residual.float().square().mean(-1, keepdim=True) + 1e-5) + gates = torch.sigmoid(F.linear(residual.float(), fn) * inv_rms * scale + base) + 1e-6 + expected = (residual.view(tokens, streams, hidden).float() * gates.unsqueeze(-1)).sum(1).bfloat16() + output = torch.full((tokens, hidden), float("nan"), device="cuda", dtype=torch.bfloat16) + actual = hc_head(residual, fn, scale, base, streams, 1e-5, 1e-6, out=output) + assert actual is output + torch.testing.assert_close(actual, expected, atol=0.04, rtol=0.02) + + +@pytest.mark.parametrize("tokens", [1, 16, 17, 129]) +def test_shared_mhc_cuda_graph_replay(tokens): + streams, hidden = 4, 4096 + residual = torch.randn(tokens, streams * hidden, device="cuda", dtype=torch.bfloat16) + fn = torch.randn(24, streams * hidden, device="cuda") * 0.005 + scale = torch.randn(3, device="cuda") + base = torch.randn(24, device="cuda") + norm = torch.randn(hidden, device="cuda", dtype=torch.bfloat16) + layer_output = torch.randn(tokens, hidden, device="cuda", dtype=torch.bfloat16) + head_fn, head_scale, head_base = fn[:streams].contiguous(), scale[:1].contiguous(), base[:streams].contiguous() + output = torch.empty(tokens, hidden, device="cuda", dtype=torch.bfloat16) + + def run(): + x, mix, post = hc_pre_norm(residual, fn, scale, base, norm, streams, 1e-6, 1e-5, 1e-6, 20) + assert x.shape == (tokens, hidden) and post.shape == (tokens, streams) + x, updated, mix, post = hc_post_pre_norm( + layer_output, residual, mix, post, fn, scale, base, norm, streams, 1e-6, 1e-5, 1e-6, 20 + ) + final = hc_post(x, updated, mix, post, streams) + return hc_head(final, head_fn, head_scale, head_base, streams, 1e-6, 1e-6, out=output) + + def reference(): + _, mix, post = _hc_pre_reference(residual, fn, scale, base, streams, 1e-6, 1e-6, 20) + updated = _hc_post_reference(layer_output, residual, mix, post, streams) + x, mix, post = _hc_pre_reference(updated, fn, scale, base, streams, 1e-6, 1e-6, 20) + x = rmsnorm_forward(x, norm, 1e-5) + final = _hc_post_reference(x, updated, mix, post, streams) + inv_rms = torch.rsqrt(final.float().square().mean(-1, keepdim=True) + 1e-6) + gates = torch.sigmoid(F.linear(final.float(), head_fn) * inv_rms * head_scale + head_base) + 1e-6 + return (final.view(tokens, streams, hidden).float() * gates.unsqueeze(-1)).sum(1).bfloat16() + + run() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = run() + for _ in range(2): + residual.normal_() + layer_output.normal_() + graph.replay() + torch.testing.assert_close(actual, reference(), atol=0.06, rtol=0.04) + + +def test_mhc_reduction_reuses_compilation_within_split_buckets(): + hashes = {} + for tokens, splits in [ + (1, 1), + (16, 1), + (2, 3), + (16, 3), + (16, 8), + (3, 3), + (17, 7), + (31, 14), + (129, 26), + (257, 33), + (513, 64), + ]: + partial = torch.randn((splits, tokens, 24), device="cuda") * 0.05 + squares = torch.rand((splits, tokens), device="cuda") * (16384 / splits) + 1 + scale = torch.randn(3, device="cuda") + base = torch.randn(24, device="cuda") + x = torch.randn((tokens, 16384), device="cuda", dtype=torch.bfloat16) + norm = torch.randn(4096, device="cuda", dtype=torch.bfloat16) + + def run(gemm, sqrsum): + post = torch.empty((tokens, 4), device="cuda") + mix = torch.empty((tokens, 4, 4), device="cuda") + output = torch.empty((tokens, 4096), device="cuda", dtype=torch.bfloat16) + block_splits = triton.next_power_of_2(gemm.shape[0]) + compiled = pre_norm._hc_pre_norm_kernel[(tokens,)]( + x, + gemm, + sqrsum, + scale, + base, + post, + mix, + norm, + output, + hidden=4096, + x_stride_m=x.stride(0), + gemm_stride_s=gemm.stride(0), + gemm_stride_m=gemm.stride(1), + sqrsum_stride_s=sqrsum.stride(0), + sqrsum_stride_m=sqrsum.stride(1), + out_stride_m=output.stride(0), + RMS_EPS=1e-5, + NORM_EPS=1e-5, + STREAMS=4, + HC_EPS=1e-6, + POST_MULTIPLIER=2.0, + SINKHORN_ITERS=20, + N_SPLITS=gemm.shape[0], + BLOCK_SPLITS=block_splits, + BLOCK_H=4096, + num_warps=4, + ) + hashes.setdefault(block_splits, set()).add(compiled.hash) + return output, mix, post + + # The fused kernel reduces splits in parallel, so compare to sequential + # FP32 accumulation with tolerance for the different reduction order. + total = torch.zeros_like(partial[:1]) + total_squares = torch.zeros_like(squares[:1]) + for index in range(splits): + total.add_(partial[index : index + 1]) + total_squares.add_(squares[index : index + 1]) + actual, expected = run(partial, squares), run(total, total_squares) + torch.testing.assert_close(actual[0], expected[0], atol=0.04, rtol=0.03) + for output, reference in zip(actual[1:], expected[1:]): + torch.testing.assert_close(output, reference, atol=1e-6, rtol=1e-5) + assert all(len(bucket_hashes) == 1 for bucket_hashes in hashes.values()) diff --git a/unit_tests/common/fused_moe/test_activation_config.py b/unit_tests/common/fused_moe/test_activation_config.py new file mode 100644 index 0000000000..c59ab2712d --- /dev/null +++ b/unit_tests/common/fused_moe/test_activation_config.py @@ -0,0 +1,132 @@ +import dataclasses +import json + +import pytest +import torch + +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.fused_moe_weight import FusedMoeWeight +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.impl.marlin_impl import FuseMoeMarlin +from lightllm.common.quantization.no_quant import NoQuantization +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.utils.envs_utils import get_env_start_args + + +@pytest.fixture(autouse=True) +def runtime(monkeypatch): + monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(dataclasses.asdict(StartArgs()))) + for name, value in { + "GLOBAL_RANK": 0, + "GLOBAL_WORLD_SIZE": 1, + "DP_WORLD_SIZE": 1, + "CURRENT_RANK_IN_DP": 0, + "CURRENT_RANK_IN_NODE": 0, + "CURRENT_DEVICE_ID": 0, + }.items(): + monkeypatch.setenv("LIGHTLLM_" + name, str(value)) + get_env_start_args.cache_clear() + yield + get_env_start_args.cache_clear() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("activation", ["silu", "clamped_silu", "clamped_silu_add_one", "gelu"]) +def test_call_parameters_reach_expert_activation(monkeypatch, activation): + monkeypatch.setattr( + "lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul.ffn_use_tanh_approximate_gelu", + lambda: activation == "gelu", + ) + dim, count = 128, 4 + config = {"norm_topk_prob": True, "num_experts_per_tok": 2, "scoring_func": "softmax"} + kwargs = { + "clamped_silu": {"alpha": 1.0, "limit": 10.0, "clamp_up_add_one": False}, + "clamped_silu_add_one": {"alpha": 1.702, "limit": 7.0}, + }.get(activation, {}) + weight = FusedMoeWeight( + gate_proj_name="gate", + up_proj_name="up", + down_proj_name="down", + e_score_correction_bias_name="", + weight_prefix="experts", + n_routed_experts=count, + hidden_size=dim, + moe_intermediate_size=dim, + data_type=torch.bfloat16, + quant_method=NoQuantization(), + network_config=config, + ) + eye = torch.eye(dim, device="cuda", dtype=torch.bfloat16) + weights = {} + for i in range(count): + weights[f"experts.{i}.gate.weight"] = eye * (1 + i / 4) + weights[f"experts.{i}.up.weight"] = eye * (2 + i / 8) + weights[f"experts.{i}.down.weight"] = eye + weight.load_hf_weights(weights) + assert weight.verify_load() + x = torch.linspace(-25, 25, 3 * dim, device="cuda", dtype=torch.bfloat16).view(3, dim) + router = torch.tensor([[1, 3, 2, 0], [4, 3, 1, 2], [1, 2, 3, 4]], device="cuda", dtype=torch.float32) + top = router.topk(2, dim=-1) + probs = top.values.softmax(-1) + expected = torch.zeros_like(x, dtype=torch.float32) + for row in range(x.shape[0]): + for choice in range(2): + i = int(top.indices[row, choice]) + gate = (x[row] * (1 + i / 4)).float() + up = (x[row] * (2 + i / 8)).float() + if kwargs: + gate = gate.clamp(max=kwargs["limit"]) + up = up.clamp(-kwargs["limit"], kwargs["limit"]) + gate = gate * torch.sigmoid(kwargs["alpha"] * gate) + if kwargs.get("clamp_up_add_one", True): + up += 1 + elif activation == "gelu": + gate = torch.nn.functional.gelu(gate, approximate="tanh") + else: + gate = torch.nn.functional.silu(gate) + expert_out = (gate.bfloat16() * up.bfloat16()).bfloat16() + expected[row] += (expert_out.float() * probs[row, choice]).bfloat16().float() + if kwargs: + default_output = weight.experts(x.clone(), router, 2, True, False, 0, 0, False) + actual = weight.experts(x.clone(), router, 2, True, False, 0, 0, False, **kwargs) + actual_with_topk = weight.experts_with_topk( + x.clone(), + topk_weights=probs, + topk_ids=top.indices, + is_prefill=False, + **kwargs, + ) + torch.testing.assert_close(actual, expected.bfloat16(), atol=0.125, rtol=0.01) + torch.testing.assert_close(actual_with_topk, expected.bfloat16(), atol=0.125, rtol=0.01) + if kwargs: + # A clamped call must not change subsequent calls on the same weight. + actual_default = weight.experts(x.clone(), router, 2, True, False, 0, 0, False) + torch.testing.assert_close(actual_default, default_output, atol=0, rtol=0) + + +def test_marlin_rejects_clamp_at_call(monkeypatch): + backend = FuseMoeMarlin + monkeypatch.setattr(backend, "create_workspace", lambda self: None) + monkeypatch.setattr(backend, "_select_experts", lambda *args, **kwargs: (None, None)) + impl = backend( + n_routed_experts=4, + num_fused_shared_experts=0, + routed_scaling_factor=1.0, + quant_method=None, + ) + with pytest.raises(NotImplementedError, match="does not support clamped SwiGLU"): + impl( + input_tensor=None, + router_logits=None, + w13=None, + w2=None, + correction_bias=None, + scoring_func="softmax", + top_k=2, + renormalize=True, + use_grouped_topk=False, + topk_group=0, + num_expert_group=0, + is_prefill=False, + alpha=1.0, + limit=10.0, + clamp_up_add_one=False, + ) diff --git a/unit_tests/common/fused_moe/test_append_shared_expert_topk.py b/unit_tests/common/fused_moe/test_append_shared_expert_topk.py new file mode 100644 index 0000000000..4a197e911f --- /dev/null +++ b/unit_tests/common/fused_moe/test_append_shared_expert_topk.py @@ -0,0 +1,222 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.fused_moe_weight import FusedMoeWeight +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.impl.triton_impl import FuseMoeTriton +from lightllm.common.basemodel.triton_kernel.fused_moe.append_shared_expert_topk import ( + append_fused_shared_experts, +) + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for Triton kernels") + + +def test_experts_with_topk_appends_the_fused_shared_expert(monkeypatch): + monkeypatch.setattr( + "lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.fused_moe_weight.get_moe_capture_callback", + lambda *args, **kwargs: None, + ) + seen = {} + + def run(**kwargs): + seen.update(kwargs) + return kwargs["input_tensor"] + + impl = SimpleNamespace( + num_fused_shared_experts=1, + n_routed_experts=4, + _fused_experts=run, + ) + impl._prepare_expert_execution = lambda **kwargs: FuseMoeTriton._prepare_expert_execution(impl, **kwargs) + weight = SimpleNamespace( + layer_num_=0, + w13=torch.ones(1, device="cuda"), + w2=torch.ones(1, device="cuda"), + fuse_moe_impl=impl, + ) + topk_ids = torch.tensor([[0, 2], [1, 3]], dtype=torch.int64, device="cuda") + topk_weights = torch.tensor([[0.2, 0.8], [0.4, 0.6]], device="cuda") + FusedMoeWeight.experts_with_topk( + weight, + torch.ones(2, 4, device="cuda"), + topk_weights, + topk_ids, + is_prefill=False, + ) + assert torch.equal(seen["topk_ids"], torch.tensor([[0, 2, 4], [1, 3, 4]], device="cuda")) + assert torch.allclose(seen["topk_weights"], torch.tensor([[0.2, 0.8, 1.0], [0.4, 0.6, 1.0]], device="cuda")) + + +def test_capture_observes_logical_ids_before_execution_mapping(): + logical_ids = torch.tensor([[0, 2]], dtype=torch.int64, device="cuda") + logical_weights = torch.tensor([[0.2, 0.8]], device="cuda") + captured = {} + executed = {} + impl = FuseMoeTriton.__new__(FuseMoeTriton) + impl.n_routed_experts = 4 + impl.num_fused_shared_experts = 1 + impl._select_experts = lambda **kwargs: (logical_weights, logical_ids) + + def fused(**kwargs): + executed.update(kwargs) + return kwargs["input_tensor"] + + impl._fused_experts = fused + hidden = torch.ones(1, 4, device="cuda") + impl( + input_tensor=hidden, + router_logits=hidden, + w13=None, + w2=None, + correction_bias=None, + scoring_func="softmax", + top_k=2, + renormalize=False, + use_grouped_topk=False, + topk_group=0, + num_expert_group=0, + is_prefill=False, + moe_capture_callback=lambda ids: captured.setdefault("ids", ids), + ) + + assert captured["ids"] is logical_ids + assert torch.equal(logical_ids, torch.tensor([[0, 2]], device="cuda")) + assert torch.equal(executed["topk_ids"], torch.tensor([[0, 2, 4]], device="cuda")) + + +def test_redundancy_repair_does_not_overwrite_logical_ids(monkeypatch): + from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.impl.deepgemm_impl import ( + FuseMoeDeepGEMM, + ) + + logical_ids = torch.tensor([[1, 3]], dtype=torch.int64, device="cuda") + seen = {} + + def repair(topk_ids, **kwargs): + seen["ptr"] = topk_ids.data_ptr() + topk_ids.fill_(-1) + + monkeypatch.setattr( + "lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.impl.deepgemm_impl.redundancy_topk_ids_repair", + repair, + ) + impl = FuseMoeDeepGEMM.__new__(FuseMoeDeepGEMM) + impl.redundancy_expert_num = 1 + impl.redundancy_expert_ids_tensor = None + impl.ep_n_routed_experts = 2 + impl.global_rank_ = 0 + impl.routed_expert_counter_tensor = None + impl.auto_update_redundancy_expert = False + impl.num_fused_shared_experts = 0 + impl.n_routed_experts = 4 + _, execution_ids = FuseMoeDeepGEMM._prepare_expert_execution( + impl, + torch.ones(1, 2, device="cuda"), + logical_ids, + is_prefill=False, + ) + assert seen["ptr"] != logical_ids.data_ptr() + assert torch.equal(logical_ids, torch.tensor([[1, 3]], device="cuda")) + assert torch.equal(execution_ids, torch.full_like(logical_ids, -1)) + + +def test_append_fused_shared_experts_without_gate(): + topk_ids = torch.tensor([[0, 2], [1, 3], [2, 0]], dtype=torch.int32, device="cuda") + topk_weights = torch.tensor([[0.2, 0.8], [0.4, 0.6], [0.7, 0.3]], dtype=torch.float32, device="cuda") + + out_weights, out_ids = append_fused_shared_experts( + topk_weights=topk_weights, + topk_ids=topk_ids, + shared_expert_start_id=4, + num_fused_shared_experts=1, + shared_expert_gate=None, + ) + + expect_ids = torch.tensor([[0, 2, 4], [1, 3, 4], [2, 0, 4]], dtype=torch.int32, device="cuda") + expect_weights = torch.tensor( + [[0.2, 0.8, 1.0], [0.4, 0.6, 1.0], [0.7, 0.3, 1.0]], dtype=torch.float32, device="cuda" + ) + assert torch.equal(out_ids, expect_ids) + assert torch.allclose(out_weights, expect_weights) + + +def test_append_fused_shared_experts_with_gate(): + topk_ids = torch.tensor([[0, 2], [1, 3], [2, 0]], dtype=torch.int32, device="cuda") + topk_weights = torch.tensor([[0.2, 0.8], [0.4, 0.6], [0.7, 0.3]], dtype=torch.float32, device="cuda") + shared_expert_gate = torch.tensor([[0.0], [2.0], [-2.0]], dtype=torch.float32, device="cuda") + + out_weights, out_ids = append_fused_shared_experts( + topk_weights=topk_weights, + topk_ids=topk_ids, + shared_expert_start_id=4, + num_fused_shared_experts=1, + shared_expert_gate=shared_expert_gate, + ) + + expect_ids = torch.tensor([[0, 2, 4], [1, 3, 4], [2, 0, 4]], dtype=torch.int32, device="cuda") + expect_weights = torch.cat([topk_weights, torch.sigmoid(shared_expert_gate)], dim=1) + assert torch.equal(out_ids, expect_ids) + assert torch.allclose(out_weights, expect_weights) + + +def test_append_fused_shared_experts_multiple_tokens_per_grid(): + token_num = 4097 + topk_ids = torch.stack( + [ + torch.arange(token_num, dtype=torch.int32, device="cuda") % 4, + (torch.arange(token_num, dtype=torch.int32, device="cuda") + 1) % 4, + ], + dim=1, + ) + topk_weights = torch.rand((token_num, 2), dtype=torch.float32, device="cuda") + shared_expert_gate = torch.randn((token_num, 1), dtype=torch.float32, device="cuda") + + out_weights, out_ids = append_fused_shared_experts( + topk_weights=topk_weights, + topk_ids=topk_ids, + shared_expert_start_id=4, + num_fused_shared_experts=1, + shared_expert_gate=shared_expert_gate, + ) + + expect_ids = torch.cat( + [topk_ids, torch.full((token_num, 1), 4, dtype=torch.int32, device="cuda")], + dim=1, + ) + expect_weights = torch.cat([topk_weights, torch.sigmoid(shared_expert_gate)], dim=1) + assert torch.equal(out_ids, expect_ids) + assert torch.allclose(out_weights, expect_weights) + + +def test_append_fused_shared_experts_multi_shared_gate(): + token_num = 7 + topk_ids = torch.stack( + [ + torch.arange(token_num, dtype=torch.int32, device="cuda") % 4, + (torch.arange(token_num, dtype=torch.int32, device="cuda") + 1) % 4, + ], + dim=1, + ) + topk_weights = torch.rand((token_num, 2), dtype=torch.float32, device="cuda") + shared_expert_gate = torch.randn((token_num, 2), dtype=torch.float32, device="cuda") + + out_weights, out_ids = append_fused_shared_experts( + topk_weights=topk_weights, + topk_ids=topk_ids, + shared_expert_start_id=4, + num_fused_shared_experts=2, + shared_expert_gate=shared_expert_gate, + ) + + expect_ids = torch.cat( + [ + topk_ids, + torch.tensor([[4, 5]], dtype=torch.int32, device="cuda").repeat(token_num, 1), + ], + dim=1, + ) + expect_weights = torch.cat([topk_weights, torch.sigmoid(shared_expert_gate)], dim=1) + assert torch.equal(out_ids, expect_ids) + assert torch.allclose(out_weights, expect_weights) diff --git a/unit_tests/common/fused_moe/test_deepep.py b/unit_tests/common/fused_moe/test_deepep.py new file mode 100644 index 0000000000..ecc92af1c0 --- /dev/null +++ b/unit_tests/common/fused_moe/test_deepep.py @@ -0,0 +1,311 @@ +import pytest + +pytest.skip(reason="need special env, install deep_ep and deep_gemm", allow_module_level=True) + +import os +import torch +import torch.distributed as dist +import deep_ep +import numpy as np +from lightllm.common.fused_moe.grouped_fused_moe_ep import fused_experts_impl +from typing import Tuple +from lightllm.utils.log_utils import init_logger + +logger = init_logger(__name__) + +seed = 42 +torch.manual_seed(seed) + +if torch.cuda.is_available(): + torch.cuda.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + + +def per_block_cast_to_fp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + assert x.dim() == 2 + m, n = x.shape + from deep_gemm import ceil_div + + x_padded = torch.zeros((ceil_div(m, 128) * 128, ceil_div(n, 128) * 128), dtype=x.dtype, device=x.device) + x_padded[:m, :n] = x + x_view = x_padded.view(-1, 128, x_padded.size(1) // 128, 128) + x_amax = x_view.abs().float().amax(dim=(1, 3), keepdim=True).clamp(1e-4) + x_scaled = (x_view * (448.0 / x_amax)).to(torch.float8_e4m3fn) + return x_scaled.view_as(x_padded)[:m, :n].contiguous(), (x_amax / 448.0).view(x_view.size(0), x_view.size(2)) + + +def init_dist(local_rank: int, num_local_ranks: int): + # NOTES: you may rewrite this function with your own cluster settings + ip = os.getenv("MASTER_ADDR", "127.0.0.1") + port = int(os.getenv("MASTER_PORT", "8361")) + num_nodes = int(os.getenv("WORLD_SIZE", 1)) + node_rank = int(os.getenv("RANK", 0)) + assert (num_local_ranks < 8 and num_nodes == 1) or num_local_ranks == 8 + + dist.init_process_group( + backend="nccl", + init_method=f"tcp://{ip}:{port}", + world_size=num_nodes * num_local_ranks, + rank=node_rank * num_local_ranks + local_rank, + ) + torch.set_default_dtype(torch.bfloat16) + torch.set_default_device("cuda") + torch.cuda.set_device(local_rank) + + return dist.get_rank(), dist.get_world_size(), dist.new_group(list(range(num_local_ranks * num_nodes))) + + +def fused_experts_impl_ref( + x: torch.Tensor, # [M, K] + w1: torch.Tensor, # [group, N, K] + w2: torch.Tensor, # [group, K, N/2] + topk_weight: torch.Tensor, # [M, topk] + topk_ids: torch.Tensor, # [M, topk] + num_experts: int, +): + N = w1.shape[1] + ep_size = torch.distributed.get_world_size() + experts_per_rank = num_experts // ep_size + + cnts = topk_ids.new_zeros((topk_ids.shape[0], num_experts)) + cnts.scatter_(1, topk_ids, 1) + tokens_per_expert = cnts.sum(dim=0) + idxs = topk_ids.view(-1).argsort() + sorted_tokens = x[idxs // topk_ids.shape[1]] + sorted_tokens_shape = sorted_tokens.shape + + if ep_size > 1: + tokens_per_ep_rank = tokens_per_expert.view(ep_size, -1).sum(dim=1) + tokens_per_expert_group = tokens_per_expert.new_empty(tokens_per_expert.shape[0]) + dist.all_to_all_single(tokens_per_expert_group, tokens_per_expert) + output_splits = tokens_per_expert_group.view(ep_size, -1).sum(1).cpu().numpy().tolist() + gathered_tokens = sorted_tokens.new_empty( + tokens_per_expert_group.sum(dim=0).cpu().item(), sorted_tokens.shape[1] + ) + input_split_sizes = tokens_per_ep_rank.cpu().numpy().tolist() + dist.all_to_all( + list(gathered_tokens.split(output_splits)), + list(sorted_tokens.split(input_split_sizes)), + ) + tokens_per_expert_post_gather = tokens_per_expert_group.view(ep_size, experts_per_rank).sum(dim=0) + gatherd_idxs = np.zeros(shape=(gathered_tokens.shape[0],), dtype=np.int32) + s = 0 + for i, k in enumerate(tokens_per_expert_group.cpu().numpy()): + gatherd_idxs[s : s + k] = i % experts_per_rank + s += k + gatherd_idxs = gatherd_idxs.argsort() + sorted_tokens = gathered_tokens[gatherd_idxs] + tokens_per_expert = tokens_per_expert_post_gather + tokens_per_expert = tokens_per_expert.cpu().numpy() + + outputs = [] + start_idx = 0 + for i, num_tokens in enumerate(tokens_per_expert): + end_idx = start_idx + num_tokens + if num_tokens == 0: + continue + tokens_for_this_expert = sorted_tokens[start_idx:end_idx] + w1_out = torch.matmul(tokens_for_this_expert, w1[i, : N // 2, :].T) + w2_out = torch.matmul(tokens_for_this_expert, w1[i, N // 2 :, :].T) + tmp = torch.nn.functional.silu(w1_out) + tmp = tmp * w2_out + expert_out = torch.matmul(tmp, w2[i].T) + outputs.append(expert_out) + start_idx = end_idx + + outs = torch.cat(outputs, dim=0) if len(outputs) else sorted_tokens.new_empty(0) + + if ep_size > 1: + new_x = torch.empty_like(outs) + new_x[gatherd_idxs] = outs + gathered_tokens = new_x.new_empty(*sorted_tokens_shape) + dist.all_to_all( + list(gathered_tokens.split(input_split_sizes)), + list(new_x.split(output_splits)), + ) + outs = gathered_tokens + + new_x = torch.empty_like(outs) + new_x[idxs] = outs + final_out = ( + new_x.view(*topk_ids.shape, -1) + .type(topk_weight.dtype) + .mul_(topk_weight.unsqueeze(dim=-1)) + .sum(dim=1) + .type(new_x.dtype) + ) + + return final_out + + +def case1(local_rank: int, num_local_ranks: int): + # Init dist + rank, num_ranks, group = init_dist(local_rank, num_local_ranks) + + torch.manual_seed(rank) + + # Construct inputs + seqlen = 16 + hidden_states = torch.randn((seqlen, 7168), device="cuda", dtype=torch.bfloat16) + w1 = torch.randn((256 // num_local_ranks, 4096, 7168), device="cuda", dtype=torch.bfloat16) + w2 = torch.randn((256 // num_local_ranks, 7168, 2048), device="cuda", dtype=torch.bfloat16) + + w1_fp8 = torch.empty_like(w1, dtype=torch.float8_e4m3fn) + w2_fp8 = torch.empty_like(w2, dtype=torch.float8_e4m3fn) + w1_scale = torch.empty((256 // num_local_ranks, 4096 // 128, 7168 // 128), device="cuda", dtype=torch.float) + w2_scale = torch.empty((256 // num_local_ranks, 7168 // 128, 2048 // 128), device="cuda", dtype=torch.float) + + for i in range(256 // num_local_ranks): + w1_fp8[i], w1_scale[i] = per_block_cast_to_fp8(w1[i]) + w2_fp8[i], w2_scale[i] = per_block_cast_to_fp8(w2[i]) + + topk_weights = torch.randn((seqlen, 8), device="cuda", dtype=torch.float32) + topk_weights = torch.softmax(topk_weights, dim=-1) # 对每行进行softmax归一化 + topk_weights = torch.tensor(topk_weights, device="cuda", dtype=torch.float32) + topk_ids = torch.zeros((seqlen, 8), device="cuda", dtype=torch.int64) + for i in range(seqlen): + topk_ids[i] = torch.randperm(254, device="cuda")[:8] + 1 + + # Init buffer + test_ll_compatibility, num_rdma_bytes = True, 0 + num_max_dispatch_tokens_per_rank = 512 + if test_ll_compatibility: + ll_num_tokens, ll_hidden, ll_num_experts, _ = num_max_dispatch_tokens_per_rank, 7168, 256, 8 + num_rdma_bytes = deep_ep.Buffer.get_low_latency_rdma_size_hint( + ll_num_tokens, ll_hidden, num_ranks, ll_num_experts + ) + + buffer = deep_ep.Buffer( + group, + int(1e9), + num_rdma_bytes, + low_latency_mode=test_ll_compatibility, + num_qps_per_rank=(ll_num_experts // num_ranks if test_ll_compatibility else 1), + ) + + # Test normal + ref_output = fused_experts_impl_ref( + x=hidden_states, w1=w1, w2=w2, topk_weight=topk_weights, topk_ids=topk_ids, num_experts=256 + ) + + output = fused_experts_impl( + hidden_states=hidden_states, + w1=w1_fp8, + w2=w2_fp8, + topk_weights=topk_weights, + topk_idx=topk_ids, + num_experts=256, + buffer=buffer, + is_prefill=True, + use_fp8_w8a8=True, + use_fp8_all2all=True, + use_int8_w8a16=False, + w1_scale=w1_scale, + w2_scale=w2_scale, + previous_event=None, + ) + + # Test ll + if test_ll_compatibility: + buffer.clean_low_latency_buffer(ll_num_tokens, ll_hidden, ll_num_experts) + ll_output = fused_experts_impl( + hidden_states=hidden_states, + w1=w1_fp8, + w2=w2_fp8, + topk_weights=topk_weights, + topk_idx=topk_ids, + num_experts=256, + buffer=buffer, + is_prefill=False, + use_fp8_w8a8=True, + use_fp8_all2all=True, + use_int8_w8a16=False, + w1_scale=w1_scale, + w2_scale=w2_scale, + previous_event=None, + ) + + # Check + norm_sim = torch.nn.functional.cosine_similarity(ref_output, output).mean() + ll_sim = torch.nn.functional.cosine_similarity(ref_output, ll_output).mean() + + assert torch.allclose(norm_sim, torch.ones(1), atol=1e-2, rtol=0) + assert torch.allclose(ll_sim, torch.ones(1), atol=1e-2, rtol=0) + logger.info(f"deepep cosine {norm_sim}") + logger.info(f"deepep ll cosine {ll_sim}") + + # Profile + profile = os.getenv("PROFILE", False) + if profile: + with torch.profiler.profile( + activities=[ + torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA, + ], + record_shapes=True, + profile_memory=True, + with_stack=True, + ) as prof: + for i in range(10): + output = fused_experts_impl( + hidden_states=hidden_states, + w1=w1_fp8, + w2=w2_fp8, + topk_weights=topk_weights, + topk_idx=topk_ids, + num_experts=256, + buffer=buffer, + is_prefill=True, + use_fp8_w8a8=True, + use_fp8_all2all=True, + use_int8_w8a16=False, + w1_scale=w1_scale, + w2_scale=w2_scale, + previous_event=None, + ) + # prof.step() + if rank == 0: + print(prof.key_averages().table(sort_by="cuda_time_total")) + prof.export_chrome_trace("normal_trace.json") + + if test_ll_compatibility: + with torch.profiler.profile( + activities=[ + torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA, + ], + record_shapes=True, + profile_memory=True, + with_stack=True, + ) as prof: + for i in range(10): + ll_output = fused_experts_impl( + hidden_states=hidden_states, + w1=w1_fp8, + w2=w2_fp8, + topk_weights=topk_weights, + topk_idx=topk_ids, + num_experts=256, + buffer=buffer, + is_prefill=False, + use_fp8_w8a8=True, + use_fp8_all2all=True, + use_int8_w8a16=False, + w1_scale=w1_scale, + w2_scale=w2_scale, + previous_event=None, + ) + if rank == 0: + print(prof.key_averages().table(sort_by="cuda_time_total")) + prof.export_chrome_trace("ll_trace.json") + + dist.barrier() + + +def test_end2end(): + num_processes = 8 + torch.multiprocessing.spawn(case1, args=(num_processes,), nprocs=num_processes) + + +if __name__ == "__main__": + pytest.main() diff --git a/unit_tests/common/fused_moe/test_deepgemm_activation.py b/unit_tests/common/fused_moe/test_deepgemm_activation.py new file mode 100644 index 0000000000..528d9a1ade --- /dev/null +++ b/unit_tests/common/fused_moe/test_deepgemm_activation.py @@ -0,0 +1,96 @@ +import pytest +import torch + +from lightllm.common.basemodel.triton_kernel.fused_moe import grouped_fused_moe_ep as ep +from lightllm.common.quantization import Quantcfg + + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or not ep.HAS_DEEPGEMM or torch.cuda.get_device_capability()[0] < 9, + reason="SM90+, DeepEP and DeepGEMM required", +) + + +@pytest.mark.parametrize("is_prefill", [False, True]) +@pytest.mark.parametrize("use_ue8m0_scales", [False, True]) +@pytest.mark.parametrize("alpha,limit,add_one", [(None, None, True), (1.0, 10.0, False), (1.702, 7.0, True)]) +def test_grouped_gemm_activation_matches_reference(monkeypatch, is_prefill, use_ue8m0_scales, alpha, limit, add_one): + monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", "0") + monkeypatch.delenv("LIGHTLLM_USE_UE8M0_SCALES") + config = {"quant_method": "fp8", "weight_block_size": [128, 128]} + if use_ue8m0_scales: + config["scale_fmt"] = "ue8m0" + Quantcfg({"n_layer": 1, "quantization_config": config}) + # Identity projections isolate activation/FP8 behavior from GEMM rounding. + dim, padded = 128, 128 + x = torch.linspace(-24, 24, 5 * dim, device="cuda").view(5, dim).to(torch.float8_e4m3fn) + eye = torch.eye(dim, device="cuda", dtype=torch.bfloat16) + w1 = torch.cat([2 * eye, 3 * eye]).unsqueeze(0).repeat(2, 1, 1).to(torch.float8_e4m3fn) + w2 = eye.unsqueeze(0).repeat(2, 1, 1).to(torch.float8_e4m3fn) + w1_scale = torch.ones(2, 2, 1, device="cuda") + w2_scale = torch.ones(2, 1, 1, device="cuda") + counts = torch.tensor([3, 2], device="cuda", dtype=torch.int32) + activation_args = dict(alpha=alpha, limit=limit, clamp_up_add_one=add_one) + + if is_prefill: + recv = torch.zeros(2, padded, dim, device="cuda", dtype=torch.float8_e4m3fn) + recv[0, :3] = x[:3] + recv[1, :2] = x[3:] + metadata = torch.zeros(5, 3, device="cuda", dtype=torch.int32) + metadata[:, 2] = torch.tensor([0, 1, 2, padded, padded + 1], device="cuda") + actual = ep.chunked_expanded_moe_forward( + num_recv_tokens_per_expert_list=[padded, padded], + num_unaligned_recv_tokens_per_expert=counts, + recv_x=(recv.view(2 * padded, dim), torch.ones(1, 2 * padded, device="cuda").T), + recv_topk_weights=torch.ones(2 * padded, device="cuda"), + recv_src_metadata=metadata, + w1=w1, + w1_scale=w1_scale, + w2=w2, + w2_scale=w2_scale, + block_size_k=128, + hidden_dtype=torch.bfloat16, + **activation_args, + ) + else: + alignment = ep.get_mk_alignment_for_contiguous_layout() + starts = [] + psum_values = [] + cursor = 0 + for count in counts.tolist(): + if starts: + cursor = (cursor + alignment - 1) // alignment * alignment + starts.append(cursor) + cursor += count + psum_values.append(cursor) + psum = torch.tensor(psum_values, device="cuda", dtype=torch.int32) + max_rows = (int(psum[-1]) + alignment - 1) // alignment * alignment + recv = torch.zeros(max_rows, dim, device="cuda", dtype=torch.float8_e4m3fn) + recv[starts[0] : starts[0] + 3] = x[:3] + recv[starts[1] : starts[1] + 2] = x[3:] + output = ep.decode_masked_group_gemm( + (recv, torch.ones(max_rows, 1, device="cuda")), + psum, + alignment, + torch.bfloat16, + w1, + w1_scale, + w2, + w2_scale, + expected_m=3, + **activation_args, + ) + actual = torch.cat([output[starts[0] : starts[0] + 3], output[starts[1] : starts[1] + 2]]) + + gate = (x.float() * 2).bfloat16().float() + up = (x.float() * 3).bfloat16().float() + if limit is not None: + gate = gate.clamp(max=limit) + up = up.clamp(-limit, limit) + int(add_one) + gate = (gate * torch.sigmoid(gate * (alpha or 1.0))).bfloat16().float() + activation = (gate * up).bfloat16().float() + scales = activation.abs().amax(-1, keepdim=True).clamp(min=1e-10) / 448 + if use_ue8m0_scales: + scales = torch.exp2(torch.ceil(torch.log2(scales))) + expected = ((activation / scales).to(torch.float8_e4m3fn).float() * scales).bfloat16() + torch.testing.assert_close(actual, expected, atol=0.05, rtol=0.02) diff --git a/unit_tests/common/fused_moe/test_ep_reduce_decode_output.py b/unit_tests/common/fused_moe/test_ep_reduce_decode_output.py new file mode 100644 index 0000000000..8b2805e3a3 --- /dev/null +++ b/unit_tests/common/fused_moe/test_ep_reduce_decode_output.py @@ -0,0 +1,36 @@ +import pytest +import torch + +from lightllm.common.basemodel.triton_kernel.fused_moe.deepep_expanded_layout_kernels import ep_reduce_decode_output + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +def test_ep_reduce_decode_output_weights_local_rows_and_compacts_metadata(): + hidden = 8 + expert_output = torch.tensor( + [[1.0] * hidden, [2.0] * hidden, [4.0] * hidden], + device="cuda", + dtype=torch.bfloat16, + ) + route_weights = torch.tensor([0.5, 0.25, 0.0], device="cuda", dtype=torch.float32) + recv_src_metadata = torch.tensor( + [[7, 9, 0, 1], [3, 4, -1, -1]], + device="cuda", + dtype=torch.int32, + ) + num_valid_recv_tokens = torch.tensor([1], device="cuda", dtype=torch.int32) + + dense_output, compact_metadata = ep_reduce_decode_output( + expert_output=expert_output, + route_weights=route_weights, + recv_src_metadata=recv_src_metadata, + num_valid_recv_tokens=num_valid_recv_tokens, + ) + + expected = expert_output[0].float() * 0.5 + expert_output[1].float() * 0.25 + torch.testing.assert_close(dense_output[0].float(), expected) + assert compact_metadata[0, :2].tolist() == [7, 9] + assert compact_metadata[0, 2].item() == 0 + assert compact_metadata[0, 3].item() == -1 diff --git a/unit_tests/common/fused_moe/test_grouped_fused_moe.py b/unit_tests/common/fused_moe/test_grouped_fused_moe.py new file mode 100644 index 0000000000..8bfa703886 --- /dev/null +++ b/unit_tests/common/fused_moe/test_grouped_fused_moe.py @@ -0,0 +1,275 @@ +import torch +import time +import pytest +import triton +from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe import ( + _moe_align_fused_atomic_token, + fused_experts_impl, + moe_align_fused, + moe_align2, + grouped_matmul, +) +from lightllm.utils.log_utils import init_logger + +logger = init_logger(__name__) + +seed = 42 +torch.manual_seed(seed) + +if torch.cuda.is_available(): + torch.cuda.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + + +def _check_moe_align_fused(topk_ids, topk_weights, expert_num, ordered=True): + expert_to_token_index = torch.empty((expert_num, topk_ids.numel()), dtype=torch.int32, device="cuda") + expert_to_weight = torch.empty((expert_num, topk_ids.numel()), dtype=torch.float32, device="cuda") + expert_token_num = torch.empty((expert_num,), dtype=torch.int32, device="cuda") + + moe_align_fused( + expert_to_token_index, + expert_to_weight, + expert_token_num, + topk_ids, + topk_weights, + ) + torch.cuda.synchronize() + + flat_topk_ids = topk_ids.flatten() + flat_topk_weights = topk_weights.flatten() + expected_token_num = torch.bincount(flat_topk_ids, minlength=expert_num).to(torch.int32) + assert torch.equal(expert_token_num, expected_token_num) + + for expert_id, token_num in enumerate(expected_token_num.tolist()): + expected_index = torch.nonzero(flat_topk_ids == expert_id, as_tuple=False).flatten() + expected_weight = flat_topk_weights[expected_index] + expected_index = expected_index.to(torch.int32) + token_index = expert_to_token_index[expert_id, :token_num] + token_weight = expert_to_weight[expert_id, :token_num] + + if not ordered: + order = torch.argsort(token_index) + token_index = token_index[order] + token_weight = token_weight[order] + + assert torch.equal(token_index, expected_index) + assert torch.allclose(token_weight, expected_weight) + + +def test_moe_align_fused_small_token(): + expert_num = 5 + small_topk_ids = torch.tensor([[0, 1, 2], [0, 3, 1], [3, 1, 4]], dtype=torch.int32, device="cuda") + small_topk_weights = torch.tensor( + [[0.3, 0.7, 0.1], [0.2, 0.8, 0.4], [0.5, 0.6, 0.9]], dtype=torch.float32, device="cuda" + ) + _check_moe_align_fused(small_topk_ids, small_topk_weights, expert_num) + + small_many_topk_ids = torch.arange(128 * 17, dtype=torch.int32, device="cuda").reshape(128, 17) % expert_num + small_many_topk_weights = torch.arange(small_many_topk_ids.numel(), dtype=torch.float32, device="cuda").reshape( + 128, 17 + ) + _check_moe_align_fused(small_many_topk_ids, small_many_topk_weights, expert_num) + + +def test_moe_align_fused_large_token(): + expert_num = 5 + + base_topk_ids = torch.tensor([[0, 1, 2], [0, 3, 1], [3, 1, 4], [2, 0, 4]], dtype=torch.int32, device="cuda") + large_topk_ids = base_topk_ids.repeat(33, 1)[:129].contiguous() + large_topk_weights = torch.arange(large_topk_ids.numel(), dtype=torch.float32, device="cuda").reshape(129, 3) + _check_moe_align_fused(large_topk_ids, large_topk_weights, expert_num, ordered=False) + + medium_topk_ids = base_topk_ids.repeat(1024, 1).contiguous() + medium_topk_weights = torch.arange(medium_topk_ids.numel(), dtype=torch.float32, device="cuda").reshape(4096, 3) + _check_moe_align_fused(medium_topk_ids, medium_topk_weights, expert_num, ordered=False) + + shared_expert_num = 257 + shared_routing = torch.arange(512 * 7, dtype=torch.int32, device="cuda").reshape(512, 7) % 256 + shared_last = torch.full((512, 1), 256, dtype=torch.int32, device="cuda") + shared_topk_ids = torch.cat([shared_routing, shared_last], dim=1).contiguous() + shared_topk_weights = torch.arange(shared_topk_ids.numel(), dtype=torch.float32, device="cuda").reshape(512, 8) + _check_moe_align_fused(shared_topk_ids, shared_topk_weights, shared_expert_num, ordered=False) + + large_atomic_topk_ids = base_topk_ids.repeat(1281, 1)[:5121].contiguous() + large_atomic_topk_weights = torch.arange(large_atomic_topk_ids.numel(), dtype=torch.float32, device="cuda").reshape( + 5121, 3 + ) + _check_moe_align_fused(large_atomic_topk_ids, large_atomic_topk_weights, expert_num, ordered=False) + + sparse_expert_num = 257 + sparse_topk_ids = base_topk_ids.repeat(1281, 1)[:5121].contiguous() + sparse_topk_weights = torch.arange(sparse_topk_ids.numel(), dtype=torch.float32, device="cuda").reshape(5121, 3) + _check_moe_align_fused(sparse_topk_ids, sparse_topk_weights, sparse_expert_num, ordered=False) + + +def test_moe_align_fused_large_token_unordered(): + expert_num = 257 + topk_ids = torch.arange(5121 * 8, dtype=torch.int32, device="cuda").reshape(5121, 8) % expert_num + topk_weights = torch.arange(topk_ids.numel(), dtype=torch.float32, device="cuda").reshape(5121, 8) + _check_moe_align_fused(topk_ids, topk_weights, expert_num, ordered=False) + + +def test_moe_align_fused_atomic_token_unordered(): + expert_num = 9 + topk_ids = torch.arange(257 * 4, dtype=torch.int32, device="cuda").reshape(257, 4) % expert_num + topk_weights = torch.arange(topk_ids.numel(), dtype=torch.float32, device="cuda").reshape(257, 4) + expert_to_token_index = torch.empty((expert_num, topk_ids.numel()), dtype=torch.int32, device="cuda") + expert_to_weight = torch.empty((expert_num, topk_ids.numel()), dtype=torch.float32, device="cuda") + expert_token_num = torch.empty((expert_num,), dtype=torch.int32, device="cuda") + + _moe_align_fused_atomic_token( + expert_to_token_index, + expert_to_weight, + expert_token_num, + topk_ids, + topk_weights, + ) + torch.cuda.synchronize() + + flat_topk_ids = topk_ids.flatten() + flat_topk_weights = topk_weights.flatten() + expected_token_num = torch.bincount(flat_topk_ids, minlength=expert_num).to(torch.int32) + assert torch.equal(expert_token_num, expected_token_num) + + for expert_id, token_num in enumerate(expected_token_num.tolist()): + expected_index = torch.nonzero(flat_topk_ids == expert_id, as_tuple=False).flatten().to(torch.int32) + expected_weight = flat_topk_weights[expected_index] + token_index = expert_to_token_index[expert_id, :token_num] + token_weight = expert_to_weight[expert_id, :token_num] + order = torch.argsort(token_index) + assert torch.equal(token_index[order], expected_index) + assert torch.allclose(token_weight[order], expected_weight) + + +def test_fused_experts_atomic_align_path_is_deterministic(): + token_num = 129 + expert_num = 9 + hidden_size = 64 + intermediate_size = 128 + topk = 4 + hidden_states = torch.randn((token_num, hidden_size), dtype=torch.bfloat16, device="cuda") / 10 + w1 = torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.bfloat16, device="cuda") / 10 + w2 = torch.randn((expert_num, hidden_size, intermediate_size // 2), dtype=torch.bfloat16, device="cuda") / 10 + topk_ids = torch.arange(token_num * topk, dtype=torch.int32, device="cuda").reshape(token_num, topk) % expert_num + topk_weights = torch.softmax(torch.randn((token_num, topk), dtype=torch.float32, device="cuda"), dim=-1) + + out_0 = fused_experts_impl(hidden_states, w1, w2, topk_weights, topk_ids) + out_1 = fused_experts_impl(hidden_states, w1, w2, topk_weights, topk_ids) + torch.cuda.synchronize() + + assert torch.equal(out_0, out_1) + + +def test_moe_align2(): + + experts_token_num = torch.zeros((4,), dtype=torch.int32, device="cuda") + experts_token_num[0] = 8 + experts_token_num[1] = 0 + experts_token_num[2] = 60 + experts_token_num[3] = 16 + + token_num_mul_topk_num = 100 + block_m = 16 + mblocks_to_tuple_info = moe_align2(token_num_mul_topk_num, experts_token_num, block_m=block_m) + expected_expert_ids = torch.tensor([0, 2, 2, 2, 2, 3, -1, -1, -1, -1], device="cuda", dtype=torch.int32) + valid_blocks = expected_expert_ids != -1 + + max_num_tokens_padded = token_num_mul_topk_num + experts_token_num.shape[0] * (block_m - 1) + expected_max_num_m_blocks = min(token_num_mul_topk_num, triton.cdiv(max_num_tokens_padded, block_m)) + assert mblocks_to_tuple_info.shape[0] == expected_max_num_m_blocks + assert torch.equal(mblocks_to_tuple_info[:, 0], expected_expert_ids) + assert torch.equal( + mblocks_to_tuple_info[valid_blocks, 1], + torch.tensor([0, 0, 1, 2, 3, 0], device="cuda", dtype=torch.int32), + ) + + # When expert_num is large relative to token_num, the padded cdiv upper bound + # exceeds token_num; allocation must be capped by token_num_mul_topk_num. + token_num_mul_topk_num = 8 + block_m = 16 + experts_token_num = torch.zeros((10,), dtype=torch.int32, device="cuda") + experts_token_num[0] = 8 + max_num_tokens_padded = token_num_mul_topk_num + experts_token_num.shape[0] * (block_m - 1) + padded_cdiv = triton.cdiv(max_num_tokens_padded, block_m) + assert padded_cdiv > token_num_mul_topk_num + + mblocks_to_tuple_info = moe_align2(token_num_mul_topk_num, experts_token_num, block_m=block_m) + assert mblocks_to_tuple_info.shape[0] == token_num_mul_topk_num + assert torch.equal( + mblocks_to_tuple_info[:1, 0], + torch.tensor([0], device="cuda", dtype=torch.int32), + ) + assert torch.all(mblocks_to_tuple_info[1:, 0] == -1) + + +def test_grouped_matmul(): + test_dtype = torch.bfloat16 + token_inputs = torch.randn((10, 512), dtype=test_dtype, device="cuda") / 10 + experts_token_num = torch.tensor([1, 9], dtype=torch.int32, device="cuda") + experts_to_token_index = torch.tensor( + [ + [0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + [1, 2, 3, 4, 5, 6, 7, 8, 9, 0], + ], + dtype=torch.int32, + device="cuda", + ) + experts_to_weights = torch.tensor( + [ + [0.5, 0, 0, 0, 0, 0, 0, 0, 0, 0], + [0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0], + ], + dtype=torch.float32, + device="cuda", + ) + expert_weights = torch.randn((2, 1024, 512), dtype=test_dtype, device="cuda") / 10 + topk_num = 1 + out = torch.empty((10, 1024), dtype=test_dtype, device="cuda") + # warm up + grouped_matmul( + 10 * 1, + token_inputs, + None, + experts_token_num, + experts_to_token_index, + experts_to_weights, + expert_weights, + None, + topk_num, + out, + mul_routed_weight=True, + use_fp8_w8a8=False, + ) + torch.cuda.synchronize() + start = time.time() + for _ in range(100): + grouped_matmul( + 10 * 1, + token_inputs, + None, + experts_token_num, + experts_to_token_index, + experts_to_weights, + expert_weights, + None, + topk_num, + out, + mul_routed_weight=True, + use_fp8_w8a8=False, + ) + torch.cuda.synchronize() + logger.info(f"grouped_matmul test cost time: {time.time() - start} s") + + ans_list = [] + ans_list.append(torch.matmul(token_inputs[0:1, :], expert_weights[0].transpose(0, 1))) + for i in range(9): + t_ans = torch.matmul(token_inputs[(i + 1) : (i + 2), :], expert_weights[1].transpose(0, 1)) + ans_list.append(t_ans) + + true_out = torch.cat(ans_list, dim=0) + + assert torch.allclose(0.5 * true_out, out, atol=1e-2, rtol=0) + + +if __name__ == "__main__": + pytest.main() diff --git a/unit_tests/common/fused_moe/test_grouped_topk.py b/unit_tests/common/fused_moe/test_grouped_topk.py new file mode 100755 index 0000000000..432e133163 --- /dev/null +++ b/unit_tests/common/fused_moe/test_grouped_topk.py @@ -0,0 +1,100 @@ +import torch +import time +import pytest +import numpy as np +from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_topk import triton_grouped_topk +from lightllm.common.basemodel.triton_kernel.fused_moe.topk_select import biased_grouped_topk as grouped_topk +from lightllm.utils.log_utils import init_logger + +logger = init_logger(__name__) + +seed = 42 +torch.manual_seed(seed) + +if torch.cuda.is_available(): + torch.cuda.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + + +@pytest.mark.parametrize( + "expert_num, topk_group, group_num, topk_num, scoring_func, token_num", + [ + (*a, b, c) + for a in [(256, 4, 8, 8), (160, 3, 8, 6)] + for b in [ + "sigmoid", + ] + for c in [1, 8, 256, 1024, 2048, 4096, 8192] + ], +) +def test_grouped_topk(expert_num, topk_group, group_num, topk_num, scoring_func, token_num): + print("test", expert_num, topk_group, group_num, topk_num, scoring_func, token_num) + dtype = torch.float32 + hidden_state = torch.randn((token_num, 1), dtype=dtype, device="cuda") + gating_output = torch.randn((token_num, expert_num), dtype=dtype, device="cuda") * 10 + correction_bias = torch.randn((expert_num,), dtype=dtype, device="cuda") + correction_bias[correction_bias <= 0.0] = 0.0 + + old_topk_weights, old_topk_ids = grouped_topk( + hidden_state, + gating_output=gating_output, + correction_bias=correction_bias, + topk=topk_num, + renormalize=True, + num_expert_group=group_num, + topk_group=topk_group, + scoring_func=scoring_func, + ) + + new_topk_weights, new_topk_ids = triton_grouped_topk( + None, + gating_output=gating_output, + correction_bias=correction_bias, + topk=topk_num, + renormalize=True, + num_expert_group=group_num, + topk_group=topk_group, + scoring_func=scoring_func, + ) + + torch.cuda.synchronize() + start = time.time() + for _ in range(60): + old_topk_weights, old_topk_ids = grouped_topk( + hidden_state, + gating_output=gating_output, + correction_bias=correction_bias, + topk=topk_num, + renormalize=True, + num_expert_group=group_num, + topk_group=topk_group, + scoring_func=scoring_func, + ) + torch.cuda.synchronize() + print(f"old cost time {time.time() - start} s") + + torch.cuda.synchronize() + start = time.time() + for _ in range(60): + new_topk_weights, new_topk_ids = triton_grouped_topk( + None, + gating_output=gating_output, + correction_bias=correction_bias, + topk=topk_num, + renormalize=True, + num_expert_group=group_num, + topk_group=topk_group, + scoring_func=scoring_func, + ) + torch.cuda.synchronize() + print(f"new cost time {time.time() - start} s") + + assert torch.equal(torch.sort(old_topk_ids, dim=1)[0], torch.sort(new_topk_ids, dim=1)[0]) + assert torch.allclose( + torch.sort(old_topk_weights, dim=1)[0], torch.sort(new_topk_weights, dim=1)[0], atol=1e-3, rtol=1e-1 + ), f"max delta {torch.max(torch.sort(old_topk_weights, dim=1)[0] - torch.sort(new_topk_weights, dim=1)[0])}" + return + + +if __name__ == "__main__": + pytest.main() diff --git a/unit_tests/common/kv_trans_kernel/test_kv_trans_v2.py b/unit_tests/common/kv_trans_kernel/test_kv_trans_v2.py new file mode 100644 index 0000000000..dd34402dd2 --- /dev/null +++ b/unit_tests/common/kv_trans_kernel/test_kv_trans_v2.py @@ -0,0 +1,123 @@ +import pytest +import torch +import random +from lightllm.common.kv_trans_kernel.kv_trans_v2 import kv_trans_v2_for_p_node, kv_trans_v2_for_d_node, kv_trans_for_dp + + +@pytest.mark.parametrize( + "token_num", + [token_num for token_num in range(5, 10)], +) +def test_kv_trans_v2_for_p_node(token_num): + dp_size_in_node = 8 + head_num = 2 + head_dim = 512 + kv_buffer_token_num = 512 + mems = [] + for _ in range(dp_size_in_node): + mems.append(torch.randn((kv_buffer_token_num, head_num, head_dim), dtype=torch.float16, device="cuda")) + input_mems = torch.tensor([e.data_ptr() for e in mems], dtype=torch.uint64, device="cuda") + input_idx = [random.randint(0, kv_buffer_token_num - 1) for _ in range(token_num)] + input_idx = torch.tensor(input_idx, dtype=torch.int32, device="cuda") + input_dp_idx = [random.randint(0, dp_size_in_node - 1) for _ in range(token_num)] + input_dp_idx = torch.tensor(input_dp_idx, dtype=torch.int32, device="cuda") + + true_output = torch.zeros((token_num, head_num, head_dim), dtype=torch.float16, device="cuda") + test_output = torch.zeros((token_num, head_num, head_dim), dtype=torch.float16, device="cuda") + output_idx = torch.arange(0, token_num, 1, dtype=torch.int32, device="cuda") + + kv_trans_v2_for_p_node(input_mems, input_idx, input_dp_idx, test_output, output_idx, dp_size_in_node) + + for dest_token_index, token_index, dp_index in zip( + list(range(token_num)), input_idx.cpu().numpy(), input_dp_idx.cpu().numpy() + ): + true_output[dest_token_index, :, :] = mems[dp_index][token_index] + + assert torch.equal(true_output, test_output) + return + + +@pytest.mark.parametrize( + "token_num", + [token_num for token_num in range(5, 10)], +) +def test_kv_trans_v2_for_d_node(token_num): + card_num = 8 + dp_size_in_node = 4 + head_num = 2 + head_dim = 512 + kv_buffer_token_num = 512 + mems = [] + for _ in range(card_num): + mems.append(torch.randn((kv_buffer_token_num, head_num, head_dim), dtype=torch.float16, device="cuda")) + output_mems = torch.tensor([e.data_ptr() for e in mems], dtype=torch.uint64, device="cuda") + output_idx = [random.randint(0, kv_buffer_token_num - 1) for _ in range(token_num)] + output_idx = torch.tensor(output_idx, dtype=torch.int32, device="cuda") + output_dp_idx = [random.randint(0, dp_size_in_node - 1) for _ in range(token_num)] + output_dp_idx = torch.tensor(output_dp_idx, dtype=torch.int32, device="cuda") + + test_input = torch.randn((token_num, head_num, head_dim), dtype=torch.float16, device="cuda") + input_idx = torch.arange(0, token_num, 1, dtype=torch.int32, device="cuda") + + kv_trans_v2_for_d_node(output_mems, output_idx, output_dp_idx, test_input, input_idx, dp_size_in_node) + + for dest_token_index, token_index, dest_token_index, dp_index in zip( + list(range(token_num)), + input_idx.cpu().numpy(), + output_idx.cpu().numpy(), + output_dp_idx.cpu().numpy(), + ): + for mem_index in range(dp_index * card_num // dp_size_in_node, (dp_index + 1) * card_num // dp_size_in_node): + torch.equal(mems[mem_index][dest_token_index, :, :], test_input[token_index, :, :]) + + return + + +@pytest.mark.parametrize( + "token_num", + [token_num for token_num in range(5, 10)], +) +def test_kv_trans_for_dp(token_num): + card_num = 8 + dp_size_in_node = 4 + layer_num = 3 + head_num = 2 + head_dim = 512 + kv_buffer_token_num = 512 + rank_in_dp = 1 + + card_num_per_d = card_num // dp_size_in_node + + # 创建多层的 mem，每个 mem 包含所有层的数据 + mems = [] + for _ in range(card_num): + mems.append( + torch.randn((layer_num, kv_buffer_token_num, head_num, head_dim), dtype=torch.float16, device="cuda") + ) + + input_mems = torch.tensor([e.data_ptr() for e in mems], dtype=torch.uint64, device="cuda") + input_idx = [random.randint(0, kv_buffer_token_num - 1) for _ in range(token_num)] + input_idx = torch.tensor(input_idx, dtype=torch.int32, device="cuda") + input_dp_idx = [random.randint(0, dp_size_in_node - 1) for _ in range(token_num)] + input_dp_idx = torch.tensor(input_dp_idx, dtype=torch.int32, device="cuda") + + true_output = torch.zeros((layer_num, kv_buffer_token_num, head_num, head_dim), dtype=torch.float16, device="cuda") + test_output = torch.zeros((layer_num, kv_buffer_token_num, head_num, head_dim), dtype=torch.float16, device="cuda") + output_idx = torch.arange(0, token_num, 1, dtype=torch.int32, device="cuda") + + kv_trans_for_dp(input_mems, input_idx, input_dp_idx, test_output, output_idx, dp_size_in_node, rank_in_dp) + + # 验证结果 + for dest_token_index, src_token_index, dp_index in zip( + list(range(token_num)), input_idx.cpu().numpy(), input_dp_idx.cpu().numpy() + ): + mem_index = rank_in_dp + dp_index * card_num_per_d + # 所有 layer 都从同一个 mem 的对应层读取 + true_output[:, dest_token_index, :, :] = mems[mem_index][:, src_token_index, :, :] + + assert torch.equal(true_output, test_output), "kv_trans_for_dp output mismatch" + return + + +if __name__ == "__main__": + pytest.main() diff --git a/unit_tests/common/kv_trans_kernel/test_nixl_kv_trans.py b/unit_tests/common/kv_trans_kernel/test_nixl_kv_trans.py new file mode 100644 index 0000000000..3ec72e15d2 --- /dev/null +++ b/unit_tests/common/kv_trans_kernel/test_nixl_kv_trans.py @@ -0,0 +1,141 @@ +import os +import pytest +import torch + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for Triton kernels") + + +def _randn_like(shape, dtype=torch.float32, device="cuda"): + torch.manual_seed(1234) + return torch.randn(shape, dtype=dtype, device=device).contiguous() + + +def _gather_along_size(t, idx): + return t.index_select(dim=1, index=idx) + + +@pytest.mark.parametrize("head_dim", [64, 96]) +@pytest.mark.parametrize( + "kv_head_num,tp_world_size", + [ + (4, 1), + (4, 2), + (8, 1), + (8, 2), + (8, 4), + (16, 2), + ], +) +def test_page_io_roundtrip_with_tp(head_dim, kv_head_num, tp_world_size): + from lightllm.common.kv_trans_kernel.nixl_kv_trans import page_io + + device = "cuda" + dtype = torch.bfloat16 + + # Shapes + layer_num = 3 + size = 32 + assert kv_head_num % 2 == 0 + page_size = 32 + page_head_num = kv_head_num * tp_world_size + + tp_indices = list(range(tp_world_size)) + + kv_buffer = _randn_like((layer_num, size, kv_head_num, head_dim), dtype=dtype, device=device) + page_tensor = torch.zeros((page_size, layer_num, page_head_num, head_dim), dtype=dtype, device=device).contiguous() + + # Select a handful of token positions to move (tid count <= page_size) + mem_indexes = torch.tensor([2, 5, 7, 9, 11], dtype=torch.int64, device=device).contiguous() + + # Write: kv_buffer -> page_tensor, done by all tp ranks to fill their partition + for tp_index in tp_indices: + page_io( + mem_indexes=mem_indexes, + page_tensor=page_tensor, + kv_buffer=kv_buffer, + tp_index=tp_index, + tp_world_size=tp_world_size, + mode="write", + ) + + # After-write expectation check + token_num = mem_indexes.numel() + k_head_num = kv_head_num // 2 + page_k_head_num = page_head_num // 2 + page_write_head_num = page_k_head_num // tp_world_size + + expected_page = torch.zeros_like(page_tensor) + for tid in range(token_num): + mem_index = int(mem_indexes[tid].item()) + for layer_index in range(layer_num): + for tp_index in tp_indices: + page_head_start = tp_index * page_write_head_num + for kv_head_id in range(page_write_head_num): + # K half + expected_page[tid, layer_index, page_head_start + kv_head_id, :] = kv_buffer[ + layer_index, mem_index, kv_head_id, : + ] + # V half + expected_page[tid, layer_index, page_k_head_num + page_head_start + kv_head_id, :] = kv_buffer[ + layer_index, mem_index, k_head_num + kv_head_id, : + ] + + assert torch.allclose(page_tensor[:token_num], expected_page[:token_num], atol=1e-3, rtol=1e-3) + + # Read back to a fresh buffer + kv_buffer_rt = torch.zeros_like(kv_buffer) + for tp_index in tp_indices: + page_io( + mem_indexes=mem_indexes, + page_tensor=page_tensor, + kv_buffer=kv_buffer_rt, + tp_index=tp_index, + tp_world_size=tp_world_size, + mode="read", + ) + + # Check equality only at selected positions along size-dim + ref = _gather_along_size(kv_buffer, mem_indexes) + out = _gather_along_size(kv_buffer_rt, mem_indexes) + assert torch.allclose(out, ref, atol=1e-3, rtol=1e-3) + + +@pytest.mark.parametrize("head_dim", [32, 80]) # include non-power-of-two +def test_mla_page_io_roundtrip(head_dim): + from lightllm.common.kv_trans_kernel.nixl_kv_trans import mla_page_io + + device = "cuda" + dtype = torch.bfloat16 + + # Shapes (single-head) + layer_num = 2 + size = 20 + kv_head_num = 1 + page_size = 10 + page_head_num = 1 + + kv_buffer = _randn_like((layer_num, size, kv_head_num, head_dim), dtype=dtype, device=device) + page_tensor = torch.zeros((page_size, layer_num, page_head_num, head_dim), dtype=dtype, device=device).contiguous() + + mem_indexes = torch.tensor([0, 3, 6, 7], dtype=torch.int64, device=device).contiguous() + + # Write kv -> page + mla_page_io(mem_indexes=mem_indexes, page_tensor=page_tensor, kv_buffer=kv_buffer, mode="write") + + # After-write expectation check + token_num = mem_indexes.numel() + expected_page = torch.zeros_like(page_tensor) + for tid in range(token_num): + mem_index = int(mem_indexes[tid].item()) + for layer_index in range(layer_num): + expected_page[tid, layer_index, 0, :] = kv_buffer[layer_index, mem_index, 0, :] + assert torch.allclose(page_tensor[:token_num], expected_page[:token_num], atol=1e-3, rtol=1e-3) + + # Read back page -> kv + kv_buffer_rt = torch.zeros_like(kv_buffer) + mla_page_io(mem_indexes=mem_indexes, page_tensor=page_tensor, kv_buffer=kv_buffer_rt, mode="read") + + ref = kv_buffer.index_select(dim=1, index=mem_indexes) + out = kv_buffer_rt.index_select(dim=1, index=mem_indexes) + assert torch.allclose(out, ref, atol=1e-3, rtol=1e-3) diff --git a/unit_tests/common/quantization/test_no_quant.py b/unit_tests/common/quantization/test_no_quant.py new file mode 100644 index 0000000000..2b66767539 --- /dev/null +++ b/unit_tests/common/quantization/test_no_quant.py @@ -0,0 +1,43 @@ +import pytest +import torch + +from lightllm.common.basemodel.layer_weights.meta_weights.mm_weight.mm_weight import MMWeightTpl +from lightllm.common.quantization.no_quant import NoQuantization +from lightllm.common.quantization.quantize_method import WeightPack + + +@pytest.fixture(autouse=True) +def device_env(monkeypatch): + monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("fp32_output", [False, True]) +def test_output_buffer_controls_matmul_precision(dtype, fp32_output): + torch.manual_seed(1593) + x = torch.randn((17, 64), dtype=dtype, device="cuda") + weight = torch.randn((32, 64), dtype=dtype, device="cuda") + out = torch.empty((17, 32), dtype=torch.float32 if fp32_output else dtype, device="cuda") + + projection = MMWeightTpl.__new__(MMWeightTpl) + projection.quant_method = NoQuantization() + projection.mm_param = WeightPack(weight=weight) + projection.bias = None + result = projection.mm(x, out=out) + + assert result is out + expected = torch.mm(x.float(), weight.float().t()).to(out.dtype) + torch.testing.assert_close(result, expected, rtol=1e-5, atol=1e-5) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_default_output_and_bias_preserve_input_dtype(): + x = torch.randn((3, 64), dtype=torch.bfloat16, device="cuda") + weight = torch.randn((32, 64), dtype=x.dtype, device=x.device) + bias = torch.randn(32, dtype=x.dtype, device=x.device) + + result = NoQuantization().apply(x, WeightPack(weight=weight), bias=bias, use_custom_tensor_mananger=False) + + assert result.dtype == x.dtype + torch.testing.assert_close(result, torch.addmm(bias, x, weight.t()), rtol=0, atol=0) diff --git a/unit_tests/common/quantization/test_ue8m0_fp8_quantization.py b/unit_tests/common/quantization/test_ue8m0_fp8_quantization.py new file mode 100644 index 0000000000..f1c9b228a2 --- /dev/null +++ b/unit_tests/common/quantization/test_ue8m0_fp8_quantization.py @@ -0,0 +1,513 @@ +import os +import pytest +import torch +import torch.nn.functional as F + +from lightllm.common.basemodel.triton_kernel.quantization import fp8act_quant_kernel as activation +from lightllm.common.basemodel.triton_kernel.quantization import fp8w8a8_block_quant_kernel as weight +from lightllm.common.quantization import Quantcfg +from lightllm.common.quantization import deepgemm + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + + +@pytest.fixture(autouse=True) +def preserve_ue8m0_env(monkeypatch): + monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", os.getenv("LIGHTLLM_USE_UE8M0_SCALES", "0")) + + +def _unpack_scales(packed, groups): + shifts = torch.arange(4, device=packed.device) * 8 + exponents = ((packed.to(torch.int64)[..., None] >> shifts) & 255).flatten(1) + return torch.exp2(exponents[:, :groups].float() - 127) + + +@pytest.fixture +def packed_weight_layout(monkeypatch): + calls = [] + transform = deepgemm.deep_gemm.transform_sf_into_required_layout + sm100 = deepgemm.is_sm100_gpu() + + def record_transform(sf, n, k, recipe, num_groups=None, is_sfa=False): + calls.append((n, k, num_groups)) + if sm100: + return transform(sf, n, k, recipe, num_groups=num_groups, is_sfa=is_sfa) + # Exercise DeepGEMM's actual pack kernel on Hopper, emulating only the SM100 layout dispatch. + assert recipe == (1, 128, 128) and not is_sfa + expanded = sf.index_select(-2, torch.arange(n, device=sf.device) // 128) + return deepgemm.deep_gemm.get_mn_major_tma_aligned_packed_ue8m0_tensor(expanded) + + monkeypatch.setattr(deepgemm.deep_gemm, "transform_sf_into_required_layout", record_transform) + return calls + + +@pytest.mark.skipif(not activation.HAS_SGL_KERNEL, reason="requires SGL kernel") +@pytest.mark.parametrize("use_ue8m0_scales,use_packed_ue8m0", [(False, False), (True, False), (True, True)]) +def test_sgl_dispatch_uses_ue8m0_flag(monkeypatch, use_ue8m0_scales, use_packed_ue8m0): + quantize = activation.sgl_ops.sgl_per_token_group_quant_fp8 + calls = [] + + def record_call(*args, **kwargs): + calls.append(kwargs["scale_ue8m0"]) + return quantize(*args, **kwargs) + + monkeypatch.setattr(activation.sgl_ops, "sgl_per_token_group_quant_fp8", record_call) + x = torch.ones((3, 512), device="cuda", dtype=torch.bfloat16) + _, scales = activation.per_token_group_quant_fp8( + x, 128, use_ue8m0_scales=use_ue8m0_scales, use_packed_ue8m0=use_packed_ue8m0 + ) + assert calls == ([use_packed_ue8m0] if not use_ue8m0_scales or use_packed_ue8m0 else []) + assert scales.dtype == (torch.int32 if use_packed_ue8m0 else torch.float32) + if use_ue8m0_scales: + actual = _unpack_scales(scales, 4) if use_packed_ue8m0 else scales + torch.testing.assert_close(actual, torch.full_like(actual, 2.0 ** -8), rtol=0, atol=0) + + +@pytest.mark.parametrize("use_packed_ue8m0", [False, True]) +def test_ue8m0_rounding_at_power_of_two_boundaries(monkeypatch, use_packed_ue8m0): + monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") + monkeypatch.setattr(activation, "HAS_SGL_KERNEL", False) + powers = torch.tensor([2.0 ** exponent for exponent in (-16, -8, 0, 8, 16, 118)], device="cuda") + above = torch.nextafter(powers, torch.full_like(powers, float("inf"))) + below = torch.nextafter(powers, torch.zeros_like(powers)) + magnitudes = torch.cat((powers, above, below, powers.new_zeros(1))) * 448.0 + x = magnitudes[:, None].expand(-1, 128).contiguous() + expected_act = torch.exp2(torch.ceil(torch.log2(magnitudes.clamp_min(1e-10) / 448.0))) + expected_weight = torch.exp2(torch.ceil(torch.log2(magnitudes.clamp_min(1e-4) / 448.0))) + + _, act_scales = activation.per_token_group_quant_fp8( + x, 128, use_ue8m0_scales=True, use_packed_ue8m0=use_packed_ue8m0 + ) + _, weight_scales = weight.weight_quant(x.repeat_interleave(128, dim=0), use_ue8m0_scales=True) + + actual = _unpack_scales(act_scales, 1) if use_packed_ue8m0 else act_scales + torch.testing.assert_close(actual[:, 0], expected_act, rtol=0, atol=0) + torch.testing.assert_close(weight_scales[:, 0], expected_weight, rtol=0, atol=0) + + +@pytest.mark.parametrize("layout", ["row", "column", "tma"]) +@pytest.mark.parametrize("group_size", [64, 128]) +@pytest.mark.parametrize("use_ue8m0_scales,use_packed_ue8m0", [(False, False), (True, False), (True, True)]) +def test_activation_scales_and_quantized_values(monkeypatch, layout, group_size, use_ue8m0_scales, use_packed_ue8m0): + monkeypatch.setattr(activation, "HAS_SGL_KERNEL", False) + torch.manual_seed(20261008) + rows, groups = 17, 4 + x = torch.randn(rows, groups * group_size, device="cuda", dtype=torch.bfloat16) + x[0].zero_() + x[1].fill_(1e-12) + q, scales = activation.per_token_group_quant_fp8( + x, + group_size, + column_major_scales=layout != "row", + scale_tma_aligned=layout == "tma", + use_ue8m0_scales=use_ue8m0_scales, + use_packed_ue8m0=use_packed_ue8m0, + ) + + amax = x.float().reshape(rows, groups, group_size).abs().amax(dim=-1) + reference_scales = amax.clamp_min(1e-10) / 448.0 + if use_ue8m0_scales: + reference_scales = torch.exp2(torch.ceil(torch.log2(reference_scales.double()))).float() + reference_q = (x.float().reshape(rows, groups, group_size) / reference_scales[..., None]).to(q.dtype) + actual_scales = _unpack_scales(scales, groups) if use_packed_ue8m0 else scales + torch.testing.assert_close(actual_scales, reference_scales, rtol=0, atol=0) + torch.testing.assert_close(q.float(), reference_q.reshape_as(x).float(), rtol=0, atol=0) + assert scales.dtype == (torch.int32 if use_packed_ue8m0 else torch.float32) + # UE8M0 always follows SGL's packed TMA layout, including with default layout flags. + if use_packed_ue8m0 or layout == "tma": + assert scales.stride() == (1, 20) + elif layout == "row": + assert scales.stride() == (groups, 1) + + +@pytest.mark.parametrize("shape", [(2, 3, 256), (2, 3, 4, 256)]) +@pytest.mark.parametrize("group_size", [64, 128]) +@pytest.mark.parametrize("use_ue8m0_scales", [False, True]) +def test_nd_activation_scales_preserve_row_major_output(monkeypatch, shape, group_size, use_ue8m0_scales): + monkeypatch.setattr(activation, "HAS_SGL_KERNEL", False) + x = torch.randn(shape, device="cuda", dtype=torch.bfloat16) + x[0, 0].zero_() + q, scales = activation.per_token_group_quant_fp8(x, group_size, use_ue8m0_scales=use_ue8m0_scales) + groups = shape[-1] // group_size + reference_scales = x.float().view(*shape[:-1], groups, group_size).abs().amax(-1).clamp_min(1e-10) / 448 + if use_ue8m0_scales: + reference_scales = reference_scales.log2().ceil().exp2() + reference_q = (x.float().view(*shape[:-1], groups, group_size) / reference_scales[..., None]).to(q.dtype) + assert scales.shape == (*shape[:-1], groups) + torch.testing.assert_close(scales, reference_scales, rtol=0, atol=0) + torch.testing.assert_close(q.float(), reference_q.view_as(x).float(), rtol=0, atol=0) + + +@pytest.mark.parametrize("use_sgl", [False, True]) +@pytest.mark.parametrize("rows", [1, 17]) +@pytest.mark.parametrize("groups", [1, 2, 3, 4, 5, 8]) +@pytest.mark.parametrize("group_size", [64, 128]) +def test_packed_ue8m0_scales(monkeypatch, rows, groups, group_size, use_sgl): + if use_sgl and not activation.HAS_SGL_KERNEL: + pytest.skip("requires SGL kernel") + monkeypatch.setattr(activation, "HAS_SGL_KERNEL", use_sgl) + torch.manual_seed(20261008) + x = torch.randn(rows, groups * group_size, device="cuda", dtype=torch.bfloat16) + x[0].zero_() + if rows > 1: + x[1].fill_(1e-12) + x[-1, -group_size:].fill_(2.0 ** 16) + q, packed_scales = activation.per_token_group_quant_fp8(x, group_size, use_ue8m0_scales=True, use_packed_ue8m0=True) + + assert packed_scales.dtype == torch.int32 + assert packed_scales.shape == (rows, (groups + 3) // 4) + assert packed_scales.stride() == (1, (rows + 3) // 4 * 4) + if groups % 4: + assert torch.count_nonzero(packed_scales[:, -1].to(torch.int64) >> (groups % 4 * 8)) == 0 + scales = _unpack_scales(packed_scales, groups) + amax = x.float().reshape(rows, groups, group_size).abs().amax(-1) + reference_scales = torch.exp2(torch.ceil(torch.log2(amax.clamp_min(1e-10) / 448.0))) + reference_q = (x.float().reshape(rows, groups, group_size) / reference_scales[..., None]).to(q.dtype) + torch.testing.assert_close(scales, reference_scales, rtol=0, atol=0) + torch.testing.assert_close(q.float(), reference_q.reshape_as(x).float(), rtol=0, atol=0) + + +@pytest.mark.parametrize("experts", [None, 2]) +@pytest.mark.parametrize("use_ue8m0_scales", [False, True]) +def test_weight_quantization_partial_blocks(monkeypatch, experts, use_ue8m0_scales): + monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") + torch.manual_seed(20261008) + shape = (129, 257) if experts is None else (experts, 129, 257) + x = torch.randn(shape, device="cuda", dtype=torch.bfloat16) + x[..., :128, :128].zero_() + x[..., 128:, 128:256].fill_(1e-12) + q, scales = weight.weight_quant(x, use_ue8m0_scales=use_ue8m0_scales) + + blocks = F.pad(x.float(), (0, 127, 0, 127)).reshape(-1, 2, 128, 3, 128) + amax = blocks.abs().amax(dim=(2, 4)).reshape_as(scales) + reference_scales = amax.clamp_min(1e-4) / 448.0 if use_ue8m0_scales else amax / 448.0 + if use_ue8m0_scales: + reference_scales = torch.exp2(torch.ceil(torch.log2(reference_scales.double()))).float() + denom = reference_scales if use_ue8m0_scales else reference_scales + 1e-6 + expanded = denom.repeat_interleave(128, dim=-2).repeat_interleave(128, dim=-1)[..., :129, :257] + torch.testing.assert_close(scales, reference_scales, rtol=0, atol=0) + torch.testing.assert_close(q.float(), (x.float() / expanded).to(q.dtype).float(), rtol=0, atol=0) + + +@pytest.mark.parametrize("env_value", ["1", "0"]) +def test_ue8m0_env_preserves_unquantized_method(monkeypatch, env_value): + monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") + monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", env_value) + method = Quantcfg({"n_layer": 1}).get_quant_method(0, "q_proj") + assert method.method_name == "none" + + +@pytest.mark.skipif(not deepgemm.HAS_DEEPGEMM, reason="requires DeepGEMM") +@pytest.mark.parametrize("sm100", [False, True]) +def test_deepgemm_selects_scale_format_for_machine(monkeypatch, sm100, packed_weight_layout): + monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", "1") + monkeypatch.setattr(deepgemm, "is_sm100_gpu", lambda: sm100) + method = Quantcfg({"n_layer": 1}, quant_type="fp8w8a8-b128-deepgemm").get_quant_method(0, "q_proj") + weight_pack, _ = method.create_weight([128], 512, torch.bfloat16, 0) + method.load_weight(torch.ones((128, 512), dtype=torch.bfloat16, device="cuda"), weight_pack) + seen = [] + + def check_gemm_inputs(a, b, out): + seen.append(a[1].dtype) + assert b[1].dtype == (torch.int32 if sm100 else torch.float32) + scales = _unpack_scales(a[1], 4) if sm100 else a[1] + torch.testing.assert_close(scales, torch.full_like(scales, 2.0 ** -8), rtol=0, atol=0) + weight_scales = _unpack_scales(b[1], 4) if sm100 else b[1] + torch.testing.assert_close(weight_scales, torch.full_like(weight_scales, 2.0 ** -8), rtol=0, atol=0) + out.zero_() + + # Exercise format selection and real quantization on Hopper; SM100 GEMM needs a separate machine. + monkeypatch.setattr(deepgemm, "_deepgemm_fp8_nt", check_gemm_inputs) + method.apply( + torch.ones((3, 512), dtype=torch.bfloat16, device="cuda"), weight_pack, use_custom_tensor_mananger=False + ) + method.apply( + torch.ones((3, 512), dtype=torch.bfloat16, device="cuda"), weight_pack, use_custom_tensor_mananger=False + ) + assert seen == [torch.int32 if sm100 else torch.float32] * 2 + assert len(packed_weight_layout) == int(sm100) + + +@pytest.mark.skipif(not deepgemm.HAS_DEEPGEMM, reason="requires DeepGEMM") +@pytest.mark.parametrize("sm100", [False, True]) +@pytest.mark.parametrize("ue8m0", [False, True]) +@pytest.mark.parametrize("prequantized", [False, True]) +@pytest.mark.parametrize("experts", [None, 2]) +def test_weight_scale_loading_preserves_fused_views( + monkeypatch, packed_weight_layout, sm100, ue8m0, prequantized, experts +): + monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", str(int(ue8m0))) + monkeypatch.setattr(deepgemm, "is_sm100_gpu", lambda: sm100) + method = Quantcfg({"n_layer": 1}, quant_type="fp8w8a8-b128-deepgemm").get_quant_method(0, "q_proj") + args = dict(out_dims=[128, 129], in_dim=640, dtype=torch.bfloat16, device_id=0) + full, parts = ( + method.create_weight(**args) if experts is None else method.create_moe_weight(**args, num_experts=experts) + ) + packed = sm100 and ue8m0 + expected_parts = [] + torch.manual_seed(20261008) + for part in parts: + raw = torch.randn_like(part.weight, dtype=torch.bfloat16) + raw[..., 128:256] *= 8 + q, scales = weight.weight_quant(raw, use_ue8m0_scales=ue8m0) + loads = ( + [(part, raw, q, scales)] + if experts is None + else [(part.get_expert(i), raw[i], q[i], scales[i]) for i in range(experts)] + ) + for dest, bf16, fp8, sf in loads: + if prequantized: + method.load_weight(fp8, dest) + method.load_weight_scale(sf.cpu(), dest) + else: + method.load_weight(bf16, dest) + assert all(dest.load_ok) + torch.testing.assert_close(part.weight.float(), q.float(), rtol=0, atol=0) + expected = ( + scales.index_select(-2, torch.arange(part.weight.shape[-2], device="cuda") // 128) if packed else scales + ) + flat = part.weight_scale.reshape(-1, part.weight_scale.shape[-1]) + actual = _unpack_scales(flat, 5).reshape_as(expected) if packed else part.weight_scale + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + expected_parts.append(expected) + + expected = torch.cat(expected_parts, dim=-2) + assert full.weight_scale.dtype == (torch.int32 if packed else torch.float32) + if packed: + assert full.weight_scale.shape[-2:] == (257, 2) + assert full.weight_scale.stride()[-2:] == (1, 260) + actual = _unpack_scales(full.weight_scale.reshape(-1, 2), 5).reshape_as(expected) + else: + actual = full.weight_scale + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + assert len(packed_weight_layout) == (2 * (experts or 1) if packed else 0) + + +@pytest.mark.skipif(not deepgemm.HAS_DEEPGEMM, reason="requires DeepGEMM") +@pytest.mark.parametrize("need_trans", [False, True]) +@pytest.mark.parametrize("k", [512, 640]) +def test_tp_moe_consumes_packed_weight_scales(monkeypatch, packed_weight_layout, need_trans, k): + from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe import grouped_matmul + + monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", "1") + monkeypatch.setattr(deepgemm, "is_sm100_gpu", lambda: True) + method = Quantcfg({"n_layer": 1}, quant_type="fp8w8a8-b128-deepgemm").get_quant_method(0, "q_proj") + full, _ = method.create_moe_weight([256], k, torch.bfloat16, 0, 2) + torch.manual_seed(20261008) + raw = torch.randn((2, 256, k), device="cuda", dtype=torch.bfloat16) + q, scales = weight.weight_quant(raw, use_ue8m0_scales=True) + method.load_weight(raw, full) + x = torch.randn((17, k), device="cuda", dtype=torch.bfloat16) + counts = torch.tensor([9, 8], device="cuda", dtype=torch.int32) + indices = torch.zeros((2, 17), device="cuda", dtype=torch.int32) + indices[0, :9] = torch.arange(9, device="cuda") + indices[1, :8] = torch.arange(9, 17, device="cuda") + outputs = [] + for sf in (scales, full.weight_scale): + out = torch.empty((17, 256), device="cuda", dtype=torch.bfloat16) + grouped_matmul( + 17, + x, + None, + counts, + indices, + torch.ones((2, 17), device="cuda"), + q, + sf, + 1, + out, + mul_routed_weight=True, + use_fp8_w8a8=True, + run_config=dict( + BLOCK_SIZE_M=16, + BLOCK_SIZE_N=32, + BLOCK_SIZE_K=64, + GROUP_SIZE_M=1, + num_warps=4, + num_stages=2, + NEED_TRANS=need_trans, + ), + ) + outputs.append(out) + torch.testing.assert_close(outputs[0], outputs[1], rtol=0, atol=0) + + +@pytest.mark.skipif(not deepgemm.HAS_DEEPGEMM, reason="requires DeepGEMM") +@pytest.mark.parametrize("prequantized", [False, True]) +@pytest.mark.parametrize("env_value", [None, "1", "0"]) +@pytest.mark.parametrize("scale_fmt", ["no_config", None, "ue8m0", "float32"]) +@pytest.mark.parametrize("rows", [1, 17, 32]) +def test_deepgemm_load_and_apply(monkeypatch, rows, scale_fmt, env_value, prequantized): + monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") + if env_value is None: + monkeypatch.delenv("LIGHTLLM_USE_UE8M0_SCALES", raising=False) + else: + monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", env_value) + torch.manual_seed(20261008) + x = torch.randn(rows, 1024, device="cuda", dtype=torch.bfloat16) + w = torch.randn(256, 1024, device="cuda", dtype=torch.bfloat16) + config = {"n_layer": 1} + if scale_fmt != "no_config": + config["quantization_config"] = {"quant_method": "fp8", "weight_block_size": [128, 128]} + if scale_fmt is not None: + config["quantization_config"]["scale_fmt"] = scale_fmt + method = Quantcfg(config, quant_type="fp8w8a8-b128-deepgemm").get_quant_method(0, "q_proj") + use_ue8m0_scales = scale_fmt == "ue8m0" or env_value == "1" + weight_pack, _ = method.create_weight([256], 1024, torch.bfloat16, 0) + assert method.use_ue8m0_scales == use_ue8m0_scales + assert method.use_packed_ue8m0 == (use_ue8m0_scales and deepgemm.is_sm100_gpu()) + monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", "0" if use_ue8m0_scales else "1") + if prequantized: + qweight, scales = weight.weight_quant(w, use_ue8m0_scales=use_ue8m0_scales) + method.load_weight(qweight, weight_pack) + method.load_weight_scale(scales, weight_pack) + else: + method.load_weight(w, weight_pack) + assert method.use_ue8m0_scales == use_ue8m0_scales + weight_scales = weight_pack.weight_scale + if method.use_packed_ue8m0: + weight_scales = _unpack_scales(weight_scales, 8)[::128] + assert torch.all(weight_scales > 0) + log_scales = torch.log2(weight_scales) + assert torch.equal(log_scales, log_scales.round()) == use_ue8m0_scales + + quantize_activation = deepgemm.per_token_group_quant_fp8 + + def check_activation_scales(*args, **kwargs): + result = quantize_activation(*args, **kwargs) + packed = use_ue8m0_scales and deepgemm.is_sm100_gpu() + assert kwargs["use_packed_ue8m0"] == packed + assert result[1].dtype == (torch.int32 if packed else torch.float32) + if not packed: + log_scales = torch.log2(result[1]) + assert torch.equal(log_scales, log_scales.round()) == use_ue8m0_scales + return result + + monkeypatch.setattr(deepgemm, "per_token_group_quant_fp8", check_activation_scales) + + out = method.apply(x, weight_pack, use_custom_tensor_mananger=False) + assert method.use_ue8m0_scales == use_ue8m0_scales + # Later calls reuse the choice made when creating the weight. + monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", "0" if use_ue8m0_scales else "1") + method.apply(x, weight_pack, out=out, use_custom_tensor_mananger=False) + reference = x.float() @ w.float().T + nrmse = (out.float() - reference).square().mean().sqrt() / reference.square().mean().sqrt() + cosine = F.cosine_similarity(out.float().flatten(), reference.flatten(), dim=0) + assert torch.isfinite(out).all() + assert nrmse.item() < 0.08 + assert cosine.item() > 0.995 + + +@pytest.mark.skipif(not deepgemm.HAS_DEEPGEMM, reason="requires DeepGEMM") +@pytest.mark.parametrize("prepacked", [False, True]) +@pytest.mark.parametrize("sm100", [False, True]) +def test_redundancy_expert_update_packs_scales(monkeypatch, packed_weight_layout, prepacked, sm100): + from types import SimpleNamespace + from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.ep_redundancy import ( + FusedMoeWeightEPAutoRedundancy, + ) + + monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", "1") + monkeypatch.setattr(deepgemm, "is_sm100_gpu", lambda: sm100) + method = Quantcfg({"n_layer": 1}, quant_type="fp8w8a8-b128-deepgemm").get_quant_method(0, "q_proj") + w13, _ = method.create_moe_weight([128, 128], 512, torch.bfloat16, 0, 4) + w2, _ = method.create_moe_weight([512], 128, torch.bfloat16, 0, 4) + update = object.__new__(FusedMoeWeightEPAutoRedundancy) + update.redundancy_expert_num = 2 + update.redundancy_expert_ids = [6, 7] + update._ep_w = SimpleNamespace( + w13=w13, + w2=w2, + quant_method=method, + redundancy_expert_ids_tensor=torch.zeros(2, device="cuda", dtype=torch.int64), + ) + expected = [] + for dest in (w13, w2): + dest.weight.zero_() + dest.weight_scale.zero_() + shape = (2,) + dest.weight.shape[-2:] + raw = torch.randn(shape, device="cuda", dtype=torch.bfloat16) + q, sf = weight.weight_quant(raw, use_ue8m0_scales=True) + if prepacked: + source, _ = method.create_moe_weight([shape[-2]], shape[-1], torch.bfloat16, 0, 2) + method.load_weight(raw, source) + sf = source.weight_scale + expected.append([q, sf]) + update.w13, update.w2 = expected + before = len(packed_weight_layout) + update.commit() + assert len(packed_weight_layout) == before + (2 if sm100 and not prepacked else 0) + for dest, (q, sf) in zip((w13, w2), expected): + assert torch.count_nonzero(dest.weight[:2].float()) == 0 + assert torch.count_nonzero(dest.weight_scale[:2]) == 0 + torch.testing.assert_close(dest.weight[2:].float(), q.float(), rtol=0, atol=0) + if sm100: + groups = (q.shape[-1] + 127) // 128 + actual = _unpack_scales(dest.weight_scale[2:].reshape(-1, dest.weight_scale.shape[-1]), groups) + expected_sf = ( + _unpack_scales(sf.reshape(-1, sf.shape[-1]), groups) + if prepacked + else sf.repeat_interleave(128, dim=-2).reshape(-1, groups) + ) + else: + actual, expected_sf = dest.weight_scale[2:], sf + torch.testing.assert_close(actual, expected_sf, rtol=0, atol=0) + assert update._ep_w.redundancy_expert_ids_tensor.tolist() == [6, 7] + + +@pytest.mark.skipif(not deepgemm.HAS_DEEPGEMM, reason="requires DeepGEMM") +@pytest.mark.parametrize("is_prefill", [False, True]) +@pytest.mark.parametrize("sm100", [False, True]) +@pytest.mark.parametrize("scale_source", ["none", "env", "config"]) +def test_ep_moe_rounds_scales_and_preserves_logical_block_size(monkeypatch, is_prefill, sm100, scale_source): + from lightllm.common.basemodel.triton_kernel.fused_moe import grouped_fused_moe_ep as ep + from lightllm.utils.envs_utils import get_env_use_ue8m0_scales + + monkeypatch.delenv("LIGHTLLM_USE_UE8M0_SCALES", raising=False) + assert not get_env_use_ue8m0_scales() + config = {"n_layer": 1} + if scale_source == "env": + monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", "1") + elif scale_source == "config": + config["quantization_config"] = {"quant_method": "fp8", "weight_block_size": [128, 128], "scale_fmt": "ue8m0"} + use_ue8m0_scales = scale_source != "none" + monkeypatch.setattr(deepgemm, "is_sm100_gpu", lambda: sm100) + method = Quantcfg(config, quant_type="fp8w8a8-b128-deepgemm").get_quant_method(0, "q_proj") + assert get_env_use_ue8m0_scales() == use_ue8m0_scales + w13, _ = method.create_moe_weight([128, 128], 512, torch.bfloat16, 0, 2) + w2, _ = method.create_moe_weight([512], 128, torch.bfloat16, 0, 2) + x = torch.randn((3, 512), device="cuda", dtype=torch.bfloat16) + q, sf = ep.quantize_fused_experts_input(x, w13, method) + assert q.shape == x.shape and sf.shape == (3, 4) + assert sf.dtype == torch.float32 + expected_scales = x.float().reshape(3, 4, 128).abs().amax(-1).clamp_min(1e-10) / 448 + if use_ue8m0_scales: + expected_scales = torch.exp2(torch.ceil(torch.log2(expected_scales))) + expected_q = (x.float().reshape(3, 4, 128) / expected_scales.unsqueeze(-1)).to(q.dtype).reshape_as(q) + torch.testing.assert_close(q.float(), expected_q.float(), atol=0, rtol=0) + torch.testing.assert_close(sf, expected_scales) + seen = [] + monkeypatch.setattr(ep.dist_group_manager, "ep_buffer", object(), raising=False) + + def check_dispatch(**kwargs): + seen.append(kwargs["block_size_k"]) + assert ( + kwargs["w1_scale"].dtype + == kwargs["w2_scale"].dtype + == (torch.int32 if sm100 and use_ue8m0_scales else torch.float32) + ) + return x + + monkeypatch.setattr(ep.prefill_decode_sync, "fused_experts_impl", check_dispatch) + ep.fused_experts( + x, + w13, + w2, + torch.ones((3, 1), device="cuda"), + torch.zeros((3, 1), device="cuda", dtype=torch.int64), + num_experts=2, + quant_method=method, + is_prefill=is_prefill, + ) + assert seen == [128] diff --git a/unit_tests/common/test_autotuner.py b/unit_tests/common/test_autotuner.py new file mode 100644 index 0000000000..547246cfe5 --- /dev/null +++ b/unit_tests/common/test_autotuner.py @@ -0,0 +1,486 @@ +import json +from contextlib import nullcontext + +import pytest +import torch +from frozendict import frozendict + +from lightllm.common.triton_utils import autotuner as autotuner_module +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, AutotuneLevel, Autotuner, autotune + + +@pytest.fixture(autouse=True) +def autotune_environment(monkeypatch): + monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", None) + monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: AutotuneLevel.ADAPTIVE_AUTOTUNE) + monkeypatch.setattr(autotuner_module.dist, "is_initialized", lambda: False) + monkeypatch.setattr(autotuner_module.KernelConfigs, "get_config_file_name", lambda params: "configs.json") + + +def make_kernel(tmp_path, monkeypatch, name, kernel_type=None, warmup_all_exist_config=True): + calls = [] + benchmarks = [] + options = {} if kernel_type is None else {"kernel_type": kernel_type} + + @autotune( + kernel_name=name, + configs_gen_func=lambda: [{"block": 1}, {"block": 2}], + static_key_func=lambda: {"dtype": "test"}, + run_key_func=lambda size: size, + warmup_all_exist_config=warmup_all_exist_config, + **options, + ) + def kernel(size, run_config=None): + calls.append((size, run_config)) + return run_config + + def bench(size, run_config): + benchmarks.append((size, run_config)) + return 1.0 / run_config["block"] + + cache_dir = tmp_path / name + cache_dir.mkdir() + kernel._cache_dir = str(cache_dir) + monkeypatch.setattr(kernel, "_bench", bench) + return kernel, calls, benchmarks, cache_dir / "configs.json" + + +@pytest.mark.parametrize("repeat_override,retry_override", [(None, None), (8, 2)]) +def test_short_kernel_benchmark_sampling_options(monkeypatch, repeat_override, retry_override): + calls, replays = [], [] + + @autotune( + kernel_name="short_kernel_sampling", + configs_gen_func=lambda: [], + static_key_func=lambda: {}, + run_key_func=lambda size: size, + benchmark_repeats=32, + benchmark_retries=5, + ) + def kernel(size, run_config=None): + calls.append(size) + + class Stream: + def synchronize(self): + pass + + class Graph: + def replay(self): + replays.append(None) + + class Event: + def record(self): + pass + + def synchronize(self): + pass + + def elapsed_time(self, other): + return 0.032 + + monkeypatch.setattr(torch.cuda, "current_stream", Stream) + monkeypatch.setattr(torch.cuda, "Stream", Stream) + monkeypatch.setattr(torch.cuda, "CUDAGraph", Graph) + monkeypatch.setattr(torch.cuda, "graph", lambda *args, **kwargs: nullcontext()) + monkeypatch.setattr(torch.cuda, "Event", lambda **kwargs: Event()) + actual = kernel._bench(16, n_repeat=repeat_override, n_retries=retry_override) + repeats = 32 if repeat_override is None else repeat_override + retries = 5 if retry_override is None else retry_override + assert len(calls) == repeats + 1 # Initial warmup plus calls captured in the graph. + assert len(replays) == retries + assert actual == pytest.approx(0.032 / repeats) + + +@pytest.mark.parametrize("repeats,retries", [(0, 3), (3, 0)]) +def test_benchmark_sampling_requires_positive_counts(repeats, retries): + with pytest.raises(ValueError, match="must be positive"): + autotune( + kernel_name="invalid_sampling", + configs_gen_func=lambda: [], + static_key_func=lambda: {}, + run_key_func=lambda size: size, + benchmark_repeats=repeats, + benchmark_retries=retries, + )(lambda size, run_config=None: None) + + +def test_key_defaults_preserve_explicit_arguments_and_required_checks(): + @autotune( + kernel_name="key_defaults", + configs_gen_func=lambda: [], + static_key_func=lambda window: {"window": window}, + run_key_func=lambda size: size, + ) + def kernel(size, window=(-1, -1), run_config=None): + return run_config + + assert kernel._static_key(8) == {"window": (-1, -1)} + assert kernel._static_key(8, (511, 0)) == {"window": (511, 0)} + assert kernel._static_key(8, window=(1023, 0)) == {"window": (1023, 0)} + assert kernel._static_key(8, window=None) == {"window": None} + with pytest.raises(KeyError, match="Missing argument 'size'"): + kernel._run_key() + + +@pytest.mark.parametrize("level", [AutotuneLevel.ADAPTIVE_AUTOTUNE, AutotuneLevel.FORCE_AUTOTUNE]) +def test_two_warmup_phases_only_persist_matching_kernel_configs(tmp_path, monkeypatch, level): + monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: level) + general, _, general_benchmarks, general_cache = make_kernel(tmp_path, monkeypatch, "general") + decode, decode_calls, decode_benchmarks, decode_cache = make_kernel( + tmp_path, monkeypatch, "decode", AutotuneKernelType.DECODE_ATTENTION + ) + + with Autotuner.autotune_warmup(): + assert general(8) == {"block": 2} + assert decode(8) is None + assert len(general_benchmarks) == 2 + assert decode_benchmarks == [] + assert decode_calls == [(8, None)] + assert not decode_cache.exists() + general_cache_before = general_cache.read_bytes() + + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + assert general(16) == {"block": 2} + assert decode(16) == {"block": 2} + assert len(general_benchmarks) == 2 + assert len(decode_benchmarks) == 2 + assert general_cache.read_bytes() == general_cache_before + assert json.loads(general_cache.read_text()) == {"8": {"block": 2}} + assert json.loads(decode_cache.read_text()) == {"16": {"block": 2}} + + +def test_other_phase_warms_history_without_tuning_and_later_uses_new_config(tmp_path, monkeypatch): + decode, calls, benchmarks, cache_file = make_kernel( + tmp_path, monkeypatch, "decode", AutotuneKernelType.DECODE_ATTENTION + ) + cache_file.write_text(json.dumps({"8": {"block": 1}, "32": {"block": 2}})) + cache_before = cache_file.read_bytes() + + with Autotuner.autotune_warmup(): + assert decode(16) == {"block": 1} + # Historical configs are warmed in either phase, but only the matching phase searches new configs. + assert calls == [(16, {"block": 1}), (16, {"block": 2}), (16, {"block": 1})] + assert benchmarks == [] + assert cache_file.read_bytes() == cache_before + + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + assert decode(16) == {"block": 2} + assert len(benchmarks) == 2 + assert decode(16) == {"block": 2} + assert json.loads(cache_file.read_text())["16"] == {"block": 2} + + +@pytest.mark.parametrize("level", [AutotuneLevel.ADAPTIVE_AUTOTUNE, AutotuneLevel.FORCE_AUTOTUNE]) +def test_excluded_kernel_does_not_enter_distributed_tuning(tmp_path, monkeypatch, level): + general, calls, benchmarks, cache_file = make_kernel(tmp_path, monkeypatch, "general") + monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: level) + monkeypatch.setattr(autotuner_module.dist, "is_initialized", lambda: True) + monkeypatch.setattr(autotuner_module, "get_global_rank", lambda: 0) + monkeypatch.setattr(autotuner_module, "get_global_world_size", lambda: 2) + + def unexpected_collective(*args, **kwargs): + pytest.fail("An excluded kernel must not enter autotuning collectives") + + monkeypatch.setattr(general, "_get_autotune_group", unexpected_collective) + monkeypatch.setattr(autotuner_module.dist, "all_gather_object", unexpected_collective) + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + assert general(16) is None + assert calls == [(16, None)] + assert benchmarks == [] + assert not cache_file.exists() + + +@pytest.mark.parametrize("level", [0, 1, 2, 3]) +@pytest.mark.parametrize("kernel_type", [AutotuneKernelType.GENERAL, AutotuneKernelType.DECODE_ATTENTION]) +def test_matching_phase_respects_kernel_autotune_policy(tmp_path, monkeypatch, level, kernel_type): + monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: level) + kernel, _, benchmarks, cache_file = make_kernel(tmp_path, monkeypatch, kernel_type.value, kernel_type) + cache_file.write_text(json.dumps({"16": {"block": 1}})) + cache_before = cache_file.read_bytes() + with Autotuner.autotune_warmup(kernel_type): + result = kernel(16) + if level == AutotuneLevel.FORCE_AUTOTUNE and kernel_type == AutotuneKernelType.GENERAL: + assert result == {"block": 2} + assert len(benchmarks) == 2 + assert json.loads(cache_file.read_text()) == {"16": {"block": 2}} + else: + assert result == (None if level == AutotuneLevel.CLOSE_AUTOTUNE else {"block": 1}) + assert benchmarks == [] + assert cache_file.read_bytes() == cache_before + + +@pytest.mark.parametrize("kernel_type", [AutotuneKernelType.GENERAL, AutotuneKernelType.DECODE_ATTENTION]) +def test_force_autotune_reuses_decode_configs_across_layers(tmp_path, monkeypatch, kernel_type): + monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: AutotuneLevel.FORCE_AUTOTUNE) + kernel, _, benchmarks, cache_file = make_kernel(tmp_path, monkeypatch, kernel_type.value, kernel_type) + with Autotuner.autotune_warmup(kernel_type): + for size in [8, 16]: + benchmarks.clear() + assert kernel(size) == {"block": 2} + assert len(benchmarks) == 2 # A new run key still needs tuning. + for _ in range(3): + assert kernel(size) == {"block": 2} + assert len(benchmarks) == (2 if kernel_type == AutotuneKernelType.DECODE_ATTENTION else 8) + assert json.loads(cache_file.read_text()) == {"8": {"block": 2}, "16": {"block": 2}} + + +def test_history_is_warmed_on_load_or_during_autotune_warmup(tmp_path, monkeypatch): + kernel, calls, benchmarks, cache_file = make_kernel(tmp_path, monkeypatch, "general") + cache_file.write_text(json.dumps({"8": {"block": 1}, "32": {"block": 2}})) + assert kernel(16) == {"block": 1} + assert calls == [(16, {"block": 1}), (16, {"block": 2}), (16, {"block": 1})] + calls.clear() + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + assert kernel(16) == {"block": 1} + assert calls == [(16, {"block": 1})] + calls.clear() + assert kernel(16) == {"block": 1} + assert calls == [(16, {"block": 1})] + assert benchmarks == [] + + +def test_repeated_configs_are_skipped_after_success(tmp_path, monkeypatch): + kernel, calls, benchmarks, cache_file = make_kernel(tmp_path, monkeypatch, "general") + config = {"block": 1, "warps": 4} + cache_file.write_text(json.dumps({"8": config, "16": config, "32": {"block": 2}})) + with Autotuner.autotune_warmup(): + assert kernel(16) == config + assert calls == [(16, config), (16, {"block": 2}), (16, config)] + + calls.clear() + with Autotuner.autotune_warmup(): + assert kernel(16) == config + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + assert kernel(16) == config + assert kernel(16) == config + assert calls == [(16, config)] * 3 + assert benchmarks == [] + + +@pytest.mark.parametrize("level", [0, 1, 2]) +@pytest.mark.parametrize("phase", [None, AutotuneKernelType.GENERAL, AutotuneKernelType.DECODE_ATTENTION]) +def test_disabled_history_warmup_loads_cache_without_extra_state_updates(tmp_path, monkeypatch, level, phase): + monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: level) + calls = [] + + @autotune( + kernel_name="stateful_kernel", + kernel_type=AutotuneKernelType.DECODE_ATTENTION, + configs_gen_func=lambda: [{"block": 1}, {"block": 2}], + static_key_func=lambda: {}, + run_key_func=lambda state: state.numel(), + warmup_all_exist_config=False, + ) + def kernel(state, run_config=None): + calls.append(run_config) + state.add_(run_config["block"]) + return state + + kernel._cache_dir = str(tmp_path) + cache_file = tmp_path / "configs.json" + cache_file.write_text(json.dumps({"4": {"block": 1}, "8": {"block": 2}})) + cache_before = cache_file.read_bytes() + + def unexpected_warmup(*args, **kwargs): + pytest.fail("Disabling history warmup must skip its execution and argument cloning") + + monkeypatch.setattr(kernel, "kernel_warmup", unexpected_warmup) + monkeypatch.setattr(kernel, "_mutate_args_clone", unexpected_warmup) + state = torch.zeros(4) + scope = nullcontext() if phase is None else Autotuner.autotune_warmup(phase) + with scope: + for _ in range(2): + assert kernel(state) is state + assert calls == [{"block": 1}, {"block": 1}] + torch.testing.assert_close(state, torch.full((4,), 2.0)) + assert cache_file.read_bytes() == cache_before + assert not kernel.warmuped_configs_set + + +@pytest.mark.parametrize("kernel_type", [AutotuneKernelType.GENERAL, AutotuneKernelType.DECODE_ATTENTION]) +def test_disabled_history_warmup_still_searches_new_configs(tmp_path, monkeypatch, kernel_type): + kernel, calls, benchmarks, cache_file = make_kernel( + tmp_path, monkeypatch, kernel_type.value, kernel_type, warmup_all_exist_config=False + ) + with Autotuner.autotune_warmup(kernel_type): + for size in [8, 16]: + calls.clear() + assert kernel(size) == {"block": 2} + assert calls == [(size, {"block": 2})] + calls.clear() + assert kernel(size) == {"block": 2} + assert calls == [(size, {"block": 2})] + assert len(benchmarks) == 4 + assert json.loads(cache_file.read_text()) == {"8": {"block": 2}, "16": {"block": 2}} + assert not kernel.warmuped_configs_set + + +def test_failed_warmup_retries_only_during_autotune_warmup(tmp_path, monkeypatch): + kernel, calls, _, cache_file = make_kernel(tmp_path, monkeypatch, "general") + cache_file.write_text(json.dumps({"8": {"block": 1}, "32": {"block": 2}})) + + def shape_sensitive_kernel(size, run_config=None): + calls.append((size, run_config)) + if size < 16 and run_config["block"] == 2: + raise RuntimeError("This config requires a larger input") + return run_config + + monkeypatch.setattr(kernel, "fn", shape_sensitive_kernel) + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + assert kernel(8) == {"block": 1} + assert calls == [(8, {"block": 1}), (8, {"block": 2}), (8, {"block": 1})] + calls.clear() + assert kernel(16) == {"block": 1} + assert calls == [(16, {"block": 1})] + calls.clear() + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + assert kernel(16) == {"block": 1} + assert calls == [(16, {"block": 2}), (16, {"block": 1})] + calls.clear() + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + assert kernel(16) == {"block": 1} + assert calls == [(16, {"block": 1})] + + +def test_new_configs_are_warmed_on_later_calls_and_state_is_not_persisted(tmp_path, monkeypatch): + kernel, calls, benchmarks, cache_file = make_kernel(tmp_path, monkeypatch, "general") + cache_file.write_text(json.dumps({"8": {"block": 1}})) + with Autotuner.autotune_warmup(): + assert kernel(8) == {"block": 1} + calls.clear() + with Autotuner.autotune_warmup(): + assert kernel(16) == {"block": 2} + assert calls == [(16, {"block": 2})] + assert len(benchmarks) == 2 + assert json.loads(cache_file.read_text()) == {"8": {"block": 1}, "16": {"block": 2}} + + calls.clear() + assert kernel(16) == {"block": 2} + assert calls == [(16, {"block": 2})] + calls.clear() + with Autotuner.autotune_warmup(): + assert kernel(16) == {"block": 2} + assert calls == [(16, {"block": 2}), (16, {"block": 2})] + calls.clear() + with Autotuner.autotune_warmup(): + assert kernel(16) == {"block": 2} + assert calls == [(16, {"block": 2})] + + # A new autotuner loading the same file must warm its configs again. + reloaded, reload_calls, _, _ = make_kernel(tmp_path, monkeypatch, "reloaded") + reloaded._cache_dir = str(cache_file.parent) + assert reloaded(16) == {"block": 2} + assert len(reload_calls) == 3 + assert {call[1]["block"] for call in reload_calls[:2]} == {1, 2} + reload_calls.clear() + with Autotuner.autotune_warmup(): + assert reloaded(16) == {"block": 2} + assert reload_calls == [(16, {"block": 2})] + + +def test_warmup_preserves_mutated_input_and_skips_repeated_execution(): + executions = [] + + @autotune( + kernel_name="mutating_kernel", + configs_gen_func=lambda: [{"block": 1}], + static_key_func=lambda: {}, + run_key_func=lambda state: state.numel(), + mutates_args=["state"], + ) + def kernel(state, run_config=None): + executions.append(run_config) + state.add_(run_config["block"]) + + state = torch.zeros(4) + static_key = frozendict({}) + kernel.kernel_warmup(static_key, state, run_config={"block": 1}) + torch.testing.assert_close(state, torch.zeros(4)) + + kernel.kernel_warmup(static_key, state, run_config={"block": 1}) + assert executions == [{"block": 1}] + torch.testing.assert_close(state, torch.zeros(4)) + + +def test_explicit_config_bypasses_tuning(tmp_path, monkeypatch): + kernel, calls, benchmarks, cache_file = make_kernel(tmp_path, monkeypatch, "general") + with Autotuner.autotune_warmup(): + assert kernel(16, run_config={"block": 3}) == {"block": 3} + assert calls == [(16, {"block": 3})] + assert benchmarks == [] + assert not cache_file.exists() + + +@pytest.mark.parametrize("keyword_input", [False, True]) +@pytest.mark.parametrize("level", [AutotuneLevel.ADAPTIVE_AUTOTUNE, AutotuneLevel.FORCE_AUTOTUNE]) +def test_rebuild_inputs_only_for_benchmarking(tmp_path, monkeypatch, keyword_input, level): + monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: level) + rebuilds = [] + executions = [] + benchmarks = [] + original = torch.zeros(2) + + def rebuild(state, size, run_config=None): + rebuilds.append(state) + return (torch.ones(size),), {"size": size} + + @autotune( + kernel_name="rebuild_inputs", + kernel_type=AutotuneKernelType.DECODE_ATTENTION, + configs_gen_func=lambda: [{"block": 1}, {"block": 2}], + static_key_func=lambda state: {"input_size": state.numel()}, + run_key_func=lambda size: size, + rebuild_input_func=rebuild, + ) + def kernel(state, size, run_config=None): + executions.append(state) + return state + + def bench(state, size, run_config): + benchmarks.append(state) + torch.testing.assert_close(state, torch.ones(8)) + assert size == 8 + return 1.0 / run_config["block"] + + kernel._cache_dir = str(tmp_path) + monkeypatch.setattr(kernel, "_bench", bench) + args, kwargs = ((), {"state": original, "size": 8}) if keyword_input else ((original, 8), {}) + + assert kernel(*args, **kwargs) is original + with Autotuner.autotune_warmup(): + assert kernel(*args, **kwargs) is original + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + assert kernel(*args, **kwargs, run_config={"block": 3}) is original + assert not rebuilds + assert kernel(*args, **kwargs) is original + assert kernel(*args, **kwargs) is original # Exact cache hit only warms historical configs. + + assert kernel(*args, **kwargs) is original + assert len(rebuilds) == 1 and rebuilds[0] is original + assert len(benchmarks) == 2 and benchmarks[0] is benchmarks[1] + assert all(state is original for state in executions) + assert kernel.cached_configs == {frozendict({"input_size": 2}): {"8": {"block": 2}}} + torch.testing.assert_close(original, torch.zeros(2)) + + +def test_default_api_and_nested_phase_restore_after_exception(): + assert not Autotuner.is_autotune_warmup() + assert not Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.GENERAL) + assert not Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.DECODE_ATTENTION) + Autotuner.start_autotune_warmup() + assert Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.GENERAL) + assert not Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.DECODE_ATTENTION) + with pytest.raises(RuntimeError, match="test failure"): + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + assert Autotuner.is_autotune_warmup() + assert Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.DECODE_ATTENTION) + assert not Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.GENERAL) + with Autotuner.autotune_warmup(): + assert Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.GENERAL) + assert Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.DECODE_ATTENTION) + raise RuntimeError("test failure") + assert Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.GENERAL) + Autotuner.end_autotune_warmup() + assert not Autotuner.is_autotune_warmup() + assert not Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.GENERAL) + assert not Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.DECODE_ATTENTION) diff --git a/unit_tests/common/test_cache_copy_control_flow.py b/unit_tests/common/test_cache_copy_control_flow.py new file mode 100644 index 0000000000..26fe317052 --- /dev/null +++ b/unit_tests/common/test_cache_copy_control_flow.py @@ -0,0 +1,98 @@ +import ast +import __future__ +from pathlib import Path +import sys +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + + +ROOT = Path(__file__).resolve().parents[2] + + +def load_methods(path, class_name, method_names, optimize=0): + tree = ast.parse((ROOT / path).read_text()) + cls = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == class_name) + cls.body = [node for node in cls.body if isinstance(node, ast.FunctionDef) and node.name in method_names] + module = ast.Module(body=[cls], type_ignores=[]) + namespace = {"logger": Mock()} + exec(compile(module, path, "exec", flags=__future__.annotations.compiler_flag, optimize=optimize), namespace) + return namespace[class_name] + + +@pytest.mark.parametrize("optimize", [0, 1, 2]) +@pytest.mark.parametrize("prefix", [0, 64]) +def test_cpu_cache_restore_allocates_before_read(optimize, prefix): + cache_class = load_methods( + "lightllm/server/router/model_infer/mode_backend/multi_level_kv_cache.py", + "MultiLevelKvCacheModule", + {"read_req_from_cpu"}, + optimize, + ) + req = SimpleNamespace(cur_kv_len=prefix, hold_kv_len=prefix) + indices = list(range(prefix)) + + def allocate(count): + indices.extend(range(req.hold_kv_len, req.hold_kv_len + count)) + req.hold_kv_len += count + return indices[-count:] + + def get_indices(start, end): + selected = indices[start:end] + return SimpleNamespace(values=selected, numel=lambda: len(selected)) + + req.kv_cache_alloc_need = Mock(side_effect=lambda length: length - req.hold_kv_len) + req.alloc_kv_mem = Mock(side_effect=allocate) + req.token_index = SimpleNamespace(get_gpu_token_indexes=get_indices) + operator = SimpleNamespace(read_req=Mock(), get_page_layouts=lambda: []) + cache = cache_class() + cache.backend = SimpleNamespace( + model=SimpleNamespace(att_state_manager=SimpleNamespace(get_cpu_cache_operator=lambda: operator)), + is_hybrid_att_model=False, + ) + cache.cpu_cache_client = SimpleNamespace(cpu_kv_cache_tensor=object()) + + cache.read_req_from_cpu(req, [0], [256], 0) + + req.alloc_kv_mem.assert_called_once_with(256 - prefix) + operator.read_req.assert_called_once() + assert operator.read_req.call_args.kwargs["token_indices"].values == list(range(256)) + assert req.cur_kv_len == req.hold_kv_len == 256 + + +@pytest.mark.parametrize("return_fail_index", [False, True]) +@pytest.mark.parametrize("success", [False, True]) +def test_native_batch_return_formats(monkeypatch, return_fail_index, success): + copier = load_methods( + "lightllm/common/cuda_wrapper.py", + "PinnedH2DBatchCopier", + {"_copy_on_stream", "_try_native_batch_h2d"}, + ) + copier._native_batch_available = True + copier._protect_pinned_sources = Mock() + error = 0 if success else SimpleNamespace(name="cudaErrorInvalidValue") + result = (error, 0) if return_fail_index else (error,) + runtime = SimpleNamespace( + cudaMemcpyAttributes=SimpleNamespace, + cudaMemcpySrcAccessOrder=SimpleNamespace(cudaMemcpySrcAccessOrderStream=1), + cudaError_t=SimpleNamespace(cudaSuccess=0), + cudaMemcpyBatchAsync=Mock(return_value=result), + ) + monkeypatch.setitem(sys.modules, "cuda", SimpleNamespace()) + monkeypatch.setitem(sys.modules, "cuda.bindings", SimpleNamespace(runtime=runtime)) + source = SimpleNamespace(data_ptr=lambda: 100, numel=lambda: 256, element_size=lambda: 4) + target = SimpleNamespace(data_ptr=lambda: 200, copy_=Mock()) + stream = SimpleNamespace(cuda_stream=7) + + copier._copy_on_stream([target], [source], stream) + + runtime.cudaMemcpyBatchAsync.assert_called_once() + assert runtime.cudaMemcpyBatchAsync.call_args.args[:4] == ([200], [100], [1024], 1) + assert copier._native_batch_available is success + if success: + target.copy_.assert_not_called() + copier._protect_pinned_sources.assert_called_once_with([source], stream) + else: + target.copy_.assert_called_once_with(source, non_blocking=True) + copier._protect_pinned_sources.assert_not_called() diff --git a/unit_tests/common/test_hybrid_req_manager_pd_sharing.py b/unit_tests/common/test_hybrid_req_manager_pd_sharing.py new file mode 100644 index 0000000000..dc3bc15ecb --- /dev/null +++ b/unit_tests/common/test_hybrid_req_manager_pd_sharing.py @@ -0,0 +1,148 @@ +import multiprocessing +import traceback +from multiprocessing.reduction import ForkingPickler + +import pytest +import torch + +from lightllm.common.attention_state.bounded_att import GDNState, KDAState +from lightllm.server.router.model_infer.mode_backend.pd.p2p_fix import reduce_tensor +from unit_tests.common.test_linear_state import make_state + + +class _SerializationFailure: + def __reduce_ex__(self, protocol): + raise RuntimeError("injected serialization failure") + + +def _assert_received_manager_roundtrip(payload, state_name, mode, expected): + resources = ForkingPickler.loads(payload) + state = getattr(resources, state_name) + assert state.big_pages_pool is state.small_pages_pool is None + assert resources.full_att_allocator is None + layout = resources.get_pd_operator().get_pd_bounded_att_page_layouts()[state_name] + block = torch.empty((1, layout.byte_end), dtype=torch.uint8, device="cuda") + state.write_to_pd_page(1, block, 0, 8, layout) + state.read_from_pd_page(block, 2, 0, 8, layout) + width = state.config.conv_width + ssm_index = 2 if mode == "replay" else 2 * (state.mtp_step + 1) + actual = (state.conv_state[:, 2, :, :width], state.ssm_state[:, ssm_index]) + for tensor, values in zip(actual, expected): + torch.testing.assert_close(tensor, torch.tensor(values, dtype=tensor.dtype, device="cuda"), rtol=0, atol=0) + assert state.mtp_state_index[2] == 0 + assert state.mtp_state_index[1] == 1 + if mode == "replay": + assert state.cursors[2] == 0 + assert state.cursors[1] == 1 + torch.cuda.synchronize() + + +def _receive_pd_managers(connection): + torch.cuda.set_device(0) + while True: + message = connection.recv() + if message is None: + return + try: + _assert_received_manager_roundtrip(*message) + connection.send(None) + except Exception: + connection.send(traceback.format_exc()) + + +@pytest.fixture(scope="module") +def pd_receiver(): + context = multiprocessing.get_context("spawn") + sender, receiver = context.Pipe() + process = context.Process(target=_receive_pd_managers, args=(receiver,)) + process.start() + receiver.close() + try: + yield sender, process + finally: + if process.is_alive(): + sender.send(None) + process.join(timeout=30) + if process.is_alive(): + process.terminate() + process.join(timeout=5) + sender.close() + assert process.exitcode == 0 + + +@pytest.mark.parametrize("cls", [GDNState, KDAState]) +@pytest.mark.parametrize("mode", ["native", "replay"]) +def test_pd_export_copy_keeps_worker_pools_and_runtime_identity(cls, mode): + state, resources, _ = make_state(cls=cls, mode=mode, device="meta") + state.big_pages_pool, state.small_pages_pool = object(), object() + before = vars(state).copy() + exported = resources.get_pd_operator().create_ipc_copy() + copy = getattr(exported, state.state_name) + assert copy.big_pages_pool is copy.small_pages_pool is None + assert copy.conv_state is state.conv_state + assert copy.ssm_state is state.ssm_state + if mode == "replay": + assert copy.cursors is state.cursors + assert copy.keys is state.keys + assert copy.mtp_state_index is state.mtp_state_index + assert all(vars(state)[name] is value for name, value in before.items()) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA IPC and pinned memory required") +@pytest.mark.parametrize("cls", [GDNState, KDAState]) +@pytest.mark.parametrize("mode", ["native", "replay"]) +@pytest.mark.parametrize("fail_serialization", [False, True]) +def test_shared_attention_state_roundtrip_preserves_worker_resources( + runtime_start_args, monkeypatch, pd_receiver, cls, mode, fail_serialization +): + # Torch 2.11 uses pidfd_getfd for expandable segments, which is unavailable in + # some containers. Disable expandable segments before allocating the exported state. + torch.cuda.memory._set_allocator_settings("expandable_segments:False") + state, resources, _ = make_state(cls=cls, mode=mode) + state.initialize_page_pools(big_page_capacity=2, small_page_capacity=2) + state.conv_state.copy_(torch.arange(5, dtype=torch.bfloat16, device="cuda")) + for slot in range(state.ssm_state.shape[1]): + state.ssm_state[:, slot].fill_(slot + 1) + state.mtp_state_index[1], state.mtp_state_index[2] = 1, 2 + if mode == "replay": + state.raw_keys.zero_() + state.raw_values.zero_() + state.gates.fill_(-0.5) + state.betas.zero_() + state.cursors[1], state.cursors[2] = 1, 3 + width = state.config.conv_width + if mode == "replay": + expected = [ + state.conv_state[:, 1, :, :width].cpu().tolist(), + state.ssm_state[:, 1].cpu().tolist(), + ] + else: + expected = [ + state.conv_state[:, 1, :, :width].cpu().tolist(), + state.ssm_state[:, state.mtp_step + 1].cpu().tolist(), + ] + before = vars(state).copy() + source_slot = 1 if mode == "replay" else state.mtp_step + 1 + source_conv, source_ssm = state.conv_state[:, 1].clone(), state.ssm_state[:, source_slot].clone() + snapshots = [] + for pool in (state.big_pages_pool, state.small_pages_pool): + for tensor in (pool.conv_state, pool.ssm_state): + tensor.fill_(3) + snapshots.append((tensor, tensor.data_ptr(), tensor.clone())) + monkeypatch.setitem(ForkingPickler._extra_reducers, torch.Tensor, reduce_tensor) + exported = resources.get_pd_operator().create_ipc_copy() + if fail_serialization: + with pytest.raises(RuntimeError, match="injected serialization failure"): + ForkingPickler.dumps((_SerializationFailure(), exported)) + else: + sender, receiver = pd_receiver + sender.send((ForkingPickler.dumps(exported).tobytes(), state.state_name, mode, expected)) + assert sender.poll(30), f"PD receiver did not respond; exit code: {receiver.exitcode}" + error = sender.recv() + assert error is None, error + assert all(vars(state)[name] is value for name, value in before.items()) + torch.testing.assert_close(state.conv_state[:, 1], source_conv, rtol=0, atol=0) + torch.testing.assert_close(state.ssm_state[:, source_slot], source_ssm, rtol=0, atol=0) + for tensor, pointer, contents in snapshots: + assert tensor.data_ptr() == pointer + torch.testing.assert_close(tensor, contents, rtol=0, atol=0) diff --git a/unit_tests/common/test_linear_state.py b/unit_tests/common/test_linear_state.py new file mode 100644 index 0000000000..e538002a17 --- /dev/null +++ b/unit_tests/common/test_linear_state.py @@ -0,0 +1,334 @@ +from dataclasses import replace +from importlib import import_module + +import pytest +import torch + +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState, KDAState + + +def make_state(cls=GDNState, mode="native", step=2, dtype=torch.float32, device="cuda", config=None): + config = config or ConvSSMStateConfig( + logical_layer_ids=(0, 1), + global_k_heads=1, + global_v_heads=1 if cls is KDAState else 2, + key_dim=32, + value_dim=32, + conv_width=3, + conv_dtype=torch.bfloat16, + ssm_dtype=dtype, + ) + if mode == "replay": + from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState + from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDAState + + cls = ReplayKDAState if cls is KDAState else ReplayGDNState + state = cls(config, 4, mtp_step=step, device=device, ssm_replay_config={"cache_len": 4}) + else: + state = cls(config, 4, mtp_step=step, device=device) + resources = AttentionStateManager(request_capacity=4) + resources.register_state("main", state) + resources.finish_state_registration() + layout = resources.get_pd_operator().get_pd_bounded_att_page_layouts()[state.state_name] + return state, resources, layout + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("cls", [GDNState, KDAState]) +@pytest.mark.parametrize("mode,step", [("native", 0), ("native", 2), ("replay", 0), ("replay", 2)]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_checkpoint_and_pd_export_same_accepted_prefix_and_restore_reused_slot( + runtime_start_args, cls, mode, step, dtype +): + state, _, layout = make_state(cls, mode, step, dtype) + state.initialize_page_pools(big_page_capacity=2, small_page_capacity=2) + src, dst, offset = 1, 2, min(1, step) + state.conv_state.copy_(torch.arange(3 + step, device="cuda").to(torch.bfloat16)) + for slot in range(state.ssm_state.shape[1]): + state.ssm_state[:, slot].fill_(slot + 1) + if state.mtp_enabled: + state.mtp_state_index[src] = offset + if mode == "native": + expected_conv = state.conv_state[:, src, ..., :3].clone() + expected_ssm = state.ssm_state[:, src * (state.mtp_step + 1)].clone() + else: + # prefill 页只保存基准窗口和 canonical SSM，不按 MTP 偏移取窗口，也不折回 history。 + expected_conv = state.conv_state[:, src, ..., :3].clone() + expected_ssm = state.ssm_state[:, src].clone() + if mode == "replay": + state.raw_keys.zero_() + state.raw_values.zero_() + state.gates.fill_(-0.5) + state.betas.zero_() + state.cursors[src] = 1 + before = state.ssm_state.clone() + block = torch.empty((1, layout.byte_end), dtype=torch.uint8, device="cuda") + state.write_to_pd_page(src, block, 0, 8, layout) + for kind in ("small", "big"): + state.save_bounded_att_runtime_to_pages([src], [8], kind, [0]) + torch.cuda.synchronize() + for pool in (state.small_pages_pool, state.big_pages_pool): + torch.testing.assert_close(pool.conv_state[0], expected_conv.cpu(), rtol=0, atol=0) + torch.testing.assert_close(pool.ssm_state[0], expected_ssm.cpu(), rtol=2e-5, atol=2e-5) + assert torch.equal(state.ssm_state, before) + if mode == "replay": + assert state.cursors[src].item() == 1 + + for restore in ( + lambda: state.restore_bounded_att_runtime_from_pages([dst], [8], "small", [0]), + lambda: state.restore_bounded_att_runtime_from_pages([dst], [8], "big", [0]), + lambda: state.read_from_pd_page(block, dst, 0, 8, layout), + lambda: state.copy_runtime_state_from_dp( + src_req_indices=[src], dst_req_indices=[dst], remote_dp_bounded_state=state + ), + ): + if state.mtp_enabled: + state.mtp_state_index[dst] = step + if mode == "replay": + state.cursors[dst] = 3 + restore() + actual_conv = state.conv_state[:, dst, ..., :3] + if mode == "replay": + actual_ssm = state.ssm_state[:, dst] + else: + actual_ssm = state.ssm_state[:, dst * (state.mtp_step + 1)] + torch.testing.assert_close(actual_conv, expected_conv, rtol=0, atol=0) + torch.testing.assert_close(actual_ssm, expected_ssm, rtol=2e-5, atol=2e-5) + if state.mtp_enabled: + assert state.mtp_state_index[dst].item() == 0 + if mode == "replay": + assert state.cursors[dst].item() == 0 + assert state.cursors[src].item() == 1 + state.reset_bounded_att_state([dst]) + assert not state.conv_state[:, dst].any() + if mode == "replay": + assert not state.ssm_state[:, dst].any() + else: + start = dst * (state.mtp_step + 1) + assert not state.ssm_state[:, start : start + state.mtp_step + 1].any() + + +def copy_on_cpu(*, copies): + for source, destination in copies: + destination.copy_(source) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("cls", [GDNState, KDAState]) +@pytest.mark.parametrize("mode", ["native", "replay"]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("destination", ["cpu", "cuda"]) +def test_checkpoint_batch_export(runtime_start_args, cls, mode, dtype, destination): + state, _, _ = make_state(cls, mode, 2, dtype) + state.initialize_page_pools(3, device=destination) + state.conv_state.copy_(torch.arange(5, device="cuda").to(torch.bfloat16)) + for slot in range(state.ssm_state.shape[1]): + state.ssm_state[:, slot].fill_(slot + 1) + state.mtp_state_index.copy_(torch.tensor([0, 2, 1, 0], device="cuda")) + if mode == "replay": + state.raw_keys.zero_() + state.raw_values.zero_() + state.gates.fill_(-0.5) + state.betas.zero_() + state.cursors.copy_(torch.tensor([1, 2, 3, 0], device="cuda")) + requests = [2, 0, 1] + width = state.config.conv_width + if mode == "replay": + expected = [(state.conv_state[:, req, :, :width].clone(), state.ssm_state[:, req].clone()) for req in requests] + else: + expected = [ + ( + state.conv_state[:, req, :, :width].clone(), + state.ssm_state[:, req * (state.mtp_step + 1)].clone(), + ) + for req in requests + ] + before = state.ssm_state.clone() + history = getattr(state, "runtime", state) + cursors = history.cursors.clone() if mode == "replay" else None + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + state.save_bounded_att_runtime_to_pages(requests, [8] * 3, "big", [1, 2, 0]) + stream.synchronize() + pool = state.big_pages_pool + for slot, (conv, ssm) in zip([1, 2, 0], expected): + torch.testing.assert_close(pool.conv_state[slot].cpu(), conv.cpu(), atol=0, rtol=0) + torch.testing.assert_close(pool.ssm_state[slot].cpu(), ssm.cpu(), atol=0, rtol=0) + assert torch.equal(state.ssm_state, before) + if cursors is not None: + assert torch.equal(state.cursors, cursors) + target, _, _ = make_state(cls, mode, 2, dtype) + if target.mtp_enabled: + target.mtp_state_index.fill_(2) + target.copy_runtime_state_from_dp(requests, [0, 1, 2], state) + for req, (conv, ssm) in enumerate(expected): + actual_conv = target.conv_state[:, req, :, : target.config.conv_width] + if mode == "replay": + actual_ssm = target.ssm_state[:, req] + else: + actual_ssm = target.ssm_state[:, req * (target.mtp_step + 1)] + torch.testing.assert_close(actual_conv, conv, atol=0, rtol=0) + torch.testing.assert_close(actual_ssm, ssm, atol=0, rtol=0) + assert torch.equal(state.ssm_state, before) + if cursors is not None: + assert torch.equal(state.cursors, cursors) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_native_checkpoint_scratch_is_bounded(runtime_start_args): + config = ConvSSMStateConfig( + logical_layer_ids=(0, 1, 2, 3), + global_k_heads=1, + global_v_heads=2, + key_dim=128, + value_dim=128, + conv_width=3, + conv_dtype=torch.bfloat16, + ssm_dtype=torch.float32, + ) + state = GDNState(config, 65, mtp_step=2, device="cuda") + state.initialize_page_pools(64) + requests = list(range(64)) + state.save_bounded_att_runtime_to_pages(requests, [128] * 64, "big", requests) + torch.cuda.synchronize() + before = torch.cuda.memory_allocated() + torch.cuda.reset_peak_memory_stats() + state.save_bounded_att_runtime_to_pages(requests, [128] * 64, "big", requests) + torch.cuda.synchronize() + assert torch.cuda.max_memory_allocated() - before < 1024 * 1024 + + +def test_pd_block_roundtrip_preserves_replicated_heads_and_accepted_mtp_state(monkeypatch): + # Exercise layout and accepted-state selection independently of the CUDA copy kernel. + monkeypatch.setattr( + import_module("lightllm.common.attention_state.bounded_att.gdn.gdn_state"), + "copy_bounded_att_tensors", + copy_on_cpu, + ) + config = ConvSSMStateConfig( + logical_layer_ids=(0,), + global_k_heads=1, + global_v_heads=2, + key_dim=2, + value_dim=3, + conv_width=3, + conv_dtype=torch.float32, + ssm_dtype=torch.float32, + tp_world_size=4, + ) + states, layouts, snapshots = [], [], [] + src, dst = 1, 2 + for rank in range(4): + state, _, layout = make_state(device="cpu", config=replace(config, tp_rank=rank)) + state.conv_state.fill_(-1) + state.ssm_state.fill_(-2) + state.mtp_state_index.fill_(1) + state.conv_state[:, src, :2, :3].fill_(10 + rank) + state.conv_state[:, src, 2:4, :3].fill_(20 + rank) + state.conv_state[:, src, 4:, :3].fill_(30 + rank) + state.ssm_state[:, src * (state.mtp_step + 1)].fill_(40 + rank) + states.append(state) + layouts.append(layout) + snapshots.append((state.conv_state.clone(), state.ssm_state.clone())) + size = layouts[0].byte_end + block = torch.full((1, size + 16), 0xA5, dtype=torch.uint8) + for state, layout in zip(states, layouts): + state.write_to_pd_page(src, block, 0, 8, layout) + expected_conv = torch.cat( + [torch.full((1, channels, 3), value) for channels, value in ((2, 10), (2, 20), (3, 30), (3, 32))], + dim=1, + ).float() + expected_ssm = torch.stack((torch.full((1, 2, 3), 40), torch.full((1, 2, 3), 42)), dim=1).float() + conv_bytes = expected_conv.numel() * 4 + torch.testing.assert_close(block[0, :conv_bytes].view(torch.float32).view_as(expected_conv), expected_conv) + torch.testing.assert_close(block[0, conv_bytes:size].view(torch.float32).view_as(expected_ssm), expected_ssm) + assert torch.all(block[0, size:] == 0xA5) + exported = block.clone() + other_requests = torch.tensor([0, 1, 3]) + for rank, (state, layout, (conv_before, ssm_before)) in enumerate(zip(states, layouts, snapshots)): + torch.testing.assert_close(state.conv_state, conv_before) + torch.testing.assert_close(state.ssm_state, ssm_before) + state.read_from_pd_page(block, dst, 0, 8, layout) + expected_local_conv = torch.cat( + (expected_conv[:, :4], expected_conv[:, 4 + rank // 2 * 3 : 7 + rank // 2 * 3]), dim=1 + ) + torch.testing.assert_close(state.conv_state[:, dst, ..., :3], expected_local_conv) + torch.testing.assert_close(state.ssm_state[:, dst * 3], expected_ssm[:, rank // 2 : rank // 2 + 1]) + assert state.mtp_state_index.tolist() == [1, 1, 0, 1] + torch.testing.assert_close(state.conv_state[:, other_requests], conv_before[:, other_requests]) + other_slots = (other_requests[:, None] * 3 + torch.arange(3)).flatten() + torch.testing.assert_close(state.ssm_state[:, other_slots], ssm_before[:, other_slots]) + assert torch.equal(block, exported) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("cls", [GDNState, KDAState]) +@pytest.mark.parametrize("step", [0, 2]) +@pytest.mark.parametrize("restore_kind", ["checkpoint", "pd", "dp"]) +def test_replay_restored_slot_decodes_like_clean_slot(runtime_start_args, cls, step, restore_kind): + from unit_tests.common.basemodel.triton_kernel.linear_att.replay_history import replay_layer + + torch.manual_seed(1690) + state, _, layout = make_state(cls=cls, mode="replay", step=step) + state.initialize_page_pools(big_page_capacity=1) + state.ssm_state.normal_(mean=0, std=0.02) + state.conv_state.normal_(mean=0, std=0.1) + block = torch.empty((1, layout.byte_end), dtype=torch.uint8, device="cuda") + state.save_bounded_att_runtime_to_pages(req_indices=[0], state_token_lens=[8], page_kind="big", page_slot_ids=[0]) + state.write_to_pd_page(req_idx=0, pd_page_buffers=block, pd_page_index=0, state_token_len=8, layout=layout) + hv, kd, vd = state.ssm_state.shape[-3:] + kh = state.num_key_heads + width = step + 1 + q = torch.randn(width, kh, kd, device="cuda", dtype=torch.bfloat16).repeat(2, 1, 1) + k = torch.randn_like(q[:width]).repeat(2, 1, 1) + v = torch.randn(width, hv, vd, device="cuda", dtype=torch.bfloat16).repeat(2, 1, 1) + a = torch.full((2 * width, hv * kd if state.kda else hv), -3.0, device="cuda") + b = torch.zeros(2 * width, hv, device="cuda") + alog = torch.zeros(hv, device="cuda") + bias = torch.zeros(a.shape[-1], device="cuda") + reqs = torch.tensor([1, 2], dtype=torch.int32, device="cuda") + cu = torch.arange(3, dtype=torch.int32, device="cuda") * width + + def decode(): + outputs = [] + for layer in state.config.logical_layer_ids: + outputs.append( + replay_layer( + state=state, + layer_idx=layer, + q=q, + k=k, + v=v, + a=a, + b=b, + a_log=alog, + bias=bias, + reqs=reqs, + cu_seqlens=cu, + run_config={"BV": 32, "num_warps": 1, "num_stages": 3}, + ) + ) + state.finish_decode_token(logical_layer_id=layer, reqs=reqs, cu_seqlens=cu) + return outputs + + decode() + assert state.cursors[1].item() == width + state.reset_bounded_att_state(req_indices=[2]) + state.restore_bounded_att_runtime_from_pages( + req_indices=[2], state_token_lens=[8], page_kind="big", page_slot_ids=[0] + ) + if restore_kind == "checkpoint": + state.restore_bounded_att_runtime_from_pages( + req_indices=[1], state_token_lens=[8], page_kind="big", page_slot_ids=[0] + ) + elif restore_kind == "pd": + state.read_from_pd_page(pd_page_buffers=block, req_idx=1, pd_page_index=0, state_token_len=8, layout=layout) + else: + state.copy_runtime_state_from_dp(src_req_indices=[0], dst_req_indices=[1], remote_dp_bounded_state=state) + assert state.cursors[1].item() == 0 + for _ in range(3): + for output in decode(): + torch.testing.assert_close(output[:width], output[width:], rtol=0, atol=0) + torch.testing.assert_close(state.ssm_state[:, 1], state.ssm_state[:, 2], rtol=0, atol=0) diff --git a/unit_tests/common/test_linear_state_contract.py b/unit_tests/common/test_linear_state_contract.py new file mode 100644 index 0000000000..fbfc4a43a9 --- /dev/null +++ b/unit_tests/common/test_linear_state_contract.py @@ -0,0 +1,241 @@ +import pytest +import torch + +from lightllm.common.attention_state.bounded_att import KDAState +from unit_tests.common.test_linear_state import make_state +from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_cursor import advance_cursor +from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState +from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ReplayKDAState +from lightllm.models.glm5_next.config import Glm5NextResourceConfig +from lightllm.server.core.objs.start_args_type import StartArgs + + +@pytest.fixture +def linear_config(): + return ConvSSMStateConfig( + logical_layer_ids=(0, 1, 3, 4, 6, 7), + global_k_heads=1, + global_v_heads=2, + key_dim=32, + value_dim=32, + conv_width=3, + conv_dtype=torch.bfloat16, + ssm_dtype=torch.float32, + ) + + +@pytest.mark.parametrize("layer,row", [(0, 0), (1, 1), (3, 2), (4, 3), (6, 4), (7, 5)]) +def test_layer_buffers_follow_packed_linear_layers(linear_config, layer, row): + owner = GDNState(linear_config, request_capacity=1, device="cpu") + for i in range(6): + owner.conv_state[i].fill_(i) + owner.ssm_state[i].fill_(i + 10) + actual_conv, actual_ssm = owner.get_layer_view(layer) + assert (actual_conv == row).all() + assert (actual_ssm == row + 10).all() + + +@pytest.mark.parametrize("layer", [-1, 2, 5, 8]) +def test_linear_mapping_rejects_full_attention_and_out_of_range_layers(linear_config, layer): + owner = GDNState(linear_config, request_capacity=1, device="cpu") + with pytest.raises(KeyError): + owner.get_layer_view(layer) + + +def test_model_configuration_retains_nondefault_kda_gate_bound(): + 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], + "gate_lower_bound": -2.5, + }, + }, + StartArgs(tp=4, data_type="bfloat16"), + ) + assert config.gate_lower_bound == -2.5 + + +def test_kda_backend_receives_gate_bound_at_construction(linear_config): + owner = ReplayKDAState(linear_config, 3, mtp_step=2, gate_lower_bound=-2.5, device="cpu") + assert owner.gate_lower_bound == -2.5 + + +@pytest.mark.parametrize("cls", [GDNState, KDAState]) +@pytest.mark.parametrize("mtp_step,expected", [(0, "inline"), (2, "precompute")]) +def test_replay_defaults_to_precompute_for_speculative_decode(linear_config, cls, mtp_step, expected): + replay_cls = ReplayKDAState if cls is KDAState else ReplayGDNState + owner = replay_cls(linear_config, request_capacity=2, mtp_step=mtp_step, device="cpu") + history = getattr(owner, "runtime", owner) + assert history.compute_mode == expected + assert history.verify_width == mtp_step + 1 + + +@pytest.mark.parametrize("cls", [GDNState, KDAState]) +def test_mtp0_decode_keeps_explicit_precompute(linear_config, cls): + replay_cls = ReplayKDAState if cls is KDAState else ReplayGDNState + owner = replay_cls( + linear_config, + request_capacity=2, + mtp_step=0, + device="cpu", + ssm_replay_config={"cache_len": 8, "compute_mode": "precompute"}, + ) + assert owner.verify_width == 1 + assert owner.compute_mode == "precompute" + assert owner.capacity == 8 + + +class _CaptureKernel: + def __init__(self): + self.args = None + + def __getitem__(self, grid): + def launch(*args, **kwargs): + self.args = args + self.kwargs = kwargs + + return launch + + +@pytest.mark.parametrize( + "module_name,cls,bv,precompute", + [ + ("replay_gdn_decode", ReplayGDNState, 32, True), + ("replay_gdn_decode", ReplayGDNState, 8, False), + ("replay_kda_decode", ReplayKDAState, 32, True), + ("replay_kda_decode", ReplayKDAState, 8, False), + ], +) +def test_mtp0_precompute_reaches_replay_kernel(monkeypatch, linear_config, module_name, cls, bv, precompute): + import importlib + + decode = importlib.import_module(f"lightllm.common.basemodel.attention.linear.{module_name}") + captured = _CaptureKernel() + monkeypatch.setattr(decode, "_decode_replay", captured) + owner = cls( + linear_config, + request_capacity=2, + mtp_step=0, + device="cpu", + ssm_replay_config={"cache_len": 8, "compute_mode": "precompute"}, + ) + heads = owner.num_key_heads + q = torch.zeros(1, heads, linear_config.key_dim) + v = torch.zeros(1, owner.ssm_state.shape[-3], linear_config.value_dim) + gate = v.shape[-3] * linear_config.key_dim if cls is ReplayKDAState else v.shape[-3] + run_config = {"BV": bv, "num_warps": 1, "num_stages": 1} + if cls is ReplayKDAState: + decode.replay( + state_ssm=owner.ssm_state[0], + state_keys=owner.keys[0], + state_deltas=owner.deltas[0], + state_gates=owner.gates[0], + state_raw_keys=owner.raw_keys[0], + state_raw_values=owner.raw_values[0], + state_betas=owner.betas[0], + state_cursors=owner.cursors, + state_tile_sync_mark=owner.tile_sync_mark, + q=q, + k=q, + v=v, + a=torch.zeros(1, gate), + b=torch.zeros(1, v.shape[-3]), + a_log=torch.zeros(v.shape[-3]), + bias=torch.zeros(gate), + reqs=torch.zeros(1, dtype=torch.int32), + cu_seqlens=torch.arange(2, dtype=torch.int32), + state_verify_width=owner.verify_width, + state_compute_mode=owner.compute_mode, + state_gate_lower_bound=owner.gate_lower_bound, + run_config=run_config, + ) + else: + decode.replay( + owner, + 0, + q, + q, + v, + torch.zeros(1, gate), + torch.zeros(1, v.shape[-3]), + torch.zeros(v.shape[-3]), + torch.zeros(gate), + torch.zeros(1, dtype=torch.int32), + torch.arange(2, dtype=torch.int32), + run_config=run_config, + ) + assert captured.kwargs["PRECOMPUTE_PROJECTIONS"] is precompute + assert captured.kwargs["KDA"] is (cls is ReplayKDAState) + + +@pytest.mark.parametrize("cls", [GDNState, KDAState]) +def test_native_commit_selects_mtp_slot(cls): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + owner, _, _ = make_state(cls, step=2) + accepted = torch.tensor([3], dtype=torch.int32, device="cuda") + owner.commit_bounded_att_state(torch.tensor([1], device="cuda"), accepted, accepted) + assert owner.mtp_state_index.tolist() == [0, 2, 0, 0] + + +@pytest.mark.parametrize("cls", [GDNState, KDAState]) +def test_replay_fused_acceptance_keeps_cursors_inside_mode(monkeypatch, cls): + import sys + + owner, _, _ = make_state(cls, mode="replay", step=3, device="cpu") + history = getattr(owner, "runtime", owner) + calls = [] + monkeypatch.setattr(sys.modules[type(history).__module__], "commit_ssm_prefix", lambda *args: calls.append(args)) + widths = torch.tensor([1, 2, 3]) + owner.commit_bounded_att_state(torch.tensor([2, 0, 1]), torch.tensor([1, 2, 3]), widths) + assert len(calls) == 1 + assert calls[0][3] is widths + assert calls[0][4] is history.cursors + assert calls[0][0] is owner.mtp_state_index + + +def _commit_prefix_cpu(indices, reqs, lengths, req_verify_width, replay_cursors) -> None: + """CPU 上按请求行核对索引和 cursor。正式提交算子只跑 CUDA。""" + + indices[reqs] = (lengths - 1).to(indices.dtype) + hold = replay_cursors.numel() - 1 + for row in range(reqs.numel()): + req = int(reqs[row]) + if req == hold: + continue + rejected = int(req_verify_width[row]) - int(lengths[row]) + replay_cursors[req] = int(replay_cursors[req]) - rejected + + +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_acceptance_keeps_request_order_with_strided_inputs(device): + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA required") + owner, _, _ = make_state(mode="replay", step=3, device=device) + owner.mtp_state_index.fill_(2) + history = getattr(owner, "runtime", owner) + history.cursors.zero_() + reqs = torch.tensor([2, -99, 0, -99, 3, -99], device=device)[::2] + lengths = torch.tensor([1, -99, 3, -99, 2, -99], dtype=torch.int32, device=device)[::2] + if device == "cuda": + flat_lengths = lengths.contiguous() + cu = torch.zeros(flat_lengths.numel() + 1, dtype=torch.int32, device=device) + cu[1:] = torch.cumsum(flat_lengths, 0) + advance_cursor(history, reqs.contiguous(), cu) + owner.commit_bounded_att_state(reqs, lengths, lengths) + else: + # 正式提交只接受 CUDA。CPU 用例只核对 stride 下的请求顺序。 + history.cursors[0] = 3 + history.cursors[2] = 1 + _commit_prefix_cpu(owner.mtp_state_index, reqs, lengths, lengths, history.cursors) + assert owner.mtp_state_index.tolist() == [2, 2, 0, 1] + assert history.cursors.tolist() == [3, 0, 1, 0] diff --git a/unit_tests/common/test_mem_manager_ipc.py b/unit_tests/common/test_mem_manager_ipc.py new file mode 100644 index 0000000000..523b61f82f --- /dev/null +++ b/unit_tests/common/test_mem_manager_ipc.py @@ -0,0 +1,114 @@ +from multiprocessing.reduction import ForkingPickler + +import pytest +import torch +import torch.multiprocessing as mp + +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.allocators import full_att_kv_allocator as allocator_module +from lightllm.common.attention_state.full_att import GQAAttentionConfig, GQAFullAttentionState + + +class _Counter: + def __init__(self, *args): + pass + + def set_value(self, value): + self.value = value + + +def _make_manager(monkeypatch, size, device="cpu"): + if device == "cuda": + # The spawn receiver reopens CUDA IPC handles. Recent torch releases use + # pidfd_getfd for expandable segments, which some containers do not expose. + torch.cuda.memory._set_allocator_settings("expandable_segments:False") + monkeypatch.setattr(allocator_module, "SharedInt", _Counter) + monkeypatch.setattr(allocator_module, "get_current_rank_in_node", lambda: 0) + manager = AttentionStateManager(request_capacity=2) + state = GQAFullAttentionState( + GQAAttentionConfig(logical_layer_groups=((0,),), global_head_num=1, head_dim=4, dtype=torch.float32), + device=device, + ) + state.allocate_storage(4) + manager.register_state("main", state) + manager.finish_state_registration() + manager.full_att_allocator = allocator_module.FullAttKVAllocator(size, pin_memory=False) + mapping = torch.arange(4, dtype=torch.int32, device=device).reshape(1, 4) + return manager, mapping + + +def _export(manager, mapping, kind): + operator = manager.get_dp_shared_operator() if kind == "dp" else manager.get_pd_operator() + return ForkingPickler.dumps((operator.create_ipc_copy(), mapping)) + + +@pytest.mark.parametrize("kind", ["dp", "pd"]) +def test_ipc_descriptor_size_does_not_scale_with_local_allocator(monkeypatch, kind): + small, small_mapping = _make_manager(monkeypatch, 256) + large, large_mapping = _make_manager(monkeypatch, 131072) + small_payload = _export(small, small_mapping, kind) + large_payload = _export(large, large_mapping, kind) + assert abs(len(large_payload) - len(small_payload)) < 128 + + +@pytest.mark.parametrize("kind", ["dp", "pd"]) +def test_ipc_roundtrip_retains_shared_buffers_and_local_owner(monkeypatch, kind): + source, mapping = _make_manager(monkeypatch, 256) + allocator = source.full_att_allocator + exported, exported_mapping = ForkingPickler.loads(_export(source, mapping, kind)) + assert source.full_att_allocator is allocator + assert allocator.alloc(1).tolist() == [0] + assert exported.full_att_allocator is None + assert exported.get_dp_shared_operator().owner is exported + assert exported.get_pd_operator().owner is exported + state = exported.get_full_attention_states()[0] + state.k_storage.fill_(17) + state.v_storage.fill_(19) + exported_mapping.fill_(23) + source_state = source.get_full_attention_states()[0] + assert torch.all(source_state.k_storage == 17) + assert torch.all(source_state.v_storage == 19) + assert torch.all(mapping == 23) + + +def _update_exported_buffers(payload, connection): + try: + exported, mapping = ForkingPickler.loads(payload) + assert exported.full_att_allocator is None + state = exported.get_full_attention_states()[0] + state.k_storage.fill_(17) + state.v_storage.fill_(19) + mapping.fill_(23) + if mapping.is_cuda: + torch.cuda.synchronize() + connection.send("complete") + finally: + connection.close() + + +@pytest.mark.parametrize("kind", ["dp", "pd"]) +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_ipc_buffers_remain_shared_between_processes(monkeypatch, device, kind): + source, mapping = _make_manager(monkeypatch, 256, device) + allocator = source.full_att_allocator + context = mp.get_context("spawn") + receiver, sender = context.Pipe(duplex=False) + process = context.Process(target=_update_exported_buffers, args=(_export(source, mapping, kind).tobytes(), sender)) + try: + process.start() + sender.close() + assert receiver.poll(60), "IPC consumer did not complete" + assert receiver.recv() == "complete" + process.join(10) + assert process.exitcode == 0 + assert source.full_att_allocator is allocator + state = source.get_full_attention_states()[0] + assert torch.all(state.k_storage == 17) + assert torch.all(state.v_storage == 19) + assert torch.all(mapping == 23) + finally: + if process.is_alive(): + process.terminate() + process.join(10) + receiver.close() + sender.close() diff --git a/unit_tests/common/test_req_manager_hold_page.py b/unit_tests/common/test_req_manager_hold_page.py new file mode 100644 index 0000000000..a5125f9326 --- /dev/null +++ b/unit_tests/common/test_req_manager_hold_page.py @@ -0,0 +1,27 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.req_manager import ReqManager + + +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) + req_manager.page_size = 4 + req_manager.bind_hold_page((100, 101, 102, 103)) + + 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() + + +def test_hold_request_indexes_require_all_requests_released(): + req_manager = ReqManager.__new__(ReqManager) + req_manager.req_list = SimpleNamespace(is_all_free=lambda: False) + + with pytest.raises(AssertionError, match="all requests are released"): + req_manager.bind_hold_page((100, 101, 102, 103)) diff --git a/unit_tests/models/deepseek3_2/triton_kernel/test_hadamard_transform.py b/unit_tests/models/deepseek3_2/triton_kernel/test_hadamard_transform.py new file mode 100644 index 0000000000..8a54d6d9fd --- /dev/null +++ b/unit_tests/models/deepseek3_2/triton_kernel/test_hadamard_transform.py @@ -0,0 +1,83 @@ +import pytest +import torch +import triton + +from lightllm.models.deepseek3_2.triton_kernel.hadamard_transform import hadamard_transform + + +TP = 8 +INDEX_N_HEADS = 64 +INDEX_HEAD_DIM = 128 +TP_INDEX_N_HEADS = INDEX_N_HEADS // TP +SCALE = INDEX_HEAD_DIM ** -0.5 + + +def _get_sgl_kernel_hadamard_transform(): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for hadamard_transform comparison") + try: + from sgl_kernel import hadamard_transform as sgl_hadamard_transform + except ImportError: + pytest.skip("sgl_kernel.hadamard_transform is not available") + return sgl_hadamard_transform + + +def _bench(fn, x): + ms = triton.testing.do_bench_cudagraph(lambda: fn(x, scale=SCALE), return_mode="median") + return ms, fn(x, scale=SCALE) + + +@pytest.mark.parametrize("tokens", [1, 16, 128, 512, 1024, 2048, 4096, 8192, 16384]) +def test_hadamard_transform_matches_sgl_kernel_deepseek_v32_shapes(tokens): + sgl_hadamard_transform = _get_sgl_kernel_hadamard_transform() + + q = torch.randn(tokens, TP_INDEX_N_HEADS, INDEX_HEAD_DIM, dtype=torch.bfloat16, device="cuda") + k = torch.randn(tokens, INDEX_HEAD_DIM, dtype=torch.bfloat16, device="cuda") + + q_expected = sgl_hadamard_transform(q, scale=SCALE) + q_actual = hadamard_transform(q, scale=SCALE) + k_expected = sgl_hadamard_transform(k, scale=SCALE) + k_actual = hadamard_transform(k, scale=SCALE) + torch.cuda.synchronize() + + assert torch.equal(q_actual, q_expected) + assert torch.equal(k_actual, k_expected) + + +def test_hadamard_transform_perf_report_deepseek_v32_shapes(): + sgl_hadamard_transform = _get_sgl_kernel_hadamard_transform() + + print( + "\nDeepSeek-V3.2 per-rank shapes with tp=8:" + "\n q: [tokens, 8, 128]" + "\n k: [tokens, 128]" + "\n\ntokens | q_diff | k_diff | sgl_q ms | tri_q ms | sgl_k ms | tri_k ms | tri(q+k) ms | slowdown q+k" + ) + + for tokens in [1, 16, 128, 512, 1024, 2048, 4096, 8192, 16384]: + q = torch.randn(tokens, TP_INDEX_N_HEADS, INDEX_HEAD_DIM, dtype=torch.bfloat16, device="cuda") + k = torch.randn(tokens, INDEX_HEAD_DIM, dtype=torch.bfloat16, device="cuda") + + q_expected = sgl_hadamard_transform(q, scale=SCALE) + q_actual = hadamard_transform(q, scale=SCALE) + k_expected = sgl_hadamard_transform(k, scale=SCALE) + k_actual = hadamard_transform(k, scale=SCALE) + torch.cuda.synchronize() + + q_diff = (q_expected.float() - q_actual.float()).abs().max().item() + k_diff = (k_expected.float() - k_actual.float()).abs().max().item() + sgl_q_ms, _ = _bench(sgl_hadamard_transform, q) + tri_q_ms, _ = _bench(hadamard_transform, q) + sgl_k_ms, _ = _bench(sgl_hadamard_transform, k) + tri_k_ms, _ = _bench(hadamard_transform, k) + sgl_sum_ms = sgl_q_ms + sgl_k_ms + tri_sum_ms = tri_q_ms + tri_k_ms + + print( + f"{tokens:6d} | {q_diff:6.1g} | {k_diff:6.1g} | " + f"{sgl_q_ms:8.4f} | {tri_q_ms:8.4f} | {sgl_k_ms:8.4f} | {tri_k_ms:8.4f} | " + f"{tri_sum_ms:11.4f} | {tri_sum_ms / sgl_sum_ms:10.2f}x" + ) + + assert q_diff == 0 + assert k_diff == 0 diff --git a/unit_tests/models/deepseek3_2/triton_kernel/test_topk_index_to_mem_index.py b/unit_tests/models/deepseek3_2/triton_kernel/test_topk_index_to_mem_index.py new file mode 100644 index 0000000000..5e3aaf38d0 --- /dev/null +++ b/unit_tests/models/deepseek3_2/triton_kernel/test_topk_index_to_mem_index.py @@ -0,0 +1,24 @@ +import torch +import pytest +from lightllm.models.deepseek3_2.triton_kernel.topk_index_to_mem_index import trans_topk_index_to_mem_index + + +def test_trans_topk_index_to_mem_index(): + """Test trans_topk_index_to_mem_index converts topk indices to memory indices correctly.""" + batch_size = 1 + topk = 2048 + + # Create topk_index tensor with some valid indices and some -1 (padding) + topk_index = torch.zeros((batch_size, topk), dtype=torch.int32, device="cuda") + topk_index[:, 0:2048] = torch.arange(0, 2048, dtype=torch.int32, device="cuda") + + # Create ragged_mem_index lookup table + ragged_mem_index = torch.arange(0, 2048, dtype=torch.int32, device="cuda") + 10 + + topk_mem_index = trans_topk_index_to_mem_index(topk_index, ragged_mem_index) + + assert torch.equal(topk_mem_index, (torch.arange(0, 2048, dtype=torch.int32, device="cuda") + 10).view(1, -1)) + + +if __name__ == "__main__": + pytest.main() diff --git a/unit_tests/models/deepseek_v4/test_c128_workspace.py b/unit_tests/models/deepseek_v4/test_c128_workspace.py new file mode 100644 index 0000000000..c7fcdca099 --- /dev/null +++ b/unit_tests/models/deepseek_v4/test_c128_workspace.py @@ -0,0 +1,548 @@ +"""C128 capacity, two-stage compression alignment and graph replay regressions.""" + +import bisect +from types import SimpleNamespace as NS, MethodType +from unittest.mock import patch +import pytest +import torch +from lightllm.common.basemodel.basemodel import TpPartBaseModel +from lightllm.common.basemodel.mtp_manager import MtpManager +from lightllm.common.basemodel.cuda_graph import CudaGraph +from lightllm.common.basemodel import cuda_graph as graph_impl +from lightllm.models.deepseek_v4.workspace import DeepseekV4Workspace +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.server.router.model_infer.mode_backend.generic_pre_process import ( + prepare_decode_inputs, + overlap_prepare_decode_inputs, +) + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +def model_for( + batch=256, + requests=256, + step=3, + tp=1, + overlap=False, + prefill_overlap=False, + graph_requests=64, + graph=True, + dynamic=False, +): + args = StartArgs( + batch_max_tokens=batch, + running_max_req_size=requests, + mtp_mode="dspark" if step else None, + mtp_step=step, + mtp_dynamic_verify=dynamic, + enable_decode_microbatch_overlap=overlap, + enable_prefill_microbatch_overlap=prefill_overlap, + enable_tpsp_mix_mode=overlap or prefill_overlap or tp > 1, + ) + mtp = MtpManager.__new__(MtpManager) + mtp.args = args + model = NS( + batch_max_tokens=batch, + max_req_num=requests, + tp_world_size_=tp, + args=args, + mtp_manager=mtp, + is_mtp_draft_model=False, + disable_cudagraph=not graph, + config={"head_dim": 512}, + req_manager=NS(HOLD_REQUEST_ID=requests), + ) + model._align_decode_batch_size = MethodType(TpPartBaseModel._align_decode_batch_size, model) + graph_requests = min(graph_requests, requests) + model.graph_max_batch_size = model._align_decode_batch_size( + (graph_requests // 2 if overlap else graph_requests) * mtp.get_decode_tokens_per_request(False) + ) + return model + + +def reqs_for(count, step): + return [ + NS( + req_idx=i, + mtp_step=step, + cur_kv_len=124, + get_cur_total_len=lambda: 125, + multimodal_params={"images": [], "audios": []}, + get_radix_cache_shared_len=lambda: 0, + shared_kv_node=None, + ) + for i in range(count) + ] + + +def inputs_for(model, count): + reqs = reqs_for(count, model.args.mtp_step) + if model.args.enable_decode_microbatch_overlap: + result = overlap_prepare_decode_inputs(reqs) + return [result[0], result[3]] + return [prepare_decode_inputs(reqs)[0]] + + +def graph_buckets(model): + with patch.object(graph_impl, "get_env_start_args", lambda: model.args): + width = model.mtp_manager.get_decode_tokens_per_request(False) + alignment = model.mtp_manager.get_decode_batch_alignment(False) + return CudaGraph.gen_cuda_graph_batch_sizes( + alignment, 32 * width, 8 * alignment, model.graph_max_batch_size, model.tp_world_size_ + ) + + +@pytest.mark.parametrize("overlap", [False, True]) +def test_reported_overflow_reproduction_is_now_covered(overlap): + m = model_for(overlap=overlap) + w = DeepseekV4Workspace(m) + w.init_c128_compress_scratch(m) + assert ( + w.c128_compress_scratch.shape == (2, 512, 512) if overlap else w.c128_compress_scratch.shape == (1, 1024, 512) + ) + for region, input in enumerate(inputs_for(m, 256)): + position = input.b_seq_len - 1 + boundaries = torch.where((position + 1) % 128 == 0)[0] + assert int(boundaries[boundaries >= 256][0]) == 259 + scratch = w.get_c128_compress_scratch(input.batch_size, region) + assert int(boundaries.max()) < scratch.shape[0] + assert w.token_capacity == 256 + + +@pytest.mark.parametrize("requests", [17, 256, 257]) +@pytest.mark.parametrize("step", [0, 3, 5]) +@pytest.mark.parametrize("tp", [1, 8]) +@pytest.mark.parametrize("overlap", [False, True]) +@pytest.mark.parametrize("graph", [False, True]) +def test_real_decode_mtp_tp_and_graph_padding_fit(requests, step, tp, overlap, graph): + m = model_for(batch=9, requests=requests, step=step, tp=tp, overlap=overlap, graph=graph) + w = DeepseekV4Workspace(m) + w.init_c128_compress_scratch(m) + buckets = graph_buckets(m) if graph else [] + for count in [0, 1, requests // 2, requests]: + for region, input in enumerate(inputs_for(m, count)): + rows = m._align_decode_batch_size(max(1, input.batch_size)) + if buckets and rows <= buckets[-1]: + rows = buckets[bisect.bisect_left(buckets, rows)] + input.input_ids = torch.zeros(input.batch_size, dtype=torch.int64) + padded = TpPartBaseModel._create_padded_decode_model_input(m, input, rows) + assert padded.batch_size == rows + assert w.get_c128_compress_scratch(rows, region).shape == (rows, 512) + for rows in buckets: + assert w.get_c128_compress_scratch(rows, 0).shape[0] == rows + assert w.token_capacity == ((9 + tp - 1) // tp) * tp + + +@pytest.mark.parametrize("dynamic", [False, True]) +def test_dynamic_verify_plans_for_maximum_width_and_odd_microbatch(dynamic): + m = model_for(batch=1, requests=17, step=5, tp=8, overlap=True, graph=True, dynamic=dynamic) + w = DeepseekV4Workspace(m) + w.init_c128_compress_scratch(m) + assert w.get_c128_compress_scratch(m._align_decode_batch_size(9 * 6), 0).shape[0] >= 54 + assert w.token_capacity == 8 + + +def test_graph_max_bucket_can_exceed_eager_bound(): + m = model_for(batch=1, requests=17, step=3, tp=8) + m.graph_max_batch_size = 256 # lower-level model API can explicitly request a larger graph bucket + w = DeepseekV4Workspace(m) + w.init_c128_compress_scratch(m) + assert w.get_c128_compress_scratch(256, 0).shape[0] == 256 + + +def test_prefill_arena_keeps_its_own_capacity(): + m = model_for(batch=4097, requests=256, step=3, tp=8, prefill_overlap=True) + w = DeepseekV4Workspace(m) + w.init_c128_compress_scratch(m) + assert w.c128_compress_scratch.shape == (2, 4104, 512) + w.init_flashmla_prefill_q(2, 4, 512, torch.bfloat16) + assert w.get_flashmla_prefill_q(4104).shape == (4104, 4, 512) + with pytest.raises(ValueError, match="capacity"): + w.get_flashmla_prefill_q(4105) + + +@pytest.mark.parametrize("overlap", [False, True]) +def test_checked_accessor_rejects_bad_rows_and_regions(overlap): + m = model_for(overlap=overlap) + w = DeepseekV4Workspace(m) + w.init_c128_compress_scratch(m) + for index in [-1, 2 if overlap else 1]: + with pytest.raises(ValueError, match="region"): + w.get_c128_compress_scratch(1, index) + for rows in [-1, w.c128_compress_scratch.shape[1] + 1]: + with pytest.raises(ValueError, match="capacity"): + w.get_c128_compress_scratch(rows, 0) + assert w.get_c128_compress_scratch(0, 0).shape == (0, 512) + + +def compress_case(workspace, batch, region=0, boundary=True): + """Four verify rows per request, with known C128 history and packed output.""" + from lightllm.models.deepseek_v4.layer_infer.compressor import fused_compress + + steps = torch.arange(4, device="cuda", dtype=torch.int32).repeat(batch) + positions = (124 if boundary else 2) + steps + request_ids = torch.arange(batch, device="cuda", dtype=torch.int32).repeat_interleave(4) + runtime = torch.randn(batch + 1, 131, 1024, device="cuda") * 0.4 + values = torch.randn(batch * 4, 1024, device="cuda") * 0.4 + packed = torch.full((batch + 1, 1168), 85, dtype=torch.uint8, device="cuda") + state = NS(config=NS(window_size=128), get_layer_view=lambda **_: runtime) + compressed = NS(config=NS(block_token_size=256), get_layer_view=lambda **_: NS(kv=packed)) + infer = NS( + position_ids=positions, + mem_index=request_ids * 256 + steps, + b_req_idx=request_ids, + b_mtp_index=steps, + b_seq_len=positions + 1, + b_ready_cache_len=None, + b_q_start_loc=None, + _dsv4_token_to_batch_idx=None, + is_prefill=False, + microbatch_index=region, + dsv4_workspace=workspace, + get_compressor_state=lambda **_: state, + get_compressed_kv_state=lambda **_: compressed, + ) + norm = torch.ones(512, device="cuda", dtype=torch.bfloat16) + cos = torch.ones(256, 32, device="cuda") + sin = torch.zeros_like(cos) + + def run(): + return fused_compress( + kv_score=values, + infer_state=infer, + layer_idx=0, + norm_weight=norm, + eps=1e-6, + head_dim=512, + qk_rope_head_dim=64, + compress_ratio=128, + cos_table=cos, + sin_table=sin, + ) + + return NS(run=run, infer=infer, runtime=runtime, values=values, packed=packed, batch=batch, region=region) + + +def check_compressed_reference(case, history, workspace): + current = case.values.view(case.batch, 4, 1024) + window = torch.cat((history[: case.batch, :124], current), dim=1) + expected = (window[..., :512] * window[..., 512:].softmax(dim=1)).sum(dim=1) + actual = workspace.get_c128_compress_scratch(case.batch * 4, case.region)[3::4] + torch.testing.assert_close(actual, expected, rtol=1e-5, atol=5e-7) + torch.testing.assert_close(case.runtime[: case.batch, 124:128], current, rtol=0, atol=0) + normalized = expected * torch.rsqrt(expected.square().mean(dim=1, keepdim=True) + 1e-6) + rope = case.packed[: case.batch, 448:576].contiguous().view(torch.bfloat16) + torch.testing.assert_close(rope.float(), normalized[:, -64:].bfloat16().float(), rtol=0.008, atol=2e-6) + # Only the first packed slot belongs to a boundary; the second one stays untouched. + assert (case.packed[: case.batch, 576:1152] == 85).all() + assert (case.packed[: case.batch, 1160:1168] == 85).all() + assert (case.packed[case.batch] == 85).all() + + +@pytest.mark.parametrize("overlap", [False, True]) +def test_c128_kernel_preserves_guards_regions_and_reference(overlap): + torch.manual_seed(1013) + model = model_for(overlap=overlap, graph=False) + workspace = DeepseekV4Workspace(model) + workspace.init_c128_compress_scratch(model) + regions, rows, head = workspace.c128_compress_scratch.shape + arena = torch.full((regions, rows + 2, head), 17.0, device="cuda") + workspace.c128_compress_scratch = arena[:, 1:-1] + batch = 128 if overlap else 256 + cases = [compress_case(workspace, batch, region) for region in range(regions)] + histories = [case.runtime.clone() for case in cases] + streams = [torch.cuda.Stream() for _ in cases] + for stream in streams: + stream.wait_stream(torch.cuda.current_stream()) + for case, stream in zip(cases, streams): + with torch.cuda.stream(stream): + case.run() + for stream in streams: + torch.cuda.current_stream().wait_stream(stream) + for case, history in zip(cases, histories): + check_compressed_reference(case, history, workspace) + assert (arena[:, 0] == 17).all() and (arena[:, -1] == 17).all() + + +def test_c128_capture_without_boundary_then_replay_with_boundary(): + torch.manual_seed(1013) + workspace = DeepseekV4Workspace(model_for(graph=True)) + workspace.init_c128_compress_scratch(model_for(graph=True)) + case = compress_case(workspace, 256, boundary=False) + for _ in range(2): + case.run() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + case.run() + history = case.runtime.clone() + case.infer.position_ids.copy_(124 + case.infer.b_mtp_index) + case.infer.b_seq_len.copy_(case.infer.position_ids + 1) + graph.replay() + check_compressed_reference(case, history, workspace) + + +@pytest.mark.parametrize("boundary", [False, True]) +def test_short_scratch_is_rejected_before_any_compression_launch(monkeypatch, boundary): + from lightllm.models.deepseek_v4.layer_infer import compressor + + model = model_for(graph=False) + workspace = DeepseekV4Workspace(model) + workspace.init_c128_compress_scratch(model) + workspace.c128_compress_scratch = torch.empty((1, 256, 512), device="cuda") + case = compress_case(workspace, 256, boundary=boundary) + + class RejectLaunch: + def __getitem__(self, grid): + pytest.fail("unsafe compression kernel was launched") + + for name in ("_compress_window_kernel", "_compress_norm_rope_pack_kernel"): + monkeypatch.setattr(compressor, name, RejectLaunch()) + with pytest.raises(ValueError, match="capacity"): + case.run() + + +def reference_window(values, state, position, seq_len, ready_len, q_start, req_idx): + """Per-channel window softmax, matching the compress stage one boundary row at a time.""" + start = position - 127 + kv_rows = [] + score_rows = [] + for pos in range(start, start + 128): + if pos < 0 or pos >= seq_len: + kv_rows.append(torch.zeros(512, device=values.device)) + score_rows.append(torch.full((512,), float("-inf"), device=values.device)) + elif pos >= ready_len: + row = values[q_start + (pos - ready_len)] + kv_rows.append(row[:512]) + score_rows.append(row[512:]) + else: + row = state[req_idx, pos % state.shape[1]] + kv_rows.append(row[:512]) + score_rows.append(row[512:]) + kv = torch.stack(kv_rows) + weight = torch.softmax(torch.stack(score_rows), dim=0) + return (kv * weight).sum(dim=0) + + +def reference_norm(compressed, weight, eps): + variance = compressed.square().sum() / compressed.numel() + return compressed * torch.rsqrt(variance + eps) * weight.float() + + +def reference_rope(normed, position, cos, sin): + rope = normed[-64:].float().view(32, 2) + even, odd = rope[:, 0], rope[:, 1] + row = (position // 128) * 128 + rotated = torch.stack((even * cos[row] - odd * sin[row], odd * cos[row] + even * sin[row]), dim=-1) + return rotated.reshape(-1).to(torch.bfloat16) + + +def packed_rope(packed, slot): + page, token_in_page = divmod(int(slot), 2) + start = token_in_page * 576 + 448 + return packed[page, start : start + 128].contiguous().view(torch.bfloat16) + + +def stage_spec(scenario): + """Decode and prefill layouts whose boundary rows exercise the window scratch and both page slots.""" + cos = torch.randn(256, 32, device="cuda") + sin = torch.randn(256, 32, device="cuda") + norm = torch.randn(512, device="cuda").to(torch.bfloat16) + prefill = scenario.startswith("prefill") + if scenario == "decode": + batch = 2 + steps = torch.arange(4, device="cuda", dtype=torch.int32).repeat(batch) + positions = 124 + steps + request_ids = torch.arange(batch, device="cuda", dtype=torch.int32).repeat_interleave(4) + mem_index = request_ids * 256 + steps + runtime = torch.randn(batch + 1, 131, 1024, device="cuda") + values = torch.randn(batch * 4, 1024, device="cuda") + b_req_idx, b_seq_len, b_mtp = request_ids, positions + 1, steps + token_to_batch = b_ready = b_q_start = None + elif scenario == "decode_second_slot": + batch = 2 + steps = torch.arange(4, device="cuda", dtype=torch.int32).repeat(batch) + positions = 124 + steps + request_ids = torch.arange(batch, device="cuda", dtype=torch.int32).repeat_interleave(4) + mem_index = request_ids * 256 + 128 + steps + runtime = torch.randn(batch + 1, 131, 1024, device="cuda") + values = torch.randn(batch * 4, 1024, device="cuda") + b_req_idx, b_seq_len, b_mtp = request_ids, positions + 1, steps + token_to_batch = b_ready = b_q_start = None + elif scenario == "prefill_current": + positions = torch.arange(128, device="cuda", dtype=torch.int32) + mem_index = positions + runtime = torch.randn(1, 131, 1024, device="cuda") + values = torch.randn(128, 1024, device="cuda") + token_to_batch = torch.zeros(128, dtype=torch.int32, device="cuda") + b_req_idx = torch.zeros(1, dtype=torch.int32, device="cuda") + b_seq_len = torch.tensor([128], dtype=torch.int32, device="cuda") + b_ready = torch.zeros(1, dtype=torch.int32, device="cuda") + b_q_start = torch.zeros(1, dtype=torch.int32, device="cuda") + b_mtp = torch.zeros(1, dtype=torch.int32, device="cuda") + else: + positions = torch.arange(64, 192, device="cuda", dtype=torch.int32) + mem_index = 256 + positions + runtime = torch.randn(1, 131, 1024, device="cuda") + values = torch.randn(128, 1024, device="cuda") + token_to_batch = torch.zeros(128, dtype=torch.int32, device="cuda") + b_req_idx = torch.zeros(1, dtype=torch.int32, device="cuda") + b_seq_len = torch.tensor([192], dtype=torch.int32, device="cuda") + b_ready = torch.tensor([64], dtype=torch.int32, device="cuda") + b_q_start = torch.zeros(1, dtype=torch.int32, device="cuda") + b_mtp = torch.zeros(1, dtype=torch.int32, device="cuda") + packed = torch.full((int(mem_index.max()) // 128 + 2, 1168), 85, dtype=torch.uint8, device="cuda") + return NS( + positions=positions, + mem_index=mem_index, + runtime=runtime, + values=values, + packed=packed, + cos=cos, + sin=sin, + norm=norm, + eps=1e-6, + prefill=prefill, + token_to_batch=token_to_batch, + b_req_idx=b_req_idx, + b_seq_len=b_seq_len, + b_mtp=b_mtp, + b_ready=b_ready, + b_q_start=b_q_start, + ) + + +def attach_stage_case(workspace, spec, runtime, packed): + from lightllm.models.deepseek_v4.layer_infer.compressor import fused_compress + + state = NS(config=NS(window_size=128), get_layer_view=lambda **_: runtime) + compressed = NS(config=NS(block_token_size=256), get_layer_view=lambda **_: NS(kv=packed)) + infer = NS( + position_ids=spec.positions, + mem_index=spec.mem_index, + b_req_idx=spec.b_req_idx, + b_mtp_index=spec.b_mtp, + b_seq_len=spec.b_seq_len, + b_ready_cache_len=spec.b_ready, + b_q_start_loc=spec.b_q_start, + _dsv4_token_to_batch_idx=spec.token_to_batch, + is_prefill=spec.prefill, + microbatch_index=0, + dsv4_workspace=workspace, + get_compressor_state=lambda **_: state, + get_compressed_kv_state=lambda **_: compressed, + ) + + def run(): + return fused_compress( + kv_score=spec.values, + infer_state=infer, + layer_idx=0, + norm_weight=spec.norm, + eps=spec.eps, + head_dim=512, + qk_rope_head_dim=64, + compress_ratio=128, + cos_table=spec.cos, + sin_table=spec.sin, + ) + + return run + + +def launch_fused_oracle(spec, runtime, packed): + import triton + from lightllm.models.deepseek_v4.triton_kernel.compressor import _fused_compress_norm_rope_kernel + + slots = torch.where((spec.positions + 1) % 128 == 0, spec.mem_index.reshape(-1) // 128, -1) + _fused_compress_norm_rope_kernel[(spec.positions.shape[0],)]( + spec.values, + spec.values.stride(0), + spec.values.stride(1), + runtime.flatten(0, 1), + spec.positions, + spec.token_to_batch, + spec.b_req_idx, + spec.b_mtp, + spec.b_seq_len, + spec.b_ready, + spec.b_q_start, + slots, + spec.norm, + spec.eps, + spec.cos, + spec.cos.stride(0), + spec.cos.stride(1), + spec.sin, + spec.sin.stride(0), + spec.sin.stride(1), + packed, + HEAD_DIM=512, + CANDIDATE_DIM=512, + STATE_DIM=1024, + COMPRESS_RATIO=128, + WINDOW_SIZE=128, + IS_C4=False, + IS_PREFILL=spec.prefill, + STATE_CAPACITY_TOKENS=runtime.shape[1], + ROPE_HEAD_DIM=64, + FP8_MAX=torch.finfo(torch.float8_e4m3fn).max, + AMAX_MIN=1e-4, + NOPE_DIM=448, + QUANT_BLOCK=64, + SCALE_BYTES=8, + PAGE_SIZE=2, + BYTES_PER_PAGE=packed.stride(0), + BLOCK=triton.next_power_of_2(512), + IS_IN_INDEXER=False, + num_warps=4, + ) + + +@pytest.mark.parametrize("scenario", ["decode", "decode_second_slot", "prefill_current", "prefill_split"]) +def test_two_stages_match_fused_kernel_and_reference(scenario): + torch.manual_seed(1013) + spec = stage_spec(scenario) + model = model_for(graph=False) + workspace = DeepseekV4Workspace(model) + workspace.init_c128_compress_scratch(model) + staged_runtime, oracle_runtime = spec.runtime.clone(), spec.runtime.clone() + staged_packed, oracle_packed = spec.packed.clone(), spec.packed.clone() + history = spec.runtime.clone() + workspace.c128_compress_scratch.fill_(17) + + attach_stage_case(workspace, spec, staged_runtime, staged_packed)() + launch_fused_oracle(spec, oracle_runtime, oracle_packed) + torch.cuda.synchronize() + assert torch.equal(staged_packed, oracle_packed) + + positions = spec.positions + boundary = (positions + 1) % 128 == 0 + rows = positions.shape[0] + compressed = workspace.get_c128_compress_scratch(rows, 0) + assert torch.equal(compressed[~boundary], torch.full_like(compressed[~boundary], 17)) + assert torch.equal(workspace.c128_compress_scratch[0, rows], torch.full((512,), 17, device="cuda")) + + for token in boundary.nonzero(as_tuple=False).flatten().tolist(): + if spec.prefill: + ready = int(spec.b_ready[int(spec.token_to_batch[token])]) + q_start = int(spec.b_q_start[int(spec.token_to_batch[token])]) + seq_len = int(spec.b_seq_len[int(spec.token_to_batch[token])]) + req_idx = int(spec.b_req_idx[int(spec.token_to_batch[token])]) + else: + ready = int(positions[token] - spec.b_mtp[token]) + q_start = token - int(spec.b_mtp[token]) + seq_len = int(spec.b_seq_len[token]) + req_idx = int(spec.b_req_idx[token]) + expected = reference_window(spec.values, history, int(positions[token]), seq_len, ready, q_start, req_idx) + torch.testing.assert_close(compressed[token], expected, rtol=1e-5, atol=1e-6) + normed = reference_norm(compressed[token], spec.norm, spec.eps) + slot = int(spec.mem_index[token] // 128) + torch.testing.assert_close( + packed_rope(staged_packed, slot).float(), + reference_rope(normed, int(positions[token]), spec.cos, spec.sin).float(), + rtol=0, + atol=0, + ) diff --git a/unit_tests/models/deepseek_v4/test_kernels.py b/unit_tests/models/deepseek_v4/test_kernels.py new file mode 100644 index 0000000000..d5afb9b219 --- /dev/null +++ b/unit_tests/models/deepseek_v4/test_kernels.py @@ -0,0 +1,342 @@ +import pytest +import torch + +from lightllm.models.deepseek_v4.triton_kernel import norm_rope, topk_softplus_sqrt, indexer_topk +from lightllm.models.deepseek_v4.triton_kernel.rotary_emb import inverse_rotary_emb +from lightllm.utils.vllm_utils import vllm_ops + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.fixture(autouse=True) +def topk_provider_fallback(monkeypatch): + if hasattr(vllm_ops, "top_k_per_row_decode"): + return + + def top_k_per_row_decode(scores, indent, lengths, out, num_rows, row_stride, stride, topk): + available = min(row_stride, scores.shape[1]) + for row in range(scores.shape[0]): + row_length = lengths[row].clamp(max=available) + count = min(available, topk) + columns = torch.arange(available, device=scores.device) + valid_scores = torch.where(columns < row_length, scores[row, :available], -torch.inf) + indices = torch.topk(valid_scores, count).indices + mask = columns[:count] < row_length + out[row, :count] = torch.where(mask, indices, -1) + out[row, count:] = -1 + + monkeypatch.setattr(vllm_ops, "top_k_per_row_decode", top_k_per_row_decode, raising=False) + + +def freqs_and_positions(batch, dtype=torch.int64): + torch.manual_seed(13) + angles = torch.randn(128, 32, device="cuda") + freqs = torch.polar(torch.ones_like(angles), angles) + positions = torch.randint(0, 128, (batch,), device="cuda", dtype=dtype) + return freqs, positions + + +def rope_reference(x, freqs, positions): + pairs = torch.view_as_complex(x[..., -64:].float().reshape(*x.shape[:-1], 32, 2).contiguous()) + rotation = freqs[positions.long()] + while rotation.ndim < pairs.ndim: + rotation = rotation.unsqueeze(1) + return torch.view_as_real(pairs * rotation).flatten(-2) + + +@pytest.mark.parametrize("tokens,heads,dim_stride", [(0, 8, 1), (1, 3, 1), (7, 16, 2), (129, 64, 1)]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +def test_inverse_rotary_emb_strides_and_unwritten_channels(tokens, heads, dim_stride, dtype): + freqs, positions = freqs_and_positions(tokens) + padded = torch.randn(tokens, heads + 2, 512 * dim_stride, device="cuda", dtype=dtype) + before = padded.clone() + attention = padded[:, :heads, ::dim_stride] + expected = rope_reference(attention, freqs.conj(), positions).to(dtype) + rotation = freqs[positions] + # Independent table strides: complex real view for cos, contiguous sin. + inverse_rotary_emb(attention[..., -64:], rotation.real, rotation.imag.contiguous()) + tolerance = 0.01 if dtype == torch.bfloat16 else 1e-5 + torch.testing.assert_close(attention[..., -64:], expected, rtol=tolerance, atol=tolerance) + torch.testing.assert_close(attention[..., :-64], before[:, :heads, ::dim_stride][..., :-64], rtol=0, atol=0) + torch.testing.assert_close(padded[:, heads:], before[:, heads:], rtol=0, atol=0) + if dim_stride == 2: + torch.testing.assert_close(padded[..., 1::2], before[..., 1::2], rtol=0, atol=0) + + +def test_inverse_rotary_emb_cuda_graph_replay(): + freqs, positions = freqs_and_positions(7) + rotation = freqs[positions] + source = torch.randn(7, 5, 512, device="cuda", dtype=torch.bfloat16) + output = torch.empty_like(source) + + def run(): + inverse_rotary_emb(output[..., -64:], rotation.real, rotation.imag) + + output.copy_(source) + run() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + for _ in range(2): + source.normal_() + positions.random_(0, 128) + rotation.copy_(freqs[positions]) + output.copy_(source) + graph.replay() + expected = source.clone() + expected[..., -64:] = rope_reference(source, freqs.conj(), positions).bfloat16() + torch.testing.assert_close(output, expected, rtol=0.01, atol=0.01) + + +@pytest.mark.parametrize("batch,heads", [(0, 8), (1, 8), (7, 16), (128, 16)]) +@pytest.mark.parametrize("position_dtype", [torch.int32, torch.int64]) +def test_q_norm_rope_strided_and_bf16_rounding(batch, heads, position_dtype): + freqs, pos = freqs_and_positions(batch, position_dtype) + q = torch.randn(batch, heads + 2, 512, device="cuda", dtype=torch.bfloat16)[:, :heads] + padded = torch.full((batch, heads + 4, 512), 9, device="cuda", dtype=torch.bfloat16) + out = padded[:, :heads] + norm_rope.fused_q_norm_rope(q, out, 1e-6, freqs, pos) + normalized = (q.float() * torch.rsqrt(q.float().square().mean(-1, keepdim=True) + 1e-6)).bfloat16() + expected = normalized.clone() + expected[..., -64:] = rope_reference(normalized, freqs, pos).bfloat16() + torch.testing.assert_close(out, expected, rtol=0.01, atol=0.02) + assert (padded[:, heads:] == 9).all() + # In-place outputs must also work without changing the normalization input early. + inplace = q.contiguous().clone() + norm_rope.fused_q_norm_rope(inplace, inplace, 1e-6, freqs, pos) + torch.testing.assert_close(inplace, expected, rtol=0.01, atol=0.02) + + +@pytest.mark.parametrize("batch", [1, 7, 129]) +@pytest.mark.parametrize("position_dtype", [torch.int32, torch.int64]) +def test_k_norm_rope_cache_layout_and_unwritten_bytes(batch, position_dtype): + freqs, pos = freqs_and_positions(batch, position_dtype) + kv = torch.randn(batch, 640, device="cuda", dtype=torch.bfloat16)[:, :512] + kv[0].zero_() + weight = torch.randn(512, device="cuda", dtype=torch.bfloat16) + slots = torch.randperm(384, device="cuda")[:batch].int() + cache = torch.full((3, 74880), 165, device="cuda", dtype=torch.uint8) + norm_rope.fused_kv_norm_rope_to_swa(kv, weight, 1e-6, freqs, pos, slots, cache, 128) + values = kv.float() * torch.rsqrt(kv.float().square().mean(-1, keepdim=True) + 1e-6) * weight.float() + grouped = values[:, :448].reshape(batch, 7, 64) + scale_raw = grouped.abs().amax(-1).clamp_min(1e-4) / 448 + bits = scale_raw.contiguous().view(torch.int32) + exponent = ((bits >> 23) & 255) + ((bits & 0x7FFFFF) != 0).int() + scales = torch.ldexp(torch.ones_like(scale_raw), exponent - 127) + expected_fp8 = (grouped / scales[..., None]).clamp(-448, 448).to(torch.float8_e4m3fn) + expected_rope = rope_reference(values, freqs, pos).bfloat16() + pages = slots.long() // 128 + offsets = slots.long() % 128 + data_offsets = offsets[:, None] * 576 + torch.arange(576, device="cuda") + data = cache[pages[:, None], data_offsets] + torch.testing.assert_close( + data[:, :448].contiguous().view(torch.float8_e4m3fn).float(), + expected_fp8.reshape(batch, 448).float(), + rtol=0.13, + atol=0.03125, + ) + torch.testing.assert_close(data[:, 448:].contiguous().view(torch.bfloat16), expected_rope, rtol=0.01, atol=0.02) + scale_offsets = 128 * 576 + offsets[:, None] * 8 + torch.arange(7, device="cuda") + torch.testing.assert_close(cache[pages[:, None], scale_offsets], exponent.byte(), rtol=0, atol=0) + written = torch.zeros_like(cache, dtype=torch.bool) + written[pages[:, None], data_offsets] = True + written[pages[:, None], scale_offsets] = True + assert (cache[~written] == 165).all() + + +@pytest.mark.parametrize("batch,heads", [(0, 64), (1, 3), (7, 64), (128, 64)]) +def test_indexer_rope_hadamard_quant_and_preallocated_outputs(batch, heads): + freqs, pos = freqs_and_positions(batch) + q = torch.randn(batch, heads, 128, device="cuda", dtype=torch.bfloat16) + weights = torch.randn(batch, heads, device="cuda", dtype=torch.bfloat16) + if batch: + q[0, 0].zero_() + output = torch.empty_like(q, dtype=torch.float8_e4m3fn) + out_weights = torch.empty(batch, heads, 1, device="cuda") + got, got_weights = norm_rope.fused_q_indexer_rope_hadamard_quant(q, weights, 0.01, freqs, pos, output, out_weights) + assert got is output and got_weights is out_weights + rotated = q.float() + rotated[..., -64:] = rope_reference(q, freqs, pos) + hadamard = torch.ones(1, 1, device="cuda") + for _ in range(7): + hadamard = torch.cat((torch.cat((hadamard, hadamard), 1), torch.cat((hadamard, -hadamard), 1)), 0) + transformed = (rotated @ hadamard) * (128 ** -0.5) + scale = transformed.abs().amax(-1, keepdim=True).clamp_min(1e-4) / 448 + expected = (transformed / scale).clamp(-448, 448).to(torch.float8_e4m3fn) + torch.testing.assert_close(got.float(), expected.float(), rtol=0.13, atol=0.03125) + torch.testing.assert_close(got_weights, weights.float().unsqueeze(-1) * 0.01 * scale, rtol=4e-6, atol=1e-8) + + +def router_reference(logits, factor, bias, tokens, table, image_bias, image_start): + scores = torch.sqrt(logits.clamp_min(0) + torch.log(1 + torch.exp(-logits.abs()))) + scores = torch.nan_to_num(scores, nan=0) + selection = scores + (bias if bias is not None else 0) + if image_bias is not None: + image = tokens >= image_start + selection = torch.where(image[:, None], scores + image_bias, selection) + indices = selection.argsort(dim=-1, descending=True, stable=True)[:, :6] + if table is not None: + hash_ids = table[tokens.clamp_max(table.shape[0] - 1)] + if image_bias is None: + indices = hash_ids + else: + indices = torch.where(image[:, None], indices, hash_ids) + weights = scores.gather(1, indices) + total = weights.sum(-1, keepdim=True) + return weights * factor / torch.where(total > 0, total, 1), indices + + +@pytest.mark.parametrize("experts", [256, 384]) +@pytest.mark.parametrize("mode", ["bias", "hash", "vision", "hash_vision"]) +@pytest.mark.parametrize("rows,distribution", [(1, "normal"), (513, "normal"), (17, "ties"), (513, "underflow")]) +def test_router_modes_scaling_ties_and_underflow(experts, mode, rows, distribution): + torch.manual_seed(23) + logits = torch.randn(rows, experts, device="cuda") + bias = torch.randn(experts, device="cuda") * 0.1 + image_bias = torch.randn_like(bias) * 0.1 if "vision" in mode else None + if distribution == "ties": + logits.zero_() + bias.zero_() + if image_bias is not None: + image_bias.zero_() + elif distribution == "underflow": + logits.fill_(-100) + tokens = torch.randint(0, 32, (rows,), device="cuda") + table = torch.stack([torch.randperm(experts, device="cuda")[:6] for _ in range(32)]) if "hash" in mode else None + if image_bias is not None: + tokens[::2] = 100_000 + torch.arange(0, rows, 2, device="cuda") + correction = None if table is not None else bias + weights = torch.empty(rows, 6, device="cuda") + indices = torch.empty(rows, 6, device="cuda", dtype=torch.int64) + topk_softplus_sqrt.topk_softplus_sqrt(weights, indices, logits, 2.5, correction, tokens, table, image_bias, 32) + ref_weights, ref_indices = router_reference(logits, 2.5, correction, tokens, table, image_bias, 32) + torch.testing.assert_close(indices, ref_indices, rtol=0, atol=0) + torch.testing.assert_close(weights, ref_weights, rtol=5e-6, atol=1e-6) + + +def assert_topk(scores, lengths, pages, out, raw, page_size): + for row, length in enumerate(lengths.tolist()): + count = min(out.shape[1], length) + chosen = raw[row, :count].long() + assert (raw[row, count:] == -1).all() and (out[row, count:] == -1).all() + assert chosen.unique().numel() == count + assert ((chosen >= 0) & (chosen < length)).all() + torch.testing.assert_close( + out[row, :count], (pages[row, chosen // page_size] * page_size + chosen % page_size).int(), rtol=0, atol=0 + ) + torch.testing.assert_close( + scores[row, chosen].sort(descending=True).values, scores[row, :length].topk(count).values, rtol=0, atol=0 + ) + + +@pytest.mark.parametrize("width", [128, 512, 640, 8192, 262144]) +@pytest.mark.parametrize("page_size", [1, 64, 128]) +def test_topk_valid_prefix_raw_indices_and_padding(width, page_size): + torch.manual_seed(31) + scores = torch.randn(8, width + 128, device="cuda")[:, :width] + lengths = torch.tensor( + [0, 1, min(511, width), min(512, width), min(513, width), width // 3, width - 3, width], + device="cuda", + dtype=torch.int32, + ) + scores.masked_fill_(torch.arange(width, device="cuda")[None, :] >= lengths[:, None], float("nan")) + pages = torch.randint(0, 4096, (8, (width + page_size - 1) // page_size), device="cuda", dtype=torch.int32) + out = torch.empty(8, 512, device="cuda", dtype=torch.int32) + raw = torch.empty_like(out) + indexer_topk.indexer_topk(scores, lengths, pages, out, page_size, raw) + assert_topk(scores, lengths, pages, out, raw, page_size) + # The allocation-free in-place translation must select the same values too. + indexer_topk.indexer_topk(scores, lengths, pages, out, page_size) + valid = out >= 0 + assert valid.sum(-1).tolist() == lengths.clamp_max(512).tolist() + + +@pytest.mark.parametrize( + "width,rows,page_size", [(513, 1, 64), (640, 2, 128), (1024, 1, 1), (1024, 4, 64), (1536, 2, 128), (2048, 2, 64)] +) +@pytest.mark.parametrize("distribution", ["normal", "ties", "concentrated"]) +def test_short_topk_rank_selection_and_optional_raw(width, rows, distribution, page_size): + torch.manual_seed(43) + scores = torch.randn(rows, width + 1, device="cuda")[:, :width] + if distribution == "ties": + scores.fill_(1) + elif distribution == "concentrated": + scores.mul_(1e-6).add_(1) + lengths = torch.full((rows,), width, device="cuda", dtype=torch.int32) + if rows > 1: + lengths[0] = 513 + scores.masked_fill_(torch.arange(width, device="cuda")[None, :] >= lengths[:, None], float("nan")) + pages = torch.randint(0, 4096, (rows, (width + page_size - 1) // page_size), device="cuda", dtype=torch.int32) + out = torch.empty(rows, 512, device="cuda", dtype=torch.int32) + raw = torch.empty_like(out) + indexer_topk.indexer_topk(scores, lengths, pages, out, page_size, raw) + assert_topk(scores, lengths, pages, out, raw, page_size) + expected = out.clone() + indexer_topk.indexer_topk(scores, lengths, pages, out, page_size) + torch.testing.assert_close(out, expected, rtol=0, atol=0) + + +@pytest.mark.parametrize("rows,width", [(8, 8192), (1, 1024), (2, 2048), (4, 1024)]) +def test_topk_cuda_graph_replay_changing_lengths(rows, width): + scores = torch.randn(rows, width, device="cuda") + lengths = torch.zeros(rows, device="cuda", dtype=torch.int32) + pages = torch.randint(0, 4096, (rows, width // 64), device="cuda", dtype=torch.int32) + out = torch.empty(rows, 512, device="cuda", dtype=torch.int32) + raw = torch.empty_like(out) + fn = lambda: indexer_topk.indexer_topk(scores, lengths, pages, out, 64, raw) + fn() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + fn() + for shift in range(8): + scores.normal_() + lengths.copy_( + torch.tensor([0, 1, 511, 512, 513, 1024, 4096, width], device="cuda", dtype=torch.int32) + .clamp_max(width) + .roll(shift)[:rows] + ) + scores.masked_fill_(torch.arange(width, device="cuda")[None, :] >= lengths[:, None], float("nan")) + graph.replay() + assert_topk(scores, lengths, pages, out, raw, 64) + + +def test_norm_cuda_graph_dependent_launch_and_replay(): + freqs, positions = freqs_and_positions(7) + source = torch.randn(7, 8, 512, device="cuda", dtype=torch.bfloat16) + first = torch.empty_like(source) + second = torch.empty_like(source) + + def run(): + norm_rope.fused_q_norm_rope(source, first, 1e-6, freqs, positions) + norm_rope.fused_q_norm_rope(first, second, 1e-6, freqs, positions) + + run() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + for _ in range(3): + source.normal_() + positions.random_(0, 128) + graph.replay() + expected = source.clone() + for _ in range(2): + expected = ( + expected.float() * torch.rsqrt(expected.float().square().mean(-1, keepdim=True) + 1e-6) + ).bfloat16() + expected[..., -64:] = rope_reference(expected, freqs, positions).bfloat16() + torch.testing.assert_close(second, expected, rtol=0.02, atol=0.03) + + +@pytest.mark.parametrize("topk", [128, 256, 1024]) +@pytest.mark.parametrize("width", [64, 640, 4096]) +def test_topk_uses_configured_output_capacity(topk, width): + scores = torch.randn(2, width, device="cuda") + lengths = torch.tensor([width // 2, width], dtype=torch.int32, device="cuda") + pages = torch.randint(0, 1000, (2, (width + 63) // 64), dtype=torch.int32, device="cuda") + out = torch.empty((2, topk), dtype=torch.int32, device="cuda") + raw = torch.empty_like(out) + indexer_topk.indexer_topk(scores, lengths, pages, out, 64, raw) + assert_topk(scores, lengths, pages, out, raw, 64) diff --git a/unit_tests/models/deepseek_v4/test_model.py b/unit_tests/models/deepseek_v4/test_model.py new file mode 100644 index 0000000000..7f6e8d5183 --- /dev/null +++ b/unit_tests/models/deepseek_v4/test_model.py @@ -0,0 +1,301 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.models.deepseek_v4.model import DeepseekV4TpPartModel +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) + model.config = dict( + qk_rope_head_dim=64, + rope_theta=10000, + compress_rope_theta=160000, + max_position_embeddings=8193, + rope_scaling=dict(type="yarn", factor=16, original_max_position_embeddings=65536, beta_fast=32, beta_slow=1), + ) + model.max_seq_length = 8193 + model.layers_infer = [SimpleNamespace(compress_ratio=ratio, index_infer=SimpleNamespace()) for ratio in (0, 4, 128)] + model._init_to_get_rotary() + + positions = torch.tensor([0, 1, 127, 1023, 8192], device="cuda") + exponents = torch.arange(0, 64, 2, device="cuda", dtype=torch.float32) / 64 + angles = torch.outer(positions.float(), 1.0 / (10000 ** exponents)) + expected = torch.complex(angles.cos(), angles.sin()) + torch.testing.assert_close(model._freqs_cis_sliding[positions], expected, rtol=1e-6, atol=1e-6) + unscaled_compress = torch.outer(positions.float(), 1.0 / (160000 ** exponents)) + assert not torch.allclose( + model._freqs_cis_compress[positions], torch.complex(unscaled_compress.cos(), unscaled_compress.sin()) + ) + assert model.layers_infer[0].freqs_cis is model._freqs_cis_sliding + for layer in model.layers_infer[1:]: + assert layer.freqs_cis is model._freqs_cis_compress + + +@pytest.mark.parametrize( + "balance,ep,tpsp,tp,redundancy,error", + [ + (True, False, False, 2, 0, "enable_dp_prefill_balance"), + (False, True, True, 2, 0, "FFN and mHC token layouts"), + (False, True, False, 2, 2, "redundant experts"), + (False, True, True, 1, 0, None), + (False, True, False, 2, 0, None), + (False, False, True, 2, 0, None), + ], +) +def test_unsupported_parallel_layouts_fail_before_inference(monkeypatch, balance, ep, tpsp, tp, redundancy, error): + monkeypatch.setattr("lightllm.models.deepseek_v4.model.get_redundancy_expert_num", lambda: redundancy) + model = DeepseekV4TpPartModel.__new__(DeepseekV4TpPartModel) + model.load_way, model.tp_world_size_ = "HF", tp + model.config = {"num_attention_heads": 64, "o_groups": 8, "index_n_heads": 64} + model.args = SimpleNamespace(enable_dp_prefill_balance=balance, enable_ep_moe=ep, enable_tpsp_mix_mode=tpsp) + if error: + with pytest.raises(ValueError, match=error): + model._verify_params() + else: + model._verify_params() + + +def test_memory_check_rejects_sequence_larger_than_full_token_pool(): + model = DeepseekV4TpPartModel.__new__(DeepseekV4TpPartModel) + 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) + + with pytest.raises(AssertionError, match="max_total_token_num must be >= max_seq_length"): + model._check_mem_size() + + +def test_memory_check_still_requires_one_prefill_batch(): + model = DeepseekV4TpPartModel.__new__(DeepseekV4TpPartModel) + 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) + + with pytest.raises(AssertionError, match="greater than batch_max_tokens"): + model._check_mem_size() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("deferred", [False, True]) +def test_draft_hidden_collector_materializes_mhc_streams(monkeypatch, deferred): + from lightllm.common.basemodel.hidden_collector import LayerHiddenCollector + from lightllm.common.basemodel.mtp_manager import MtpManager + from lightllm.models.deepseek_v4.model import DeepseekV4LayerHiddenCollector + + monkeypatch.setattr(LayerHiddenCollector, "_load_layer_ids", lambda self: frozenset([0])) + model = DeepseekV4TpPartModel.__new__(DeepseekV4TpPartModel) + model.config = {"hc_mult": 4, "hidden_size": 8} + model.layers_num = 2 + model.pre_infer = SimpleNamespace(_tpsp_allgather=lambda input, infer_state: input) + model.mtp_manager = MtpManager.__new__(MtpManager) + model.mtp_manager.args = SimpleNamespace(mtp_mode="dspark") + model._init_hidden_collector() + assert isinstance(model.hidden_collector_prototype, DeepseekV4LayerHiddenCollector) + collector = model.hidden_collector_prototype.new_instance() + + torch.manual_seed(6103) + residual = torch.randn((2, 4, 8), dtype=torch.bfloat16, device="cuda") + output = torch.randn((2, 8), dtype=torch.bfloat16, device="cuda") + post_mix = torch.randn((2, 4), device="cuda") + residual_mix = torch.randn((2, 4, 4), device="cuda") + streams = ( + torch.einsum("bsh,bst->bth", residual.float(), residual_mix) + output.float()[:, None] * post_mix[..., None] + ) + streams = streams.bfloat16() + hidden = (output, residual.flatten(1), residual_mix, post_mix) if deferred else streams.flatten(1) + + collector.add(layer_index=1, hidden=None) + collector.add(layer_index=0, hidden=hidden) + actual = collector.finish_output(SimpleNamespace(need_dp_prefill_balance=False)).spec_hidden + torch.testing.assert_close(actual, streams.mean(dim=1), rtol=0.02, atol=0.02) + assert not model.hidden_collector_prototype.layer_hiddens + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("tokens", [1, 17]) +@pytest.mark.parametrize("prefill", [False, True], ids=["decode", "prefill"]) +@pytest.mark.parametrize("truncated", [False, True], ids=["full", "truncated"]) +def test_mhc_layers_and_head_cuda_graph_replay(monkeypatch, tokens, prefill, truncated): + import triton + + from lightllm.models.deepseek_v4.layer_infer.post_layer_infer import DeepseekV4PostLayerInfer + from lightllm.models.llama.layer_infer.post_layer_infer import LlamaPostLayerInfer + from unit_tests.common.basemodel.triton_kernel.mhc.test_mhc import _hc_post_reference, _hc_pre_reference + + torch.manual_seed(6104) + triton.set_allocator(lambda size, alignment, stream: torch.empty(size, device="cuda", dtype=torch.int8)) + streams, hidden = 4, 4096 + cfg = {"hc_mult": streams, "hidden_size": hidden, "rms_norm_eps": 1e-5, "hc_eps": 1e-6} + residual = torch.randn(tokens, streams * hidden, device="cuda", dtype=torch.bfloat16) + layers = [] + weights = [] + for index in range(2): + layer = DeepseekV4TransformerLayerInfer.__new__(DeepseekV4TransformerLayerInfer) + layer.hc_mult, layer.embed_dim_ = streams, hidden + layer.eps_, layer.hc_eps, layer.sinkhorn_iters = cfg["rms_norm_eps"], cfg["hc_eps"], 20 + layer.is_last_layer = index == 1 and not truncated + layer.context_attention_forward = layer.token_attention_forward = lambda x, *_: x * 0.125 + layer._ffn = lambda x, *_: x * 0.25 + layers.append(layer) + weight = SimpleNamespace() + for prefix in ("attn", "ffn"): + setattr( + weight, + f"hc_{prefix}_fn_", + SimpleNamespace(weight=torch.randn(24, streams * hidden, device="cuda") * 0.005), + ) + setattr(weight, f"hc_{prefix}_scale_", SimpleNamespace(weight=torch.randn(3, device="cuda"))) + setattr(weight, f"hc_{prefix}_base_", SimpleNamespace(weight=torch.randn(24, device="cuda"))) + setattr( + weight, + f"{prefix}_norm_", + SimpleNamespace(weight=torch.randn(hidden, device="cuda", dtype=torch.bfloat16)), + ) + weights.append(weight) + + post = DeepseekV4PostLayerInfer.__new__(DeepseekV4PostLayerInfer) + post.alloc_tensor = torch.empty + post_weight = SimpleNamespace( + network_config_=cfg, + hc_head_fn_=SimpleNamespace(weight=torch.randn(streams, streams * hidden, device="cuda") * 0.005), + hc_head_scale_=SimpleNamespace(weight=torch.randn(1, device="cuda")), + hc_head_base_=SimpleNamespace(weight=torch.randn(streams, device="cuda")), + ) + # Keep real mHC calls and replace the logits path with its input. + monkeypatch.setattr(LlamaPostLayerInfer, "token_forward", lambda self, x, *_: x) + + def run(): + state = residual + for layer, weight in zip(layers, weights): + forward = layer.context_forward if prefill else layer.token_forward + state = forward(state, None, weight) + if isinstance(state, tuple): + x, current_residual, residual_mix, post_mix = state + assert x.shape == (tokens, hidden) + assert current_residual.shape == (tokens, streams * hidden) + assert residual_mix.shape == (tokens, streams, streams) + assert post_mix.shape == (tokens, streams) + assert isinstance(state, tuple) == truncated + return post.token_forward(state, None, post_weight) + + def reference(): + state = residual + for weight in weights: + for prefix, factor in (("attn", 0.125), ("ffn", 0.25)): + x, mix, post_mix = _hc_pre_reference( + state, + getattr(weight, f"hc_{prefix}_fn_").weight, + getattr(weight, f"hc_{prefix}_scale_").weight, + getattr(weight, f"hc_{prefix}_base_").weight, + streams, + cfg["rms_norm_eps"], + cfg["hc_eps"], + 20, + ) + normalized = x.float() * torch.rsqrt(x.float().square().mean(-1, keepdim=True) + cfg["rms_norm_eps"]) + normalized = (normalized * getattr(weight, f"{prefix}_norm_").weight.float()).bfloat16() + state = _hc_post_reference(normalized * factor, state, mix, post_mix, streams) + inv_rms = torch.rsqrt(state.float().square().mean(-1, keepdim=True) + cfg["rms_norm_eps"]) + gates = ( + torch.sigmoid( + (state.float() @ post_weight.hc_head_fn_.weight.T) * inv_rms * post_weight.hc_head_scale_.weight + + post_weight.hc_head_base_.weight + ) + + cfg["hc_eps"] + ) + return (state.view(tokens, streams, hidden).float() * gates[..., None]).sum(1).bfloat16() + + torch.testing.assert_close(run(), reference(), atol=0.06, rtol=0.04) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = run() + for _ in range(2): + residual.normal_() + graph.replay() + torch.testing.assert_close(actual, reference(), atol=0.06, rtol=0.04) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("tp_size,rank", [(1, 0), (2, 0), (2, 1)]) +def test_output_projection_matches_bf16_reference(monkeypatch, tp_size, rank): + from lightllm.common.basemodel.layer_weights.base_layer_weight import BaseLayerWeight + from lightllm.common.basemodel.layer_weights.meta_weights import platform_op + from lightllm.models.deepseek_v4.layer_weights.transformer_layer_weight import DeepseekV4TransformerLayerWeight + + monkeypatch.setattr( + platform_op, + "get_env_start_args", + lambda: SimpleNamespace(hardware_platform="cuda", enable_torch_fallback=False, enable_triton_fallback=False), + ) + monkeypatch.setenv("LIGHTLLM_DP_WORLD_SIZE", str(tp_size)) + monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_DP", str(rank)) + monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") + torch.manual_seed(6102) + layer = DeepseekV4TransformerLayerWeight.__new__(DeepseekV4TransformerLayerWeight) + BaseLayerWeight.__init__(layer) + layer.prefix, layer.data_type_ = "layers.0", torch.bfloat16 + layer.hidden, layer.n_heads, layer.head_dim = 128, 8, 64 + layer.q_lora_rank, layer.o_lora_rank, layer.o_groups = 128, 128, 2 + layer.num_fused_shared_experts = 0 + layer.get_quant_method = lambda name: None + layer._init_qkvo() + layer.experts = SimpleNamespace(quant_method=SimpleNamespace(weight_scale_suffix=None)) + + weight_name = "layers.0.attn.wo_a.weight" + scale_name = "layers.0.attn.wo_a.scale" + packed = torch.randn(256, 256, device="cuda").to(torch.float8_e4m3fn) + scales = torch.tensor([[0.125, 0.25], [0.5, 1]], device="cuda").to(torch.float8_e8m0fnu) + expected_weight = packed.float().unflatten(0, (-1, 128)).unflatten(-1, (-1, 128)) * scales[:, None, :, None].float() + expected_weight = expected_weight.flatten(2, 3).flatten(0, 1).bfloat16() + weight_b = torch.randn(128, 256, device="cuda", dtype=torch.bfloat16) * 0.1 + layer.load_hf_weights({weight_name: packed, scale_name: scales, "layers.0.attn.wo_b.weight": weight_b}) + assert layer.wo_a_.verify_load() + local_groups = layer.o_groups // tp_size + grouped_weight = expected_weight.view(2, 128, 256)[rank * local_groups : (rank + 1) * local_groups] + heads = layer.n_heads // tp_size + padded = torch.full((5, 64, 64), 9, device="cuda", dtype=torch.bfloat16) + attention = padded[:, :heads] + attention.normal_(std=0.1) + angles = torch.randn(5, 32, device="cuda") + cos, sin = angles.cos(), angles.sin() + pairs = torch.view_as_complex(attention.float().reshape(5, heads, 32, 2).contiguous()) + rotated = torch.view_as_real(pairs * torch.complex(cos, sin).conj()[:, None]).flatten(-2).bfloat16() + expected = torch.einsum("tgd,grd->tgr", rotated.reshape(5, local_groups, 256), grouped_weight).flatten(1) + expected = expected @ weight_b[:, rank * local_groups * 128 : (rank + 1) * local_groups * 128].T + + infer_layer = DeepseekV4TransformerLayerInfer.__new__(DeepseekV4TransformerLayerInfer) + infer_layer.qk_rope_head_dim, infer_layer.tp_groups, infer_layer.compress_ratio = 64, local_groups, 0 + infer_layer._tpsp_reduce = lambda input, infer_state: input + state = SimpleNamespace(position_cos_sliding=cos, position_sin_sliding=sin) + torch.testing.assert_close(infer_layer._get_o(attention, state, layer), expected, rtol=0.01, atol=0.01) + assert padded[:, heads:].eq(9).all() diff --git a/unit_tests/models/deepseek_v4/test_shared_expert_overlap.py b/unit_tests/models/deepseek_v4/test_shared_expert_overlap.py new file mode 100644 index 0000000000..c088d811eb --- /dev/null +++ b/unit_tests/models/deepseek_v4/test_shared_expert_overlap.py @@ -0,0 +1,247 @@ +"""Shared-expert overlap on the EP prefill and decode paths.""" + +from types import SimpleNamespace +import pytest +import torch +from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe_ep import ( + fused_experts_impl, + prefill_decode_sync as sync, +) +from lightllm.models.deepseek_v4.layer_infer import transformer_layer_infer as layer_module +from lightllm.models.deepseek_v4.layer_infer.transformer_layer_infer import DeepseekV4TransformerLayerInfer +from lightllm.models.deepseek_v4.layer_weights.transformer_layer_weight import DeepseekV4TransformerLayerWeight + + +def _layer(ep_mode): + layer = DeepseekV4TransformerLayerInfer.__new__(DeepseekV4TransformerLayerInfer) + layer.embed_dim_ = 4 + layer.layer_num_ = 0 + layer.enable_ep_moe = ep_mode + layer.swiglu_limit = 7 + layer._tpsp_allgather = lambda input, infer_state: input + ids = torch.zeros(3, 2, dtype=torch.int64) + layer._select_experts = lambda *args: (torch.ones(3, 2), ids) + return layer, ids + + +def _shared_recorder(events, scale=2): + def shared(*args, **kwargs): + x = kwargs["input"] if "input" in kwargs else args[0] + events.append("shared") + return x * scale + + return shared + + +def _moe_inputs(): + hidden = torch.ones(3, 4) + topk_weights = torch.ones(3, 2) + topk_idx = torch.zeros(3, 2, dtype=torch.int64) + w1 = torch.ones(2, 8, 4) + w2 = torch.ones(2, 4, 4) + return hidden, topk_weights, topk_idx, w1, w2 + + +def test_ep_ffn_passes_the_shared_expert_callback(monkeypatch): + monkeypatch.setattr(layer_module, "use_sm100_mega_moe", lambda method: False) + layer, _ids = _layer(True) + events = [] + layer._ffn_shared_expert = _shared_recorder(events) + + def routed(**kwargs): + events.append("routed") + return torch.full((3, 4), 3.0) + kwargs["shared_expert_func"]() + + experts = SimpleNamespace(quant_method=None, experts_with_topk=routed) + result = layer._ffn( + torch.ones(3, 4), + SimpleNamespace(is_prefill=False), + SimpleNamespace(experts=experts), + ) + assert events == ["routed", "shared"] + assert torch.all(result == 5) + + +def test_decode_runs_shared_expert_between_combine_send_and_recv(monkeypatch): + events = [] + hidden, topk_weights, topk_idx, w1, w2 = _moe_inputs() + import deep_ep + + monkeypatch.setattr(deep_ep.ElasticBuffer, "capture", staticmethod(lambda: events.append("event") or "event")) + monkeypatch.setattr(sync, "per_token_group_quant_fp8", lambda *args, **kwargs: (hidden, hidden)) + monkeypatch.setattr(sync, "get_mk_alignment_for_contiguous_layout", lambda: 128) + monkeypatch.setattr(sync, "get_deepep_num_max_dispatch_tokens_per_rank_decode", lambda: 16) + monkeypatch.setattr(sync, "get_ep_num_sms", lambda overlap_with_compute: 8) + + def dispatch(*args, **kwargs): + events.append("dispatch") + handle = SimpleNamespace( + psum_num_recv_tokens_per_expert=torch.tensor([3, 3]), + expert_alignment=128, + recv_src_metadata=torch.zeros(3, 4, dtype=torch.int32), + psum_num_recv_tokens_per_scaleup_rank=torch.tensor([3]), + ) + return (hidden, hidden), None, topk_weights, handle, None + + def gemm(**kwargs): + assert kwargs["expected_m"] == 3 # ceil(3 tokens * 8 ranks * topk 2 / 16 experts) + events.append("gemm") + return hidden + + def reduce(**kwargs): + events.append("reduce") + return hidden, kwargs["recv_src_metadata"] + + class _Event: + def current_stream_wait(self): + events.append("combine_recv") + + def combine(*args, **kwargs): + assert kwargs["previous_event"] == "event" + assert kwargs["async_with_compute_stream"] is True + events.append("combine_send") + return torch.full((3, 4), 3.0), None, _Event() + + monkeypatch.setattr(sync, "decode_masked_group_gemm", gemm) + monkeypatch.setattr(sync, "ep_reduce_decode_output", reduce) + result = fused_experts_impl( + hidden, + w1, + w2, + topk_weights, + topk_idx, + num_experts=16, + buffer=SimpleNamespace(num_ranks=8, dispatch=dispatch, combine=combine), + is_prefill=False, + use_fp8_w8a8=True, + use_fp8_all2all=True, + shared_expert_func=lambda: events.append("shared") or torch.full((3, 4), 2.0), + ) + assert events == ["dispatch", "gemm", "reduce", "event", "combine_send", "shared", "combine_recv"] + assert torch.all(result == 5) + + +def test_prefill_overlaps_shared_expert_with_combine(monkeypatch): + events = [] + hidden, topk_weights, topk_idx, w1, w2 = _moe_inputs() + import deep_ep + + monkeypatch.setattr(deep_ep.ElasticBuffer, "capture", staticmethod(lambda: events.append("event") or "event")) + monkeypatch.setattr(sync, "get_deepep_num_max_dispatch_tokens_per_rank_prefill", lambda: 16) + monkeypatch.setattr(sync, "per_token_group_quant_fp8", lambda *args, **kwargs: (hidden, hidden)) + monkeypatch.setattr(sync, "get_mk_alignment_for_contiguous_layout", lambda: 128) + monkeypatch.setattr(sync, "get_ep_num_sms", lambda overlap_with_compute: 8) + + def dispatch(*args, **kwargs): + events.append("dispatch") + handle = SimpleNamespace( + num_recv_tokens_per_expert_list=[0], + recv_src_metadata=torch.zeros(3, 4), + ) + return (hidden, hidden), None, topk_weights, handle, None + + class _Event: + def current_stream_wait(self): + events.append("combine_recv") + + def combine(*args, **kwargs): + assert kwargs["previous_event"] == "event" + assert kwargs["async_with_compute_stream"] is True + events.append("combine_send") + return torch.full((3, 4), 3.0), None, _Event() + + result = fused_experts_impl( + hidden, + w1, + w2, + topk_weights, + topk_idx, + num_experts=16, + buffer=SimpleNamespace(dispatch=dispatch, combine=combine), + is_prefill=True, + use_fp8_w8a8=True, + use_fp8_all2all=True, + shared_expert_func=lambda: events.append("shared") or torch.full((3, 4), 2.0), + ) + assert events == ["dispatch", "event", "combine_send", "shared", "combine_recv"] + assert torch.all(result == 5) + + +def test_rename_shared_expert_as_the_last_routed_expert(): + weight = DeepseekV4TransformerLayerWeight.__new__(DeepseekV4TransformerLayerWeight) + weight.prefix = "layers.0" + weight.n_routed_experts = 4 + weight.experts = SimpleNamespace( + quant_method=SimpleNamespace(weight_scale_suffix="weight_scale_inv", weight_zero_point_suffix=None) + ) + tensors = { + "layers.0.ffn.shared_experts.w1.weight": torch.ones(2), + "layers.0.ffn.shared_experts.w1.scale": torch.ones(1), + "layers.0.ffn.shared_experts.w2.weight": torch.ones(2), + "layers.0.ffn.shared_experts.w3.weight": torch.ones(2), + } + weight._rename_shared_experts(tensors) + assert torch.equal(tensors["layers.0.ffn.experts.4.w1.weight"], torch.ones(2)) + assert torch.equal(tensors["layers.0.ffn.experts.4.w1.scale"], torch.ones(1)) + assert "layers.0.ffn.experts.4.w2.weight" in tensors + assert "layers.0.ffn.experts.4.w3.weight" in tensors + + +def test_auto_enable_fused_shared_experts_for_v4_tp(monkeypatch): + from lightllm.utils import config_utils + + monkeypatch.setattr(config_utils, "get_model_type", lambda path: "deepseek_v4") + args = SimpleNamespace(model_dir="/model", enable_ep_moe=False, enable_fused_shared_experts=False) + config_utils.auto_set_fused_shared_experts(args) + assert args.enable_fused_shared_experts is True + + args.enable_ep_moe = True + args.enable_fused_shared_experts = False + config_utils.auto_set_fused_shared_experts(args) + assert args.enable_fused_shared_experts is False + + +def test_tp_fused_shared_expert_skips_the_standalone_ffn(monkeypatch): + monkeypatch.setattr(layer_module, "use_sm100_mega_moe", lambda method: False) + layer, _ids = _layer(False) + events = [] + layer._ffn_shared_expert = _shared_recorder(events) + reductions = [] + layer._tpsp_reduce = lambda input, infer_state: reductions.append(input.clone()) or input + + def routed(**kwargs): + events.append("routed") + return torch.full((3, 4), 3.0) + + experts = SimpleNamespace(quant_method=None, experts_with_topk=routed) + result = layer._ffn( + torch.ones(3, 4), + SimpleNamespace(is_prefill=False), + SimpleNamespace(experts=experts, num_fused_shared_experts=1), + ) + assert events == ["routed"] + assert torch.all(result == 3) + assert len(reductions) == 1 + assert torch.all(reductions[0] == 3) + + +@pytest.mark.parametrize("ep_mode", [False, True]) +def test_inplace_backends_compute_shared_expert_first(monkeypatch, ep_mode): + monkeypatch.setattr(layer_module, "use_sm100_mega_moe", lambda method: ep_mode) + layer, _ids = _layer(ep_mode) + events = [] + x = torch.ones(3, 4) + layer._ffn_shared_expert = _shared_recorder(events) + + def routed(**kwargs): + events.append("routed") + kwargs["input_tensor"].fill_(9) + return torch.full((3, 4), 3.0) + + reductions = [] + layer._tpsp_reduce = lambda input, infer_state: reductions.append(input) or input + experts = SimpleNamespace(quant_method=None, experts_with_topk=routed) + result = layer._ffn(x, SimpleNamespace(is_prefill=False), SimpleNamespace(experts=experts)) + assert events == ["shared", "routed"] + assert torch.all(result == 5) + assert len(reductions) == (0 if ep_mode else 1) diff --git a/unit_tests/models/deepseek_v4/test_tokenizer.py b/unit_tests/models/deepseek_v4/test_tokenizer.py new file mode 100644 index 0000000000..b96db74cab --- /dev/null +++ b/unit_tests/models/deepseek_v4/test_tokenizer.py @@ -0,0 +1,79 @@ +import copy +from types import SimpleNamespace + +import pytest + +from lightllm.models.deepseek_v4.tokenizer import DeepSeekV4Tokenizer + + +IMAGE_TOKEN = "<｜deepseek_image｜>" + + +class _Tokenizer: + def convert_tokens_to_ids(self, token): + assert token == IMAGE_TOKEN + return 99 + + def encode(self, prompt, **kwargs): + ids = [] + for index, text in enumerate(prompt.split(IMAGE_TOKEN)): + if index: + ids.append(99) + ids.extend([7] * len(text)) + return ids + + +@pytest.fixture +def tokenizer(tmp_path): + # An empty model directory selects the bundled encoder used by default. + return DeepSeekV4Tokenizer(_Tokenizer(), str(tmp_path), {"vision_n_layers": 1}) + + +@pytest.mark.parametrize("tokenize", [False, True]) +def test_mixed_content_preserves_image_order_and_expands_placeholders(tokenizer, tokenize): + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "before"}, + {"type": "image_url", "image_url": {"url": "first"}}, + {"type": "text", "text": "between"}, + {"type": "image_url", "image_url": {"url": "second"}}, + {"type": "text", "text": "after"}, + ], + } + ] + original = copy.deepcopy(messages) + prompt = tokenizer.apply_chat_template(messages=messages, tokenize=tokenize) + if not tokenize: + assert f"before{IMAGE_TOKEN}between{IMAGE_TOKEN}after" in prompt + images = [SimpleNamespace(token_id=1000, token_num=13), SimpleNamespace(token_id=2000, token_num=17)] + ids = tokenizer.encode(prompt, SimpleNamespace(images=images)) + for image in images: + expected = list(range(image.token_id + image.start_idx % 4, image.token_id + image.token_num)) + assert ids[image.start_idx : image.start_idx + len(expected)] == expected + assert images[0].start_idx < images[1].start_idx + assert 99 not in ids + assert messages == original + + +@pytest.mark.parametrize( + "kwargs,thinking", + [ + ({"reasoning_effort": "none"}, False), + ({"reasoning_effort": "none", "thinking": False}, False), + ({"reasoning_effort": "none", "enable_thinking": False}, False), + ({"reasoning_effort": "none", "thinking": True}, True), + ({"reasoning_effort": "high"}, True), + ({"reasoning_effort": "max"}, True), + ({}, True), + ], +) +def test_reasoning_effort_uses_encoder_modes(tokenizer, kwargs, thinking): + prompt = tokenizer.apply_chat_template(messages=[{"role": "user", "content": "hello"}], **kwargs) + assert prompt.endswith("" if thinking else "") + + +def test_unknown_reasoning_effort_is_not_silently_replaced(tokenizer): + with pytest.raises(AssertionError, match="Invalid reasoning effort: low"): + tokenizer.apply_chat_template(messages=[{"role": "user", "content": "hello"}], reasoning_effort="low") diff --git a/unit_tests/models/deepseek_v4/test_vision.py b/unit_tests/models/deepseek_v4/test_vision.py new file mode 100644 index 0000000000..ad931ac2b5 --- /dev/null +++ b/unit_tests/models/deepseek_v4/test_vision.py @@ -0,0 +1,350 @@ +import importlib.util +import io +import json +import os +import sys +from types import SimpleNamespace + +import numpy as np +import pytest +import torch +from PIL import Image +from torch import nn +from torch.nn.attention import SDPBackend, sdpa_kernel + +from lightllm.common.basemodel.attention_vit.sdpa.fp import SdpaVitAttBackend +from lightllm.models.deepseek_v4 import deepseek_v4_visual as visual +from lightllm.models.deepseek_v4 import image_processor as ours +from lightllm.server import visualserver + + +REFERENCE_DIR = "/mtc/models/DeepSeek-V4-Flash-Vision-Exp/inference" +REFERENCE_PROCESSOR_PATH = os.path.join(REFERENCE_DIR, "image_processor.py") +REFERENCE_VISION_PATH = os.path.join(REFERENCE_DIR, "vision.py") + +PATCH_SIZE = 14 +DOWNSAMPLE_RATIO = 3 +MAX_N_TOKEN = 384 +MIN_PIXELS = 147456 +MAX_WH_RATIO = 8 + +REFERENCE_ARGS = SimpleNamespace( + vision_patch_size=PATCH_SIZE, + vision_downsample_ratio=DOWNSAMPLE_RATIO, + vision_max_n_token=MAX_N_TOKEN, + vision_min_pixels=MIN_PIXELS, + vision_max_wh_ratio=MAX_WH_RATIO, +) + + +def _load_reference_processor(): + spec = importlib.util.spec_from_file_location("deepseek_v4_reference_image_processor", REFERENCE_PROCESSOR_PATH) + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def _load_reference_vision(): + spec = importlib.util.spec_from_file_location("deepseek_v4_reference_vision", REFERENCE_VISION_PATH) + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +@pytest.fixture(scope="module") +def reference_processor(): + if not os.path.exists(REFERENCE_PROCESSOR_PATH): + pytest.skip("DeepSeek-V4 vision reference processor is not available") + return _load_reference_processor() + + +@pytest.fixture(scope="module") +def reference_vision(): + if not os.path.exists(REFERENCE_VISION_PATH): + pytest.skip("DeepSeek-V4 reference vision tower is not available") + return _load_reference_vision() + + +@pytest.mark.parametrize( + "height,width", + [(14, 14), (15, 30), (41, 512), (100, 100), (410, 756), (756, 42)], +) +def test_grid_tokens_matches_reference(reference_processor, height, width): + assert ours.grid_tokens(height, width, PATCH_SIZE, DOWNSAMPLE_RATIO) == reference_processor.grid_tokens( + height, width, PATCH_SIZE, DOWNSAMPLE_RATIO + ) + + +@pytest.mark.parametrize("start_pos", range(8)) +@pytest.mark.parametrize("n_llm_h,n_llm_w", [(1, 1), (2, 3), (3, 2), (4, 5), (7, 8)]) +def test_image_block_and_permutation_match_reference(reference_processor, n_llm_h, n_llm_w, start_pos): + expected_types, expected_perm = reference_processor.build_image_block(n_llm_h, n_llm_w, start_pos) + types, perm = ours.build_image_block(n_llm_h, n_llm_w, start_pos) + assert torch.equal(types, expected_types) + assert torch.equal(perm, expected_perm) + + +@pytest.mark.parametrize("prefix_len", [0, 1, 2, 3, 511, 512, 513, 514]) +@pytest.mark.parametrize("text_gap", range(4)) +@pytest.mark.parametrize("reuse_image", [False, True]) +def test_tokenizer_spans_match_reference_blocks(reference_processor, prefix_len, text_gap, reuse_image): + from lightllm.models.deepseek_v4.tokenizer import DeepSeekV4Tokenizer + from lightllm.server.multimodal_params import MultimodalParams + from lightllm.utils.multimodal_chunk_utils import adjust_mm_chunk_end, clamp_mm_cache_len, get_mm_input_spans + + tokenizer = DeepSeekV4Tokenizer( + SimpleNamespace(convert_tokens_to_ids=lambda token: 99), "/unused", {"vision_n_layers": 1} + ) + grids = [(1, 1), (3, 2), (1, 1)] + params = MultimodalParams(images=[{"type": "url", "data": "unused"} for _ in grids]) + prompt = [7] * prefix_len + expected_ids = prompt.copy() + expected_spans = [] + expected_cache = [] + for index, (image, (height, width)) in enumerate(zip(params.images, grids)): + canonical_types, _ = reference_processor.build_image_block(height, width, start_pos=0) + image.token_id = 2 ** 62 + (0 if reuse_image and index == 2 else index * 1024) + image.token_num = len(canonical_types) + image.start_index_in_embed_cache = 100 + (0 if reuse_image and index == 2 else index * 1024) + expected_cache.append((image.token_id, image.token_num, image.start_index_in_embed_cache)) + start = len(expected_ids) + actual_types, _ = reference_processor.build_image_block(height, width, start_pos=start) + suffix_start = len(canonical_types) - len(actual_types) + assert torch.equal(canonical_types[suffix_start:], actual_types) + expected_ids.extend(range(image.token_id + suffix_start, image.token_id + image.token_num)) + expected_spans.append((start, len(expected_ids))) + prompt.extend([99] + [8] * text_gap) + expected_ids.extend([8] * text_gap) + + input_ids = tokenizer.encode(prompt, params) + assert input_ids == expected_ids + assert [ + (image.token_id, image.token_num, image.start_index_in_embed_cache) for image in params.images + ] == expected_cache + spans = get_mm_input_spans(np.asarray(input_ids, dtype=np.int64), params.to_dict()) + assert spans == expected_spans + for candidate in range(len(input_ids) + 1): + expected_end = next((start for start, end in expected_spans if start < candidate < end), candidate) + assert clamp_mm_cache_len(candidate, spans) == expected_end + + for chunk_size in (1, 4, 16, 256, 512): + start = 0 + while start < len(input_ids): + end = adjust_mm_chunk_end(start, min(start + chunk_size, len(input_ids)), spans) + assert start < end <= len(input_ids) + assert not any(left < end < right for left, right in expected_spans) + start = end + + +@pytest.mark.parametrize( + "width,height", + [(800, 600), (100, 2000), (2000, 100), (50, 50), (13, 7), (384, 384)], +) +def test_image_geometry_and_pixels_match_reference(reference_processor, width, height): + rng = np.random.default_rng(width * 10000 + height) + image = Image.fromarray(rng.integers(0, 256, (height, width, 3), dtype=np.uint8)) + image_bytes = io.BytesIO() + image.save(image_bytes, format="PNG") + + expected = reference_processor.load_image({"data": image_bytes.getvalue()}, REFERENCE_ARGS) + actual = ours.load_image(image, REFERENCE_ARGS) + assert actual[1:] == expected[1:] + assert torch.equal(actual[0], expected[0]) + + n_vit_h, n_vit_w, n_llm_h, n_llm_w, token_num = ours.get_image_grid( + height, + width, + PATCH_SIZE, + DOWNSAMPLE_RATIO, + MAX_N_TOKEN, + MIN_PIXELS, + MAX_WH_RATIO, + ) + assert (n_vit_h, n_vit_w, n_llm_h, n_llm_w) == actual[1:] + reference_types, _ = reference_processor.build_image_block(n_llm_h, n_llm_w, start_pos=0) + assert token_num == len(reference_types) + assert token_num <= MAX_N_TOKEN + + +def _tiny_model(): + return visual.DeepseekV4VisionModel( + {"weight_dir": "/unused"}, + hidden_size=8, + vision_n_layers=1, + vision_dim=8, + vision_n_heads=2, + vision_inter_dim=12, + vision_patch_size=2, + vision_rope_theta=10000.0, + vision_downsample_ratio=1, + vision_max_n_token=64, + vision_min_pixels=0, + vision_max_wh_ratio=8, + ) + + +def _tiny_vision_args(): + return SimpleNamespace( + hidden_size=12, + dim=12, + vision_n_layers=2, + vision_dim=16, + vision_n_heads=4, + vision_inter_dim=24, + vision_patch_size=2, + vision_rope_theta=10000.0, + vision_downsample_ratio=2, + ) + + +@pytest.mark.parametrize("n_vit_h,n_vit_w", [(4, 4), (5, 3)]) +def test_vit_matches_reference(reference_vision, n_vit_h, n_vit_w, monkeypatch): + monkeypatch.setattr(visualserver, "VIT_ATTN_BACKEND", SdpaVitAttBackend) + torch.manual_seed(0) + args = _tiny_vision_args() + model = visual.ViT(args).eval() + previous_dtype = torch.get_default_dtype() + torch.set_default_dtype(torch.bfloat16) + try: + reference = reference_vision.ViT(args).eval() + finally: + torch.set_default_dtype(previous_dtype) + reference.load_state_dict(model.state_dict()) + patches = torch.randn( + n_vit_h * n_vit_w, + 3, + args.vision_patch_size, + args.vision_patch_size, + dtype=torch.bfloat16, + ) + with torch.inference_mode(), sdpa_kernel(SDPBackend.MATH): + actual = model(patches, n_vit_h, n_vit_w) + expected = reference(patches, n_vit_h, n_vit_w) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +@pytest.mark.parametrize("n_vit_h,n_vit_w", [(4, 4), (5, 3)]) +def test_aligner_matches_reference(reference_vision, n_vit_h, n_vit_w): + torch.manual_seed(0) + args = _tiny_vision_args() + model = visual.Aligner(args).eval() + previous_dtype = torch.get_default_dtype() + torch.set_default_dtype(torch.bfloat16) + try: + reference = reference_vision.Aligner(args).eval() + finally: + torch.set_default_dtype(previous_dtype) + reference.load_state_dict(model.state_dict()) + hidden = torch.randn(n_vit_h * n_vit_w, args.vision_dim, dtype=torch.bfloat16) + with torch.inference_mode(): + actual = model(hidden, n_vit_h, n_vit_w) + expected = reference(hidden, n_vit_h, n_vit_w) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +def test_visual_parameter_dtypes(): + model = _tiny_model() + for module in model.modules(): + if isinstance(module, nn.Linear): + assert module.weight.dtype == torch.bfloat16 + if module.bias is not None: + assert module.bias.dtype == torch.bfloat16 + elif isinstance(module, visual.RMSNorm): + assert module.weight.dtype == torch.float32 + assert model.image_start.dtype == torch.bfloat16 + assert model.image_pad.dtype == torch.bfloat16 + + +def test_load_model_reads_only_indexed_visual_weights(tmp_path, monkeypatch): + model = _tiny_model() + expected = { + name: torch.full_like(tensor, index + 1) for index, (name, tensor) in enumerate(model.state_dict().items()) + } + weight_map = {name: "visual.safetensors" for name in expected} + weight_map["layers.0.self_attn.weight"] = "language.safetensors" + (tmp_path / "model.safetensors.index.json").write_text(json.dumps({"weight_map": weight_map})) + + requested = [] + opened = [] + + class FakeSafetensors: + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, traceback): + return False + + def get_tensor(self, key): + requested.append(key) + return expected[key] + + def fake_safe_open(path, framework, device): + opened.append(os.path.basename(path)) + assert framework == "pt" + assert device == "cpu" + return FakeSafetensors() + + monkeypatch.setattr(visual, "safe_open", fake_safe_open) + model.load_model(str(tmp_path)) + + assert opened == ["visual.safetensors"] + assert set(requested) == set(expected) + for name, tensor in model.state_dict().items(): + assert torch.equal(tensor, expected[name]) + + +def test_encode_returns_canonical_cache_block(monkeypatch): + model = _tiny_model() + with torch.no_grad(): + model.image_start.fill_(10) + model.image_pad.fill_(20) + model.image_newline.fill_(30) + model.image_end.fill_(40) + + class FakeVision(nn.Module): + def forward(self, patches, n_h, n_w): + return torch.zeros((n_h * n_w, 1), dtype=torch.bfloat16) + + class FakeAligner(nn.Module): + def forward(self, x, n_h, n_w): + values = torch.arange(n_h * n_w, dtype=torch.bfloat16) + return values[:, None].expand(-1, model.config.hidden_size) + + model.vision = FakeVision() + model.aligner = FakeAligner() + + image_bytes = io.BytesIO() + Image.new("RGB", (6, 4), color=(1, 2, 3)).save(image_bytes, format="PNG") + shm_names = [] + monkeypatch.setattr(visual, "get_shm_name_data", lambda uuid: f"{uuid}-data") + + def fake_read_shm(name): + shm_names.append(name) + return image_bytes.getvalue() + + monkeypatch.setattr(visual, "read_shm", fake_read_shm) + embeds, uuids, valid_ids = model.encode([SimpleNamespace(uuid="image-1")]) + + types, perm = ours.build_canonical_image_block(2, 3) + sentinel_table = torch.stack( + [ + model.image_start, + model.image_pad, + model.image_pad, + model.image_newline, + model.image_end, + ] + ) + expected = sentinel_table[types] + row_major = torch.arange(6, dtype=torch.bfloat16)[:, None].expand(-1, model.config.hidden_size) + expected[types == ours.IMAGE] = row_major[perm] + + assert shm_names == ["image-1-data"] + assert uuids == ["image-1"] + assert valid_ids == [[0, len(types)]] + assert torch.equal(embeds, expected) + assert torch.equal(types[:3], torch.full((3,), ours.IMAGE_PAD, dtype=torch.int64)) + assert types[3] == ours.IMAGE_START + assert types[-1] == ours.IMAGE_END diff --git a/unit_tests/models/deepseek_v4/test_vision_integration.py b/unit_tests/models/deepseek_v4/test_vision_integration.py new file mode 100644 index 0000000000..571f993266 --- /dev/null +++ b/unit_tests/models/deepseek_v4/test_vision_integration.py @@ -0,0 +1,321 @@ +from types import SimpleNamespace + +import pytest +import torch + + +VISION_CONFIG = { + "vision_n_layers": 32, + "vision_patch_size": 14, + "vision_downsample_ratio": 3, + "vision_max_n_token": 384, + "vision_min_pixels": 147456, + "vision_max_wh_ratio": 8, +} + + +class _Tokenizer: + def convert_tokens_to_ids(self, token): + assert token == "<｜deepseek_image｜>" + return 99 + + def encode(self, prompt, **kwargs): + return prompt + + +@pytest.mark.parametrize("prefix_len", range(8)) +def test_tokenizer_uses_position_dependent_canonical_suffix(prefix_len): + from lightllm.models.deepseek_v4.tokenizer import DeepSeekV4Tokenizer + from lightllm.server.multimodal_params import MultimodalParams + + tokenizer = DeepSeekV4Tokenizer(_Tokenizer(), "/unused", VISION_CONFIG) + params = MultimodalParams(images=[{"type": "url", "data": "unused"}]) + image = params.images[0] + image.token_id, image.token_num = 100_000, 13 + prompt = [7] * prefix_len + [99, 8] + + result = tokenizer.encode(prompt, params) + + cache_skip = prefix_len % 4 + assert result == [7] * prefix_len + list(range(image.token_id + cache_skip, image.token_id + 13)) + [8] + assert image.start_idx == prefix_len + image_end = prefix_len + image.token_num - cache_skip + assert result[image.start_idx : image_end] == list( + range(image.token_id + cache_skip, image.token_id + image.token_num) + ) + + from lightllm.utils.multimodal_chunk_utils import get_mm_input_spans + + assert get_mm_input_spans(result, params.to_dict()) == [(image.start_idx, image_end)] + + +def test_tokenizer_keeps_dsv4_on_one_dimensional_positions(): + from lightllm.models.deepseek_v4.tokenizer import DeepSeekV4Tokenizer + + tokenizer = DeepSeekV4Tokenizer(_Tokenizer(), "/unused", VISION_CONFIG) + image = SimpleNamespace(image_h=64, image_w=64, grid_thwd=None) + + assert tokenizer.get_image_token_length(image) > 0 + assert image.grid_thwd is None + + +def test_chunk_boundary_moves_before_complete_image_block(): + from lightllm.server.router.model_infer.infer_batch import InferReq + + req = InferReq.__new__(InferReq) + req.args = SimpleNamespace(chunked_prefill_size=8192, batch_max_tokens=8192) + req.cur_kv_len = 0 + req.cur_output_len = 0 + req.shm_req = SimpleNamespace(input_len=10_000) + req.mm_input_spans = [(8100, 8400)] + + assert req.get_chuncked_input_token_len() == 8100 + + req.cur_kv_len = 8100 + assert req.get_chuncked_input_token_len() == 10_000 + + +def test_chunk_boundary_between_adjacent_images_never_splits_the_previous_image(): + from lightllm.server.router.model_infer.infer_batch import InferReq + + req = InferReq.__new__(InferReq) + req.args = SimpleNamespace(chunked_prefill_size=640, batch_max_tokens=640) + req.cur_kv_len = 0 + req.cur_output_len = 0 + req.shm_req = SimpleNamespace(input_len=768) + req.mm_input_spans = [(0, 384), (384, 768)] + + assert req.get_chuncked_input_token_len() == 384 + + req.cur_kv_len = 384 + assert req.get_chuncked_input_token_len() == 768 + + +def test_atomic_image_block_can_exceed_chunk_size(): + from lightllm.server.router.model_infer.infer_batch import InferReq + + req = InferReq.__new__(InferReq) + req.args = SimpleNamespace(chunked_prefill_size=256, batch_max_tokens=1024) + req.cur_kv_len = 0 + req.cur_output_len = 0 + req.shm_req = SimpleNamespace(input_len=1000) + req.mm_input_spans = [(100, 800)] + + assert req.get_chuncked_input_token_len() == 100 + + req.cur_kv_len = 100 + assert req.get_chuncked_input_token_len() == 800 + + +def test_vision_model_allows_cpu_cache(): + from lightllm.models.deepseek_v4.model import DeepseekV4TpPartModel + from lightllm.server.core.objs import StartArgs + + model = SimpleNamespace( + load_way="HF", + tp_world_size_=1, + config={ + "num_attention_heads": 8, + "o_groups": 8, + "index_n_heads": 8, + "vision_n_layers": 1, + }, + args=StartArgs(enable_cpu_cache=True), + ) + + DeepseekV4TpPartModel._verify_params(model) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("is_hash", [False, True]) +def test_router_uses_bias_vl_only_for_image_tokens(is_hash): + from lightllm.models.deepseek_v4.layer_infer.transformer_layer_infer import DeepseekV4TransformerLayerInfer + + router = DeepseekV4TransformerLayerInfer.__new__(DeepseekV4TransformerLayerInfer) + router.has_vision = True + router.is_hash = is_hash + router.vocab_size = 32 + router.num_experts_per_tok = 6 + router.routed_scaling_factor = 1.5 + router.alloc_tensor = torch.empty + + torch.manual_seed(0) + logits = torch.randn((2, 256), dtype=torch.float32, device="cuda") + hash_table = torch.zeros((router.vocab_size, 6), dtype=torch.int64, device="cuda") + hash_table[2] = torch.tensor([1, 4, 7, 10, 13, 16], device="cuda") + text_bias = torch.zeros(256, dtype=torch.float32, device="cuda") + text_bias[20:26] = torch.tensor([60.0, 55.0, 50.0, 45.0, 40.0, 35.0], device="cuda") + vision_bias = torch.zeros(256, dtype=torch.float32, device="cuda") + vision_bias[200:206] = torch.tensor([60.0, 55.0, 50.0, 45.0, 40.0, 35.0], device="cuda") + layer_weight = SimpleNamespace( + n_routed_experts=256, + gate_weight_=SimpleNamespace(mm=lambda x, out: out.copy_(x)), + gate_tid2eid_=SimpleNamespace(weight=hash_table), + gate_bias_=SimpleNamespace(weight=text_bias), + gate_bias_vl_=SimpleNamespace(weight=vision_bias), + ) + infer_state = SimpleNamespace(is_prefill=True, input_ids=torch.tensor([2, 100_000], device="cuda")) + + weights, indices = router._select_experts(logits, infer_state, layer_weight) + + scores = torch.sqrt(torch.nn.functional.softplus(logits)) + text_indices = hash_table[infer_state.input_ids[0]] if is_hash else (scores[0] + text_bias).topk(6).indices + image_indices = (scores[1] + vision_bias).topk(6).indices + expected_indices = torch.stack((text_indices, image_indices)) + torch.testing.assert_close(indices, expected_indices) + expected = scores.gather(1, expected_indices) + expected = expected / expected.sum(dim=-1, keepdim=True) * 1.5 + torch.testing.assert_close(weights, expected, rtol=2e-5, atol=1e-6) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_build_image_visibility_scatter(): + from lightllm.models.deepseek_v4.triton_kernel.build_swa_indices import build_image_visibility + + image_spans = torch.tensor([[0, 1, 4], [0, 7, 3], [1, 6, 3]], dtype=torch.int32, device="cuda") + b_q_start_loc = torch.tensor([0, 6], dtype=torch.int32, device="cuda") + b_ready_cache_len = torch.tensor([2, 5], dtype=torch.int32, device="cuda") + b_q_seq_len = torch.tensor([6, 5], dtype=torch.int32, device="cuda") + image_left = torch.zeros(11, dtype=torch.int32, device="cuda") + image_right = torch.zeros(11, dtype=torch.int32, device="cuda") + + build_image_visibility( + image_spans, + b_q_start_loc, + b_ready_cache_len, + b_q_seq_len, + image_left, + image_right, + ) + + assert image_left.cpu().tolist() == [1, 2, 3, 0, 0, 0, 0, 0, 1, 2, 0] + assert image_right.cpu().tolist() == [2, 1, 0, 0, 0, 2, 0, 2, 1, 0, 0] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_swa_index_adds_bidirectional_image_visibility(): + from lightllm.models.deepseek_v4.triton_kernel.build_swa_indices import build_swa_indices + + positions = torch.tensor([5, 8, 10], dtype=torch.int32, device="cuda") + req_idx = torch.zeros(3, dtype=torch.int32, device="cuda") + prefill_state = SimpleNamespace( + prefill_token_ranges=torch.tensor([[0, 128]], dtype=torch.int32, device="cuda"), + swa_prefill_workspace_offsets=torch.tensor([128], dtype=torch.int32, device="cuda"), + HOLD_SWA_SLOT=256, + ) + requests = SimpleNamespace( + config=SimpleNamespace(page_token_size=128), + swa_capacity_tokens=128, + HOLD_SWA_SLOT=128, + HOLD_REQUEST_ID=1, + ) + write_slots = torch.empty_like(positions) + output = torch.empty((3, 8), dtype=torch.int32, device="cuda") + lengths = torch.empty(3, dtype=torch.int32, device="cuda") + image_left = torch.tensor([0, 2, 4], dtype=torch.int32, device="cuda") + image_right = torch.tensor([0, 2, 0], dtype=torch.int32, device="cuda") + + build_swa_indices( + req_idx, + positions, + requests, + output, + lengths, + write_slots, + sliding_window=4, + image_left=image_left, + image_right=image_right, + prefill_state=prefill_state, + ) + + assert output.cpu().tolist() == [ + [133, 132, 131, 130, -1, -1, -1, -1], + [133, 134, 135, 136, 137, 138, -1, -1], + [134, 135, 136, 137, 138, -1, -1, -1], + ] + assert lengths.cpu().tolist() == [4, 6, 5] + + +@pytest.mark.parametrize( + "loaded_start,load_end,spans,expected", + [ + (0, 4096, [], 4096), + (0, 2048, [(2048, 2500)], 2048), + (0, 4096, [(3500, 4096)], 4096), + (0, 4096, [(3500, 4500)], 2048), + (0, 4096, [(1000, 2500), (3500, 4500)], 0), + (2304, 4096, [(3000, 4500)], 2304), + ], +) +def test_common_cpu_cache_load_never_splits_an_image(monkeypatch, loaded_start, load_end, spans, expected): + from lightllm.server.router.model_infer.mode_backend import multi_level_kv_cache as common + + ends = list(range(2048, load_end + 1, 2048)) + page_ids = list(range(10, 10 + len(ends))) + released, loaded = [], [] + req = SimpleNamespace( + req_idx=0, + cur_kv_len=loaded_start, + hold_kv_len=loaded_start, + mm_input_spans=spans, + sampling_param=SimpleNamespace(shm_param=SimpleNamespace(prompt_logprobs=-1)), + shm_req=SimpleNamespace( + input_len=8192, + disk_prompt_cache_len=0, + cpu_cache_match_page_indexes=SimpleNamespace(get_all=lambda: page_ids), + token_hash_page_len_list=SimpleNamespace(get_all=lambda: ends), + ), + ) + table = torch.arange(8192, dtype=torch.int32).view(1, -1) + + def allocate(alloc_token_num): + if alloc_token_num == 0: + return None + start = req.cur_kv_len + req.hold_kv_len += alloc_token_num + return torch.arange(start, start + alloc_token_num, dtype=torch.int32) + + req.args = SimpleNamespace(page_size=256) + req.token_index = SimpleNamespace(get_gpu_token_indexes=lambda start, end: table[0, start:end]) + req.kv_cache_alloc_need = lambda target: max(0, target - req.hold_kv_len) + req.alloc_kv_mem = allocate + + from collections import deque + + module = common.MultiLevelKvCacheModule.__new__(common.MultiLevelKvCacheModule) + module.load_tasks = deque() + module.args = SimpleNamespace(page_size=256) + module.backend = SimpleNamespace( + is_master_in_dp=True, + dp_world_size=1, + is_hybrid_att_model=False, + model=SimpleNamespace( + req_manager=SimpleNamespace( + req_to_token_indexs=table, alloc_token_indexes=allocate, wait_token_indexes=lambda: None + ), + att_state_manager=SimpleNamespace( + get_cpu_cache_operator=lambda: SimpleNamespace( + get_page_layouts=lambda: {}, + read_req=lambda **kwargs: loaded.append(req.cur_kv_len + kwargs["token_indices"].numel()), + ), + full_att_allocator=None, + ), + ), + ) + module.cpu_cache_client = SimpleNamespace( + 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) + monkeypatch.setattr(common.dist, "barrier", lambda **kwargs: None) + monkeypatch.setattr(torch.Tensor, "cuda", lambda self, **kwargs: self) + module.load_cpu_cache_to_reqs([req]) + for task in getattr(module, "load_tasks", ()): + task.event.synchronize() + module.update_cpu_cache_load_task_states() + assert req.cur_kv_len == req.hold_kv_len == expected + assert loaded == ([expected] if expected > loaded_start else []) + assert released == page_ids diff --git a/unit_tests/models/deepseek_v4/test_workspace.py b/unit_tests/models/deepseek_v4/test_workspace.py new file mode 100644 index 0000000000..2bc4a56719 --- /dev/null +++ b/unit_tests/models/deepseek_v4/test_workspace.py @@ -0,0 +1,97 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.models.deepseek_v4.model import DeepseekV4TpPartModel +from lightllm.models.deepseek_v4.workspace import DeepseekV4C4PrefillWorkspace, DeepseekV4Workspace +from lightllm.models.deepseek_v4_dspark.model import DeepseekV4DSparkModel +from lightllm.models.deepseek_v4_mtp.model import DeepseekV4MTPModel + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + + +@pytest.mark.parametrize("draft_class", [DeepseekV4MTPModel, DeepseekV4DSparkModel]) +def test_draft_initialization_preserves_captured_prefill_buffer(monkeypatch, draft_class): + args = SimpleNamespace(llm_kv_type="fp8kv_dsa") + monkeypatch.setattr("lightllm.models.deepseek_v4.model.get_env_start_args", lambda: args) + + def backend_class(**kwargs): + return lambda model: SimpleNamespace(tp_q_head_num=2, padded_q_head_num=4) + + monkeypatch.setattr("lightllm.models.deepseek_v4.model.get_nsa_prefill_att_backend_class", backend_class) + monkeypatch.setattr("lightllm.models.deepseek_v4.model.get_nsa_decode_att_backend_class", backend_class) + main = DeepseekV4TpPartModel.__new__(DeepseekV4TpPartModel) + main.batch_max_tokens, main.tp_world_size_, main.data_type = 8, 1, torch.bfloat16 + main.config = {"head_dim": 512, "block_size": 3, "dspark_noise_token_id": 1} + main.layers_infer = [SimpleNamespace(compress_ratio=0)] + main.dsv4_workspace = DeepseekV4Workspace(main) + for name in ( + "_freqs_cis_sliding", + "_freqs_cis_compress", + "_cos_cached_sliding", + "_sin_cached_sliding", + "_cos_cached_compress", + "_sin_cached_compress", + ): + setattr(main, name, torch.ones((1, 32), device="cuda")) + main._init_att_backend() + buffer = main.dsv4_workspace.flashmla_prefill_q + buffer[:, :2].fill_(2) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + result = main.dsv4_workspace.flashmla_prefill_q + 1 + + draft = draft_class.__new__(draft_class) + draft.main_model, draft.config, draft.data_type = main, main.config, main.data_type + draft.layers_infer = [SimpleNamespace(compress_ratio=0)] + draft.context_wkv_weight = object() + draft._init_custom() + draft._init_att_backend() + + assert draft.dsv4_workspace is main.dsv4_workspace + assert main.dsv4_workspace.flashmla_prefill_q is buffer + assert draft.layers_infer[0].flashmla_q_head_num_ == 4 + main.dsv4_workspace.flashmla_prefill_q[:, :2].fill_(4) + graph.replay() + torch.testing.assert_close(result[:, :2], torch.full_like(result[:, :2], 5), rtol=0, atol=0) + torch.testing.assert_close(result[:, 2:], torch.ones_like(result[:, 2:]), rtol=0, atol=0) + + +@pytest.mark.parametrize("batch_max_tokens,tp,capacity", [(4097, 8, 4104), (4096, 8, 4096), (1, 8, 8), (3, 1, 3)]) +def test_workspace_covers_tp_padding_and_checks_q_capacity(batch_max_tokens, tp, capacity): + from lightllm.models.deepseek_v4.triton_kernel.norm_rope import fused_q_norm_rope + + model = SimpleNamespace( + batch_max_tokens=batch_max_tokens, + tp_world_size_=tp, + config={"index_n_heads": 2, "index_head_dim": 128}, + ) + workspace = DeepseekV4Workspace(model) + assert workspace.token_capacity == capacity + workspace.init_flashmla_prefill_q(2, 4, 512, torch.bfloat16) + workspace.init_c4_prefill_aux(model) + q = workspace.get_flashmla_prefill_q(capacity) + q_input = torch.ones((capacity, 2, 512), dtype=torch.bfloat16, device="cuda") + freqs = torch.ones((1, 32), dtype=torch.complex64, device="cuda") + positions = torch.zeros(capacity, dtype=torch.int32, device="cuda") + fused_q_norm_rope(q_input, q[:, :2], 1e-6, freqs, positions) + torch.testing.assert_close(q[:, :2], q_input, rtol=0, atol=0) + assert q[:, 2:].count_nonzero() == 0 + assert workspace.c4_prefill_aux.indexer_q_outputs(capacity)[0].shape[0] == capacity + with pytest.raises(ValueError, match="capacity"): + workspace.get_flashmla_prefill_q(capacity + 1) + + +@pytest.mark.parametrize("overlap", [False, True]) +def test_c4_aux_is_one_workspace(overlap): + model = SimpleNamespace( + batch_max_tokens=4, + tp_world_size_=1, + config={"index_n_heads": 2, "index_head_dim": 8}, + args=SimpleNamespace(enable_prefill_microbatch_overlap=overlap), + ) + workspace = DeepseekV4Workspace(model) + workspace.init_c4_prefill_aux(model) + assert isinstance(workspace.c4_prefill_aux, DeepseekV4C4PrefillWorkspace) diff --git a/unit_tests/models/glm5_next/test_cache.py b/unit_tests/models/glm5_next/test_cache.py new file mode 100644 index 0000000000..6da740353d --- /dev/null +++ b/unit_tests/models/glm5_next/test_cache.py @@ -0,0 +1,270 @@ +from unit_tests.common.att_resources.model_helpers import build_hybrid_resources +import dataclasses +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.attention_state.allocators import FullAttKVAllocator + +from lightllm.common.attention_state.full_att import ( + PackedMLAFullAttentionState, + PackedMLAAttentionConfig, + PackedMLALogicalConfig, +) +from lightllm.common.req_manager import ReqManager +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 + + +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], + }, + }, + StartArgs(tp=4, data_type="bfloat16"), + ) + # 模型配置保留原始 latent 宽度；packed 行宽由 State 创建入口计算。 + assert config.full_att_head_dim == 512 + + +@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("mode", ["native", "replay"]) +def test_hybrid_checkpoint_restore_and_packed_kv_copy( + monkeypatch, + small_page, + tp_world_size, + mtp_step, + mode, +): + 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, + data_type="bfloat16", + linear_att_hash_page_size=4, + linear_att_page_block_num=2, + cpu_cache_token_page_size=8, + mtp_step=mtp_step, + mtp_mode="eagle_with_att" if mtp_step else None, + ) + set_unique_server_name(args) + get_env_start_args.cache_clear() + set_env_start_args(dataclasses.asdict(args)) + config = Glm5NextResourceConfig( + tp_world_size=tp_world_size, + full_att_all_num_kv_heads=1, + full_att_head_dim=512, + global_linear_k_heads=2 * tp_world_size, + global_linear_v_heads=2 * tp_world_size, + head_linear_k_dim=128, + head_linear_v_dim=128, + conv_kernel_size=4, + conv_state_dtype=torch.bfloat16, + ssm_state_dtype=torch.float32, + full_attention_interval=4, + all_layer_num=4, + ) + full_layer_ids = config.get_full_att_layer_indices() + if mtp_step: + full_layer_ids += (config.all_layer_num,) + full = PackedMLAFullAttentionState( + PackedMLAAttentionConfig( + 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, + ) + full.allocate_storage(16) + req = ReqManager(3, 16) + 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 = build_hybrid_resources( + config=config, + request_capacity=4, + mtp_step=mtp_step, + ssm_state_mode=mode, + 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, + ) + full = main_full + linear = resources.main_kda_conv_bf16_ssm_fp32_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) + 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( + 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) + 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_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]) + linear.conv_state[:, 0].normal_() + linear.ssm_state[:, 0].normal_() + conv = linear.conv_state[:, 0, ..., :3].clone() + ssm = linear.ssm_state[:, 0].clone() + if mode == "replay": + linear.raw_keys.zero_() + linear.raw_values.zero_() + linear.gates.fill_(-0.5) + linear.betas.zero_() + linear.cursors[0] = 1 + linear.cursors[2] = 3 + if mtp_step: + linear.mtp_state_index[2] = mtp_step + resources.save_bounded_att_runtime_to_pages( + req_indices=[0], + state_token_lens=[8], + page_kind=kind, + page_slot_ids=[slot], + ) + torch.cuda.synchronize() + linear.conv_state[:, 0].zero_() + linear.ssm_state[:, 0].zero_() + 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], + ) + torch.cuda.synchronize() + assert torch.equal(linear.conv_state[:, 2, ..., :3], conv) + restored_slot = 2 if mode == "replay" else 2 * (linear.mtp_step + 1) + assert torch.equal(linear.ssm_state[:, restored_slot], ssm) + if mode == "replay": + assert linear.cursors[2] == 0 + if mtp_step: + assert linear.mtp_state_index[2] == 0 + # 整组边界不依赖旧窗口，恢复大小页无需改动无重叠窗口。 + assert torch.equal(tail.buffer, unchanged) + resources.reset_bounded_att_state([2]) + assert not linear.conv_state[:, 2].any() + if mode == "replay": + assert not linear.ssm_state[:, 2].any() + else: + assert not linear.ssm_state[:, 2 * (linear.mtp_step + 1) : 3 * (linear.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 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) + packed_bytes[:, 0].random_(0, 256) + full.copy_full_att_tokens(torch.tensor([0], device="cuda"), torch.tensor([7], device="cuda")) + assert torch.equal(packed_bytes[:, 0], packed_bytes[:, 7]) + + +@pytest.mark.parametrize("gate_type", ["silu", "sigmoid"]) +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( + 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=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=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( + 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( + 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( + actual=out, + expected=expected.bfloat16(), + atol=0.008, + rtol=0.008, + ) diff --git a/unit_tests/models/glm5_next/test_compact_topk_modes.py b/unit_tests/models/glm5_next/test_compact_topk_modes.py new file mode 100644 index 0000000000..973ef62c07 --- /dev/null +++ b/unit_tests/models/glm5_next/test_compact_topk_modes.py @@ -0,0 +1,65 @@ +import pytest +import torch +from triton.runtime.jit import JITFunction + +from lightllm.models.glm5_next.triton_kernel.fp8_prefill import _compact_topk_indices_kernel + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.mark.parametrize("index_dtype", [torch.int32, torch.int64]) +@pytest.mark.parametrize("aligned_chunks", [False, True]) +def test_compact_topk_preserves_layout_and_reuses_kernel(monkeypatch, index_dtype, aligned_chunks): + torch.manual_seed(1525) + misses, hashes = [], set() + original = JITFunction._do_compile + + def observe(self, *args, **kwargs): + if self is _compact_topk_indices_kernel: + misses.append(1) + return original(self, *args, **kwargs) + + monkeypatch.setattr(JITFunction, "_do_compile", observe) + # Exercise aligned production chunks and unaligned chunk boundaries separately. + chunk_starts = [0, 2097152, 4194304, 0] if aligned_chunks else [127, 8193, 8193, 127] + for index in range(4): + batch = [1, 3, 33, 65][index] + low = chunk_starts[index] + high = low + [1, 2305, 16385, 30000][index] + topk = torch.randint(-7, 20000, (batch, 2176), device="cuda", dtype=index_dtype) + topk[:, ::17] = -1 + base = max(0, low - 32) + starts = torch.arange(batch, device="cuda", dtype=torch.int32) * 31 + base + # Include positive indices outside either end of the current chunk. + topk[0, 1:6] = torch.tensor( + [low - 2 - base, low - 1 - base, low - base, high - 1 - base, high - base], + device="cuda", + dtype=torch.int32, + ) + output = torch.full_like(topk, -99999) + lengths = torch.full((batch,), -987, device="cuda", dtype=torch.int32) + before_misses = len(misses) + compiled = _compact_topk_indices_kernel[(batch,)]( + topk, + starts, + output, + lengths, + low, + high, + topk.stride(0), + topk.stride(1), + TOPK=2176, + BLOCK=4096, + ) + if hashes: + assert len(misses) == before_misses + hashes.add(compiled.hash) + ragged = topk + starts[:, None] + valid = (topk >= 0) & (ragged >= low) & (ragged < high) + assert torch.equal(lengths, valid.sum(1).int()) + for row in range(batch): + expected = ragged[row][valid[row]] - low + assert torch.equal(output[row, : expected.numel()], expected) + assert (output[row, expected.numel() :] == -99999).all() + assert len(hashes) == 1 diff --git a/unit_tests/models/glm5_next/test_direct_nsa_indices.py b/unit_tests/models/glm5_next/test_direct_nsa_indices.py new file mode 100644 index 0000000000..d0e6dcb9bd --- /dev/null +++ b/unit_tests/models/glm5_next/test_direct_nsa_indices.py @@ -0,0 +1,257 @@ +import weakref +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.basemodel.attention.base_att import AttControl +from lightllm.common.basemodel.attention.nsa.glm5_next import ( + Glm5NextSparseAttBackend, + Glm5NextSparseDecodeState, +) +from lightllm.common.basemodel.attention.nsa.fp8_glm5_next import ( + Fp8Glm5NextSparseAttBackend, + Fp8Glm5NextSparseDecodeState, +) +from lightllm.common.basemodel.infer_struct import InferStateInfo +from lightllm.common.attention_state.full_att.packed_mla.fp8_e8m0_packed_mla_write import write_fp8_e8m0_packed_kv +from lightllm.common.basemodel.attention.nsa import fp8_glm5_next +from lightllm.models.glm5_next.triton_kernel.kpool import expand_topk, prefill_nsa_params + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +def tensor(values): + return torch.tensor(values, device="cuda", dtype=torch.int32) + + +@pytest.mark.parametrize("counts", [(3, 0, 1, 2), (257, 0, 1, 513), (1,) * 129]) +def test_prefill_metadata_reuses_existing_prefixes_during_graph_replay(counts): + batch, queries = len(counts), sum(counts) + positions = torch.empty(queries, device="cuda", dtype=torch.int32) + req = torch.empty(batch, device="cuda", dtype=torch.int32) + cu_q, cu_kv = torch.empty(batch + 1, device="cuda", dtype=torch.int32), torch.empty( + batch + 1, device="cuda", dtype=torch.int32 + ) + + def update(counts, phase): + q_prefix, kv_prefix, visible, starts, requests = [0], [0], [], [], [] + req_values = [(batch - i + phase) % (batch + 1) for i in range(batch)] + for i, count in enumerate(counts): + prefix = 1000000 + i * 7 + phase + visible.extend(range(prefix + 1, prefix + count + 1)) + starts.extend([kv_prefix[-1]] * count) + requests.extend([req_values[i]] * count) + q_prefix.append(q_prefix[-1] + count) + kv_prefix.append(kv_prefix[-1] + prefix + count) + positions.copy_(tensor(visible) - 1) + req.copy_(tensor(req_values)) + cu_q.copy_(tensor(q_prefix)) + cu_kv.copy_(tensor(kv_prefix)) + return tensor(visible), tensor(requests), tensor(starts) + + expected = update(counts, 0) + actual = prefill_nsa_params(positions, req, cu_q, cu_kv) + for observed, reference in zip(actual, expected): + torch.testing.assert_close(observed, reference, rtol=0, atol=0) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = prefill_nsa_params(positions, req, cu_q, cu_kv) + expected = update(counts[::-1], 11) + graph.replay() + for observed, reference in zip(actual, expected): + torch.testing.assert_close(observed, reference, rtol=0, atol=0) + + +@pytest.mark.parametrize("dense", [False, True]) +def test_topk_direct_request_lookup_preserves_mtp_hold_and_strided_tables(dense): + table = torch.arange(8 * 70, device="cuda", dtype=torch.int32).reshape(8, 70)[::2, ::2] + req = tensor([2, 2, 0, 3, 3]) + lengths = tensor([5, 7, 9, 2, 2] if dense else [17, 18, 9, 2, 2]) + groups = tensor([[3, 1], [1, 0], [0, 1], [0, 0], [0, 0]]) + + def check(actual): + mem, relative = actual + expected_mem, expected_relative = torch.full_like(mem, -1), torch.full_like(relative, -1) + for i, (request, length) in enumerate(zip(req.tolist(), lengths.tolist())): + tokens = ( + list(range(length)) + if dense + else [group * 4 + j for group in groups[i].tolist()[: min(length // 4, 2)] for j in range(4)] + + list(range(length // 4 * 4, length)) + ) + expected_relative[i, : len(tokens)] = tensor(tokens) + expected_mem[i, : len(tokens)] = table[request, tensor(tokens).long()] + torch.testing.assert_close(mem, expected_mem, rtol=0, atol=0) + torch.testing.assert_close(relative, expected_relative, rtol=0, atol=0) + + check(expand_topk(groups, lengths, req, table, 8, dense=dense)) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = expand_topk(groups, lengths, req, table, 8, dense=dense) + req.copy_(tensor([0, 0, 2, 3, 3])) + table.add_(10000) + lengths[:2].add_(1) + graph.replay() + check(actual) + + +@pytest.mark.parametrize("fp8", [False, True]) +def test_decode_aliases_infer_state_without_ragged_allocation(monkeypatch, fp8): + import lightllm.common.basemodel.triton_kernel.gen_nsa_ks_ke as legacy + + def forbidden(*args, **kwargs): + raise AssertionError("GLM must not generate or copy a full ragged index") + + monkeypatch.setattr(legacy, "gen_nsa_ks_ke", forbidden) + monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", str(torch.cuda.current_device())) + model = type( + "Model", + (), + { + "config": {"index_kpool": 4, "index_topk": 2048}, + "batch_max_tokens": 4, + "max_req_num": 1, + "max_seq_length": 32, + "att_state_role": "decode_test", + "args": SimpleNamespace(enable_prefill_microbatch_overlap=False, run_mode="decode"), + }, + )() + backend = (Fp8Glm5NextSparseAttBackend if fp8 else Glm5NextSparseAttBackend)(model) + assert not hasattr(backend, "ragged_mem_buffers") + infer = InferStateInfo() + infer.b_req_idx = tensor([2, 2, 0, 3, 3]) + infer.b_seq_len = tensor([999999, 1000000, 65536, 2, 2]) + cls = Fp8Glm5NextSparseDecodeState if fp8 else Glm5NextSparseDecodeState + state = cls(backend=backend, _infer_state_weak_ref=weakref.ref(infer)) + state.init_state() + assert state.q_kv_seq_len is infer.b_seq_len + assert state.q_req_idx is infer.b_req_idx + assert not hasattr(state, "ragged_mem_index") and not hasattr(state, "ks") + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + state.init_state() + infer.b_seq_len[:3].sub_(100) + infer.b_req_idx[:3].copy_(tensor([0, 0, 1])) + graph.replay() + assert state.q_kv_seq_len is infer.b_seq_len and state.q_req_idx is infer.b_req_idx + + +@pytest.mark.parametrize( + "queries,heads,total_kv,chunk_size", + [ + (1, 64, 65536, 2097152), + (2, 64, 65536, 2097152), + (3, 8, 65536, 2097152), + (9, 64, 512, 2097152), + (9, 64, 131073, 2097152), + (4, 16, 40, 3), + (4, 64, 40, 7), + ], +) +def test_fp8_prefill_reads_request_table_without_ragged_buffer(monkeypatch, queries, heads, total_kv, chunk_size): + monkeypatch.setattr(fp8_glm5_next, "KV_GATHER_CHUNK_SIZE", chunk_size) + + torch.manual_seed(1685) + batch = min(queries, 2) + counts = [queries] if batch == 1 else [queries // 2, queries - queries // 2] + seqs = [total_kv] if batch == 1 else [total_kv // 2, total_kv - total_kv // 2] + req = tensor([2, 0][:batch]) + table = torch.full((4, 2 * (max(seqs) + 7)), -1, device="cuda", dtype=torch.int32)[:, ::2] + physical = torch.randperm(total_kv + 32, device="cuda", dtype=torch.int32)[:total_kv] + offset = 0 + for request, length in zip(req.tolist(), seqs): + table[request, :length] = physical[offset : offset + length] + offset += length + q = torch.randn(queries, heads, 512, device="cuda", dtype=torch.bfloat16) + source = torch.randn(total_kv + 32, 1, 512, device="cuda", dtype=torch.bfloat16) + packed = torch.zeros(total_kv + 32, 1, 720, device="cuda", dtype=torch.uint8) + write_fp8_e8m0_packed_kv(source, torch.arange(total_kv + 32, device="cuda", dtype=torch.int32), packed) + prefill_kv = torch.randn(queries, 1, 512, device="cuda", dtype=torch.bfloat16) + visible, query_req, current_positions, q_prefix, kv_prefix = [], [], [], [0], [0] + for request, length, count in zip(req.tolist(), seqs, counts): + visible.extend(range(length - count + 1, length + 1)) + query_req.extend([request] * count) + current_positions.extend(range(length - count, length)) + q_prefix.append(q_prefix[-1] + count) + kv_prefix.append(kv_prefix[-1] + length) + query_req, current_positions = tensor(query_req), tensor(current_positions) + mem_index = table[query_req.long(), current_positions.long()] + topk = torch.full((queries, 2176), -1, device="cuda", dtype=torch.int32) + for i, length in enumerate(visible): + topk[i, :7] = tensor([0, length - 1, 0, 1, length // 2, length - 1, -1]) + topk_mem = torch.empty_like(topk) + + def refresh_physical(): + topk_mem.copy_(table[query_req.long()[:, None], topk.clamp_min(0).long()]) + topk_mem.masked_fill_(topk < 0, -1) + mem_index.copy_(table[query_req.long(), current_positions.long()]) + + refresh_physical() + infer = InferStateInfo() + infer.mem_index, infer.b_req_idx = mem_index, req + infer.position_ids = tensor(visible) - 1 + infer.b1_cu_q_seq_len, infer.b1_cu_kv_seq_len = tensor(q_prefix), tensor(kv_prefix) + infer.total_token_num = total_kv + infer.req_manager = SimpleNamespace(req_to_token_indexs=table) + monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", str(torch.cuda.current_device())) + model = type( + "Model", + (), + { + "batch_max_tokens": queries, + "prefill_graph": None, + "max_req_num": batch, + "max_seq_length": max(seqs), + "att_state_role": "indices_test", + "config": {"index_topk": 2048}, + "args": SimpleNamespace(enable_prefill_microbatch_overlap=False, run_mode="normal"), + }, + )() + backend = Fp8Glm5NextSparseAttBackend(model) + state = backend.create_att_prefill_state(infer) + state.init_state() + assert not hasattr(state, "ragged_mem_index") and not hasattr(backend, "ragged_mem_buffers") + control = AttControl( + nsa_prefill=True, + nsa_prefill_dict={ + "topk_indices": topk, + "topk_mem_indices": topk_mem, + "prefill_cache_kv": prefill_kv, + "softmax_scale": 512 ** -0.5, + "kv_lora_rank": 512, + }, + ) + + def run(): + return state.prefill_att(q, packed[..., :584], None, control) + + def check(output): + scales = packed[..., 576:583].view(torch.float8_e8m0fnu).float().reshape(-1, 7, 1) + values = packed[..., :448].view(torch.float8_e4m3fn).float().reshape(-1, 7, 64) * scales + kv = torch.cat((values.reshape(-1, 1, 448).bfloat16(), packed[..., 448:576].view(torch.bfloat16)), -1) + kv[mem_index.long()] = prefill_kv + expected = [] + for i in range(queries): + keys = kv[topk_mem[i][topk_mem[i] >= 0].long(), 0].float() + expected.append((q[i].float() @ keys.T * 512 ** -0.5).softmax(-1) @ keys) + torch.testing.assert_close(output.float(), torch.stack(expected), atol=0.012, rtol=0.015) + + check(run()) + workspace = backend._prefill_workspaces[0] + if workspace is not None: + assert set(workspace) == {"gathered_kv", "chunk_topk_indices", "chunk_cache_seqlens", "cu_seqlens_q"} + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = run() + graph.replay() + check(output) + q.normal_() + prefill_kv.normal_() + for request, length in zip(req.tolist(), seqs): + table[request, :length].copy_(table[request, :length].roll(1)) + topk[0].fill_(-1) + refresh_physical() + graph.replay() + check(output) diff --git a/unit_tests/models/glm5_next/test_fp8_compile_cache.py b/unit_tests/models/glm5_next/test_fp8_compile_cache.py new file mode 100644 index 0000000000..6a07b19811 --- /dev/null +++ b/unit_tests/models/glm5_next/test_fp8_compile_cache.py @@ -0,0 +1,94 @@ +import pytest +import torch + +from lightllm.models.glm5_next.triton_kernel import ( + fp8_prefill as fp8, +) + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +def record_kernel_hashes(monkeypatch, kernel): + hashes = [] + original = kernel.run + + def run(*args, **kwargs): + compiled = original(*args, **kwargs) + hashes.append(compiled.hash) + return compiled + + monkeypatch.setattr(kernel, "run", run) + return hashes + + +@pytest.mark.parametrize("transposed_lse", [False, True]) +def test_online_attention_merge_reuses_kernel_across_query_counts(monkeypatch, transposed_lse): + hashes = record_kernel_hashes(monkeypatch, fp8._merge_att_output_kernel) + for num_q_tokens in [1, 3, 8, 16, 17, 33, 64]: + output = torch.randn(num_q_tokens, 4, 512, dtype=torch.float32, device="cuda") + lse = torch.randn(4, num_q_tokens, dtype=torch.float32, device="cuda") + chunk_output = torch.randn_like(output) + chunk_lse = ( + torch.randn(num_q_tokens, 4, dtype=torch.float32, device="cuda").T + if transposed_lse + else torch.randn_like(lse) + ) + expected_lse = torch.logaddexp(lse, chunk_lse) + expected = output * torch.exp(lse - expected_lse).T.unsqueeze(-1) + chunk_output * torch.exp( + chunk_lse - expected_lse + ).T.unsqueeze(-1) + fp8._merge_att_output_kernel[(num_q_tokens, 1)]( + output, + lse, + chunk_output, + chunk_lse, + chunk_lse.stride(1), + chunk_lse.stride(0), + q_head_num=4, + q_token_num=num_q_tokens, + BLOCK_HEADS=4, + num_warps=4, + ) + torch.testing.assert_close(output, expected, rtol=2e-5, atol=2e-6) + torch.testing.assert_close(lse, expected_lse, rtol=2e-5, atol=2e-6) + assert len(hashes) == 7 and len(set(hashes)) == 1 + + +def test_prefill_kv_gather_reuses_kernel_across_token_counts(monkeypatch): + hashes = record_kernel_hashes(monkeypatch, fp8._gather_prefill_kv_kernel) + cached_kv = torch.randn(64, 512, dtype=torch.bfloat16, device="cuda") + packed_kv = torch.zeros(64, 1, 584, dtype=torch.uint8, device="cuda") + packed_kv[:, 0, :448] = cached_kv[:, :448].to(torch.float8_e4m3fn).view(torch.uint8) + packed_kv[:, 0, 448:576] = cached_kv[:, 448:].contiguous().view(torch.uint8) + packed_kv[:, 0, 576:583] = 127 + cached_kv[:, :448] = cached_kv[:, :448].to(torch.float8_e4m3fn).to(torch.bfloat16) + table = torch.arange(64, dtype=torch.int32, device="cuda").view(1, -1) + requests = torch.zeros(1, dtype=torch.int32, device="cuda") + cu_q = torch.tensor([0, 1], dtype=torch.int32, device="cuda") + for num_tokens in [1, 3, 8, 16, 17, 32, 33]: + cu_kv = torch.tensor([0, num_tokens], dtype=torch.int32, device="cuda") + prefill_kv = torch.randn(1, 1, 512, dtype=torch.bfloat16, device="cuda") + output = torch.empty(num_tokens, 1, 512, dtype=torch.bfloat16, device="cuda") + fp8._gather_prefill_kv_kernel[(num_tokens // 4 + 1,)]( + packed_kv, + prefill_kv, + output, + table, + requests, + cu_kv, + cu_q, + 0, + num_tokens, + packed_kv.stride(0), + prefill_kv.stride(0), + prefill_kv.stride(2), + req_to_token_indexs_stride_b=table.stride(0), + req_to_token_indexs_stride_s=table.stride(1), + batch_size=1, + BLOCK_ROWS=4, + num_warps=4, + ) + expected = cached_kv[:num_tokens].clone() + expected[-1] = prefill_kv[0, 0] + torch.testing.assert_close(output[:, 0], expected, rtol=0, atol=0) + assert len(hashes) == 7 and len(set(hashes)) == 1 diff --git a/unit_tests/models/glm5_next/test_fp8_kv.py b/unit_tests/models/glm5_next/test_fp8_kv.py new file mode 100644 index 0000000000..a6a5c7c9f2 --- /dev/null +++ b/unit_tests/models/glm5_next/test_fp8_kv.py @@ -0,0 +1,1073 @@ +import weakref +from types import SimpleNamespace + +import pytest +import torch +import triton + +from lightllm.common.basemodel.attention.base_att import AttControl +from lightllm.common.basemodel.infer_struct import InferStateInfo +from lightllm.common.basemodel.attention.nsa.glm5_next import ( + Glm5NextSparsePrefillState, + Glm5NextSparseDecodeState, +) +from lightllm.common.basemodel.attention.nsa.fp8_glm5_next import ( + Fp8Glm5NextSparseAttBackend, + Fp8Glm5NextSparsePrefillState, + Fp8Glm5NextSparseDecodeState, +) +from lightllm.models.glm5_next.triton_kernel.fp8_sparse_decode import ( + can_use_fp8_sparse_direct_decode, + 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.packed_mla.fp8_e8m0_packed_mla_write import write_fp8_e8m0_packed_kv +from lightllm.common.basemodel.attention.nsa import fp8_glm5_next +from lightllm.models.glm5_next.triton_kernel.fp8_prefill import ( + _gather_prefill_kv_kernel, +) + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.fixture(autouse=True) +def current_cuda_device(monkeypatch): + monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", str(torch.cuda.current_device())) + + +def _make_infer_state(**attributes): + state = InferStateInfo() + vars(state).update(attributes) + return state + + +def _make_workspace_backend(batch_max_tokens, overlap=False): + backend = object.__new__(Fp8Glm5NextSparseAttBackend) + backend.model = SimpleNamespace( + batch_max_tokens=batch_max_tokens, + prefill_graph=None, + config={"index_topk": 2048, "index_kpool": 4}, + max_req_num=1, + max_seq_length=fp8_glm5_next.KV_GATHER_CHUNK_SIZE, + att_state_role="unit_test", + args=SimpleNamespace(enable_prefill_microbatch_overlap=overlap, run_mode="normal"), + ) + backend._prefill_workspaces = backend.allocate_prefill_workspaces(backend.model) + return backend + + +class _PrefillCase: + def __init__(self, seq_lens, q_lens, heads, width=2176, capacity=None, backend=None, softmax_scale=0.0625): + device = "cuda" + tensor = lambda values: torch.tensor(values, device=device, dtype=torch.int32) + self.seq_lens, self.q_lens = seq_lens, q_lens + self.requests = list(range(2 * len(seq_lens) - 1, 0, -2)) + total, queries = sum(seq_lens), sum(q_lens) + capacity = capacity or total + 17 + self.req_to_token_indexs = torch.full( + (2 * len(seq_lens) + 1, max(seq_lens) + 7), -1, device=device, dtype=torch.int32 + ) + slots = torch.randperm(capacity, device=device, dtype=torch.int32)[:total] + cu_q, cu_kv, q_req_idx, positions = [0], [0], [], [] + for req, length, count in zip(self.requests, seq_lens, q_lens): + self.req_to_token_indexs[req, :length] = slots[cu_kv[-1] : cu_kv[-1] + length] + q_req_idx.extend([req] * count) + positions.extend(range(length - count, length)) + cu_q.append(cu_q[-1] + count) + cu_kv.append(cu_kv[-1] + length) + self.q_req_idx, self.position_ids = tensor(q_req_idx).long(), tensor(positions).long() + self.q = torch.randn(queries, heads, 512, device=device, dtype=torch.bfloat16) + self.source_kv = torch.randn(capacity, 1, 512, device=device, dtype=torch.bfloat16) + self.packed_kv = torch.zeros(capacity, 1, 720, device=device, dtype=torch.uint8) + self.all_slots = torch.arange(capacity, device=device, dtype=torch.int32) + write_fp8_e8m0_packed_kv(self.source_kv, self.all_slots, self.packed_kv) + self.prefill_cache_kv = torch.randn(queries, 1, 1024, device=device, dtype=torch.bfloat16)[..., ::2] + visible = tensor(positions) + 1 + self.topk_indices = torch.randint(max(seq_lens), (queries, width), device=device, dtype=torch.int32) + self.topk_indices.remainder_(visible[:, None]) + self.topk_indices.masked_fill_( + torch.arange(width, device=device)[None, :] >= visible.clamp_max(2048)[:, None], -1 + ) + self.topk_indices[:, ::19] = -1 + self.topk_indices[:, 0] = visible - 1 + self.topk_indices[:, 1] = visible - 1 + self.topk_mem_indices = torch.empty(queries, width * 2, device=device, dtype=torch.int32)[:, ::2] + self.infer_state = _make_infer_state( + is_prefill=True, + req_manager=SimpleNamespace(req_to_token_indexs=self.req_to_token_indexs), + b_req_idx=tensor(self.requests), + b_seq_len=tensor(seq_lens), + b1_cu_q_seq_len=tensor(cu_q), + b1_cu_kv_seq_len=tensor(cu_kv), + position_ids=tensor(positions), + total_token_num=total, + mem_index=torch.empty(queries, device=device, dtype=torch.int32), + ) + self.backend = backend or _make_workspace_backend(queries) + _, self.att_state = self.new_state() + self.control = AttControl( + nsa_prefill=True, + nsa_prefill_dict={ + "topk_indices": self.topk_indices, + "topk_mem_indices": self.topk_mem_indices, + "prefill_cache_kv": self.prefill_cache_kv, + "softmax_scale": softmax_scale, + "kv_lora_rank": 512, + }, + ) + + def new_state(self, **overrides): + if overrides: + infer = _make_infer_state(**vars(self.infer_state)) + vars(infer).update(overrides) + else: + infer = self.infer_state + state = self.backend.create_att_prefill_state(infer) + state.init_state() + return infer, state + + def refresh(self): + self.infer_state.mem_index.copy_(self.req_to_token_indexs[self.q_req_idx, self.position_ids]) + self.topk_mem_indices.copy_( + self.req_to_token_indexs[self.q_req_idx[:, None], self.topk_indices.clamp_min(0).long()] + ) + self.topk_mem_indices.masked_fill_(self.topk_indices < 0, -1) + self.control.nsa_prefill_dict["topk_indices"] = self.topk_indices + + def run(self, state=None): + self.refresh() + active = self.att_state if state is None else state + return active.prefill_att(self.q, self.packed_kv[..., :584], None, self.control) + + def check(self, actual): + values = _dequantize_kv_cache(self.packed_kv) + values[self.infer_state.mem_index.long()] = self.prefill_cache_kv + physical = self.req_to_token_indexs[self.q_req_idx[:, None], self.topk_indices.clamp_min(0).long()].masked_fill( + self.topk_indices < 0, -1 + ) + expected = _torch_sparse_attention(self.q, values, physical, self.control.nsa_prefill_dict["softmax_scale"]) + torch.testing.assert_close(actual.float(), expected, atol=0.012, rtol=0.015) + + def mutate(self): + self.q.normal_() + self.prefill_cache_kv.normal_() + self.source_kv.normal_() + write_fp8_e8m0_packed_kv(self.source_kv, self.all_slots, self.packed_kv) + for req, length in zip(self.requests, self.seq_lens): + self.req_to_token_indexs[req, :length].copy_(self.req_to_token_indexs[req, :length].roll(1)) + self.topk_indices.copy_(self.topk_indices.roll(1, dims=1)) + + +@pytest.mark.parametrize("draft", [False, True]) +@pytest.mark.parametrize("kv_type", ["None", "fp8kv_dsa"]) +def test_main_and_native_mtp_backend_selection(monkeypatch, draft, kv_type): + from lightllm.models.glm5_next.model import Glm5NextTpPartModel + from lightllm.models.glm5_next_mtp.model import Glm5NextMTPModel + + monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") + model = object.__new__(Glm5NextMTPModel if draft else Glm5NextTpPartModel) + model.args = SimpleNamespace(llm_kv_type=kv_type, enable_prefill_microbatch_overlap=False, run_mode="normal") + model.config = {"index_topk": 2048, "index_kpool": 4} + model.graph_max_batch_size = 4 + model.max_seq_length = 128 + model.max_req_num = 2 + model.batch_max_tokens = 4 + model.mtp_previous_draft_models = [] + model._init_att_backend() + infer_state = InferStateInfo() + prefill = model.prefill_att_backend.create_att_prefill_state(infer_state) + decode = model.decode_att_backend.create_att_decode_state(infer_state) + assert type(prefill) is (Fp8Glm5NextSparsePrefillState if kv_type == "fp8kv_dsa" else Glm5NextSparsePrefillState) + assert type(decode) is (Fp8Glm5NextSparseDecodeState if kv_type == "fp8kv_dsa" else Glm5NextSparseDecodeState) + infer_state.prefill_att_state = prefill + infer_state.decode_att_state = decode + infer_state.check_att_states() + assert prefill.infer_state is infer_state + assert decode.infer_state is infer_state + owner_ref = weakref.ref(infer_state) + del infer_state + assert owner_ref() is None + with pytest.raises(ReferenceError): + prefill.infer_state + with pytest.raises(ReferenceError): + decode.infer_state + + +def _dequantize_kv_cache(packed_kv): + values = packed_kv[..., :448].view(torch.float8_e4m3fn).float().reshape(-1, 7, 64) + scales = packed_kv[..., 576:583].view(torch.float8_e8m0fnu).float().reshape(-1, 7, 1) + tail = packed_kv[..., 448:576].view(torch.bfloat16) + return torch.cat(((values * scales).reshape(-1, 1, 448).bfloat16(), tail), dim=-1) + + +def _torch_sparse_attention(q, kv, indices, softmax_scale=0.0625): + outputs = [] + for row in range(q.shape[0]): + keys = kv[indices[row][indices[row] >= 0].long(), 0].float() + outputs.append((q[row].float() @ keys.T * softmax_scale).softmax(-1) @ keys) + return torch.stack(outputs) + + +def _capture_cuda_graph(run): + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + run() + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = run() + return graph, output + + +@pytest.mark.parametrize("index_dtype", [torch.int32, torch.int64]) +@pytest.mark.parametrize("source_d_stride", [1, 2]) +def test_fp8_scatter_quantization_and_gather_prefill(index_dtype, source_d_stride): + torch.manual_seed(53) + # Noncontiguous KV from a fused projection, and a strided packed cache. + raw_source = torch.randn(11, 1, 512 * source_d_stride, device="cuda", dtype=torch.bfloat16) + source = raw_source if source_d_stride == 1 else raw_source[..., ::2] + source[0].zero_() + 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) + 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) + actual_scales = packed_kv[..., 576:583].view(torch.float8_e8m0fnu).float() + torch.testing.assert_close(actual_scales[mem_index], scales, rtol=0, atol=0) + assert torch.equal(packed_kv[mem_index, :, :448], values.view(torch.uint8).reshape(11, 1, 448)) + assert torch.equal(packed_kv[..., 448:576].view(torch.bfloat16)[mem_index], source[..., 448:]) + assert not packed_kv[mem_index, :, 583].any() + assert (packed_kv[..., 584:] == 71).all() + unused = torch.ones(31, device="cuda", dtype=torch.bool) + unused[mem_index] = False + assert (packed_kv[unused] == 71).all() + + # The packer is graph-safe for both source layouts and mutable destination mapping. + 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) + mem_index.copy_(mem_index.roll(1)) + graph.replay() + changed_groups = source[..., :448].float().reshape(11, 1, 7, 64) + changed_scales = torch.exp2(torch.ceil(torch.log2((changed_groups.abs().amax(-1) / 448).clamp_min(1e-4)))) + changed_values = (changed_groups / changed_scales.unsqueeze(-1)).to(torch.float8_e4m3fn) + assert torch.equal(packed_kv[mem_index, :, :448], changed_values.view(torch.uint8).reshape(11, 1, 448)) + + prefill_kv = source[8:].clone().add_(0.25) + gathered_kv = torch.empty((8, 1, 512), device="cuda", dtype=torch.bfloat16) + _gather_prefill_kv_kernel[(3,)]( + packed_kv[..., :584], + prefill_kv, + gathered_kv, + mem_index[None], + torch.zeros(1, device="cuda", dtype=torch.int32), + torch.tensor([0, 11], device="cuda", dtype=torch.int32), + torch.tensor([0, 3], device="cuda", dtype=torch.int32), + 2, + 8, + packed_kv.stride(0), + prefill_kv.stride(0), + prefill_kv.stride(2), + req_to_token_indexs_stride_b=11, + req_to_token_indexs_stride_s=1, + batch_size=1, + BLOCK_ROWS=4, + ) + expected = _dequantize_kv_cache(packed_kv) + expected[mem_index[8:].long()] = prefill_kv + assert torch.equal(gathered_kv, expected[mem_index[2:10].long()]) + + +@pytest.mark.parametrize("num_heads", [8, 16, 64, 128]) +def test_fp8_nope_decode_and_cuda_graph(num_heads): + if torch.cuda.get_device_capability()[0] != 9: + pytest.skip("Validated FlashMLA sparse decode target is Hopper") + import flash_mla + + torch.manual_seed(53) + 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) + 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) + indices = torch.randint(0, 4096, (9, 2176), device="cuda", dtype=torch.int32) + lanes = torch.arange(2176, device="cuda") + indices.masked_fill_(lanes[None, :] >= lengths[:, None], -1) + control = AttControl( + nsa_decode=True, nsa_decode_dict={"topk_mem_indices": indices, "softmax_scale": 0.0625, "kv_lora_rank": 512} + ) + infer_state = _make_infer_state() + state = Fp8Glm5NextSparseDecodeState( + _infer_state_weak_ref=weakref.ref(infer_state), flashmla_sched_meta=flash_mla.get_mla_metadata()[0] + ) + + def run(): + return state.decode_att((q, q[..., :0]), packed_kv[..., :584], None, control) + + def check(actual): + expected = _torch_sparse_attention(q, _dequantize_kv_cache(packed_kv), indices) + torch.testing.assert_close(actual.float(), expected, atol=0.012, rtol=0.015) + + check(run()) + graph, out = _capture_cuda_graph(run) + graph.replay() + check(out) + q.normal_() + source.normal_() + 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) + graph.replay() + check(out) + + +@pytest.mark.parametrize("cached", [False, True]) +@pytest.mark.parametrize("num_heads", [8, 16, 64, 128]) +@pytest.mark.parametrize("chunk_size", [3, 7, 32768]) +def test_fp8_nope_batched_prefill_and_graph(monkeypatch, cached, num_heads, chunk_size): + monkeypatch.setattr(fp8_glm5_next, "KV_GATHER_CHUNK_SIZE", chunk_size) + torch.manual_seed(53) + prefix = 4 if cached else 0 + case = _PrefillCase([prefix + 3] * 2, [3, 3], num_heads, width=128, capacity=64, softmax_scale=512 ** -0.5) + case.q = torch.randn(num_heads, 6, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1) + case.topk_indices[0].fill_(-1) + case.topk_indices[1, :2] = prefix + case.topk_indices[3].fill_(-1) + case.topk_indices[3, 0] = prefix + case.check(case.run()) + graph, output = _capture_cuda_graph(case.run) + graph.replay() + case.check(output) + case.mutate() + graph.replay() + case.check(output) + + +@pytest.mark.parametrize("num_heads", [8, 64]) +@pytest.mark.parametrize("chunk_size", [3, 32768]) +def test_fp8_uncached_prefill_reuses_current_kv(monkeypatch, num_heads, chunk_size): + from sgl_kernel.flash_attn import flash_attn_with_kvcache + + monkeypatch.setattr(fp8_glm5_next, "KV_GATHER_CHUNK_SIZE", chunk_size) + case = _PrefillCase([3, 3], [3, 3], num_heads, width=128) + case.prefill_cache_kv = torch.randn(6, 1, 1536, device="cuda", dtype=torch.bfloat16)[..., :512] + case.control.nsa_prefill_dict["prefill_cache_kv"] = case.prefill_cache_kv + + def check_kv_storage(*args, **kwargs): + assert kwargs["v_cache"].untyped_storage().data_ptr() == case.prefill_cache_kv.untyped_storage().data_ptr() + return flash_attn_with_kvcache(*args, **kwargs) + + monkeypatch.setattr("sgl_kernel.flash_attn.flash_attn_with_kvcache", check_kv_storage) + case.check(case.run()) + graph, output = _capture_cuda_graph(case.run) + case.mutate() + graph.replay() + case.check(output) + + +@pytest.mark.parametrize("num_heads", [8, 16, 64, 128]) +@pytest.mark.parametrize("softmax_scale", [0.0625, 512 ** -0.5]) +def test_fp8_chunked_prefill_sparse_rows_and_graph(monkeypatch, num_heads, softmax_scale): + monkeypatch.setattr(fp8_glm5_next, "KV_GATHER_CHUNK_SIZE", 129) + torch.manual_seed(53) + case = _PrefillCase([1025], [9], num_heads, capacity=2048, softmax_scale=softmax_scale) + lengths = torch.tensor([0, 1, 3, 127, 128, 511, 2048, 2049, 2051], device="cuda") + case.topk_indices.masked_fill_(torch.arange(2176, device="cuda")[None, :] >= lengths[:, None], -1) + case.topk_indices[1, 0] = 1016 + case.topk_indices[2, :3] = torch.tensor([2, 3, 4], device="cuda") + case.topk_indices[3, :127] = torch.randint(1016, 1025, (127,), device="cuda", dtype=torch.int32) + case.check(case.run()) + graph, output = _capture_cuda_graph(case.run) + graph.replay() + case.check(output) + case.mutate() + graph.replay() + case.check(output) + + +@pytest.mark.parametrize("num_q", [64, 129]) +@pytest.mark.parametrize("num_heads", [8, 64]) +def test_fp8_prefill_uses_flashattention_and_request_table_graph(monkeypatch, num_q, num_heads): + from sgl_kernel.flash_attn import flash_attn_with_kvcache + + torch.manual_seed(53100 + num_q) + case = _PrefillCase([32768, 32768], [num_q // 2, num_q - num_q // 2], num_heads, capacity=131072) + observed = [] + + def observed_flashattention(*args, **kwargs): + observed.append(kwargs["qv"].shape) + return flash_attn_with_kvcache(*args, **kwargs) + + def unexpected_flashmla(*args, **kwargs): + raise AssertionError("FP8 prefill must use its single FlashAttention backend") + + monkeypatch.setattr("sgl_kernel.flash_attn.flash_attn_with_kvcache", observed_flashattention) + monkeypatch.setattr("sgl_kernel.flash_mla.flash_mla_sparse_fwd", unexpected_flashmla) + case.check(case.run()) + assert observed == [(num_q, num_heads, 512)] + graph, output = _capture_cuda_graph(case.run) + case.mutate() + graph.replay() + case.check(output) + case.topk_indices.fill_(-1) + graph.replay() + assert torch.count_nonzero(output) == 0 + + +@pytest.mark.parametrize( + "num_q,num_heads", + [(1, 8), (2, 8), (8, 8), (1, 64), (2, 64), (3, 64), (4, 64), (5, 64), (6, 64), (7, 64), (8, 64)], +) +def test_fp8_cached_prefill_strided_holes_and_current(num_q, num_heads): + torch.manual_seed(53016 + num_q) + seqs = [65536 + value for value in [127, 521, 1031, 2053, 4099, 8191, 12289, 16381][:num_q]] + case = _PrefillCase(seqs, [1] * num_q, num_heads) + case.q = torch.randn(num_heads, num_q, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1) + backing = torch.full((num_q, 4352), -1, device="cuda", dtype=torch.int32) + backing[:, ::2] = case.topk_indices + case.topk_indices = backing[:, ::2] + output = case.run() + case.check(output) + assert output.is_contiguous() and torch.count_nonzero(output) > 0 + graph, output = _capture_cuda_graph(case.run) + graph.replay() + case.check(output) + case.mutate() + graph.replay() + case.check(output) + + +@pytest.mark.parametrize("num_q", [2, 3, 5, 8, 16, 32, 128]) +def test_fp8_cached_prefill_multi_query_and_graph(num_q): + torch.manual_seed(53200 + num_q) + case = _PrefillCase([131072, 131072], [num_q // 2, num_q - num_q // 2], 64) + case.check(case.run()) + for _ in range(2): + infer, state = case.new_state() + case.check(case.run(state)) + fresh_infer, fresh = case.new_state(mem_index=case.infer_state.mem_index.clone()) + case.check(case.run(fresh)) + graph, output = _capture_cuda_graph(case.run) + case.mutate() + graph.replay() + case.check(output) + case.topk_indices.fill_(-1) + graph.replay() + assert torch.count_nonzero(output) == 0 + + +def test_fp8_prefill_graph_with_empty_requests_and_changing_boundaries(): + case = _PrefillCase([65536] * 4, [0, 4, 0, 3], 64) + case.check(case.run()) + graph, output = _capture_cuda_graph(case.run) + replacement = _PrefillCase([65536] * 4, [2, 0, 5, 0], 64) + for name in ( + "q", + "packed_kv", + "prefill_cache_kv", + "req_to_token_indexs", + "q_req_idx", + "position_ids", + "topk_indices", + "topk_mem_indices", + ): + getattr(case, name).copy_(getattr(replacement, name)) + case.infer_state.b1_cu_q_seq_len.copy_(replacement.infer_state.b1_cu_q_seq_len) + case.infer_state.position_ids.copy_(replacement.infer_state.position_ids) + for name in ("q_kv_seq_len", "q_req_idx", "q_kv_start_loc"): + getattr(case.att_state, name).copy_(getattr(replacement.att_state, name)) + graph.replay() + case.check(output) + + +def test_fp8_one_query_prefill_and_decode_read_their_own_kv_sources(): + case = _PrefillCase([65536], [1], 64) + case.q.zero_() + case.prefill_cache_kv.fill_(1) + case.source_kv.fill_(2) + write_fp8_e8m0_packed_kv(case.source_kv, case.all_slots, case.packed_kv) + case.topk_indices.fill_(-1) + case.topk_indices[:, 0] = 65535 + case.refresh() + decode_infer = _make_infer_state(**vars(case.infer_state)) + decode_infer.is_prefill = False + decode_infer.is_cuda_graph_capture = True + decode_state = case.backend.create_att_decode_state(decode_infer) + decode_state.init_state() + control = AttControl( + nsa_decode=True, + nsa_decode_dict={ + "topk_mem_indices": case.topk_mem_indices, + "softmax_scale": 0.0625, + "kv_lora_rank": 512, + }, + ) + + def run(): + prefill = case.run() + decode = decode_state.decode_att((case.q, case.q[..., :0]), case.packed_kv[..., :584], None, control) + return prefill, decode + + prefill, decode = run() + torch.testing.assert_close(prefill, torch.ones_like(prefill), rtol=0, atol=0) + torch.testing.assert_close(decode, torch.full_like(decode, 2), rtol=0, atol=0) + graph, (prefill, decode) = _capture_cuda_graph(run) + case.prefill_cache_kv.fill_(3) + case.source_kv.fill_(4) + write_fp8_e8m0_packed_kv(case.source_kv, case.all_slots, case.packed_kv) + graph.replay() + torch.testing.assert_close(prefill, torch.full_like(prefill, 3), rtol=0, atol=0) + torch.testing.assert_close(decode, torch.full_like(decode, 4), rtol=0, atol=0) + + +def test_fp8_h64_b1_decode_uses_native_eager_and_direct_graph(): + if torch.cuda.get_device_capability()[0] != 9: + pytest.skip("Validated FlashMLA sparse decode target is Hopper") + import flash_mla + + torch.manual_seed(53300) + capacity = 4096 + 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) + 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) + control = AttControl( + nsa_decode=True, nsa_decode_dict={"topk_mem_indices": indices, "softmax_scale": 0.0625, "kv_lora_rank": 512} + ) + infer = _make_infer_state(is_cuda_graph_capture=False) + state = Fp8Glm5NextSparseDecodeState( + _infer_state_weak_ref=weakref.ref(infer), + decode_topk_bound=2051, + flashmla_sched_meta=flash_mla.get_mla_metadata()[0], + ) + + def run(): + return state.decode_att((q, q[..., :0]), packed[..., :584], None, control) + + expected = _torch_sparse_attention(q, _dequantize_kv_cache(packed), indices) + eager = run() + torch.testing.assert_close(eager.float(), expected, atol=0.012, rtol=0.015) + infer.is_cuda_graph_capture = True + graph, output = _capture_cuda_graph(run) + infer.is_cuda_graph_capture = False + graph.replay() + torch.testing.assert_close(output.float(), expected, atol=0.012, rtol=0.015) + + +@pytest.mark.parametrize("num_heads", [8, 16, 64, 128]) +def test_fp8_decode_qpad_kernel_strides_and_zero_padding(num_heads): + """The capture padding kernel reads a true D-stride and writes every padded head as zero.""" + torch.manual_seed(53016 + num_heads) + batch = 2 + # This has BMM-style B/H strides and a non-unit final dimension stride. + q = torch.randn(num_heads, batch, 1024, device="cuda", dtype=torch.bfloat16)[..., ::2].transpose(0, 1) + padded_heads = ((num_heads + 63) // 64) * 64 + head_dim = q.shape[-1] + block_size = triton.next_power_of_2(2 * head_dim) + padded = torch.empty((batch, 1, padded_heads, head_dim), device="cuda", dtype=torch.bfloat16) + _fused_strided_qpad_kernel[(batch, triton.cdiv(padded_heads * head_dim, block_size))]( + q, + padded, + q.stride(0), + q.stride(1), + q.stride(2), + num_heads, + padded_heads, + head_dim=head_dim, + BLOCK_SIZE=block_size, + num_warps=4, + ) + expected = torch.zeros_like(padded) + expected[:, 0, :num_heads] = q + assert torch.equal(padded, expected) + assert torch.count_nonzero(padded[:, 0, num_heads:]) == 0 + + +@pytest.mark.parametrize("num_heads", [8, 16, 64, 128]) +def test_fp8_decode_infer_state_graph_qpad_and_mutation(num_heads): + """Exercise the actual backend's eager and graph flags with strided BMM Q.""" + if torch.cuda.get_device_capability()[0] != 9: + pytest.skip("Validated FlashMLA sparse decode target is Hopper") + import flash_mla + + torch.manual_seed(53100 + num_heads) + batch, capacity = 2, 4096 + 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) + 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) + control = AttControl( + nsa_decode=True, + nsa_decode_dict={"topk_mem_indices": indices, "softmax_scale": 0.0625, "kv_lora_rank": 512}, + ) + infer_state = _make_infer_state(is_cuda_graph_capture=False) + state = Fp8Glm5NextSparseDecodeState( + _infer_state_weak_ref=weakref.ref(infer_state), flashmla_sched_meta=flash_mla.get_mla_metadata()[0] + ) + + def run(): + return state.decode_att((q, q[..., :0]), packed[..., :584], None, control) + + def reference(): + return _torch_sparse_attention(q, _dequantize_kv_cache(packed), indices) + + eager = run() + torch.testing.assert_close(eager.float(), reference(), atol=0.012, rtol=0.015) + assert eager.shape == q.shape + infer_state.is_cuda_graph_capture = True + graph, output = _capture_cuda_graph(run) + infer_state.is_cuda_graph_capture = False + graph.replay() + torch.testing.assert_close(output.float(), reference(), atol=0.012, rtol=0.015) + q.copy_(q.roll(1, dims=-1)) + saved = indices.clone() + indices[:, :2049].copy_(saved[:, :2049].roll(1, dims=1)) + assert torch.equal(indices[:, 2049:], saved[:, 2049:]) + graph.replay() + torch.testing.assert_close(output.float(), reference(), atol=0.012, rtol=0.015) + + +def test_fp8_sparse_direct_decode_factory_bound_guard(): + backend = object.__new__(Fp8Glm5NextSparseAttBackend) + backend.model = SimpleNamespace(config={"index_topk": 2048, "index_kpool": 4}) + infer_state = InferStateInfo() + assert backend.create_att_decode_state(infer_state).decode_topk_bound == 2051 + backend.model = SimpleNamespace(config={"index_topk": 2048, "index_kpool": 8}) + assert backend.create_att_decode_state(infer_state).decode_topk_bound is None + backend.model = SimpleNamespace(config={"index_topk": 2047, "index_kpool": 4}) + assert backend.create_att_decode_state(infer_state).decode_topk_bound is None + + +@pytest.mark.parametrize("batch", [1, 8]) +def test_fp8_sparse_direct_decode_actual_entry_edges_and_graph(monkeypatch, batch): + if torch.cuda.get_device_capability()[0] != 9: + pytest.skip("Validated FlashMLA sparse decode target is Hopper") + import flash_mla + + torch.manual_seed(53200 + batch) + capacity = 4096 + 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) + 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] + for row in range(batch): + count = counts[row] + if count: + indices[row, :count] = torch.randint(0, capacity, (count,), device="cuda", dtype=torch.int32) + if batch == 8: + indices[5, :64] = -1 # holes + indices[6, 64:128] = indices[6, :64] # duplicates + assert torch.equal(indices[:, 2051:], torch.full_like(indices[:, 2051:], -1)) + control = AttControl( + nsa_decode=True, nsa_decode_dict={"topk_mem_indices": indices, "softmax_scale": 0.0625, "kv_lora_rank": 512} + ) + state_owner = _make_infer_state(is_cuda_graph_capture=False) + state = Fp8Glm5NextSparseDecodeState( + _infer_state_weak_ref=weakref.ref(state_owner), + flashmla_sched_meta=flash_mla.get_mla_metadata()[0], + decode_topk_bound=2051, + ) + + def native_must_not_run(*args, **kwargs): + raise AssertionError("supported GLM-5.3 decode took native FlashMLA fallback") + + def run(): + return state.decode_att((q, q[..., :0]), packed[..., :584], None, control) + + assert not can_use_fp8_sparse_direct_decode(q, packed, indices.unsqueeze(1), 512, 2051) + expected = _torch_sparse_attention(q, _dequantize_kv_cache(packed), indices) + eager = run() + torch.testing.assert_close(eager.float(), expected, atol=0.02, rtol=0.02) + if batch == 1: + assert torch.count_nonzero(eager) > 0 + monkeypatch.setattr(flash_mla, "flash_mla_with_kvcache", native_must_not_run) + state.infer_state.is_cuda_graph_capture = True + graph, output = _capture_cuda_graph(run) + state.infer_state.is_cuda_graph_capture = False + graph.replay() + torch.testing.assert_close(output.float(), expected, atol=0.02, rtol=0.02) + q.copy_(q.roll(1, dims=-1)) + saved = indices.clone() + nonempty = (indices >= 0).sum(dim=1) + for row, count in enumerate(nonempty.tolist()): + if count: + indices[row, :count].copy_(saved[row, :count].roll(1)) + assert torch.equal(indices[:, 2051:], saved[:, 2051:]) + graph.replay() + expected = _torch_sparse_attention(q, _dequantize_kv_cache(packed), indices) + torch.testing.assert_close(output.float(), expected, atol=0.02, rtol=0.02) + indices.fill_(-1) + graph.replay() + assert torch.count_nonzero(output) == 0 + + +@pytest.mark.parametrize( + "num_heads,batch", + [ + (1, 1), + (4, 3), + (7, 5), + (8, 2), + (8, 3), + (8, 4), + (8, 5), + (16, 1), + (16, 3), + (16, 4), + (17, 3), + (32, 1), + (32, 2), + (32, 3), + (64, 1), + (64, 2), + (128, 1), + ], +) +def test_fp8_sparse_direct_decode_general_graph_factory_entry(monkeypatch, num_heads, batch): + """General head/token shapes use the direct path selected by the central predicate.""" + if torch.cuda.get_device_capability()[0] != 9: + pytest.skip("Validated FlashMLA sparse decode target is Hopper") + import flash_mla + + torch.manual_seed(53300 + batch) + capacity = 4096 + 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) + 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) + indices = torch.full((batch, 2176), -1, device="cuda", dtype=torch.int32) + counts = [([2051, 0, 2048, 2050, 1, 3, 127, 128][row % 8]) for row in range(batch)] + for row, count in enumerate(counts): + if count: + indices[row, :count] = torch.randint(0, capacity, (count,), device="cuda", dtype=torch.int32) + if batch >= 3: + indices[2, :32] = -1 # holes + if batch >= 4: + indices[3, 64:128] = indices[3, :64] # duplicates + assert torch.equal(indices[:, 2051:], torch.full_like(indices[:, 2051:], -1)) + control = AttControl( + nsa_decode=True, + nsa_decode_dict={"topk_mem_indices": indices, "softmax_scale": 0.0625, "kv_lora_rank": 512}, + ) + backend = object.__new__(Fp8Glm5NextSparseAttBackend) + backend.model = SimpleNamespace(config={"index_topk": 2048, "index_kpool": 4}) + infer_state = _make_infer_state(is_cuda_graph_capture=False) + state = backend.create_att_decode_state(infer_state) + state.flashmla_sched_meta = flash_mla.get_mla_metadata()[0] + assert state.decode_topk_bound == 2051 + native_state_owner = _make_infer_state(is_cuda_graph_capture=False) + native_state = Fp8Glm5NextSparseDecodeState( + _infer_state_weak_ref=weakref.ref(native_state_owner), + flashmla_sched_meta=flash_mla.get_mla_metadata()[0], + decode_topk_bound=None, + ) + + def run(current_state): + return current_state.decode_att((q, q[..., :0]), packed[..., :584], None, control) + + native_calls = [] + original_native = flash_mla.flash_mla_with_kvcache + + def observe_native(*args, **kwargs): + native_calls.append(True) + return original_native(*args, **kwargs) + + monkeypatch.setattr(flash_mla, "flash_mla_with_kvcache", observe_native) + eager = run(state) + native_output = run(native_state) + assert native_calls + assert torch.equal(eager, native_output) + expected = _torch_sparse_attention(q, _dequantize_kv_cache(packed), indices) + torch.testing.assert_close(eager.float(), expected, atol=0.02, rtol=0.02) + + def native_must_not_run(*args, **kwargs): + raise AssertionError("supported graph shape did not use the direct packed-KV path") + + monkeypatch.setattr(flash_mla, "flash_mla_with_kvcache", native_must_not_run) + state.infer_state.is_cuda_graph_capture = True + graph, output = _capture_cuda_graph(lambda: run(state)) + state.infer_state.is_cuda_graph_capture = False + graph.replay() + torch.testing.assert_close(output.float(), expected, atol=0.02, rtol=0.02) + q.copy_(q.roll(1, dims=-1)) + saved = indices.clone() + for row, count in enumerate((indices >= 0).sum(dim=1).tolist()): + if count: + indices[row, :count].copy_(saved[row, :count].roll(1)) + assert torch.equal(indices[:, 2051:], saved[:, 2051:]) + graph.replay() + torch.testing.assert_close( + output.float(), _torch_sparse_attention(q, _dequantize_kv_cache(packed), indices), atol=0.02, rtol=0.02 + ) + indices.fill_(-1) + graph.replay() + assert torch.count_nonzero(output) == 0 + + +def test_fp8_sparse_direct_decode_unknown_bound_fallback_and_strided_q(monkeypatch): + if torch.cuda.get_device_capability()[0] != 9: + pytest.skip("Validated FlashMLA sparse decode target is Hopper") + import flash_mla + + torch.manual_seed(53250) + capacity = 4096 + 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) + 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 + control = AttControl( + nsa_decode=True, nsa_decode_dict={"topk_mem_indices": late, "softmax_scale": 0.0625, "kv_lora_rank": 512} + ) + calls = [] + original = flash_mla.flash_mla_with_kvcache + + def native(*args, **kwargs): + calls.append(True) + return original(*args, **kwargs) + + monkeypatch.setattr(flash_mla, "flash_mla_with_kvcache", native) + unknown_owner = _make_infer_state(is_cuda_graph_capture=False) + unknown = Fp8Glm5NextSparseDecodeState( + _infer_state_weak_ref=weakref.ref(unknown_owner), flashmla_sched_meta=flash_mla.get_mla_metadata()[0] + ) + actual = unknown.decode_att((q, q[..., :0]), packed[..., :584], None, control) + torch.testing.assert_close( + actual.float(), _torch_sparse_attention(q, _dequantize_kv_cache(packed), late), atol=0.02, rtol=0.02 + ) + assert calls and not can_use_fp8_sparse_direct_decode( + q, packed[..., :584], late.unsqueeze(1), 512, None, is_cuda_graph_capture=True + ) + q_d2 = torch.randn(1, 8, 1024, device="cuda", dtype=torch.bfloat16)[..., ::2] + good = torch.full((1, 2176), -1, device="cuda", dtype=torch.int32) + good[:, :2051] = torch.randint(0, capacity, (1, 2051), device="cuda", dtype=torch.int32) + assert can_use_fp8_sparse_direct_decode( + q_d2, packed[..., :584], good.unsqueeze(1), 512, 2051, is_cuda_graph_capture=True + ) + assert not can_use_fp8_sparse_direct_decode( + q, packed[..., :584], good.unsqueeze(1), 256, 2051, is_cuda_graph_capture=True + ) + + +@pytest.mark.parametrize("num_heads,batch", [(1, 3), (7, 5), (16, 16), (32, 7), (65, 3), (128, 8)]) +@pytest.mark.parametrize("index_dtype", [torch.int32, torch.int64]) +def test_fp8_sparse_direct_decode_strided_inputs_and_packed_row_padding(monkeypatch, num_heads, batch, index_dtype): + """The kernel handles partial head tiles, larger batches and actual Q/index strides.""" + torch.manual_seed(53600 + num_heads + batch) + capacity = 4096 + source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) + packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) + 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) + indices[0].fill_(-1) + + def run(): + return fp8_sparse_direct_decode(q, packed, indices, 0.0625) + + def expected(): + return _torch_sparse_attention(q, _dequantize_kv_cache(packed), indices[:, 0]) + + if (num_heads, batch) in ((1, 3), (7, 5)): + infer = _make_infer_state(is_cuda_graph_capture=True) + state = Fp8Glm5NextSparseDecodeState(_infer_state_weak_ref=weakref.ref(infer), decode_topk_bound=2051) + + def native_must_not_run(*args, **kwargs): + raise AssertionError("supported strided decode took native fallback") + + monkeypatch.setattr("flash_mla.flash_mla_with_kvcache", native_must_not_run) + torch.testing.assert_close( + state.decode_att( + (q, q[..., :0]), + packed, + None, + AttControl( + nsa_decode=True, + nsa_decode_dict={"topk_mem_indices": indices, "softmax_scale": 0.0625, "kv_lora_rank": 512}, + ), + ).float(), + expected(), + atol=0.02, + rtol=0.02, + ) + + torch.testing.assert_close(run().float(), expected(), atol=0.02, rtol=0.02) + graph, output = _capture_cuda_graph(run) + graph.replay() + torch.testing.assert_close(output.float(), expected(), atol=0.02, rtol=0.02) + q.normal_() + indices[..., :2051].copy_(indices[..., :2051].roll(1, dims=0)) + graph.replay() + torch.testing.assert_close(output.float(), expected(), atol=0.02, rtol=0.02) + + +@pytest.mark.parametrize("index_dtype", [torch.int32, torch.int64]) +def test_fp8_sparse_direct_decode_central_predicate_accepts_strides_and_rejects_unaligned_kv(index_dtype): + q = torch.empty(1, 8, 1024, device="cuda", dtype=torch.bfloat16)[..., ::2] + packed = torch.empty(8, 1, 720, device="cuda", dtype=torch.uint8) + topk = torch.empty(1, 1, 4352, device="cuda", dtype=index_dtype)[..., ::2] + assert can_use_fp8_sparse_direct_decode(q, packed, topk, 512, 2051, is_cuda_graph_capture=True) + odd_rows = torch.empty(8, 1, 721, device="cuda", dtype=torch.uint8)[..., :584] + odd_base = torch.empty(8, 1, 722, device="cuda", dtype=torch.uint8)[..., 1:585] + strided_columns = torch.empty(8, 1, 1440, device="cuda", dtype=torch.uint8)[..., ::2] + for invalid in (odd_rows, odd_base, strided_columns): + assert not can_use_fp8_sparse_direct_decode(q, invalid, topk, 512, 2051, is_cuda_graph_capture=True) + + +@pytest.mark.parametrize("strided_input", ["q", "topk"]) +def test_fp8_sparse_direct_decode_offsets_beyond_int32(strided_input): + batch, heads, capacity = 4, 1, 64 + stride = 800_000_000 + shape = (batch, heads, 512) if strided_input == "q" else (batch, 1, 2176) + dtype = torch.bfloat16 if strided_input == "q" else torch.int32 + storage_size = (batch - 1) * stride + shape[-1] + required_bytes = storage_size * torch.empty((), dtype=dtype).element_size() + torch.cuda.empty_cache() + if torch.cuda.mem_get_info()[0] < required_bytes + 512 * 1024 ** 2: + pytest.skip("Large-stride regression requires sufficient free GPU memory") + # Allocate address space without initializing the gaps between logical rows. + storage = torch.empty(storage_size, device="cuda", dtype=dtype) + strided = storage.as_strided(shape, (stride, shape[-1], 1)) + q = strided if strided_input == "q" else torch.empty(batch, heads, 512, device="cuda", dtype=torch.bfloat16) + topk = strided if strided_input == "topk" else torch.empty(batch, 1, 2176, device="cuda", dtype=torch.int32) + q.copy_(torch.randn(q.shape, device="cuda", dtype=q.dtype)) + topk.fill_(-1) + 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) + 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]) + torch.testing.assert_close(actual.float(), expected, atol=0.02, rtol=0.02) + + +@pytest.mark.parametrize("num_heads,batch", [(16, 16), (32, 8), (64, 4), (128, 2)]) +def test_fp8_sparse_direct_decode_large_workload_retains_native_graph(monkeypatch, num_heads, batch): + if torch.cuda.get_device_capability()[0] != 9: + pytest.skip("Validated FlashMLA sparse decode target is Hopper") + import flash_mla + + 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) + 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) + assert not can_use_fp8_sparse_direct_decode(q, packed, topk, 512, 2051, is_cuda_graph_capture=True) + + def direct_must_not_run(*args, **kwargs): + raise AssertionError("larger workload took the slower direct path") + + monkeypatch.setattr( + "lightllm.models.glm5_next.triton_kernel.fp8_sparse_decode.fp8_sparse_direct_decode", direct_must_not_run + ) + infer = _make_infer_state(is_cuda_graph_capture=True) + state = Fp8Glm5NextSparseDecodeState( + _infer_state_weak_ref=weakref.ref(infer), + decode_topk_bound=2051, + flashmla_sched_meta=flash_mla.get_mla_metadata()[0], + ) + control = AttControl( + nsa_decode=True, nsa_decode_dict={"topk_mem_indices": topk, "softmax_scale": 0.0625, "kv_lora_rank": 512} + ) + run = lambda: state.decode_att((q, q[..., :0]), packed[..., :584], None, control) + expected = _torch_sparse_attention(q, _dequantize_kv_cache(packed), topk[:, 0]) + torch.testing.assert_close(run().float(), expected, atol=0.02, rtol=0.02) + graph, output = _capture_cuda_graph(run) + graph.replay() + torch.testing.assert_close(output.float(), expected, atol=0.02, rtol=0.02) + + +@pytest.mark.parametrize("num_heads", [8, 64]) +def test_fp8_prefill_workspace_reuse_mutable_inputs_and_graph_replay(num_heads): + torch.manual_seed(53450) + backend = _make_workspace_backend(66) + case = _PrefillCase([32768, 16384, 16384], [11, 11, 11], num_heads, backend=backend) + case.check(case.run()) + workspace = backend._prefill_workspaces[0] + assert set(workspace) == {"gathered_kv", "chunk_topk_indices", "chunk_cache_seqlens", "cu_seqlens_q"} + pointers = {name: value.data_ptr() for name, value in workspace.items()} + for _ in range(10): + infer, state = case.new_state() + case.check(case.run(state)) + case.mutate() + case.check(case.run()) + fresh = _PrefillCase([16384, 32768, 16384], [9, 11, 13], num_heads, capacity=131072, backend=backend) + fresh.check(fresh.run()) + assert {name: value.data_ptr() for name, value in backend._prefill_workspaces[0].items()} == pointers + for overrides in ({"prefill_cuda_graph_exe_list": []}, {"is_cuda_graph_capture": True}): + infer, state = case.new_state(**overrides) + case.check(case.run(state)) + graph, output = _capture_cuda_graph(case.run) + case.mutate() + graph.replay() + case.check(output) + + +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 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.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", + ) + state = create_packed_mla_state(config=config, 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_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_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() + 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/glm5_next/test_fused_mtp_norm.py b/unit_tests/models/glm5_next/test_fused_mtp_norm.py new file mode 100644 index 0000000000..c92ae5b86a --- /dev/null +++ b/unit_tests/models/glm5_next/test_fused_mtp_norm.py @@ -0,0 +1,126 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.basemodel.triton_kernel.norm.rmsnorm import rmsnorm_forward +from lightllm.models.glm5_next_mtp.layer_infer import pre_layer_infer as mtp +from lightllm.models.glm5_next_mtp.triton_kernel.fused_eh_norm import fused_mtp_norm + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +def reference( + embeddings: torch.Tensor, + previous_hidden: torch.Tensor, + weights: tuple[torch.Tensor, torch.Tensor, torch.Tensor], + eps: float, +) -> torch.Tensor: + rmsnorm_forward(x=previous_hidden, weight=weights[0], eps=eps, out=previous_hidden) + rmsnorm_forward(x=embeddings, weight=weights[1], eps=eps, out=embeddings) + rmsnorm_forward(x=previous_hidden, weight=weights[2], eps=eps, out=previous_hidden) + return torch.cat(tensors=(embeddings, previous_hidden), dim=-1) + + +@pytest.mark.parametrize("rows", [1, 4, 16, 64, 512, 8192]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("scale", [1e-4, 1.0, 100.0]) +def test_norm_rounding_and_inplace_updates(rows: int, dtype: torch.dtype, scale: float) -> None: + torch.manual_seed(seed=5514) + hidden = 4096 + e = (torch.randn(size=(rows, hidden + 16), dtype=dtype, device="cuda") * scale)[:, :hidden] + p = (torch.randn(size=(rows, hidden + 16), dtype=dtype, device="cuda") * scale)[:, :hidden] + er, pr = e.clone(), p.clone() + weights = tuple(torch.randn(size=(hidden,), dtype=dtype, device="cuda") * 0.2 + 1 for _ in range(3)) + expected = reference(embeddings=er, previous_hidden=pr, weights=weights, eps=1e-5) + actual = fused_mtp_norm( + embeddings=e, + previous_hidden=p, + main_weight=weights[0], + embedding_weight=weights[1], + hidden_weight=weights[2], + eps=1e-5, + ) + torch.testing.assert_close(actual=actual, expected=expected, atol=0, rtol=0) + torch.testing.assert_close(actual=e, expected=er, atol=0, rtol=0) + torch.testing.assert_close(actual=p, expected=pr, atol=0, rtol=0) + + +def test_graph_replay_changes_both_inputs_and_preserves_strided_storage() -> None: + torch.manual_seed(seed=5515) + rows, hidden = 16, 4096 + e_storage = torch.randn(size=(rows, hidden * 2), device="cuda", dtype=torch.bfloat16) + p_storage = torch.randn(size=(rows, hidden * 2), device="cuda", dtype=torch.bfloat16) + e, p = e_storage[:, ::2], p_storage[:, ::2] + weights = tuple(torch.randn(size=(hidden,), device="cuda", dtype=e.dtype) * 0.2 + 1 for _ in range(3)) + fused_mtp_norm( + embeddings=e, + previous_hidden=p, + main_weight=weights[0], + embedding_weight=weights[1], + hidden_weight=weights[2], + eps=1e-5, + ) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(cuda_graph=graph): + actual = fused_mtp_norm( + embeddings=e, + previous_hidden=p, + main_weight=weights[0], + embedding_weight=weights[1], + hidden_weight=weights[2], + eps=1e-5, + ) + for scale in (0.01, 10.0): + e_storage.normal_().mul_(scale) + p_storage.normal_().mul_(scale) + untouched_e, untouched_p = e_storage[:, 1::2].clone(), p_storage[:, 1::2].clone() + er, pr = e.clone(), p.clone() + expected = reference(embeddings=er, previous_hidden=pr, weights=weights, eps=1e-5) + graph.replay() + torch.testing.assert_close(actual=actual, expected=expected, atol=0, rtol=0) + torch.testing.assert_close(actual=e, expected=er, atol=0, rtol=0) + torch.testing.assert_close(actual=p, expected=pr, atol=0, rtol=0) + assert torch.equal(input=e_storage[:, 1::2], other=untouched_e) + assert torch.equal(input=p_storage[:, 1::2], other=untouched_p) + + +@pytest.mark.parametrize("hidden", [4096, 1024]) +def test_model_integration(hidden: int) -> None: + torch.manual_seed(seed=5516) + e = torch.randn(size=(7, hidden), dtype=torch.bfloat16, device="cuda") + p = torch.randn_like(input=e) + er, pr = e.clone(), p.clone() + weights = tuple(torch.randn(size=(hidden,), dtype=e.dtype, device="cuda") * 0.2 + 1 for _ in range(3)) + + class Norm: + def __init__(self, weight: torch.Tensor) -> None: + self.weight = weight + + def __call__(self, input: torch.Tensor, eps: float, out: torch.Tensor) -> torch.Tensor: + return rmsnorm_forward(x=input, weight=self.weight, eps=eps, out=out) + + class Projection: + def __init__(self) -> None: + self.weight = torch.randn(size=(2 * hidden, 128), dtype=e.dtype, device=e.device) * 0.1 + + def mm(self, input_tensor: torch.Tensor) -> torch.Tensor: + return torch.mm(input=input_tensor, mat2=self.weight) + + layer = mtp.Glm5NextMTPPreLayerInfer.__new__(mtp.Glm5NextMTPPreLayerInfer) + layer.eps_ = 1e-5 + layer_weights = SimpleNamespace( + main_norm_weight_=Norm(weight=weights[0]), + enorm_weight_=Norm(weight=weights[1]), + hnorm_weight_=Norm(weight=weights[2]), + eh_proj_weight_=Projection(), + ) + expected = layer_weights.eh_proj_weight_.mm( + input_tensor=reference(embeddings=er, previous_hidden=pr, weights=weights, eps=1e-5) + ) + actual = layer._fuse_hidden( + input_embeddings=e, infer_state=SimpleNamespace(mtp_draft_input_hiddens=p), layer_weight=layer_weights + ) + torch.testing.assert_close(actual=actual, expected=expected, atol=0, rtol=0) + torch.testing.assert_close(actual=e, expected=er, atol=0, rtol=0) + torch.testing.assert_close(actual=p, expected=pr, atol=0, rtol=0) diff --git a/unit_tests/models/glm5_next/test_indexer_topk.py b/unit_tests/models/glm5_next/test_indexer_topk.py new file mode 100644 index 0000000000..1a4ca0f184 --- /dev/null +++ b/unit_tests/models/glm5_next/test_indexer_topk.py @@ -0,0 +1,161 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.models.glm5_next.indexer import HAS_VLLM, Glm5NextNsaInfer + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.fixture +def indexer(): + return Glm5NextNsaInfer( + 0, {"index_topk": 2048, "index_n_heads": 32, "index_head_dim": 128, "rms_norm_eps": 1e-5, "index_kpool": 4}, 1 + ) + + +def _require_vllm_topk(): + if not HAS_VLLM: + pytest.skip("vLLM top_k_per_row_decode required") + + +def _assert_topk(logits, lengths, indices): + topk = indices.shape[1] + valid = torch.arange(topk, device=logits.device)[None, :] < lengths[:, None] + assert (indices[~valid] == -1).all() + assert (indices[valid] >= 0).all() + assert (indices < lengths[:, None]).all() + for row, length in zip(indices, lengths.tolist()): + count = min(length, topk) + assert row[:count].unique().numel() == count + + got = logits.gather(1, indices.long().clamp_min(0)).masked_fill(~valid, -float("inf")) + positions = torch.arange(logits.shape[1], device=logits.device) + masked = logits.masked_fill(positions[None, :] >= lengths[:, None], -float("inf")) + # Compare values, since different equal-valued candidates are valid top-k. + torch.testing.assert_close(got.sort(descending=True).values, masked.topk(topk).values, rtol=0, atol=0) + + +@pytest.mark.parametrize("pools", [640, 8192, 65536, 262144]) +@pytest.mark.parametrize("distribution", ["normal", "concentrated", "ties"]) +def test_topk_variable_lengths(indexer, pools, distribution): + _require_vllm_topk() + torch.manual_seed(42) + # DeepGEMM logits can have a padded row stride, and output is a query chunk. + logits = torch.randn(8, pools + 128, device="cuda")[:, :pools] + if distribution == "concentrated": + logits.mul_(0.001).add_(1) + elif distribution == "ties": + logits.copy_(torch.randint(0, 4, logits.shape, device="cuda")) + lengths = torch.tensor([0, 1, 511, 512, 513, pools // 3, pools - 3, pools], device="cuda", dtype=torch.int32) + positions = torch.arange(pools, device="cuda") + logits.masked_fill_(positions[None, :] >= lengths[:, None], float("nan")) + output = torch.full((10, 512), -2, device="cuda", dtype=torch.int32) + indices = output[1:-1] + indexer.select_topk_indices(logits, lengths, indices) + _assert_topk(logits, lengths, indices) + assert (output[[0, -1]] == -2).all() + + +@pytest.mark.parametrize("pools", [8192, 262144]) +def test_topk_cuda_graph_with_changing_lengths(indexer, pools): + _require_vllm_topk() + logits = torch.zeros(8, pools, device="cuda") + lengths = torch.zeros(8, device="cuda", dtype=torch.int32) + indices = torch.empty(8, 512, device="cuda", dtype=torch.int32) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + indexer.select_topk_indices(logits, lengths, indices) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + indexer.select_topk_indices(logits, lengths, indices) + + for shift in range(3): + logits.normal_().mul_(0.001).add_(1) + lengths.copy_( + torch.tensor([0, 1, 511, 512, 513, pools // 3, pools - 3, pools], device="cuda", dtype=torch.int32).roll( + shift + ) + ) + logits.masked_fill_(torch.arange(pools, device="cuda")[None, :] >= lengths[:, None], float("nan")) + graph.replay() + _assert_topk(logits, lengths, indices) + + +def test_topk_without_vllm(indexer, monkeypatch): + from lightllm.models.glm5_next import indexer as indexer_module + + monkeypatch.setattr(indexer_module, "HAS_VLLM", False) + logits = torch.randn(5, 640, device="cuda") + lengths = torch.tensor([0, 1, 511, 512, 639], device="cuda", dtype=torch.int32) + indices = torch.empty(5, 512, device="cuda", dtype=torch.int32) + indexer.select_topk_indices(logits, lengths, indices) + _assert_topk(logits, lengths, indices) + + +@pytest.mark.parametrize("mtp_size", [1, 3, 6]) +@pytest.mark.parametrize("max_pools", [1024, 262144]) +def test_paged_decode_matches_nonpaged_with_graph_and_mtp(indexer, monkeypatch, mtp_size, max_pools): + pytest.importorskip("deep_gemm") + _require_vllm_topk() + torch.manual_seed(27) + rows = 5 * mtp_size + storage = torch.empty(1024, 1, 584, device="cuda", dtype=torch.bfloat16) + packed = storage.view(torch.uint8)[:, :, -132:] + packed[:, 0, :128] = torch.randn(1024, 128, device="cuda").to(torch.float8_e4m3fn).view(torch.uint8) + scales = torch.pow(2.0, torch.randint(-3, 2, (1024,), device="cuda").float()) + packed[:, 0, 128:] = scales.view(torch.uint8).view(-1, 4) + table = torch.full((5, max_pools * 4), -1, device="cuda", dtype=torch.int32) + table[:, 3::4] = torch.randint(0, 1024, (5, max_pools), device="cuda", dtype=torch.int32) + req_idx = torch.tensor([1, 3, 0, 2, 4], device="cuda", dtype=torch.int32).repeat_interleave(mtp_size) + lengths = torch.zeros(rows, device="cuda", dtype=torch.int32) + q = torch.randn(rows, 32, 128, device="cuda").to(torch.float8_e4m3fn) + weights = torch.randn(rows, 32, device="cuda") * 0.05 + infer_state = SimpleNamespace( + b_req_idx=req_idx, + b_seq_len=lengths, + b1_cu_q_seq_len=torch.arange(rows + 1, device="cuda", dtype=torch.int32), + max_q_seq_len=1, + req_manager=SimpleNamespace(req_to_token_indexs=table), + ) + att_state = SimpleNamespace(q_kv_seq_len=lengths) + logits_outputs = [] + select_topk = indexer.select_topk_indices + + def capture_logits(logits, pool_lengths, indices): + logits_outputs.append(logits) + select_topk(logits, pool_lengths, indices) + + monkeypatch.setattr(indexer, "select_topk_indices", capture_logits) + + def run(): + return indexer._get_decode_indices(q, weights, packed, infer_state, att_state, max_pools) + + # Match startup capture with only empty HOLD rows, then replay real work. + run() + logits_outputs.clear() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = run() + graph_logits = logits_outputs.pop() + for bases in ([0, 2047, 2051, max_pools * 4 - mtp_size + 1, 0], [7, 65, 511, 2053, 0]): + seqs = torch.tensor(bases, device="cuda", dtype=torch.int32)[:, None] + seqs = seqs + torch.arange(mtp_size, device="cuda", dtype=torch.int32)[None, :] + seqs[-1].zero_() + lengths.copy_(seqs.flatten()) + graph.replay() + indexer._get_prefill_indices(q, weights, packed, infer_state, att_state, max_pools) + reference_logits = logits_outputs.pop() + pool_lengths = lengths // 4 + valid = torch.arange(max_pools, device="cuda")[None, :] < pool_lengths[:, None] + torch.testing.assert_close(graph_logits[valid], reference_logits[valid], rtol=0, atol=0) + # Fragmented tables deliberately repeat keys; equal-score top-k ties + # may choose different indices. Check exact selected values and uniqueness. + _assert_topk(reference_logits, pool_lengths, actual) + assert (actual[-mtp_size:] == -1).all() + # Exercise request-slot reuse and speculative rollback on the same graph. + req_idx.copy_(req_idx.roll(mtp_size)) diff --git a/unit_tests/models/glm5_next/test_indexer_workspace.py b/unit_tests/models/glm5_next/test_indexer_workspace.py new file mode 100644 index 0000000000..8360f184c4 --- /dev/null +++ b/unit_tests/models/glm5_next/test_indexer_workspace.py @@ -0,0 +1,78 @@ +import weakref +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.models.glm5_next.indexer import Glm5NextNsaInfer + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires the GLM GPU indexer kernels") +@pytest.mark.parametrize( + "max_pools,q_lengths,kv_lengths,expected_chunks", + [(640, [2047, 2048, 2049] * 4, [2304, 2432, 2560] * 4, 1), (262144, [31, 1026], [65536, 1048576], 3)], +) +def test_workspace_preserves_mixed_request_pool_selection( + monkeypatch, max_pools, q_lengths, kv_lengths, expected_chunks +): + """Chunked prefill bounds padded allocations and preserves top-k results.""" + import deep_gemm + + indexer = Glm5NextNsaInfer( + 0, {"index_kpool": 4, "index_topk": 2048, "index_n_heads": 32, "index_head_dim": 128, "rms_norm_eps": 1e-5}, 1 + ) + fp8_mqa_logits = deep_gemm.fp8_mqa_logits + select_topk = indexer.select_topk_indices + topk_scores = [] + previous_logits = None + chunk_count = 0 + + def track_logits_lifetime(*args, **kwargs): + nonlocal previous_logits, chunk_count + if previous_logits is not None: + assert previous_logits() is None, "the previous logits chunk is still alive" + logits = fp8_mqa_logits(*args, **kwargs) + allocated_bytes = logits.untyped_storage().nbytes() + assert allocated_bytes <= 512 * 1024 ** 2 + previous_logits = weakref.ref(logits) + chunk_count += 1 + return logits + + def capture_topk_scores(logits, lengths, indices): + select_topk(logits, lengths, indices) + scores = logits.gather(1, indices.long().clamp_min(0)).masked_fill(indices < 0, -float("inf")) + topk_scores.append(scores.sort(dim=1).values) + + monkeypatch.setattr(deep_gemm, "fp8_mqa_logits", track_logits_lifetime) + monkeypatch.setattr(indexer, "select_topk_indices", capture_topk_scores) + torch.manual_seed(20261029) + batch, tokens = len(q_lengths), sum(q_lengths) + packed_buffer = torch.empty((batch * max_pools, 1, 132), dtype=torch.uint8, device="cuda") + packed_buffer[:, 0, :128] = ( + torch.randn(batch * max_pools, 128, device="cuda").to(torch.float8_e4m3fn).view(torch.uint8) + ) + packed_buffer[:, 0, 128:] = torch.ones(batch * max_pools, device="cuda").view(torch.uint8).view(-1, 4) + req_to_token_indexs = torch.full((batch, max_pools * 4), -1, dtype=torch.int32, device="cuda") + req_to_token_indexs[:, 3::4] = torch.arange(batch * max_pools, dtype=torch.int32, device="cuda").view(batch, -1) + q_fp8 = torch.randn(tokens, 32, 128, device="cuda").to(torch.float8_e4m3fn) + weights = torch.randn(tokens, 32, device="cuda") * 0.05 + lengths = torch.cat([torch.arange(k - q + 1, k + 1, device="cuda") for q, k in zip(q_lengths, kv_lengths)]).int() + infer_state = SimpleNamespace( + b_req_idx=torch.arange(batch, dtype=torch.int32, device="cuda"), + b_seq_len=torch.tensor(kv_lengths, dtype=torch.int32, device="cuda"), + b1_cu_q_seq_len=torch.tensor([0] + q_lengths, device="cuda").cumsum(0, dtype=torch.int32), + max_q_seq_len=max(q_lengths), + req_manager=SimpleNamespace(req_to_token_indexs=req_to_token_indexs), + ) + att_state = SimpleNamespace(q_kv_seq_len=lengths) + indexer._get_prefill_indices(q_fp8, weights, packed_buffer, infer_state, att_state, max_pools) + assert chunk_count == expected_chunks + expected_scores = torch.cat(topk_scores) + topk_scores.clear() + chunk_size = min(indexer._get_prefill_chunk_size(max_pools) // 2, tokens // 2) + monkeypatch.setattr(indexer, "_get_prefill_chunk_size", lambda _: chunk_size) + actual = indexer._get_prefill_indices(q_fp8, weights, packed_buffer, infer_state, att_state, max_pools) + torch.testing.assert_close(torch.cat(topk_scores), expected_scores, rtol=0, atol=0) + valid = actual >= 0 + assert (valid.sum(dim=1) == (lengths // 4).clamp_max(indexer.index_topk // 4)).all().item() + assert ((actual == -1) | (valid & (actual < (lengths // 4)[:, None]))).all().item() diff --git a/unit_tests/models/glm5_next/test_kda_conv_qkv.py b/unit_tests/models/glm5_next/test_kda_conv_qkv.py new file mode 100644 index 0000000000..b677477519 --- /dev/null +++ b/unit_tests/models/glm5_next/test_kda_conv_qkv.py @@ -0,0 +1,186 @@ +import weakref +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.basemodel.attention.base_att import AttControl +from lightllm.common.basemodel.attention.linear import kda as backend_module +from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import causal_conv1d_fn +from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_qkv import ( + causal_conv1d_qkv, +) +from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda import chunk_kda_with_fused_gate +from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.l2norm import l2norm_fwd, USE_DEFAULT_FLA_NORM + + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (9, 0) or USE_DEFAULT_FLA_NORM, + reason="requires Hopper and the LightLLM L2 norm", +) + + +def reference(x, weight, states, starts, slots, warm, heads): + mixed = causal_conv1d_fn( + x.clone().T, + weight, + query_start_loc=starts, + cache_indices=slots, + has_initial_state=warm, + conv_states=states, + ).T + q, k, v = [part.view(1, x.shape[0], heads, 128) for part in mixed.chunk(3, dim=-1)] + return l2norm_fwd(q.contiguous()), l2norm_fwd(k.contiguous()), v.contiguous() + + +@pytest.mark.parametrize("heads", [1, 2, 8, 64]) +@pytest.mark.parametrize("state_width", [3, 4, 7]) +@pytest.mark.parametrize("lengths", [(2049, 2, 1024), (257, 1, 1), (1, 2, 3), (0, 1, 0, 2), (1,)]) +def test_conv_norm_preserves_outputs_and_history(heads, state_width, lengths): + torch.manual_seed(42) + channels = heads * 128 * 3 + opts = dict(device="cuda", dtype=torch.bfloat16) + projection = torch.randn(sum(lengths), channels + heads + 256, **opts) + original_projection = projection.clone() + x = projection[:, :channels] + weight = torch.randn(channels, 4, **opts) * 0.125 + states = torch.randn(5, channels, state_width, **opts) + expected_states = states.clone() + starts = torch.tensor([0, *lengths], device="cuda", dtype=torch.int32).cumsum(0, dtype=torch.int32) + slots = torch.tensor([1, 2, -1, 3][: len(lengths)], device="cuda", dtype=torch.int32) + cache_seqlens = torch.tensor([0, 7, 0, 1][: len(lengths)], device="cuda", dtype=torch.int32) + expected = reference(x, weight, expected_states[..., :3], starts, slots, cache_seqlens > 0, heads) + actual = causal_conv1d_qkv(x, weight, states[..., :3], starts, slots, cache_seqlens, max(lengths)) + for a, b in zip(actual, expected): + assert a.is_contiguous() + assert torch.equal(a.view(torch.uint8), b.view(torch.uint8)) + assert len({a.untyped_storage().data_ptr() for a in actual}) == 3 + assert torch.equal(states, expected_states) + assert torch.equal(projection, original_projection) + + +@pytest.mark.parametrize("lengths", [(3, 2, 4), (7, 1, 1)]) +def test_conv_norm_cuda_graph_changing_boundaries(lengths): + torch.manual_seed(42) + heads, channels = 2, 3 * 2 * 128 + opts = dict(device="cuda", dtype=torch.bfloat16) + x = torch.randn(9, channels + 258, **opts)[:, :channels] + weight = torch.randn(channels, 4, **opts) * 0.125 + initial = torch.randn(5, channels, 4, **opts) + states = initial.clone() + initial_boundaries = [0, lengths[0], lengths[0] + lengths[1], sum(lengths)] + starts = torch.tensor(initial_boundaries, device="cuda", dtype=torch.int32) + slots = torch.tensor([1, 2, -1], device="cuda", dtype=torch.int32) + cache_seqlens = torch.tensor([4, 0, 0], device="cuda", dtype=torch.int32) + causal_conv1d_qkv(x, weight, states[..., :3], starts, slots, cache_seqlens, max(lengths)) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = causal_conv1d_qkv(x, weight, states[..., :3], starts, slots, cache_seqlens, max(lengths)) + for boundaries, cached in ((initial_boundaries, [4, 0, 0]), ([0, 4, 8, 9], [0, 7, 0])): + starts.copy_(torch.tensor(boundaries, device="cuda", dtype=torch.int32)) + cache_seqlens.copy_(torch.tensor(cached, device="cuda", dtype=torch.int32)) + states.copy_(initial) + expected_states = initial.clone() + expected = reference(x, weight, expected_states[..., :3], starts, slots, cache_seqlens > 0, heads) + graph.replay() + for a, b in zip(actual, expected): + assert torch.equal(a.view(torch.uint8), b.view(torch.uint8)) + assert torch.equal(states, expected_states) + + +@pytest.mark.parametrize("heads", [2, 8]) +@pytest.mark.parametrize("state_width", [3, 7]) +def test_conv_norm_chained_prefills(heads, state_width): + torch.manual_seed(42) + channels = 3 * heads * 128 + opts = dict(device="cuda", dtype=torch.bfloat16) + weight = torch.randn(channels, 4, **opts) * 0.125 + states = torch.randn(5, channels, state_width, **opts) + expected_states = states.clone() + slots = torch.tensor([1, 2, -1, 3], device="cuda", dtype=torch.int32) + cached = torch.tensor([0, 7, 0, 1], device="cuda", dtype=torch.int32) + for lengths in [(1, 2, 0, 3), (0, 1, 2, 0), (7, 0, 1, 2), (2, 1, 0, 1)]: + x = torch.randn(sum(lengths), channels + heads + 256, **opts)[:, :channels] + starts = torch.tensor([0, *lengths], device="cuda", dtype=torch.int32).cumsum(0, dtype=torch.int32) + expected = reference(x, weight, expected_states[..., :3], starts, slots, cached > 0, heads) + actual = causal_conv1d_qkv(x, weight, states[..., :3], starts, slots, cached, max(lengths)) + for a, b in zip(actual, expected): + assert torch.equal(a.view(torch.uint8), b.view(torch.uint8)) + assert torch.equal(states, expected_states) + cached += torch.tensor(lengths, device="cuda", dtype=torch.int32) + + +@pytest.mark.parametrize("lengths", [(129, 2), (129, 1, 1), (1,)]) +def test_backend_fused_conv_output_and_state(lengths): + torch.manual_seed(42) + heads, channels, tokens = 2, 3 * 2 * 128, sum(lengths) + opts = dict(device="cuda", dtype=torch.bfloat16) + projection = torch.randn(tokens, channels + heads + 256, **opts) + weight = torch.randn(channels, 4, **opts) * 0.125 + conv_states = torch.randn(4, channels, 4, **opts) + ssm_states = torch.randn(4, heads, 128, 128, device="cuda") * 0.01 + linear_state = SimpleNamespace(get_layer_view=lambda **kwargs: (conv_states, ssm_states)) + infer_state = type("State", (), {})() + infer_state.att_state_manager = SimpleNamespace(get_state=lambda **kwargs: linear_state) + infer_state.att_state_role = "main" + infer_state.b1_cu_q_seq_len = torch.tensor([0, *lengths], device="cuda", dtype=torch.int32).cumsum( + 0, dtype=torch.int32 + ) + infer_state.b_ready_cache_len = torch.arange(len(lengths), device="cuda", dtype=torch.int32) % 2 * 7 + infer_state.max_q_seq_len = max(lengths) + backend = SimpleNamespace( + head_dim=128, + tp_num_heads=heads, + tp_hidden_size=heads * 128, + conv_kernel_size=4, + lower_bound=-5.0, + ) + att = backend_module.KDAPrefillAttState(backend=backend, _infer_state_weak_ref=weakref.ref(infer_state)) + att.b_conv_buffer_idx = torch.arange(1, len(lengths) + 1, device="cuda", dtype=torch.int32) + att.b_ssm_buffer_idx = att.b_conv_buffer_idx + layer_weight = SimpleNamespace( + get_merged_kda_conv_weight=lambda: weight, + linear_A_log=SimpleNamespace(weight=torch.zeros(heads, device="cuda")), + linear_dt_bias=SimpleNamespace(weight=torch.zeros(heads * 128, device="cuda")), + ) + control = AttControl( + linear_att_prefill=True, + linear_att_prefill_dict=dict( + layer_weight=layer_weight, + layer_num=0, + mixed_qkv=projection[:, :channels], + raw_gate=torch.randn(tokens, heads * 128, **opts) * 0.1, + raw_beta=torch.randn(tokens, heads, **opts), + ), + ) + expected_conv, expected_ssm = conv_states.clone(), ssm_states.clone() + q, k, v = reference( + projection[:, :channels], + weight, + expected_conv[..., :3], + infer_state.b1_cu_q_seq_len, + att.b_conv_buffer_idx, + infer_state.b_ready_cache_len > 0, + heads, + ) + params = control.linear_att_prefill_dict + expected, final_state = chunk_kda_with_fused_gate( + q=q, + k=k, + v=v, + raw_g=params["raw_gate"].view(1, tokens, heads, 128), + beta=params["raw_beta"].float().sigmoid().view(1, tokens, heads), + A_log=layer_weight.linear_A_log.weight, + g_bias=layer_weight.linear_dt_bias.weight, + initial_state=ssm_states[att.b_ssm_buffer_idx].contiguous(), + output_final_state=True, + use_qk_l2norm_in_kernel=False, + cu_seqlens=infer_state.b1_cu_q_seq_len, + safe_gate=True, + lower_bound=-5.0, + ) + expected_ssm[att.b_ssm_buffer_idx] = final_state.to(ssm_states.dtype) + actual = att.prefill_att(None, None, None, att_control=control) + assert torch.equal(actual.view(torch.uint8), expected.view(torch.uint8)) + assert torch.equal(conv_states, expected_conv) + assert torch.equal(ssm_states, expected_ssm) diff --git a/unit_tests/models/glm5_next/test_kernels.py b/unit_tests/models/glm5_next/test_kernels.py new file mode 100644 index 0000000000..49b2a08eb6 --- /dev/null +++ b/unit_tests/models/glm5_next/test_kernels.py @@ -0,0 +1,665 @@ +import weakref +import dataclasses +import math +from types import SimpleNamespace + +import pytest +import torch +import triton + +from lightllm.common.basemodel.infer_struct import InferStateInfo +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.utils.envs_utils import set_env_start_args +from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda import ( + chunk_kda_with_fused_gate, + fused_kda_gate_chunk_cumsum, +) +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, + gather_pools, + gather_paged_pools, + get_pool_ranges, + expand_topk, +) +from lightllm.models.glm5_next.triton_kernel.index_quant import hadamard_transform_quant_fp8 + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.fixture(autouse=True) +def setup(): + torch.manual_seed(1525) + set_env_start_args(dataclasses.asdict(StartArgs())) + triton.set_allocator(lambda size, alignment, stream: torch.empty(size, device="cuda", dtype=torch.int8)) + + +@pytest.mark.parametrize("tokens", [1, 3, 65, 129]) +def test_kda_chunk_and_decode_match_recurrence(tokens): + heads, dim = 2, 128 + rand = lambda *shape: torch.randn(shape, device="cuda", dtype=torch.bfloat16) + q, k, v, gate = [rand(1, tokens, heads, dim) for _ in range(4)] + beta = rand(1, tokens, heads) + a = torch.randn(heads, device="cuda") + bias = torch.randn(heads, dim, device="cuda") + initial = torch.randn(1, heads, dim, dim, device="cuda") * 0.1 + expected = [] + state = initial[0].clone() + for i in range(tokens): + out, state = _reference_kda(q[0, i], k[0, i], v[0, i], gate[0, i], beta[0, i], a, bias, state) + expected.append(out) + expected = torch.stack(expected).unsqueeze(0) + actual, final = chunk_kda_with_fused_gate( + q=q, + k=k, + v=v.clone(), + raw_g=gate, + beta=beta.float().sigmoid(), + A_log=a, + g_bias=bias.flatten(), + initial_state=initial, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + cu_seqlens=torch.tensor([0, tokens], dtype=torch.int32, device="cuda"), + safe_gate=True, + ) + torch.testing.assert_close(actual.float(), expected, atol=4e-3, rtol=3e-2) + torch.testing.assert_close(final[0], state, atol=8e-3, rtol=3e-2) + # Reuse a nonzero request slot; the neighboring requests must stay intact. + states = torch.randn(4, heads, dim, dim, device="cuda") + states[2] = initial[0] + unchanged = states[[0, 1, 3]].clone() + for i in range(tokens): + out, _ = fused_recurrent_kda( + q[:, i : i + 1], + k[:, i : i + 1], + v[:, i : i + 1], + gate[:, i : i + 1].reshape(1, 1, -1), + beta[:, i : i + 1], + a, + bias.flatten(), + states, + torch.tensor([2], device="cuda", dtype=torch.int32), + ) + torch.testing.assert_close(out[0, 0].float(), expected[0, i], atol=2e-3, rtol=1e-2) + torch.testing.assert_close(states[2], state, atol=2e-5, rtol=2e-4) + assert torch.equal(states[[0, 1, 3]], unchanged) + + +@pytest.mark.parametrize("seq_lens", [(65,), (3, 65, 129)]) +@pytest.mark.parametrize("safe_gate", [False, True]) +@pytest.mark.parametrize("strided", [False, True]) +def test_kda_gate_cumsum_packed_shape(seq_lens, safe_gate, strided): + heads, dim, chunk_size = 2, 128, 64 + if strided: + # All three input strides differ from contiguous [T, H, D]; empty_like + # allocates a contiguous output for this sliced, non-dense input view. + storage = torch.randn(sum(seq_lens) * 2, heads * 2, dim * 2, device="cuda", dtype=torch.bfloat16) + raw_g = storage[::2, ::2, ::2] + else: + raw_g = torch.randn(sum(seq_lens), heads, dim, device="cuda", dtype=torch.bfloat16) + a_log = torch.randn(heads, device="cuda") + bias = torch.randn(heads, dim, device="cuda") if len(seq_lens) > 1 else None + cu_seqlens = torch.tensor([0, *seq_lens], device="cuda", dtype=torch.int32).cumsum(0, dtype=torch.int32) + lower_bound = -3.0 + + actual = fused_kda_gate_chunk_cumsum( + raw_g, + A_log=a_log, + g_bias=bias.flatten() if bias is not None else None, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + safe_gate=safe_gate, + lower_bound=lower_bound, + ) + + gate_input = raw_g.float() + (bias if bias is not None else 0) + amplitude = a_log.exp()[None, :, None] + if safe_gate: + log_gate = lower_bound * torch.sigmoid(amplitude * gate_input) + else: + log_gate = -amplitude * torch.nn.functional.softplus(gate_input) + expected = torch.empty_like(log_gate) + seq_start = 0 + for seq_len in seq_lens: + for offset in range(0, seq_len, chunk_size): + start = seq_start + offset + end = seq_start + min(offset + chunk_size, seq_len) + expected[start:end] = log_gate[start:end].cumsum(0) / math.log(2) + seq_start += seq_len + + assert actual.shape == raw_g.shape + torch.testing.assert_close(actual, expected, atol=2e-5, rtol=1e-5) + + +@pytest.mark.parametrize( + "q_lens", + [ + (3, 2, 4), + (0, 1, 0, 1, 0), + (1,) * 64, + (257, 1, 513, 0), + (8129,) + (1,) * 63, + (3, 0, 1, 2) * 31 + (3, 0, 1), + (3, 0, 1, 2) * 32, + (3, 0, 1, 2) * 32 + (3,), + ], +) +def test_kpool_ranges_cuda_graph_with_changing_query_boundaries(q_lens): + # Move the longest query across the batch while replaying the same graph. + # Empty queries must not use the one-query-per-request decode shortcut. + tensor = lambda values: torch.tensor(values, device="cuda", dtype=torch.int32) + lengths = torch.empty(sum(q_lens), device="cuda", dtype=torch.int32) + cu_q_lens = torch.empty(len(q_lens) + 1, device="cuda", dtype=torch.int32) + max_pools = 4096 + + def set_inputs(counts, phase): + cu, visible_lengths, pool_starts = [0], [], [] + for batch, count in enumerate(counts): + prefix = batch * 7 + phase + visible_lengths.extend(range(prefix + 1, prefix + count + 1)) + pool_starts.extend([batch * max_pools] * count) + cu.append(cu[-1] + count) + lengths.copy_(tensor(visible_lengths)) + cu_q_lens.copy_(tensor(cu)) + expected_lengths = tensor([length // 4 for length in visible_lengths]) + expected_starts = tensor(pool_starts) + return expected_starts, expected_starts + expected_lengths, expected_lengths + + expected = set_inputs(q_lens, 0) + actual = get_pool_ranges(lengths, cu_q_lens, max(q_lens), max_pools) + for output, reference in zip(actual, expected): + assert torch.equal(output, reference) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = get_pool_ranges(lengths, cu_q_lens, max(q_lens), max_pools) + expected = set_inputs(q_lens[::-1], 3) + graph.replay() + for output, reference in zip(actual, expected): + assert torch.equal(output, reference) + + +@pytest.mark.parametrize("cuda_graph", [False, True]) +def test_kpool_chunk_boundaries_and_fragmented_token_kv(cuda_graph): + # Unaligned chunks, reordered requests, and pools spanning old/new tails. + seqs = [9, 7] + ragged = torch.randperm(40, device="cuda", dtype=torch.int32)[: sum(seqs)] + req_idx = torch.tensor([1, 3], device="cuda", dtype=torch.int32) + table = torch.zeros(5, 12, device="cuda", dtype=torch.int32) + table[1, :9], table[3, :7] = ragged[:9], ragged[9:] + source = torch.randn(sum(seqs), 256, device="cuda", dtype=torch.bfloat16) + ape = torch.randn(4, 128, device="cuda") + packed_storage = torch.zeros(40, 1, 584, device="cuda", dtype=torch.bfloat16) + packed = packed_storage.view(torch.uint8)[:, :, -132:] + tail = torch.randn(5, 4, 256, device="cuda", dtype=torch.bfloat16) + unchanged = tail[[0, 2, 4]].clone() + tensor = lambda x: torch.tensor(x, device="cuda", dtype=torch.int32) + for chunks in [[(1, 0, 3), (3, 0, 2)], [(3, 2, 3), (1, 3, 9)], [(3, 3, 7)]]: + raw, lengths, starts, locations = [], [], [], [] + cu_q_lens = [0] + for req, first, end in chunks: + offset = 0 if req == 1 else 9 + raw.append(source[offset + first : offset + end]) + lengths.extend(range(first + 1, end + 1)) + starts.extend([len(locations)] * (end - first)) + locations.extend(table[req, :end].tolist()) + cu_q_lens.append(cu_q_lens[-1] + end - first) + args = ( + torch.cat(raw), + tail, + packed, + ape, + tensor(lengths), + tensor(locations)[tensor(starts) + tensor(lengths) - 1], + tensor([r for r, _, _ in chunks]), + tensor(cu_q_lens), + tensor([end for _, _, end in chunks]), + ) + max_q_len = max(end - first for _, first, end in chunks) + if cuda_graph: + # Compile before capture, then restore this chunk's input cache state. + saved_tail, saved_packed = tail.clone(), packed.clone() + compress_pools(*args, max_q_seq_len=max_q_len) + tail.copy_(saved_tail) + packed.copy_(saved_packed) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + compress_pools(*args, max_q_seq_len=max_q_len) + graph.replay() + else: + compress_pools(*args, max_q_seq_len=max_q_len) + for req, _, end in chunks: + offset = 0 if req == 1 else 9 + assert torch.equal(tail[req, : end % 4], source[offset + end // 4 * 4 : offset + end]) + assert torch.equal(tail[[0, 2, 4]], unchanged) + # Model the full KV copy performed by cache offload/load or request move. + packed_storage = packed_storage.clone() + packed = packed_storage.view(torch.uint8)[:, :, -132:] + keys, scales = gather_pools(packed, table, req_idx, torch.tensor(seqs, device="cuda", dtype=torch.int32), 3) + for batch, start in enumerate([0, 9]): + 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) + 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) + starts = torch.tensor([0, 9], device="cuda", dtype=torch.int32) + groups = torch.tensor([[1, 0], [0, -1]], device="cuda", dtype=torch.int32) + out, relative = expand_topk(groups, lengths, req_idx, table, index_topk=8) + for batch, expected in enumerate([[4, 5, 6, 7, 0, 1, 2, 3, 8], list(range(7))]): + assert relative[batch, : len(expected)].tolist() == expected + assert ( + out[batch, : len(expected)].tolist() + == ragged[starts[batch].item() + torch.tensor(expected, device="cuda")].tolist() + ) + assert (out[batch, len(expected) :] == -1).all() + + +def test_kpool_decode_cuda_graph_and_padding(): + # Replay one graph across all four tail phases, with reordered requests + # and repeated padding slots. Padding must not race on the hold state. + source = torch.randn(2, 8, 256, device="cuda", dtype=torch.bfloat16) + ape = torch.randn(4, 128, device="cuda") + tail = torch.zeros(5, 4, 256, device="cuda", dtype=torch.bfloat16) + storage = torch.zeros(32, 1, 584, device="cuda", dtype=torch.bfloat16) + packed = storage.view(torch.uint8)[:, :, -132:] + table = torch.randperm(32, device="cuda", dtype=torch.int32).view(4, 8) + raw = torch.zeros(4, 256, device="cuda", dtype=torch.bfloat16) + req_idx = torch.tensor([1, 3, 4, 4], device="cuda", dtype=torch.int32) + lengths = torch.ones(4, device="cuda", dtype=torch.int32) + mem_index = torch.zeros(4, device="cuda", dtype=torch.int32) + cu_q_lens = torch.arange(5, device="cuda", dtype=torch.int32) + + def run(): + compress_pools(raw, tail, packed, ape, lengths, mem_index, req_idx, cu_q_lens, lengths, max_q_seq_len=1) + + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + run() + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + tail.zero_() + for pos in range(8): + order = [pos % 2, 1 - pos % 2] + req_idx[:2] = torch.tensor([1 if i == 0 else 3 for i in order], device="cuda") + raw[:2] = source[order, pos] + lengths.fill_(pos + 1) + mem_index[:2] = table[order, pos] + graph.replay() + for batch, req in enumerate([1, 3]): + pending = (pos + 1) % 4 + assert torch.equal(tail[req, :pending], source[batch, (pos + 1) // 4 * 4 : pos + 1]) + 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) + 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()) + assert not tail[[0, 2, 4]].any() + + +@pytest.mark.parametrize("max_pools", [640, 262144]) +def test_gather_paged_pools_valid_pages_and_graph_replay(max_pools): + storage = torch.zeros(64, 1, 584, device="cuda", dtype=torch.bfloat16) + packed = storage.view(torch.uint8)[:, :, -132:] + source_keys = torch.randn(64, 128, device="cuda").to(torch.float8_e4m3fn).view(torch.uint8) + source_scales = torch.rand(64, device="cuda") + 0.1 + packed[:, 0, :128] = source_keys + packed[:, 0, 128:] = source_scales.view(torch.uint8).view(64, 4) + table = torch.full((4, max_pools * 4), -1, device="cuda", dtype=torch.int32) + locations = torch.randint(0, 64, (4, max_pools), device="cuda", dtype=torch.int32) + table[:, 3::4] = locations + req_idx = torch.tensor([2, 0, 3], device="cuda", dtype=torch.int32) + lengths = torch.tensor([max_pools, 3, 0], device="cuda", dtype=torch.int32) + + def check(pages, block_table): + page_bytes = pages.view(pages.shape[0], -1) + for row, (req, length) in enumerate(zip(req_idx.tolist(), lengths.tolist())): + page_ids = block_table[row, : triton.cdiv(length, 64)].long() + keys = page_bytes[page_ids, : 64 * 128].reshape(-1, 128) + scales = page_bytes[page_ids, 64 * 128 :].contiguous().view(torch.float32).flatten() + locs = locations[req, :length].long() + assert torch.equal(keys[:length], source_keys[locs]) + assert torch.equal(scales[:length], source_scales[locs]) + assert not keys[length:].any() + assert (scales[length:] == 1).all() + + check(*gather_paged_pools(packed, table, req_idx, lengths, max_pools)) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + pages, block_table = gather_paged_pools(packed, table, req_idx, lengths, max_pools) + # Empty HOLD rows initialize one dummy page, even with a 1M graph capacity. + pages.fill_(127) + lengths.zero_() + graph.replay() + page_bytes = pages.view(3, -1, 64 * 132) + assert not page_bytes[:, 0, : 64 * 128].any() + assert (page_bytes[:, 0, 64 * 128 :].contiguous().view(torch.float32) == 1).all() + assert (page_bytes[:, 1:] == 127).all() + for counts in ([513, 65, max_pools], [1, 0, 3]): + lengths.copy_(torch.tensor(counts, device="cuda", dtype=torch.int32)) + req_idx.copy_(torch.tensor([1, 3, 0], device="cuda", dtype=torch.int32)) + graph.replay() + check(pages, block_table) + + +@pytest.mark.parametrize("max_pools", [65535, 65536, 262144]) +def test_gather_pools_long_context_cuda_graph(max_pools): + # 1M-token graph warmup pads to 262144 pools even for empty hold requests. + # Also gather real pools at the end of that range from fragmented token KV. + storage = torch.zeros(64, 1, 584, device="cuda", dtype=torch.bfloat16) + packed = storage.view(torch.uint8)[:, :, -132:] + source_keys = torch.randn(64, 128, device="cuda").to(torch.float8_e4m3fn).view(torch.uint8) + source_scales = torch.rand(64, device="cuda") + 0.1 + packed[:, 0, :128] = source_keys + packed[:, 0, 128:] = source_scales.view(torch.uint8).view(64, 4) + table = torch.full((4, max_pools * 4), -1, device="cuda", dtype=torch.int32) + locations = torch.randint(0, 64, (4, max_pools), device="cuda", dtype=torch.int32) + table[:, 3::4] = locations + req_idx = torch.tensor([2, 0, 3], device="cuda", dtype=torch.int32) + seq_len = torch.tensor([max_pools * 4, 14, 2], device="cuda", dtype=torch.int32) + + def check(keys, scales): + keys = keys.view(torch.uint8).view(3, max_pools, 128) + scales = scales.view(3, max_pools) + for batch, (req, count) in enumerate([(2, max_pools), (0, 3), (3, 0)]): + locs = locations[req, :count].long() + assert torch.equal(keys[batch, :count], source_keys[locs]) + assert torch.equal(scales[batch, :count], source_scales[locs]) + assert not keys[batch, count:].any() + assert (scales[batch, count:] == 1).all() + + check(*gather_pools(packed, table, req_idx, seq_len, max_pools)) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + keys, scales = gather_pools(packed, table, req_idx, seq_len, max_pools) + # Replay with the same short hold lengths used by model startup. + original_lengths = seq_len.clone() + seq_len.fill_(2) + graph.replay() + assert not keys.view(torch.uint8).any() + assert (scales == 1).all() + seq_len.copy_(original_lengths) + graph.replay() + 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 + 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) + tail = torch.zeros(2, 4, 256, device="cuda", dtype=torch.bfloat16) + ape = torch.randn(4, 128, device="cuda") + lengths = torch.arange(1, 5, device="cuda", dtype=torch.int32) + zero = torch.zeros(4, device="cuda", dtype=torch.int32) + ragged = torch.tensor([0, 1, 2, loc], device="cuda", dtype=torch.int32) + cu_q_lens = torch.tensor([0, 4], device="cuda", dtype=torch.int32) + compress_pools(raw, tail, packed, ape, lengths, ragged, zero[:1], cu_q_lens, lengths[-1:], max_q_seq_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) + assert torch.equal(keys.view(torch.uint8), expected_key.view(torch.uint8)) + assert torch.equal(scales, expected_scale.flatten()) + + +@pytest.mark.parametrize("tp_world_size", [1, 4]) +def test_mhc_keeps_streams_through_decode_autotuning(monkeypatch, tp_world_size): + from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType + from lightllm.models.glm5_next.layer_infer.transformer_layer_infer import Glm5NextTransformerLayerInfer + from lightllm.models.glm5_next.model import Glm5NextTpPartModel + + hidden = 4096 + monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_DP", "0") + monkeypatch.setenv("LIGHTLLM_DP_WORLD_SIZE", str(tp_world_size)) + pre_infer = Glm5NextTpPartModel.pre_layer_infer_class({"hc_mult": 4}) + embeddings = torch.randn(1, hidden, device="cuda", dtype=torch.bfloat16) + input_ids = torch.zeros(1, device="cuda", dtype=torch.long) + pre_weight = SimpleNamespace(wte_weight_=lambda input_ids, alloc_func: embeddings[input_ids]) + pre_weight.wte_weight_.weight = embeddings + pre_weight.wte_weight_.tp_vocab_start_id = 0 + pre_weight.wte_weight_.tp_vocab_end_id = 1 + infer_state = SimpleNamespace(dist_group=None, multimodal_params=[]) + + def all_reduce(input_embeddings, **kwargs): + # TP communication must operate on the original embedding width. + assert input_embeddings.shape == (1, hidden) + input_embeddings.mul_(tp_world_size) + + monkeypatch.setattr("lightllm.models.llama.layer_infer.pre_layer_infer.all_reduce", all_reduce) + monkeypatch.setattr("lightllm.models.qwen_vl.layer_infer.pre_layer_infer.all_reduce", all_reduce) + expected_streams = (embeddings * tp_world_size).unsqueeze(1).expand(-1, 4, -1) + weight = SimpleNamespace( + att_norm_weight_=SimpleNamespace(weight=torch.ones(hidden, device="cuda", dtype=torch.bfloat16)), + ffn_norm_weight_=SimpleNamespace(weight=torch.ones(hidden, device="cuda", dtype=torch.bfloat16)), + ) + for prefix in ("attn", "ffn"): + setattr(weight, f"hc_{prefix}_fn", SimpleNamespace(weight=torch.randn(24, 4 * hidden, device="cuda") * 0.005)) + setattr(weight, f"hc_{prefix}_base", SimpleNamespace(weight=torch.zeros(24, device="cuda"))) + setattr(weight, f"hc_{prefix}_scale", SimpleNamespace(weight=torch.ones(3, device="cuda"))) + layer = object.__new__(Glm5NextTransformerLayerInfer) + layer.use_mhc = True + layer.embed_dim_, layer.mhc_streams = hidden, 4 + layer.num_hidden_layers, layer.autotune_layer_num = 5, 4 + layer.eps_, layer.hc_eps, layer.hc_sinkhorn_iters = 1e-5, 1e-6, 20 + layer.token_attention_forward = layer.context_attention_forward = layer._ffn = lambda x, *_: x * 0.1 + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + x = pre_infer.token_forward(input_ids, infer_state, pre_weight) + torch.testing.assert_close(x.view(1, 4, hidden), expected_streams, atol=0, rtol=0) + for i in range(5): + layer.layer_num_ = i + x = layer.token_forward(x, None, weight) + assert x.shape == (1, hidden if i == 4 else 4 * hidden) + with Autotuner.autotune_warmup(): + x = pre_infer.context_forward(input_ids, infer_state, pre_weight) + torch.testing.assert_close(x.view(1, 4, hidden), expected_streams, atol=0, rtol=0) + for i in range(4): + layer.layer_num_ = i + x = layer.context_forward(x, None, weight) + assert x.shape == (1, hidden if i == 3 else 4 * hidden) + + +def test_shared_chunk_kernel_preserves_gdn_natural_log_decay(): + from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops import chunk_gated_delta_rule + + tokens, heads, dim = 67, 2, 128 + q, k, v = [torch.randn(1, tokens, heads, dim, device="cuda", dtype=torch.bfloat16) for _ in range(3)] + q = torch.nn.functional.normalize(q.float(), dim=-1).bfloat16() + k = torch.nn.functional.normalize(k.float(), dim=-1).bfloat16() + gate = -torch.rand(1, tokens, heads, device="cuda") + beta = torch.rand_like(gate) + state = torch.randn(1, heads, dim, dim, device="cuda") * 0.1 + initial = state.clone() + expected = [] + for i in range(tokens): + 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) + out, final = chunk_gated_delta_rule( + q, + k, + v, + gate, + beta, + initial_state=initial, + output_final_state=True, + cu_seqlens=torch.tensor([0, tokens], device="cuda", dtype=torch.int32), + ) + torch.testing.assert_close(out.float(), torch.stack(expected, 1), atol=4e-3, rtol=3e-2) + torch.testing.assert_close(final, state, atol=8e-3, rtol=3e-2) + + +@pytest.mark.parametrize("heads", [16, 64, 128]) +def test_nope_attention_native_512_and_cuda_graph(heads): + from lightllm.common.basemodel.attention.base_att import AttControl + from lightllm.common.basemodel.attention.nsa.glm5_next import Glm5NextSparsePrefillState, Glm5NextSparseDecodeState + + if torch.cuda.get_device_capability()[0] != 9: + pytest.skip("FA3 NoPE requires Hopper") + # The 72-element tail stores indexer data, not a zero RoPE embedding. + packed = torch.randn(4096, 1, 584, dtype=torch.bfloat16, device="cuda") + kv = packed[:, :, :512] + # Match the head-major projection's non-contiguous query layout. + q = torch.randn(heads, 9, 512, dtype=torch.bfloat16, device="cuda").transpose(0, 1) + lengths = torch.tensor([0, 1, 3, 127, 128, 511, 2048, 2049, 2051], dtype=torch.int32, device="cuda") + indexes = torch.randint(0, 4096, (9, 2176), dtype=torch.int32, device="cuda") + offsets = torch.arange(indexes.shape[1], device="cuda") + indexes.masked_fill_(offsets[None, :] >= lengths[:, None], -1) + + def reference(): + expected = [] + for i, length in enumerate(lengths.tolist()): + keys = kv[indexes[i, :length].long(), 0].float() + expected.append((q[i].float() @ keys.T * 0.0625).softmax(-1) @ keys) + return torch.stack(expected) + + expected = reference() + control = AttControl( + nsa_prefill_dict={"topk_mem_indices": indexes, "softmax_scale": 0.0625, "kv_lora_rank": kv.shape[-1]} + ) + prefill_state = Glm5NextSparsePrefillState() + prefill = prefill_state._nsa_prefill_att(q, kv, control) + torch.testing.assert_close(prefill.float(), expected, atol=0.012, rtol=0.015) + infer_state = InferStateInfo() + infer_state.b1_cu_q_seq_len = torch.arange(10, dtype=torch.int32, device="cuda") + infer_state.max_q_seq_len = 1 + decode = Glm5NextSparseDecodeState( + _infer_state_weak_ref=weakref.ref(infer_state), + nsa_cache_seqlens=lengths, + nsa_cu_seqlens_k_new=torch.nn.functional.pad(lengths.cumsum(0, dtype=torch.int32), (1, 0)), + ) + control.nsa_decode_dict = control.nsa_prefill_dict + out = decode._nsa_decode_att((q, q[..., :0]), kv, control) + torch.testing.assert_close(out.float(), expected, atol=0.012, rtol=0.015) + + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + prefill_state._nsa_prefill_att(q, kv, control) + decode._nsa_decode_att((q, q[..., :0]), kv, control) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_prefill = prefill_state._nsa_prefill_att(q, kv, control) + graph_decode = decode._nsa_decode_att((q, q[..., :0]), kv, control) + graph.replay() + torch.testing.assert_close(graph_prefill.float(), expected, atol=0.012, rtol=0.015) + torch.testing.assert_close(graph_decode.float(), expected, atol=0.012, rtol=0.015) + + q.normal_() + packed.normal_() + lengths.copy_(lengths.roll(1)) + decode.nsa_cu_seqlens_k_new[1:].copy_(lengths.cumsum(0, dtype=torch.int32)) + indexes.random_(0, 4096) + indexes.masked_fill_(offsets[None, :] >= lengths[:, None], -1) + graph.replay() + expected = reference() + torch.testing.assert_close(graph_prefill.float(), expected, atol=0.012, rtol=0.015) + torch.testing.assert_close(graph_decode.float(), expected, atol=0.012, rtol=0.015) + + +@pytest.mark.parametrize("max_kv_seq_len", [2065, 1 << 20]) +def test_kpool_indexer_long_prefill_and_cached_decode(max_kv_seq_len): + from lightllm.models.glm5_next.indexer import Glm5NextNsaInfer + + tokens, heads, dim = 2065, 32, 128 + hidden = torch.randn(tokens, dim, device="cuda", dtype=torch.bfloat16) + q_weight = torch.randn(dim, heads * dim, device="cuda", dtype=torch.bfloat16) * 0.1 + ape = torch.randn(4, dim, device="cuda") + weights = SimpleNamespace( + wk_proj_=SimpleNamespace(mm=lambda x: x), + k_norm_=lambda x, eps: x, + index_kpool_compress_gate=SimpleNamespace(mm=lambda x: x * 0.25), + index_kpool_compress_ape=SimpleNamespace(weight=ape), + wq_b_proj_=SimpleNamespace(mm=lambda x: x @ q_weight), + weights_proj_=SimpleNamespace(mm=lambda x: torch.ones(x.shape[0], heads, device=x.device)), + ) + 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] + full_att_state = SimpleNamespace( + get_layer_view=lambda logical_layer_id: SimpleNamespace(indexer_k=storage.view(torch.uint8)[:, :, -132:]), + ) + infer = SimpleNamespace( + 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 "compressor_indexer" in state_type_filter + else 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]), + 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), + max_q_seq_len=tokens, + ) + state = SimpleNamespace( + q_kv_seq_len=torch.arange(1, tokens + 1, device="cuda", dtype=torch.int32), + q_req_idx=torch.zeros(tokens, device="cuda", dtype=torch.int32), + ) + indexer = Glm5NextNsaInfer( + 0, + {"index_topk": 2048, "index_n_heads": heads, "index_head_dim": dim, "rms_norm_eps": 1e-5, "index_kpool": 4}, + 1, + ) + _, 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() + 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. Stale window data is irrelevant at a complete group boundary. + prefix = 2060 + storage = storage.clone() + storage[ragged[prefix:].long()] = 0 + tail.fill_(19) + infer.is_prefill = False + infer.b1_cu_q_seq_len[1] = 1 + infer.max_q_seq_len = 1 + for pos in range(prefix, tokens): + infer.mem_index = ragged[pos : pos + 1] + infer.max_kv_seq_len = pos + 1 + infer.b_seq_len.fill_(pos + 1) + state.q_kv_seq_len = torch.tensor([pos + 1], device="cuda", dtype=torch.int32) + state.q_req_idx = infer.b_req_idx + _, decoded = indexer._get_indices(hidden[pos : pos + 1], hidden[pos : pos + 1], infer, state, weights) + valid = 2048 + (pos + 1) % 4 + assert set(decoded[0, :valid].tolist()) == set(full[pos, :valid].tolist()) + + 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) + query = (hidden[-1:] @ q_weight).view(heads, dim) + 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) + 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_kpool_gather.py b/unit_tests/models/glm5_next/test_kpool_gather.py new file mode 100644 index 0000000000..9c25401a3f --- /dev/null +++ b/unit_tests/models/glm5_next/test_kpool_gather.py @@ -0,0 +1,141 @@ +"""Byte-exact gather and graph-replay coverage for GLM K-pool staging.""" + +import pytest +import torch +import triton + +from lightllm.models.glm5_next.triton_kernel import kpool +from lightllm.models.glm5_next.triton_kernel.kpool import gather_paged_pools, gather_pools + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +def _inputs(max_pools, requests=8): + # The packed view has both a nonzero byte offset and a larger row stride. + storage = torch.empty((128, 1, 584), dtype=torch.bfloat16, device="cuda") + packed = storage.view(torch.uint8)[:, :, -132:] + keys = torch.randn(128, 128, device="cuda").to(torch.float8_e4m3fn).view(torch.uint8) + scales = torch.pow(2.0, torch.randint(-4, 3, (128,), device="cuda")).float() + packed[:, 0, :128] = keys + packed[:, 0, 128:] = scales.view(torch.uint8).view(128, 4) + table_storage = torch.full((requests, max_pools * 4 + 16), -1, device="cuda", dtype=torch.int32) + table = table_storage[:, 4:-12] + locations = torch.randint(0, 128, (requests, max_pools), device="cuda", dtype=torch.int32) + table[:, 3::4] = locations + return packed, table, locations, keys, scales + + +@pytest.mark.parametrize("pools", [17, 33, 640, 65536, 262144]) +def test_gather_pools_bytes_scales_and_changing_graph_inputs(pools): + torch.manual_seed(29) + packed, table, locations, source_keys, source_scales = _inputs(pools) + req = torch.tensor([3, 1, 0, 2], device="cuda", dtype=torch.int32) + lengths = torch.zeros(4, device="cuda", dtype=torch.int32) + gather_pools(packed, table, req, lengths, pools) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + keys, scales = gather_pools(packed, table, req, lengths, pools) + for values in ([0, 3, pools * 4 - 3, pools * 4], [pools * 4, 0, 7, 65], [2, 2, 2, 2]): + lengths.copy_(torch.tensor(values, device="cuda", dtype=torch.int32).clamp_max(pools * 4)) + req.copy_(req.roll(1)) + graph.replay() + valid = torch.arange(pools, device="cuda")[None, :] < (lengths // 4)[:, None] + loc = locations[req.long()].long() + expected_keys = source_keys[loc].masked_fill(~valid[:, :, None], 0) + expected_scales = source_scales[loc].masked_fill(~valid, 1.0) + torch.testing.assert_close(keys.view(torch.uint8).view(4, pools, 128), expected_keys, rtol=0, atol=0) + torch.testing.assert_close(scales.view(4, pools), expected_scales, rtol=0, atol=0) + + +def _check_paged(pages, block_table, req, lengths, locations, source_keys, source_scales): + page_bytes = pages.view(pages.shape[0], -1) + for row, (slot, count) in enumerate(zip(req.tolist(), lengths.tolist())): + # Shared pages may contain the owner's later keys. Only the original + # per-query causal length is readable by the indexer. + page_ids = block_table[row, : triton.cdiv(count, 64)].long() + actual_keys = page_bytes[page_ids, : 64 * 128].reshape(-1, 128)[:count] + actual_scales = page_bytes[page_ids, 64 * 128 :].contiguous().view(torch.float32).flatten()[:count] + loc = locations[slot, :count].long() + torch.testing.assert_close(actual_keys, source_keys[loc], rtol=0, atol=0) + torch.testing.assert_close(actual_scales, source_scales[loc], rtol=0, atol=0) + + +@pytest.mark.parametrize("max_pools", [640, 4096, 262144]) +@pytest.mark.parametrize("reuse_group", [1, 2, 4, 8]) +def test_paged_reuse_mixed_requests_causal_lengths_and_graph_slot_recycling(max_pools, reuse_group): + torch.manual_seed(30) + packed, table, locations, source_keys, source_scales = _inputs(max_pools) + batch = reuse_group * 2 + 1 + shared = [1] * reuse_group + [2 + i % 2 for i in range(reuse_group)] + [4] + mixed = [1 + i % 3 for i in range(batch)] + req = torch.tensor(shared, device="cuda", dtype=torch.int32) + lengths = torch.tensor([max(0, max_pools - i * 63) for i in range(batch)], device="cuda", dtype=torch.int32) + gather_paged_pools(packed, table, req, lengths, max_pools, reuse_group=reuse_group) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + pages, block_table = gather_paged_pools(packed, table, req, lengths, max_pools, reuse_group=reuse_group) + identity = torch.arange(block_table.numel(), device="cuda", dtype=torch.int32).view_as(block_table) + scenarios = [ + (shared, [max(0, max_pools - i * 63) for i in range(batch)]), + (mixed, [max_pools if i % 2 else min(65, max_pools) for i in range(batch)]), + (shared, [min(i * 17, max_pools) for i in range(batch)]), + ([0] * batch, [0] * batch), + ([5] * reuse_group + mixed[reuse_group:], [max_pools] * batch), + ] + for slots, counts in scenarios: + req.copy_(torch.tensor(slots, device="cuda", dtype=torch.int32)) + lengths.copy_(torch.tensor(counts, device="cuda", dtype=torch.int32)) + # Recycle physical KV locations without changing any captured pointer. + locations.copy_(locations.roll(1, dims=1)) + table[:, 3::4] = locations + original_lengths = lengths.clone() + if not any(counts): + pages.fill_(127) + graph.replay() + torch.testing.assert_close(lengths, original_lengths, rtol=0, atol=0) + _check_paged(pages, block_table, req, lengths, locations, source_keys, source_scales) + if max(counts) < 4096 or reuse_group == 1: + torch.testing.assert_close(block_table, identity, rtol=0, atol=0) + elif slots == shared and reuse_group > 1: + valid_pages = triton.cdiv(max(counts[:reuse_group]), 64) + for row in range(1, reuse_group): + torch.testing.assert_close(block_table[row, :valid_pages], block_table[0, :valid_pages]) + if not any(counts): + dummy = pages.view(pages.shape[0], -1)[block_table[:, 0].long()] + assert not dummy[:, : 64 * 128].any() + assert (dummy[:, 64 * 128 :].contiguous().view(torch.float32) == 1).all() + + +@pytest.mark.parametrize("batch", [1, 3, 5]) +def test_paged_partial_groups_keep_valid_per_query_pages(batch): + packed, table, locations, source_keys, source_scales = _inputs(4097) + req = torch.ones(batch, device="cuda", dtype=torch.int32) + lengths = torch.tensor([4097 - i for i in range(batch)], device="cuda", dtype=torch.int32) + pages, block_table = gather_paged_pools(packed, table, req, lengths, 4097, reuse_group=4) + _check_paged(pages, block_table, req, lengths, locations, source_keys, source_scales) + if batch < 4: + expected = torch.arange(block_table.numel(), device="cuda", dtype=torch.int32).view_as(block_table) + torch.testing.assert_close(block_table, expected, rtol=0, atol=0) + + +def test_paged_batch_sizes_reuse_compiled_kernel(): + packed, table, locations, source_keys, source_scales = _inputs(4097) + req = torch.ones(16, device="cuda", dtype=torch.int32) + lengths = torch.arange(4097, 4081, -1, device="cuda", dtype=torch.int32) + kernel = kpool._gather_paged_pools + cache = kernel.device_caches[torch.cuda.current_device()][0] + # Warm the shared-group specialization, then vary both complete and partial + # groups. Grid dimensions must not create additional compiled variants. + gather_paged_pools(packed, table, req[:4], lengths[:4], 4097) + variants = len(cache) + for batch in (5, 8, 12, 16): + pages, block_table = gather_paged_pools(packed, table, req[:batch], lengths[:batch], 4097) + _check_paged(pages, block_table, req[:batch], lengths[:batch], locations, source_keys, source_scales) + assert len(cache) == variants + + +@pytest.mark.parametrize("reuse_group", [0, 3, 16]) +def test_reject_unsupported_reuse_group_before_allocating(reuse_group): + with pytest.raises(ValueError, match="reuse_group"): + gather_paged_pools(None, None, None, None, 4096, reuse_group=reuse_group) diff --git a/unit_tests/models/glm5_next/test_layers.py b/unit_tests/models/glm5_next/test_layers.py new file mode 100644 index 0000000000..4af14b0149 --- /dev/null +++ b/unit_tests/models/glm5_next/test_layers.py @@ -0,0 +1,451 @@ +import weakref +import dataclasses +import json +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F + +from lightllm.common.basemodel.infer_struct import InferStateInfo +from lightllm.common.basemodel.attention.nsa.glm5_next import Glm5NextSparsePrefillState, Glm5NextSparseDecodeState +from lightllm.common.quantization import Quantcfg +from lightllm.models.glm5_next.layer_infer.transformer_layer_infer import Glm5NextTransformerLayerInfer +from lightllm.models.glm5_next.layer_weights.transformer_layer_weight import Glm5NextTransformerLayerWeight +from lightllm.models.glm5_next.model import Glm5NextTpPartModel +from lightllm.models.glm5_next_mtp.model import Glm5NextMTPModel +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.utils.envs_utils import get_env_start_args + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.fixture(autouse=True) +def runtime(monkeypatch): + monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(dataclasses.asdict(StartArgs()))) + for name, value in { + "GLOBAL_RANK": 0, + "GLOBAL_WORLD_SIZE": 1, + "DP_WORLD_SIZE": 1, + "CURRENT_RANK_IN_DP": 0, + "CURRENT_RANK_IN_NODE": 0, + "CURRENT_DEVICE_ID": 0, + }.items(): + monkeypatch.setenv("LIGHTLLM_" + name, str(value)) + get_env_start_args.cache_clear() + torch.manual_seed(53) + yield + get_env_start_args.cache_clear() + + +@pytest.fixture +def config(): + return { + "hidden_size": 128, + "n_embed": 128, + "intermediate_size": 128, + "moe_intermediate_size": 128, + "num_hidden_layers": 4, + "n_layer": 4, + "num_attention_heads": 4, + "num_key_value_heads": 4, + "vocab_size": 256, + "rms_norm_eps": 1e-5, + "q_lora_rank": 128, + "kv_lora_rank": 512, + "qk_nope_head_dim": 128, + "qk_rope_head_dim": 0, + "v_head_dim": 128, + "index_n_heads": 4, + "index_head_dim": 128, + "index_topk": 2048, + "index_kpool": 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, + "n_shared_experts": 1, + "first_k_dense_replace": 1, + "num_experts_per_tok": 2, + "norm_topk_prob": True, + "scoring_func": "sigmoid", + "n_group": 1, + "topk_group": 1, + "swiglu_limit": 10.0, + "num_nextn_predict_layers": 1, + } + + +def random_weight(*shape): + return torch.randn(*shape, device="cuda", dtype=torch.bfloat16) * 0.05 + + +def rmsnorm(x, eps): + return (x.float() * torch.rsqrt(x.float().square().mean(-1, keepdim=True) + eps)).to(x.dtype) + + +def test_kda_projections_need_no_sparse_attention_config(config): + for key in ("q_lora_rank", "kv_lora_rank", "qk_nope_head_dim", "qk_rope_head_dim", "v_head_dim"): + del config[key] + for key in list(config): + if key.startswith("index_"): + del config[key] + layer = Glm5NextTransformerLayerInfer(0, config) + weight = Glm5NextTransformerLayerWeight(0, torch.bfloat16, config, Quantcfg(config)) + projections = { + name: random_weight(size, 128) + for name, size in {"q": 512, "k": 512, "v": 512, "b": 4, "f_a": 128, "g_a": 128, "f_b": 512, "g_b": 512}.items() + } + weight.load_hf_weights( + {f"model.language_model.layers.0.self_attn.{name}_proj.weight": value for name, value in projections.items()} + ) + hidden = random_weight(3, 128) + actual = layer._kda_projections(hidden, SimpleNamespace(), weight) + expected = ( + torch.cat([F.linear(hidden, projections[name]) for name in ("q", "k", "v")], dim=-1), + F.linear(F.linear(hidden, projections["f_a"]), projections["f_b"]), + F.linear(hidden, projections["b"]), + F.linear(F.linear(hidden, projections["g_a"]), projections["g_b"]), + ) + for result, reference in zip(actual, expected): + torch.testing.assert_close(result, reference, atol=1e-3, rtol=0.01) + + +@pytest.mark.parametrize("rank", [0, 1]) +@pytest.mark.parametrize("is_prefill", [False, True]) +def test_mla_loading_preserves_replicated_qkv_and_sharded_bmm(config, monkeypatch, rank, is_prefill): + monkeypatch.setenv("LIGHTLLM_DP_WORLD_SIZE", "2") + monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_DP", str(rank)) + layer = Glm5NextTransformerLayerInfer(3, config) + weight = Glm5NextTransformerLayerWeight(3, torch.bfloat16, config, Quantcfg(config)) + q_a, kv_a = random_weight(128, 128), random_weight(512, 128) + q_b, kv_b, o = random_weight(512, 128), random_weight(1024, 512), random_weight(128, 512) + tensors = { + "q_a_proj.weight": q_a, + "kv_a_proj_with_mqa.weight": kv_a, + "q_b_proj.weight": q_b, + "kv_b_proj.weight": kv_b, + "o_proj.weight": o, + "q_a_layernorm.weight": torch.ones(128, device="cuda", dtype=torch.bfloat16), + "kv_a_layernorm.weight": torch.ones(512, device="cuda", dtype=torch.bfloat16), + } + weight.load_hf_weights( + {f"model.language_model.layers.3.self_attn.{name}": value for name, value in tensors.items()} + ) + hidden = random_weight(3, 128) + state = SimpleNamespace(is_prefill=is_prefill, need_dp_prefill_balance=False) + q, kv = layer._get_qkv(hidden, state, weight) + q_lora = rmsnorm(F.linear(hidden, q_a), config["rms_norm_eps"]) + expected_q = F.linear(q_lora, q_b).view(3, 4, 128)[:, rank * 2 : (rank + 1) * 2] + expected_kv = rmsnorm(F.linear(hidden, kv_a), config["rms_norm_eps"]).unsqueeze(1) + torch.testing.assert_close(q, expected_q, atol=0.015, rtol=0.015) + torch.testing.assert_close(kv, expected_kv, atol=0.015, rtol=0.015) + head_weights = kv_b.view(4, 256, 512)[rank * 2 : (rank + 1) * 2] + projected = weight.k_b_proj_.bmm(q.transpose(0, 1)) + torch.testing.assert_close(projected, torch.bmm(q.transpose(0, 1), head_weights[:, :128])) + latent_output = random_weight(3, 2, 512) + monkeypatch.setattr(layer, "_tpsp_reduce", lambda input, infer_state: input) + actual = layer._get_o(latent_output, state, weight) + values = torch.bmm(latent_output.transpose(0, 1), head_weights[:, 128:].transpose(1, 2)) + expected = F.linear(values.transpose(0, 1).reshape(3, 256), o[:, rank * 256 : (rank + 1) * 256]) + torch.testing.assert_close(actual, expected) + + +def test_sparse_layer_prefill_and_decode_match_dense_nope_attention(config): + if torch.cuda.get_device_capability()[0] != 9: + pytest.skip("GLM sparse attention kernels require Hopper") + layer = Glm5NextTransformerLayerInfer(3, config) + weight = Glm5NextTransformerLayerWeight(3, torch.bfloat16, config, Quantcfg(config)) + kv_b = random_weight(1024, 512) + weight.load_hf_weights({"model.layers.3.self_attn.kv_b_proj.weight": kv_b}) + q = random_weight(3, 4, 128) + # Use noncontiguous token slots and a padded sparse index table. + packed = random_weight(8, 1, 584) + kv = packed[:, :, :512] + slots = torch.tensor([7, 1, 5, 0, 3], device="cuda", dtype=torch.int32) + lengths = torch.tensor([2, 3, 5], device="cuda", dtype=torch.int32) + indices = torch.full((3, 128), -1, device="cuda", dtype=torch.int32) + for row, length in enumerate(lengths.tolist()): + indices[row, :length] = slots[:length] + layer.indexer = SimpleNamespace(_get_indices=lambda **kwargs: (indices, indices)) + state = InferStateInfo() + 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)) + ) + 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 + state.decode_att_state = Glm5NextSparseDecodeState( + _infer_state_weak_ref=weakref.ref(state), + nsa_cache_seqlens=lengths, + nsa_cu_seqlens_k_new=F.pad(lengths.cumsum(0, dtype=torch.int32), (1, 0)), + ) + projected = torch.bmm(q.transpose(0, 1), kv_b.view(4, 256, 512)[:, :128]).transpose(0, 1) + expected = [] + for row, length in enumerate(lengths.tolist()): + keys = kv[slots[:length].long(), 0].float() + scores = projected[row].float() @ keys.T * config["qk_nope_head_dim"] ** -0.5 + expected.append(scores.softmax(-1) @ keys) + expected = torch.stack(expected) + prefill = layer._context_attention_kernel(q, None, state, weight) + state.get_topk_indices_params = {"hidden_states": None, "q_lora": None} + decode = layer._token_attention_kernel(q, state, weight) + for output in (prefill, decode): + torch.testing.assert_close(output.float(), expected, atol=1e-3, rtol=0.015) + + +@pytest.mark.parametrize("mode", ["dense", "moe", "fused_moe"]) +def test_ffn_dispatch_preserves_clamp_and_shared_expert(config, monkeypatch, mode): + args = StartArgs(enable_fused_shared_experts=mode == "fused_moe") + monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(dataclasses.asdict(args))) + get_env_start_args.cache_clear() + index = 0 if mode == "dense" else 3 + layer = Glm5NextTransformerLayerInfer(index, config) + weight = Glm5NextTransformerLayerWeight(index, torch.bfloat16, config, Quantcfg(config)) + prefix = f"model.language_model.layers.{index}.mlp" + eye = torch.eye(128, device="cuda", dtype=torch.bfloat16) + tensors = {} + + def mlp_weights(name, gate_scale, up_scale): + tensors[f"{name}.gate_proj.weight"] = eye * gate_scale + tensors[f"{name}.up_proj.weight"] = eye * up_scale + tensors[f"{name}.down_proj.weight"] = eye + + if mode == "dense": + mlp_weights(prefix, 1.5, 0.75) + else: + mlp_weights(f"{prefix}.shared_experts", 1.5, 0.75) + for expert in range(4): + mlp_weights(f"{prefix}.experts.{expert}", 1 + expert / 4, 2 + expert / 8) + router = torch.zeros(4, 128, device="cuda", dtype=torch.float32) + router[:, :4] = torch.eye(4, device="cuda") * 0.2 + tensors[f"{prefix}.gate.weight"] = router + tensors[f"{prefix}.gate.e_score_correction_bias"] = torch.zeros(4, device="cuda") + weight.load_hf_weights(tensors) + x = torch.linspace(-25, 25, 3 * 128, device="cuda", dtype=torch.bfloat16).view(3, 128) + + def reference_mlp(gate_scale, up_scale): + gate = (x * gate_scale).float().clamp(max=10) + up = (x * up_scale).float().clamp(-10, 10) + return (F.silu(gate).bfloat16().float() * up).bfloat16() + + expected = reference_mlp(1.5, 0.75) + if mode != "dense": + scores = F.linear(x.float(), router).sigmoid() + probabilities, ids = scores.topk(2, dim=-1) + probabilities /= probabilities.sum(-1, keepdim=True) + routed = torch.zeros_like(x) + for expert in range(4): + coefficient = (probabilities * (ids == expert)).sum(-1, keepdim=True) + routed += (reference_mlp(1 + expert / 4, 2 + expert / 8).float() * coefficient).bfloat16() + expected += routed + actual = layer._ffn(x.clone(), SimpleNamespace(is_prefill=True), weight) + torch.testing.assert_close(actual, expected, atol=0.25, rtol=0.015) + + +@pytest.mark.parametrize("is_prefill", [False, True]) +def test_ep_ffn_uses_combined_output_and_replicated_shared_expert(config, monkeypatch, is_prefill): + args = StartArgs(tp=2, enable_ep_moe=True, enable_fused_shared_experts=True) + monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(dataclasses.asdict(args))) + monkeypatch.setenv("LIGHTLLM_DP_WORLD_SIZE", "2") + monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_DP", "1") + monkeypatch.setenv("LIGHTLLM_GLOBAL_WORLD_SIZE", "2") + monkeypatch.setenv("LIGHTLLM_GLOBAL_RANK", "1") + get_env_start_args.cache_clear() + layer = Glm5NextTransformerLayerInfer(3, config) + weight = Glm5NextTransformerLayerWeight(3, torch.bfloat16, config, Quantcfg(config)) + assert weight.num_fused_shared_experts == 0 + prefix = "model.language_model.layers.3.mlp.shared_experts" + eye = torch.eye(128, device="cuda", dtype=torch.bfloat16) + weight.load_hf_weights({f"{prefix}.{name}.weight": eye for name in ("gate_proj", "up_proj", "down_proj")}) + x = torch.linspace(-25, 25, 3 * 128, device="cuda", dtype=torch.bfloat16).view(3, 128) + original = x.clone() + monkeypatch.setattr(weight.moe_gate, "mm", lambda x: torch.zeros(3, 4, device="cuda")) + monkeypatch.setattr(layer, "_tpsp_reduce", lambda **kwargs: pytest.fail("EP output must not be TP-reduced again")) + + def combined_output(hidden_states, **kwargs): + assert kwargs["is_prefill"] is is_prefill + assert kwargs["alpha"] == 1.0 and kwargs["limit"] == 10.0 + assert kwargs["clamp_up_add_one"] is False + return torch.full_like(hidden_states, 17) + + monkeypatch.setattr(weight.experts, "experts", combined_output) + actual = layer._ffn(x, SimpleNamespace(is_prefill=is_prefill), weight) + shared = (F.silu(x.float().clamp(max=10)).bfloat16().float() * x.float().clamp(-10, 10)).bfloat16() + torch.testing.assert_close(actual, shared + 17, atol=0.25, rtol=0.01) + torch.testing.assert_close(x, original, atol=0, rtol=0) + + +def test_fp8_shared_expert_scales_and_bf16_kv_b_load_together(config, monkeypatch): + monkeypatch.setenv( + "LIGHTLLM_START_ARGS", json.dumps(dataclasses.asdict(StartArgs(enable_fused_shared_experts=True))) + ) + get_env_start_args.cache_clear() + config["quantization_config"] = {"quant_method": "fp8", "weight_block_size": [128, 128]} + weight = Glm5NextTransformerLayerWeight(3, torch.bfloat16, config, Quantcfg(config)) + prefix = "model.language_model.layers.3" + tensors = {} + for expert in range(5): + name = f"experts.{expert}" if expert < 4 else "shared_experts" + for projection in ("gate_proj", "up_proj", "down_proj"): + tensors[f"{prefix}.mlp.{name}.{projection}.weight"] = random_weight(128, 128).to(torch.float8_e4m3fn) + tensors[f"{prefix}.mlp.{name}.{projection}.weight_scale_inv"] = torch.ones(1, 1, device="cuda") + tensors[f"{prefix}.mlp.gate.e_score_correction_bias"] = torch.zeros(4, device="cuda") + kv_b = random_weight(1024, 512) + tensors[f"{prefix}.self_attn.kv_b_proj.weight"] = kv_b + # There is deliberately no kv_b scale: this matrix stays BF16 in the FP8 checkpoint. + weight.load_hf_weights(tensors) + assert weight.experts.verify_load() + assert weight.k_b_proj_.verify_load() and weight.v_b_proj_.verify_load() + q = random_weight(4, 2, 128) + actual = weight.k_b_proj_.bmm(q) + expected = torch.bmm(q, kv_b.view(4, 256, 512)[:, :128]) + torch.testing.assert_close(actual, expected) + + +@pytest.mark.parametrize("method", ["context_forward", "token_forward"]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_mtp_pre_layer_applies_main_norm_before_hidden_fusion(config, monkeypatch, method, dtype): + from lightllm.common.basemodel.layer_weights.meta_weights import RMSNormWeight + from lightllm.models.glm5_next_mtp.layer_infer.pre_layer_infer import Glm5NextMTPPreLayerInfer + from lightllm.models.qwen_vl.layer_infer.pre_layer_infer import LlamaMultimodalPreLayerInfer + + hidden_size, eps = config["hidden_size"], config["rms_norm_eps"] + norms = [RMSNormWeight(hidden_size, name, dtype) for name in ("main_norm", "enorm", "hnorm")] + for norm in norms: + norm.weight.copy_(torch.randn_like(norm.weight)) + embeddings = torch.randn(5, hidden_size, device="cuda", dtype=dtype) + hidden = torch.randn_like(embeddings) + projection = torch.randn(2 * hidden_size, hidden_size, device="cuda", dtype=dtype) * 0.1 + weight = SimpleNamespace( + main_norm_weight_=norms[0], + enorm_weight_=norms[1], + hnorm_weight_=norms[2], + eh_proj_weight_=SimpleNamespace(mm=lambda *, input_tensor: input_tensor @ projection), + ) + state = SimpleNamespace(mtp_draft_input_hiddens=hidden.clone()) + monkeypatch.setattr(LlamaMultimodalPreLayerInfer, method, lambda self, ids, state, weight: embeddings.clone()) + actual = getattr(Glm5NextMTPPreLayerInfer(config), method)(None, state, weight) + + def reference_norm(value, norm): + return F.rms_norm(value.float(), (hidden_size,), norm.weight.float(), eps).to(dtype) + + normalized_hidden = reference_norm(reference_norm(hidden, norms[0]), norms[2]) + expected = torch.cat((reference_norm(embeddings, norms[1]), normalized_hidden), dim=-1) @ projection + torch.testing.assert_close(actual, expected, atol=0.015 if dtype == torch.bfloat16 else 1e-5, rtol=0.015) + + +@pytest.mark.parametrize("model_class", [Glm5NextTpPartModel, Glm5NextMTPModel]) +def test_glm_weight_loading_requires_hf(config, model_class): + model = object.__new__(model_class) + model.config, model.tp_world_size_ = config, 1 + model.args, model.load_way = StartArgs(), "HF" + model._verify_params() + model.load_way = "DS" + with pytest.raises(AssertionError, match="only support HF format weights"): + model._verify_params() + + +@pytest.mark.parametrize( + "kv_type, data_type, kv_lora_rank", + [ + ("fp8kv_sph", torch.bfloat16, 512), + ("fp8kv_spt", torch.bfloat16, 512), + ("fp8kv_dsa", torch.float16, 512), + ("fp8kv_dsa", torch.bfloat16, 256), + ], +) +def test_glm_rejects_unsupported_kv_config(config, kv_type, data_type, kv_lora_rank): + model = object.__new__(Glm5NextTpPartModel) + model.config = dict(config, kv_lora_rank=kv_lora_rank) + model.tp_world_size_, model.data_type = 1, data_type + model.args, model.load_way = StartArgs(llm_kv_type=kv_type), "HF" + with pytest.raises(ValueError, match="GLM-5.3 Flash"): + model._verify_params() + + +@pytest.mark.parametrize("tp_world_size", [1, 2]) +def test_glm_rejects_tpsp(config, tp_world_size): + model = object.__new__(Glm5NextTpPartModel) + model.config, model.tp_world_size_ = config, tp_world_size + model.args, model.load_way = StartArgs(), "HF" + model._verify_params() + + model.args.enable_tpsp_mix_mode = True + with pytest.raises(AssertionError, match="GLM-5.3 Flash does not support TP/SP mixed mode"): + model._verify_params() + + +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(data_type="bfloat16", mtp_mode="vanilla_with_att", mtp_step=2), "HF" + main._verify_params() + main._init_some_value() + 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 = 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.tp_rank_, draft.data_type = 1, 0, torch.bfloat16 + draft.args = main.args + draft.load_way = "HF" + draft._init_config() + draft._verify_params() + draft.quant_cfg = Quantcfg(draft.config) + draft._init_weights() + draft._init_infer_layer() + draft._init_some_value() + 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"] + assert draft.layers_num == 1 and draft.head_dim_ == 512 + assert draft.layers_infer[0].layer_num_ == 4 + step + assert draft.trans_layers_weight[0].layer_num_ == 4 + 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 + drafts.append(draft) + + 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() + assert main.att_state_manager.main_kda_conv_bf16_ssm_fp32_state.conv_state.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 + draft._init_att_backend1() + 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.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}_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): + state = model.infer_state_class() + state.is_prefill = True + state.input_ids = torch.zeros(3, device="cuda", dtype=torch.int64) + state.b_ready_cache_len = torch.tensor([2], device="cuda", dtype=torch.int32) + state.b_seq_len = torch.tensor([5], device="cuda", dtype=torch.int32) + state.init_some_extra_state(model) + torch.testing.assert_close( + state.position_ids, torch.arange(2, 5, device="cuda", dtype=state.position_ids.dtype) + ) diff --git a/unit_tests/models/glm5_next/test_mtp.py b/unit_tests/models/glm5_next/test_mtp.py new file mode 100644 index 0000000000..30c4fc22c4 --- /dev/null +++ b/unit_tests/models/glm5_next/test_mtp.py @@ -0,0 +1,308 @@ +import weakref +import pytest +import torch + +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.common.basemodel.triton_kernel.linear_att.replayssm_fold import fold_prefix +from lightllm.models.glm5_next.triton_kernel.kpool import compress_pools + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.mark.parametrize("mtp_step", [1, 2, 3, 5]) +@pytest.mark.parametrize("state_dtype", [torch.float32, torch.bfloat16]) +def test_kda_mtp_matches_single_token_decode_and_restores_accepted_state(mtp_step, state_dtype): + torch.manual_seed(53) + width, heads, dim = mtp_step + 1, 2, 128 + # Nonconsecutive requests plus a zero-length graph-padding sequence. + req_ids = torch.tensor([3, 1, 5], device="cuda", dtype=torch.int32) + slots = req_ids[:, None] * width + torch.arange(width, device="cuda", dtype=torch.int32) + states = torch.randn(6 * width, heads, dim, dim, device="cuda", dtype=state_dtype) * 0.1 + accepted = torch.tensor([width, 1, 1], device="cuda", dtype=torch.int32) + for lengths in ([width, 1, 0], [2, width, 0], [1, 1, 0]): + cu = torch.tensor( + [0, lengths[0], sum(lengths), sum(lengths)], + device="cuda", + dtype=torch.int32, + ) + tokens = sum(lengths) + q, k, v, gate = [torch.randn(1, tokens, heads, dim, device="cuda", dtype=torch.bfloat16) for _ in range(4)] + beta = torch.randn(1, tokens, heads, device="cuda", dtype=torch.bfloat16) + a = torch.randn(heads, device="cuda") + bias = torch.randn(heads * dim, device="cuda") + reference = states.clone() + before = states.clone() + expected = torch.empty_like(v) + for request, length in enumerate(lengths): + if not length: + continue + start = sum(lengths[:request]) + current = reference[slots[request, accepted[request] - 1]].clone().unsqueeze(0) + for offset in range(length): + token = start + offset + out, _ = fused_recurrent_kda( + q[:, token : token + 1], + k[:, token : token + 1], + v[:, token : token + 1], + gate[:, token : token + 1].reshape(1, 1, -1), + beta[:, token : token + 1], + a, + bias, + current, + torch.zeros(1, device="cuda", dtype=torch.int32), + ) + expected[:, token : token + 1] = out + reference[slots[request, offset]] = current[0] + actual, _ = fused_recurrent_kda( + q, + k, + v, + gate.reshape(1, tokens, -1), + beta, + a, + bias, + states, + slots, + cu_seqlens=cu, + num_accepted_tokens=accepted, + ) + torch.testing.assert_close(actual, expected, atol=2e-3, rtol=1e-2) + torch.testing.assert_close(states, reference, atol=2e-3, rtol=1e-2) + assert torch.equal(states[slots[2]], before[slots[2]]) + # The next verify must begin at the accepted candidate, not the last written slot. + accepted = torch.tensor([min(2, lengths[0]), 1, 1], device="cuda", dtype=torch.int32) + + +@pytest.mark.parametrize("mtp_step", [1, 2, 3, 5]) +@pytest.mark.parametrize("prefix", [1, 2, 3, 4, 7, 15]) +def test_kpool_verify_and_draft_rewind_match_full_prefill(mtp_step, prefix): + torch.manual_seed(53) + width, capacity = mtp_step + 1, 128 + raw = torch.randn(capacity, 256, device="cuda", dtype=torch.bfloat16) + ape = torch.randn(4, 128, device="cuda", dtype=torch.bfloat16) + ring = torch.zeros(2, 4 + mtp_step, 256, device="cuda", dtype=torch.bfloat16) + packed = torch.zeros(capacity, 132, device="cuda", dtype=torch.uint8) + ragged = torch.arange(capacity, device="cuda", dtype=torch.int32) + + def prefill(values, destination, tail): + count = values.shape[0] + compress_pools( + values, + tail, + destination, + ape, + torch.arange(1, count + 1, device="cuda", dtype=torch.int32), + ragged[:count], + torch.zeros(1, device="cuda", dtype=torch.int32), + torch.tensor([0, count], device="cuda", dtype=torch.int32), + torch.tensor([count], device="cuda", dtype=torch.int32), + count, + ) + + prefill(raw[:prefix], packed, ring) + # Alternate expanded verification, rewound single-token drafting, and re-verification. + 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[start : start + count], + ring, + packed, + ape, + lengths, + ragged[start : start + count], + torch.zeros(count, device="cuda", dtype=torch.int32), + torch.arange(count + 1, device="cuda", dtype=torch.int32), + lengths, + 1, + b_mtp_index=torch.arange(count, device="cuda", dtype=torch.int32), + ) + expected = torch.zeros_like(packed) + prefill( + raw[: start + count], + expected, + torch.zeros(2, 4, 256, device="cuda", dtype=torch.bfloat16), + ) + closing = torch.arange(start, start + count, device="cuda") + closing = closing[(closing + 1) % 4 == 0] + assert torch.equal(packed[closing], expected[closing]) + + +@pytest.mark.parametrize("dynamic", [False, True]) +@pytest.mark.parametrize("role", ["main", "draft0"]) +@pytest.mark.parametrize("mode,mtp_step", [("native", 2), ("replay", 2), ("native", 0), ("replay", 0)]) +def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str, mode, mtp_step): + from types import SimpleNamespace + + from lightllm.common.basemodel.attention.base_att import AttControl + from lightllm.common.basemodel.attention.linear.kda import KDADecodeAttState + from lightllm.common.basemodel.attention.linear.replay_kda import ( + ReplayKDADecodeAttState, + ) + from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ( + ReplayKDAState, + ) + from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import ( + causal_conv1d_update, + ) + + torch.manual_seed(53) + heads, dim, width = 2, 128, mtp_step + 1 + hidden = heads * dim + lengths = [2, 1, 3] if dynamic and mtp_step else [width] * 3 + requests = [2, 0, 1] + request_rows = [req for req, length in zip(requests, lengths) for _ in range(length)] + offsets = [offset for length in lengths for offset in range(length)] + if dynamic: + request_rows += [3, 3] + offsets += [0, 0] + tokens = len(request_rows) + conv = torch.randn(4, 3 * hidden, 3 + mtp_step, device="cuda", dtype=torch.bfloat16) + ssm = torch.randn(4 * width, heads, dim, dim, device="cuda") * 0.01 + old_conv, expected_ssm = conv.clone(), ssm.clone() + accepted_offsets = torch.tensor([1, 2, 0, 0] if mtp_step else [0, 0, 0, 0], dtype=torch.int32, device="cuda") + mixed = torch.randn(tokens, 3 * hidden, device="cuda", dtype=torch.bfloat16) + gate = torch.randn(tokens, hidden, device="cuda", dtype=torch.bfloat16) + beta = torch.randn(tokens, heads, device="cuda", dtype=torch.bfloat16) + conv_weight = torch.randn(3 * hidden, 4, device="cuda", dtype=torch.bfloat16) * 0.1 + a, bias = torch.randn(heads, device="cuda"), torch.randn(hidden, device="cuda") + expected = [] + row = 0 + zero = torch.zeros(1, dtype=torch.int32, device="cuda") + for req, length in zip(requests, lengths): + offset = int(accepted_offsets[req]) + history = old_conv[req : req + 1, :, offset : offset + 3].contiguous() + state = expected_ssm[req * width + offset].clone().unsqueeze(0) + for token in range(length): + convolved = causal_conv1d_update( + mixed[row : row + 1].clone(), + history, + conv_weight, + bias=None, + activation="silu", + conv_state_indices=zero, + ) + q, k, v = [part.reshape(1, 1, heads, dim) for part in convolved.split(hidden, -1)] + output, _ = fused_recurrent_kda( + q, + k, + v, + gate[row : row + 1].reshape(1, 1, hidden), + beta[row : row + 1].reshape(1, 1, heads), + a, + bias, + state, + zero, + ) + expected.append(output.reshape(1, heads, dim)) + expected_ssm[req * width + token] = state[0] + row += 1 + backend = SimpleNamespace( + mtp_step=mtp_step, + tp_num_heads=heads, + head_dim=dim, + tp_hidden_size=hidden, + lower_bound=-5.0, + uses_dynamic_spec_verify_layout=lambda: dynamic, + split_qkv=lambda x: x.split(hidden, -1), + ) + 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, + ) + selected = None + for model_role in ("main", "draft0"): + state_cls = ReplayKDAState if mode == "replay" else KDAState + linear = state_cls(config, request_capacity=4, mtp_step=mtp_step, device="cuda") + if mtp_step: + linear.mtp_state_index.copy_(accepted_offsets if model_role == role else torch.zeros_like(accepted_offsets)) + if model_role == role: + linear.conv_state[0].copy_(conv) + if mode == "native": + linear.ssm_state[0].copy_(ssm) + else: + indexes = torch.arange(4, device="cuda") * width + accepted_offsets + linear.ssm_state[0].copy_(ssm[indexes]) + selected = linear + 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_state_manager = resources + infer.att_state_role = role + weight = SimpleNamespace( + get_merged_kda_conv_weight=lambda: conv_weight, + linear_A_log=SimpleNamespace(weight=a), + linear_dt_bias=SimpleNamespace(weight=bias), + ) + decode_cls = ReplayKDADecodeAttState if mode == "replay" else KDADecodeAttState + state = decode_cls(backend=backend, _infer_state_weak_ref=weakref.ref(infer)) + state.init_state() + actual = state.decode_att( + None, + None, + None, + AttControl( + linear_att_decode=True, + linear_att_decode_dict=dict( + layer_weight=weight, + layer_num=0, + mixed_qkv=mixed, + raw_gate=gate, + raw_beta=beta, + ), + ), + ) + torch.testing.assert_close( + actual.reshape(tokens, heads, dim)[:row], + torch.cat(expected), + atol=2e-3, + rtol=1e-2, + ) + if mode == "native": + checked_slots = len(requests) if dynamic and not mtp_step else expected_ssm.shape[0] + torch.testing.assert_close( + selected.ssm_state[0, :checked_slots], + expected_ssm[:checked_slots], + atol=2e-3, + rtol=1e-2, + ) + else: + for req, length in zip(requests, lengths): + saved = selected.ssm_state + clone = saved.clone() + selected.ssm_state = clone + try: + fold_prefix(selected, torch.tensor([req], dtype=torch.int32, device="cuda")) + finally: + selected.ssm_state = saved + torch.testing.assert_close( + clone[0, req], + expected_ssm[req * width + length - 1], + atol=2e-3, + rtol=1e-2, + ) + if dynamic and mtp_step: + assert torch.equal(selected.conv_state[0, 3], old_conv[3]) diff --git a/unit_tests/models/glm5_next/test_mtp_prefill_cache.py b/unit_tests/models/glm5_next/test_mtp_prefill_cache.py new file mode 100644 index 0000000000..5b60fece98 --- /dev/null +++ b/unit_tests/models/glm5_next/test_mtp_prefill_cache.py @@ -0,0 +1,278 @@ +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch + +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.bounded_att import WindowCompressorState, WindowCompressorStateConfig +from lightllm.common.attention_state.full_att.packed_mla.fp8_e8m0_packed_mla_state import ( + FP8E8M0PackedMLAConfig, + FP8E8M0PackedMLAState, +) +from lightllm.common.attention_state.full_att.packed_mla.packed_mla_state import ( + PackedMLAAttentionConfig, + PackedMLAFullAttentionState, + PackedMLALogicalConfig, +) +from lightllm.common.basemodel.batch_objs import ModelInput, PostLayerOutput +from lightllm.common.basemodel.hidden_collector import FinalHiddenCollector +from lightllm.common.basemodel.infer_struct import InferStateInfo +from lightllm.common.basemodel.triton_kernel.norm.rmsnorm import rmsnorm_forward +from lightllm.models.glm5_next.indexer import Glm5NextNsaInfer +from lightllm.models.glm5_next_mtp.layer_infer.post_layer_infer import Glm5NextMTPPostLayerInfer +from lightllm.models.glm5_next_mtp.layer_infer.transformer_layer_infer import Glm5NextMTPTransformerLayerInfer +from lightllm.models.glm5_next_mtp.model import Glm5NextMTPModel +from lightllm.models.llama.layer_infer.post_layer_infer import LlamaPostLayerInfer +from lightllm.server.core.objs.start_args_type import StartArgs + + +def make_model(): + model = Glm5NextMTPModel.__new__(Glm5NextMTPModel) + model.args = StartArgs() + model.args.mtp_mode = "eagle_with_att" + model.args.mtp_step = 1 + model.args.llm_kv_type = "fp8kv_dsa" + model.tp_world_size_ = model.layers_num = 1 + model.prefill_graph = None + model.layers_infer = [SimpleNamespace(is_linear_attention_layer=False, use_mhc=False)] + model.trans_layers_weight = [object()] + return model + + +@pytest.mark.parametrize( + "field,value", + [ + ("mtp_mode", "eagle_no_att"), + ], +) +def test_unsupported_execution_uses_full_forward(field, value): + model = make_model() + assert model._supports_prefill_kv_only() + setattr(model.args, field, value) + assert not model._supports_prefill_kv_only() + + +@pytest.mark.parametrize("field,value", [("mtp_mode", "eagle_no_att")]) +def test_unsupported_layout_uses_full_forward(field, value): + model = make_model() + setattr(model.args, field, value) + assert not model._supports_prefill_kv_only() + + +def test_eagle_multistep_and_vanilla_multistep_capabilities(): + model = make_model() + model.args.mtp_step = 3 + assert model._supports_prefill_kv_only() + model.args.mtp_mode = "vanilla_with_att" + assert not model._supports_prefill_kv_only() + + +@pytest.mark.parametrize("kv_type", ["fp8kv_dsa", "bf16"]) +@pytest.mark.parametrize("tp_world_size", [1, 4]) +def test_cache_only_capability_does_not_depend_on_kv_format_or_tp(kv_type, tp_world_size): + model = make_model() + model.args.llm_kv_type = kv_type + model.tp_world_size_ = tp_world_size + assert model._supports_prefill_kv_only() + + +@pytest.mark.parametrize("tokens", [0, 1, 8192]) +@pytest.mark.parametrize("mode", ["eagle_with_att", "vanilla_with_att"]) +def test_forward_uses_common_prefill(monkeypatch, tokens, mode): + from lightllm.common.basemodel import basemodel + + model = make_model() + model.args.mtp_mode = mode + model.mtp_previous_draft_models = [] + model.req_manager = SimpleNamespace(HOLD_REQUEST_ID=3) + model.att_state_manager = object() + model.hidden_collector_prototype = FinalHiddenCollector() + model.return_all_prompt_logics = False + model.prefill_att_backend = model.prefill_att_backend1 = None + model._select_mem_indexes = lambda model_input: torch.zeros_like(model_input.input_ids) + batch = int(tokens > 0) + model_input = ModelInput( + is_prefill=True, + batch_size=batch, + total_token_num=tokens, + max_q_seq_len=tokens, + max_kv_seq_len=tokens, + max_cache_len=0, + input_ids=torch.zeros(tokens, dtype=torch.int64), + b_req_idx=torch.zeros(batch, dtype=torch.int32), + b_mtp_index=torch.zeros(batch, dtype=torch.int32), + b_seq_len=torch.full((batch,), tokens, dtype=torch.int32), + b_ready_cache_len=torch.zeros(batch, dtype=torch.int32), + b_prefill_start_loc=torch.zeros(batch, dtype=torch.int32), + b_is_decode_req=torch.zeros(batch, dtype=torch.bool), + b_prefill_has_output_cpu=[True] * batch, + multimodal_params=[{"images": [], "audios": []} for _ in range(batch)], + mtp_draft_input_hiddens=torch.ones(tokens, 2), + ) + model.pre_infer = SimpleNamespace( + context_forward=Mock(side_effect=lambda ids, *_: torch.zeros(ids.shape[0], 2)), + _tpsp_sp_split=lambda input, **_: input, + ) + model.post_infer = Glm5NextMTPPostLayerInfer.__new__(Glm5NextMTPPostLayerInfer) + model.post_infer.tp_world_size_ = 1 + model.post_infer._tpsp_allgather = Mock(side_effect=lambda input, **_: input) + post_forward = Mock(side_effect=lambda hidden, state, _: PostLayerOutput(torch.ones(state.batch_size, 3))) + monkeypatch.setattr(LlamaPostLayerInfer, "token_forward", post_forward) + model.pre_post_weight = object() + layer = Glm5NextMTPTransformerLayerInfer.__new__(Glm5NextMTPTransformerLayerInfer) + layer.is_linear_attention_layer = layer.use_mhc = False + layer.layer_num_, layer.num_hidden_layers = 45, 1 + layer.embed_dim_ = 2 + layer._att_norm = layer._ffn_norm = lambda hidden, *_: hidden + layer.context_attention_forward = Mock(side_effect=lambda hidden, *_: torch.zeros_like(hidden)) + layer._ffn = Mock(side_effect=lambda hidden, *_: torch.zeros_like(hidden)) + layer._project_sparse_qkv_a = lambda hidden, *_: (hidden, None, hidden) + layer._normalize_sparse_kv = lambda kv, _: kv + layer.indexer = SimpleNamespace(write_k_pools=Mock()) + writer = Mock() + layer._post_cache_kv = writer + model.layers_infer = [layer] + allocator = SimpleNamespace(cache_env_in=Mock(), cache_env_out=Mock()) + monkeypatch.setattr(basemodel, "g_cache_manager", allocator) + monkeypatch.setattr(basemodel.dist_group_manager, "get_group", Mock()) + monkeypatch.setattr(basemodel.dist_group_manager, "clear_deepep_buffer", Mock()) + monkeypatch.setattr(ModelInput, "to_cuda", lambda self: self.check_input()) + monkeypatch.setattr(InferStateInfo, "init_some_extra_state", Mock()) + event = Mock() + monkeypatch.setattr(torch.cuda, "Event", Mock(return_value=event)) + output = model.forward(model_input) + assert output.prefill_mem_indexes_ready_event is event + assert output.logits.numel() == 0 + assert output.mtp_collector.spec_hidden is None + ids, infer, _ = model.pre_infer.context_forward.call_args.args + assert ids.shape == (max(1, tokens),) + assert infer.b_req_idx.tolist() == ([0] if tokens else [model.req_manager.HOLD_REQUEST_ID]) + if mode == "eagle_with_att": + assert infer.mtp_draft_input_hiddens is None + writer.assert_called_once() + layer.context_attention_forward.assert_not_called() + layer._ffn.assert_not_called() + post_forward.assert_not_called() + placeholder = model.post_infer._tpsp_allgather.call_args.kwargs["input"] + assert placeholder.shape == (0, layer.embed_dim_) + assert placeholder.untyped_storage().nbytes() == 0 + layer.indexer.write_k_pools.assert_called_once() + else: + writer.assert_called_once() + layer.context_attention_forward.assert_not_called() + layer._ffn.assert_not_called() + post_forward.assert_not_called() + allocator.cache_env_in.assert_called_once() + allocator.cache_env_out.assert_called_once() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("kv_type", ["fp8kv_dsa", "bf16"]) +@pytest.mark.parametrize( + "requests", + [((3, 0, 1),), ((0, 0, 8192),), ((0, 131071, 129), (2, 1048573, 3)), ((1, 3, 1), (0, 5, 7))], +) +def test_cache_only_matches_full_prefill_cache_writes(monkeypatch, kv_type, requests): + from lightllm.models.glm5_next import indexer as indexer_module + + torch.manual_seed(53) + layer_id, tokens = 45, sum(q for _, _, q in requests) + layer = Glm5NextMTPTransformerLayerInfer.__new__(Glm5NextMTPTransformerLayerInfer) + layer.layer_num_, layer.num_hidden_layers = layer_id, 1 + layer.use_mhc = layer.is_linear_attention_layer = False + layer.tp_world_size_ = layer.tp_q_head_num_ = 1 + layer.embed_dim_ = layer.q_lora_rank = 64 + layer.kv_lora_rank = layer.qk_nope_head_dim = 512 + layer.eps_ = 1e-5 + layer.indexer = Glm5NextNsaInfer( + layer_id, dict(index_topk=2048, index_n_heads=1, index_head_dim=128, rms_norm_eps=1e-5, index_kpool=4), 1 + ) + + def projection(rows, cols): + weight = torch.randn(rows, cols, device="cuda", dtype=torch.bfloat16) * 0.1 + return SimpleNamespace(mm=Mock(side_effect=lambda x: x @ weight)) + + def norm(input, eps, **_): + return rmsnorm_forward(input, None, eps) + + weights = SimpleNamespace( + att_norm_weight_=norm, + qkv_a_proj_with_mqa_=projection(64, 576), + q_a_layernorm_=SimpleNamespace(weight=torch.ones(64, device="cuda", dtype=torch.bfloat16)), + kv_a_layernorm_=SimpleNamespace(weight=torch.ones(512, device="cuda", dtype=torch.bfloat16)), + q_b_proj_=projection(64, 512), + wk_proj_=projection(64, 128), + k_norm_=norm, + index_kpool_compress_gate=projection(64, 128), + index_kpool_compress_ape=SimpleNamespace(weight=torch.randn(4, 128, device="cuda", dtype=torch.bfloat16)), + ) + resources = AttentionStateManager(request_capacity=4) + stores = [] + for role in ("main", "draft0"): + if kv_type == "fp8kv_dsa": + full = FP8E8M0PackedMLAState(FP8E8M0PackedMLAConfig(logical_layer_groups=((layer_id,),))) + else: + full = PackedMLAFullAttentionState( + PackedMLAAttentionConfig( + PackedMLALogicalConfig( + logical_layer_groups=((layer_id,),), + latent_dim=512, + format="packed_bf16", + indexer_dim=128, + ) + ), + device="cuda", + ) + full.allocate_storage(2 * tokens + 2) + full.storage.fill_(0x5A) + tail = WindowCompressorState( + WindowCompressorStateConfig((layer_id,), 4, 4, 128, "indexer", torch.bfloat16), 4, mtp_step=1 + ) + tail.buffer.normal_() + resources.register_state(role, full) + resources.register_state(role, tail) + stores.extend((full.storage, tail.buffer)) + resources.finish_state_registration() + initial = [s.clone() for s in stores] + q_lengths = [q for _, _, q in requests] + lengths = torch.cat([torch.arange(past + 1, past + q + 1, device="cuda") for _, past, q in requests]).int() + infer = SimpleNamespace( + is_prefill=True, + is_draft_model=True, + prefill_kv_only=True, + need_dp_prefill_balance=False, + att_state_role="draft0", + att_state_manager=resources, + mem_index=2 * torch.arange(tokens, device="cuda", dtype=torch.int32) + 1, + b_req_idx=torch.tensor([r for r, _, _ in requests], device="cuda", dtype=torch.int32), + b_seq_len=torch.tensor([p + q for _, p, q in requests], device="cuda", dtype=torch.int32), + b1_cu_q_seq_len=torch.tensor([0, *q_lengths], device="cuda", dtype=torch.int32).cumsum(0).int(), + max_q_seq_len=max(q_lengths), + max_kv_seq_len=2048, + prefill_att_state=SimpleNamespace(q_kv_seq_len=lengths, q_req_idx=None), + req_manager=SimpleNamespace(req_to_token_indexs=None), + ) + # Top-k outputs are discarded here; run the normal projections and cache writers. + monkeypatch.setattr(indexer_module, "expand_topk", Mock(return_value=(None, None))) + hidden = torch.randn(tokens, 64, device="cuda", dtype=torch.bfloat16) + _, kv = layer._get_qkv(layer._att_norm(hidden, infer, weights), infer, weights) + layer._post_cache_kv(kv, infer, weights) + layer.indexer._get_indices( + **infer.get_topk_indices_params, infer_state=infer, att_state=infer.prefill_att_state, layer_weight=weights + ) + expected = [s.clone() for s in stores] + assert not torch.equal(initial[2], expected[2]) + for store, before in zip(stores, initial): + store.copy_(before) + weights.q_b_proj_.mm.reset_mock() + placeholder = layer.context_forward(hidden, infer, weights) + assert placeholder.shape == (0, layer.embed_dim_) + assert placeholder.untyped_storage().nbytes() == 0 + weights.q_b_proj_.mm.assert_not_called() + for actual, reference in zip(stores, expected): + assert torch.equal(actual, reference) + for actual, before in zip(stores[:2], initial[:2]): + assert torch.equal(actual, before) + assert torch.equal(stores[2][:, ::2], initial[2][:, ::2]) + assert torch.equal(stores[3][:, 3], initial[3][:, 3]) diff --git a/unit_tests/models/glm5_next/test_pd_cache.py b/unit_tests/models/glm5_next/test_pd_cache.py new file mode 100644 index 0000000000..15dcb4a499 --- /dev/null +++ b/unit_tests/models/glm5_next/test_pd_cache.py @@ -0,0 +1,447 @@ +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 +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.attention_state.allocators import FullAttKVAllocator +from unit_tests.common.att_resources.model_helpers import create_full_attention_roles + +from lightllm.common.attention_state.full_att import ( + PackedMLAFullAttentionState, + 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 +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"), + pytest.mark.usefixtures("kv_allocator_environment"), +] + + +@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 = request.param + args = StartArgs( + data_type="bfloat16", + pd_kv_page_size=2048, + linear_att_hash_page_size=4, + linear_att_page_block_num=2, + mtp_step=mtp_step, + mtp_mode="eagle_with_att" if mtp_step else None, + ) + set_unique_server_name(args) + get_env_start_args.cache_clear() + set_env_start_args(dataclasses.asdict(args)) + + def create(tp, target_mtp_step=None, mode="native"): + 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_head_dim=512, + global_linear_k_heads=8, + global_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, + full_attention_interval=4, + all_layer_num=8, + ) + 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_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, + ) + full.allocate_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, + ssm_state_mode=mode, + 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, + ) + mems.append( + SimpleNamespace( + 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_compressor_indexer_w4_s4_bf16_state, + tails=window_states(resources), + resources=resources, + ) + ) + return mems + + return create + + +@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]) +@pytest.mark.parametrize("source_mode", ["native", "replay"]) +@pytest.mark.parametrize("dest_mode", ["native", "replay"]) +def test_pd_kv_and_runtime_state_roundtrip( + make_mems, + prefill_tp, + decode_tp, + remainder, + decode_mtp_step, + source_mode, + dest_mode, +): + source = make_mems(prefill_tp, mode=source_mode) + dest = make_mems(decode_tp, target_mtp_step=decode_mtp_step, mode=dest_mode) + dst_ssm_stride = 1 if dest_mode == "replay" else dest[0].linear.mtp_step + 1 + mtp_size = get_env_start_args().mtp_step + 1 + 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() + dst_indexes = torch.randperm(32)[:length].tolist() + 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( + 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 + for state in mem.full_states: + 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)) + src_slot = src_req if source_mode == "replay" else src_req * (mem.linear.mtp_step + 1) + mem.linear.ssm_state[:, src_slot].copy_(ssm[:, heads]) + if source_mode == "replay": + backend = mem.linear + backend.raw_keys.zero_() + backend.raw_values.zero_() + backend.gates.fill_(-0.5) + backend.betas.zero_() + backend.cursors[src_req] = 1 + offset = 0 + for state in mem.tails: + end = offset + len(state.config.logical_layer_ids) + state.buffer[:, src_req].copy_(tail[offset:end]) + offset = end + + untouched = [] + original_tails = [] + for mem in dest: + for state in mem.full_states: + state.storage.zero_() + mem.linear.conv_state.normal_() + mem.linear.ssm_state.normal_() + if dest_mode == "replay": + mem.linear.cursors[dst_req] = 3 + if mem.linear.mtp_enabled: + mem.linear.mtp_state_index[dst_req] = decode_mtp_step + for state in mem.tails: + 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_ssm_stride), + *((state.buffer, 1) for state in mem.tails), + ) + ] + ) + + 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_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, + 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"), + task_kind=kind, + req_idx=src_req, + state_token_len=length, + write_to_pd_page=True, + ) + dest_buffer.copy_(source_buffer) + copy_page( + 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"), + task_kind=kind, + req_idx=dst_req, + state_token_len=length, + write_to_pd_page=False, + ) + torch.cuda.synchronize() + + for rank, mem in enumerate(dest): + 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 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), + ) + 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_ssm_stride], ssm[:, heads]) + 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]] + assert torch.equal(restored_tail, expected_tail) + if dest_mode == "replay": + assert mem.linear.cursors[dst_req] == 0 + if mem.linear.mtp_enabled: + assert mem.linear.mtp_state_index[dst_req] == 0 + for (state, stride), expected in zip( + ( + (mem.linear.conv_state, 1), + (mem.linear.ssm_state, dst_ssm_stride), + *((state.buffer, 1) for state in mem.tails), + ), + untouched[rank], + ): + indexes = torch.tensor([0, 1, 3], device="cuda") * stride + 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. + count = 4 - remainder + 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 = [] + for mem, req_idx, indexes in ((source[0], src_req, src_indexes), (dest[0], dst_req, dst_indexes)): + new_indexes = [i for i in range(32) if i not in indexes][:count] + compress_pools( + k_and_gate=raw, + indexer_tail=mem.tail.buffer[0], + indexer_k=mem.full.get_layer_view(3).indexer_k, + index_kpool_compress_ape=ape, + q_kv_seq_len=tensor(range(length + 1, length + count + 1)), + mem_index=tensor(new_indexes), + b_req_idx=tensor([req_idx]), + b1_cu_q_seq_len=tensor([0, count]), + b_seq_len=tensor([length + count]), + max_q_seq_len=count, + ) + 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, monkeypatch, caplog): + from lightllm.common.attention_state import pd_operator + + monkeypatch.setattr(get_env_start_args(), "pd_kv_page_size", 8) + caplog.set_level(30, logger=pd_operator.logger.name) + monkeypatch.setattr(pd_operator.logger, "handlers", [caplog.handler]) + 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() + ) + runtime_bytes = max( + 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_operator().get_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)]) +@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.bounded_att import WindowCompressorStateConfig, create_window_compressor_state + from lightllm.common.attention_state.types import StateTransferLayout + + torch.manual_seed(53) + config = WindowCompressorStateConfig( + logical_layer_ids=(0,), + window_size=4, + compress_stride=4, + candidate_dim=128, + target="indexer", + dtype=torch.bfloat16, + ) + source = create_window_compressor_state( + config=config, + request_capacity=2, + mtp_step=source_step, + device="cuda", + ) + target = create_window_compressor_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( + k_and_gate=values, + indexer_tail=tail, + indexer_k=destination, + index_kpool_compress_ape=ape, + q_kv_seq_len=torch.arange( + start=1, + end=count + 1, + device="cuda", + dtype=torch.int32, + ), + mem_index=ragged[:count], + b_req_idx=torch.zeros(1, device="cuda", dtype=torch.int32), + b1_cu_q_seq_len=torch.tensor([0, count], device="cuda", dtype=torch.int32), + b_seq_len=torch.tensor([count], device="cuda", dtype=torch.int32), + max_q_seq_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( + 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_page( + req_idx=0, + pd_page_buffers=pages, + pd_page_index=0, + state_token_len=prefix, + layout=layout, + ) + target.buffer.fill_(17) + target.read_from_pd_page( + pd_page_buffers=pages, + req_idx=0, + pd_page_index=0, + state_token_len=prefix, + layout=layout, + ) + 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)]: + 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( + k_and_gate=raw[start : start + count], + indexer_tail=target.get_layer_view(0), + indexer_k=packed, + index_kpool_compress_ape=ape, + q_kv_seq_len=lengths, + mem_index=ragged[start : start + count], + b_req_idx=torch.zeros(count, device="cuda", dtype=torch.int32), + b1_cu_q_seq_len=torch.arange(count + 1, device="cuda", dtype=torch.int32), + b_seq_len=lengths, + max_q_seq_len=1, + b_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/glm5_next/test_prefill_workspace.py b/unit_tests/models/glm5_next/test_prefill_workspace.py new file mode 100644 index 0000000000..0d3a85d965 --- /dev/null +++ b/unit_tests/models/glm5_next/test_prefill_workspace.py @@ -0,0 +1,116 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.basemodel.basemodel import TpPartBaseModel +from lightllm.common.basemodel.attention.nsa import fp8_glm5_next +from lightllm.models.glm5_next.model import Glm5NextTpPartModel +from lightllm.models.glm5_next_mtp.model import Glm5NextMTPModel +from unit_tests.models.glm5_next.test_fp8_kv import _PrefillCase, _make_workspace_backend, _capture_cuda_graph + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.fixture(autouse=True) +def current_cuda_device(monkeypatch): + monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", str(torch.cuda.current_device())) + + +@pytest.mark.parametrize("overlap", [False, True]) +@pytest.mark.parametrize("run_mode", ["normal", "decode"]) +def test_prefill_workspace_is_reserved_once_and_reused_at_runtime(monkeypatch, overlap, run_mode): + model = SimpleNamespace( + args=SimpleNamespace(enable_prefill_microbatch_overlap=overlap, run_mode=run_mode), + batch_max_tokens=8, + max_req_num=2, + max_seq_length=13, + config={"index_topk": 2048}, + att_state_role=f"workspace_test_{run_mode}_{overlap}", + ) + backend_class = fp8_glm5_next.Fp8Glm5NextSparseAttBackend + reserved = backend_class.allocate_prefill_workspaces(model) + pointers = [{key: value.data_ptr() for key, value in buffers.items()} for buffers in reserved] + backend = object.__new__(backend_class) + backend.__init__(model) + assert [ + {key: value.data_ptr() for key, value in buffers.items()} for buffers in backend._prefill_workspaces + ] == pointers + assert len(reserved) == 1 + int(overlap) + capacity = 8 if run_mode == "decode" else 34 + assert reserved[0]["gathered_kv"].shape == (capacity, 1, 512) + assert reserved[0]["chunk_topk_indices"].shape == (8, 2176) + + def forbidden(*args, **kwargs): + raise AssertionError("Runtime must reuse the workspace reserved before KV capacity planning") + + monkeypatch.setattr(backend, "allocate_prefill_workspaces", forbidden) + monkeypatch.setattr(backend, "get_gpu_workspace_buffer", forbidden) + for slot in range(1 + int(overlap)): + for q_tokens in (1, 8, 3): + topk = torch.empty((q_tokens, 2176), dtype=torch.int32, device="cuda") + views = backend.get_prefill_workspace(SimpleNamespace(microbatch_index=slot), topk) + assert {key: value.data_ptr() for key, value in views.items()} == pointers[slot] + torch.testing.assert_close( + views["cu_seqlens_q"], torch.arange(q_tokens + 1, dtype=torch.int32, device="cuda") + ) + + +def test_prefill_microbatches_use_independent_scratch_during_graph_replay(monkeypatch): + monkeypatch.setattr(fp8_glm5_next, "KV_GATHER_CHUNK_SIZE", 7) + backend = _make_workspace_backend(8, overlap=True) + first = _PrefillCase([15], [3], 8, backend=backend) + second = _PrefillCase([17], [3], 8, backend=backend) + second.infer_state.microbatch_index = 1 + first.check(first.run()) + second.check(second.run()) + pointers = [{key: value.data_ptr() for key, value in buffers.items()} for buffers in backend._prefill_workspaces] + assert all(pointers[0][key] != pointers[1][key] for key in pointers[0]) + + def forbidden(*args, **kwargs): + raise AssertionError("Attention must not allocate a new prefill workspace") + + monkeypatch.setattr(backend, "allocate_prefill_workspaces", forbidden) + first_graph, first_output = _capture_cuda_graph(first.run) + second_graph, second_output = _capture_cuda_graph(second.run) + first_stream, second_stream = torch.cuda.Stream(), torch.cuda.Stream() + first.mutate() + second.mutate() + for _ in range(3): + first_stream.wait_stream(torch.cuda.current_stream()) + second_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(first_stream): + first_graph.replay() + with torch.cuda.stream(second_stream): + second_graph.replay() + torch.cuda.current_stream().wait_stream(first_stream) + torch.cuda.current_stream().wait_stream(second_stream) + first.check(first_output) + second.check(second_output) + assert [ + {key: value.data_ptr() for key, value in buffers.items()} for buffers in backend._prefill_workspaces + ] == pointers + + +@pytest.mark.parametrize("model_class", [Glm5NextTpPartModel, Glm5NextMTPModel]) +@pytest.mark.parametrize("kv_type", ["None", "fp8kv_dsa"]) +def test_glm_reserves_prefill_workspace_during_resource_preparation(monkeypatch, model_class, kv_type): + model = object.__new__(model_class) + model.args = SimpleNamespace(llm_kv_type=kv_type) + model.mtp_previous_draft_models = [] + model.att_state_manager = SimpleNamespace(get_full_attention_states=lambda role: [object()]) + calls = [] + + def prepare(self): + assert self is model + calls.append("prepare_states") + + def reserve(cls, prepared_model): + assert prepared_model is model + calls.append("reserve_workspace") + + monkeypatch.setattr(TpPartBaseModel, "prepare_attention_states", prepare) + monkeypatch.setattr(fp8_glm5_next.Fp8Glm5NextSparseAttBackend, "allocate_prefill_workspaces", classmethod(reserve)) + model.prepare_attention_states() + assert calls == (["prepare_states", "reserve_workspace"] if kv_type == "fp8kv_dsa" else ["prepare_states"]) diff --git a/unit_tests/models/glm5_next/test_token_major_mla_bmm.py b/unit_tests/models/glm5_next/test_token_major_mla_bmm.py new file mode 100644 index 0000000000..fa9e727f87 --- /dev/null +++ b/unit_tests/models/glm5_next/test_token_major_mla_bmm.py @@ -0,0 +1,92 @@ +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch + +from lightllm.common.basemodel.layer_weights.meta_weights.mm_weight.mm_weight import BMMWeightTpl +from lightllm.models.glm5_next.layer_infer.transformer_layer_infer import Glm5NextTransformerLayerInfer + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +def make_projection(layer, heads: int, inner: int, cols: int, dtype: torch.dtype): + projection = BMMWeightTpl.__new__(BMMWeightTpl) + projection.weight = torch.randn(heads, inner, cols, device="cuda", dtype=dtype) * 0.1 + layer.tp_q_head_num_ = heads + if inner == 256: + layer.kv_lora_rank = cols + layer.layer_num_ = 0 + layer.softmax_scale = inner ** -0.5 + layer.indexer = SimpleNamespace(_get_indices=lambda **kwargs: (None, None)) + state = SimpleNamespace( + att_state_role="main", + att_state_manager=Mock(), + prefill_att_state=SimpleNamespace(prefill_att=lambda **kwargs: kwargs["q"]), + ) + weight = SimpleNamespace(k_b_proj_=projection) + + def run(x): + state.get_topk_indices_params = {"hidden_states": None, "q_lora": None} + return layer._context_attention_kernel(x, None, state, weight) + + else: + layer.v_head_dim = cols + layer._tpsp_reduce = lambda input, infer_state: input + state = SimpleNamespace(is_prefill=False, need_dp_prefill_balance=False) + weight = SimpleNamespace(v_b_proj_=projection, o_weight_=SimpleNamespace(mm=lambda x: x)) + + def run(x): + return layer._get_o(x, state, weight).view(x.shape[0], heads, cols) + + return projection, run + + +@pytest.mark.parametrize("rows", [1, 128, 1024, 1659, 4096, 8192]) +@pytest.mark.parametrize("heads", [16, 64]) +@pytest.mark.parametrize("inner,cols", [(256, 512), (512, 256)]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +def test_token_major_projection_matches_reference_and_keeps_guard_rows(rows, heads, inner, cols, dtype): + torch.manual_seed(272) + layer = Glm5NextTransformerLayerInfer.__new__(Glm5NextTransformerLayerInfer) + projection, run = make_projection(layer, heads, inner, cols, dtype) + x = torch.randn(rows, heads, inner, device="cuda", dtype=dtype) + expected = projection.bmm(x.transpose(0, 1)).transpose(0, 1).contiguous() + slab = torch.full((rows + 2, heads, cols), 113.0, device="cuda", dtype=dtype) + output = slab[1:-1] + layer.alloc_tensor = Mock(return_value=output) + actual = run(x) + assert actual.data_ptr() == output.data_ptr() and actual.is_contiguous() + assert torch.equal(actual, expected) + assert (slab[0] == 113).all() and (slab[-1] == 113).all() + + +@pytest.mark.parametrize("inner,cols", [(256, 512), (512, 256)]) +def test_token_major_projection_accepts_head_padding(inner, cols): + # TP4 FlashMLA crops a 64-head output to 16 heads without copying it. + layer = Glm5NextTransformerLayerInfer.__new__(Glm5NextTransformerLayerInfer) + projection, run = make_projection(layer, 16, inner, cols, torch.bfloat16) + x = torch.randn(128, 64, inner, device="cuda", dtype=torch.bfloat16)[:, :16] + assert not x.is_contiguous() + expected = projection.bmm(x.transpose(0, 1)).transpose(0, 1) + actual = run(x) + assert actual.is_contiguous() + assert torch.equal(actual, expected) + + +@pytest.mark.parametrize("inner,cols", [(256, 512), (512, 256)]) +def test_token_major_projection_replays_new_inputs_in_cuda_graph(inner, cols): + layer = Glm5NextTransformerLayerInfer.__new__(Glm5NextTransformerLayerInfer) + projection, run = make_projection(layer, 16, inner, cols, torch.bfloat16) + x = torch.randn(128, 16, inner, device="cuda", dtype=torch.bfloat16) + for _ in range(3): + run(x) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = run(x) + x.normal_() + graph.replay() + expected = projection.bmm(x.transpose(0, 1)).transpose(0, 1) + assert torch.equal(actual, expected) diff --git a/unit_tests/models/glm5_next/test_visual.py b/unit_tests/models/glm5_next/test_visual.py new file mode 100644 index 0000000000..e4e38812c4 --- /dev/null +++ b/unit_tests/models/glm5_next/test_visual.py @@ -0,0 +1,210 @@ +import dataclasses +import json +from types import SimpleNamespace + +import numpy as np +import pytest +import torch +from PIL import Image + +from lightllm.models.glm5_next.tokenizer import Glm5NextTokenizer +from lightllm.models.glm5_next.vision_process import Glm5NextImageProcessor, smart_resize +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.server.multimodal_params import ImageItem, MultimodalParams +from lightllm.utils.envs_utils import set_env_start_args + + +@pytest.mark.parametrize("width,height,canvas", [(7, 11, (168, 112)), (173, 89, (112, 196)), (448, 448, (448, 448))]) +def test_image_canvas_and_token_expansion(tmp_path, width, height, canvas): + assert smart_resize(height, width) == canvas + (tmp_path / "processor_config.json").write_text(json.dumps({"image_processor": {}})) + tokenizer = Glm5NextTokenizer( + None, {"image_start_token_id": 100, "image_end_token_id": 101, "image_token_id": 102}, str(tmp_path) + ) + image = ImageItem(type="base64", data="") + image.image_w, image.image_h = width, height + image.token_id = 1000 + image.token_num = tokenizer.get_image_token_length(image) + assert image.token_num == canvas[0] * canvas[1] // 28 ** 2 + multi = MultimodalParams() + multi.images = [image] + ids = tokenizer.encode([1, 100, 102, 101, 2], multi) + assert ids == [1, 100, *range(1000, 1000 + image.token_num), 101, 2] + assert image.start_idx == 2 + assert image.grid_thwd[-1] == 0 + assert tokenizer.encode([1, 2, 3]) == [1, 2, 3] + assert tokenizer.encode([100, 102, 101]) == [100, 101] + + +def test_image_patch_order_temporal_repeat_and_black_padding(): + pixels = np.zeros((29, 31, 3), dtype=np.uint8) + pixels[:, :, 0] = np.arange(31)[None, :] + pixels[:, :, 1] = np.arange(29)[:, None] + pixels[:, :, 2] = 127 + processor = Glm5NextImageProcessor(min_image_tokens=1, do_rescale=False, do_normalize=False) + patches, grid = processor._preprocess_bydevice(Image.fromarray(pixels), device="cpu") + assert grid.tolist() == [[1, 4, 4]] + patches = patches.reshape(16, 3, 2, 14, 14) + assert torch.equal(patches[:, :, 0], patches[:, :, 1]) + assert torch.equal(patches[0, :, 0], torch.from_numpy(pixels[:14, :14]).permute(2, 0, 1)) + assert torch.equal(patches[1, :, 0], torch.from_numpy(pixels[:14, 14:28]).permute(2, 0, 1)) + assert torch.equal(patches[2, :, 0], torch.from_numpy(pixels[14:28, :14]).permute(2, 0, 1)) + assert not patches[-1].any() + + +def test_vision_attention_uses_triton_rope_with_fa3_backend(monkeypatch): + import lightllm.models.glm5_next.glm5_next_visual as glm_visual + import lightllm.server.visualserver as visualserver + from lightllm.common.basemodel.attention_vit.fa3.fp import Fa3VitAttBackend + from lightllm.server.visualserver import set_vit_att_backend + + rope_inputs = [] + norm_inputs = [] + attn_inputs = [] + + def fake_qk_norm(qkv, q_weight, k_weight, eps): + norm_inputs.append((qkv, q_weight, k_weight, eps)) + return qkv[:, 0].contiguous(), qkv[:, 1].contiguous() + + def fake_rope(x, cos, sin): + rope_inputs.append((x, cos, sin)) + return x + + def fake_fa3(q, k, v, o, cu_seqlens, max_seqlen): + attn_inputs.append((q, k, v, cu_seqlens, max_seqlen)) + o.copy_(q) + return o + + monkeypatch.setattr(glm_visual, "qk_rms_norm", fake_qk_norm) + monkeypatch.setattr(glm_visual, "apply_rotary_pos_emb_triton", fake_rope) + monkeypatch.setattr(Fa3VitAttBackend, "_vit_att_fwd", staticmethod(fake_fa3)) + monkeypatch.setattr(visualserver, "VIT_ATTN_BACKEND", visualserver.VIT_ATTN_BACKEND) + set_vit_att_backend("fa3") + + attention = glm_visual.Glm5NextVisionAttention(hidden_size=16, num_heads=2, eps=1e-5, bias=True) + x = torch.randn(3, 16) + cos = torch.randn(3, 4) + sin = torch.randn(3, 4) + cu_seqlens = torch.tensor([0, 3], dtype=torch.int32) + result = attention(x, cu_seqlens, 3, cos, sin) + + assert result.shape == x.shape + assert len(norm_inputs) == 1 + assert norm_inputs[0][0].shape == (3, 3, 2, 8) + assert [(q.shape, cos_.shape, sin_.shape) for q, cos_, sin_ in rope_inputs] == [ + ((3, 2, 8), (3, 4), (3, 4)), + ((3, 2, 8), (3, 4), (3, 4)), + ] + assert len(attn_inputs) == 1 + assert attn_inputs[0][0].shape == attn_inputs[0][1].shape == attn_inputs[0][2].shape == (3, 2, 8) + + +def test_vision_rms_norm_preserves_glm_rounding(monkeypatch): + import lightllm.models.glm5_next.glm5_next_visual as glm_visual + + norm_calls = [] + + def fake_rms_norm(x, weight, eps, round_norm_before_weight): + norm_calls.append((x, weight, eps, round_norm_before_weight)) + return x + + monkeypatch.setattr(glm_visual, "rms_norm", fake_rms_norm) + norm = glm_visual.Glm5NextVisionRMSNorm(hidden_size=8, eps=1e-5) + x = torch.randn(2, 8) + assert norm(x) is x + assert len(norm_calls) == 1 + assert norm_calls[0][0] is x + assert norm_calls[0][1] is norm.weight + assert norm_calls[0][2:] == (1e-5, True) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_triton_glm_vision_norms_match_reference(): + from lightllm.models.vit.triton_kernel.rms_norm_vit import qk_rms_norm, rms_norm + + def glm_rms_norm(x, weight, eps): + normalized = x.float() * torch.rsqrt(x.float().square().mean(-1, keepdim=True) + eps) + return normalized.to(x.dtype) * weight + + torch.manual_seed(53) + eps = 1e-5 + qkv = torch.randn(7, 3, 16, 64, dtype=torch.bfloat16, device="cuda") + q_weight = torch.randn(64, dtype=torch.bfloat16, device="cuda") + k_weight = torch.randn(64, dtype=torch.bfloat16, device="cuda") + q, k = qk_rms_norm(qkv, q_weight, k_weight, eps) + + torch.testing.assert_close(q, glm_rms_norm(qkv[:, 0], q_weight, eps), atol=1e-2, rtol=1e-2) + torch.testing.assert_close(k, glm_rms_norm(qkv[:, 1], k_weight, eps), atol=1e-2, rtol=1e-2) + + x = torch.randn(7, 1024, dtype=torch.bfloat16, device="cuda") + weight = torch.randn(1024, dtype=torch.bfloat16, device="cuda") + torch.testing.assert_close( + rms_norm(x, weight, eps, round_norm_before_weight=True), + glm_rms_norm(x, weight, eps), + atol=1e-2, + rtol=1e-2, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_vision_batch_preserves_independent_image_attention(): + from lightllm.models.glm5_next.glm5_next_visual import Glm5NextVisionTransformer + from lightllm.server.visualserver import set_vit_att_backend + + torch.manual_seed(53) + set_vit_att_backend("sdpa") + model = ( + Glm5NextVisionTransformer( + {"data_type": "float32"}, + hidden_size=128, + out_hidden_size=256, + depth=2, + intermediate_size=256, + projection_intermediate_size=512, + num_heads=2, + ) + .cuda() + .eval() + ) + grid = torch.tensor([[1, 4, 6], [1, 6, 4]]) + x = torch.randn(48, 3 * 2 * 14 ** 2, device="cuda") + with torch.inference_mode(): + batch = model(x, grid) + individual = torch.cat([model(x[:24], grid[:1]), model(x[24:], grid[1:])]) + assert batch.shape == (12, 256) + torch.testing.assert_close(batch, individual, atol=2e-5, rtol=2e-5) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_target_and_mtp_prefill_embed_images_before_hidden_fusion(monkeypatch): + from lightllm.models.glm5_next.layer_infer.pre_layer_infer import Glm5NextPreLayerInfer + from lightllm.models.glm5_next_mtp.layer_infer.pre_layer_infer import Glm5NextMTPPreLayerInfer + from lightllm.server.router.model_infer.infer_batch import g_infer_context + + set_env_start_args(dataclasses.asdict(StartArgs(mtp_mode="eagle_with_att", mtp_step=2))) + monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_DP", "0") + monkeypatch.setenv("LIGHTLLM_DP_WORLD_SIZE", "1") + hidden = 128 + config = {"hidden_size": hidden, "rms_norm_eps": 1e-5, "hc_mult": 4} + image_cache = torch.randn(3, 1, hidden, pin_memory=True) + monkeypatch.setattr(g_infer_context, "cpu_embed_cache_client", SimpleNamespace(cpu_embed_cache_tensor=image_cache)) + text = torch.randn(8, hidden, device="cuda") + ids = torch.tensor([1, 1000, 1001, 2], device="cuda") + metadata = {"token_id": 1000, "token_num": 2, "start_index_in_embed_cache": 1} + state = SimpleNamespace(multimodal_params=[{"images": [metadata], "audios": []}]) + weight = SimpleNamespace(wte_weight_=SimpleNamespace(weight=text, tp_vocab_start_id=0, tp_vocab_end_id=8)) + expected = torch.cat((text[1:2], image_cache[1:, 0].cuda(), text[2:3])) + expanded = Glm5NextPreLayerInfer(config).context_forward(ids, state, weight) + torch.testing.assert_close(expanded.reshape(4, 4, hidden), expected[:, None].expand(-1, 4, -1)) + + state.mtp_draft_input_hiddens = torch.randn(4, hidden, device="cuda") + old_hidden = state.mtp_draft_input_hiddens.clone() + projection = torch.randn(2 * hidden, hidden, device="cuda") * 0.1 + weight.enorm_weight_ = weight.hnorm_weight_ = weight.main_norm_weight_ = SimpleNamespace( + weight=torch.ones(hidden, device="cuda") + ) + weight.eh_proj_weight_ = SimpleNamespace(mm=lambda *, input_tensor: input_tensor @ projection) + actual = Glm5NextMTPPreLayerInfer(config).context_forward(ids, state, weight) + old_hidden = torch.nn.functional.rms_norm(old_hidden, (hidden,), weight.main_norm_weight_.weight, 1e-5) + normalized = [x * torch.rsqrt(x.square().mean(-1, keepdim=True) + 1e-5) for x in (expected, old_hidden)] + torch.testing.assert_close(actual, torch.cat(normalized, -1) @ projection) diff --git a/unit_tests/models/qwen2_vl/test_mrope.py b/unit_tests/models/qwen2_vl/test_mrope.py new file mode 100644 index 0000000000..e98989179a --- /dev/null +++ b/unit_tests/models/qwen2_vl/test_mrope.py @@ -0,0 +1,66 @@ +import torch +import pytest + +# Import the Triton kernel function under test. Adjust the import path as needed. +from lightllm.models.qwen2_vl.triton_kernel.mrope import mrope_triton_fused + +# Reference Python implementation for multimodal rotary positional embeddings + + +def rotate_half(x): + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def apply_multimodal_rotary_pos_emb(q, k, cos, sin, mrope_section, unsqueeze_dim=1): + mrope_section = mrope_section * 2 + cos = torch.cat([m[i % 3] for i, m in enumerate(cos.split(mrope_section, dim=-1))], dim=-1).unsqueeze(unsqueeze_dim) + sin = torch.cat([m[i % 3] for i, m in enumerate(sin.split(mrope_section, dim=-1))], dim=-1).unsqueeze(unsqueeze_dim) + + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + + return q_embed, k_embed + + +@pytest.mark.parametrize( + "B,H_q,H_k,L,D,mrope_section", + [ + (1, 4, 2, 32, 16, [4, 2, 2]), + ], +) +def test_mrope_triton_correctness(B, H_q, H_k, L, D, mrope_section): + """ + Test that the Triton kernel matches the reference PyTorch implementation. + """ + torch.manual_seed(0) + device = "cuda" + HALF = D // 2 + + q = torch.rand((B, H_q, L, D), dtype=torch.float32, device=device) + k = torch.rand((B, H_k, L, D), dtype=torch.float32, device=device) + + cos_half = torch.rand((1, L, HALF), dtype=torch.float32, device=device).expand(3, -1, -1) + sin_half = torch.rand((1, L, HALF), dtype=torch.float32, device=device).expand(3, -1, -1) + + cos_full = torch.cat([cos_half, cos_half], dim=-1) + sin_full = torch.cat([sin_half, sin_half], dim=-1) + + cos_ref = cos_full.unsqueeze(1) + sin_ref = sin_full.unsqueeze(1) + + ref_q, ref_k = apply_multimodal_rotary_pos_emb(q, k, cos_ref, sin_ref, mrope_section, unsqueeze_dim=1) + + q = q.transpose(1, 2).contiguous().view(L, H_q, D) + k = k.transpose(1, 2).contiguous().view(L, H_k, D) + mrope_section = torch.tensor(mrope_section, dtype=torch.int32, device="cuda") + mrope_triton_fused(q, k, cos_half, sin_half, mrope_section, is_interleaved=False) + q = q.transpose(0, 1).contiguous().view(B, H_q, L, D) + k = k.transpose(0, 1).contiguous().view(B, H_k, L, D) + assert torch.allclose(q, ref_q, rtol=1e-3, atol=1e-3) + assert torch.allclose(k, ref_k, rtol=1e-3, atol=1e-3) + + +if __name__ == "__main__": + pytest.main() diff --git a/unit_tests/models/qwen3-vl/test_deepstack_emb.py b/unit_tests/models/qwen3-vl/test_deepstack_emb.py new file mode 100644 index 0000000000..f629a16352 --- /dev/null +++ b/unit_tests/models/qwen3-vl/test_deepstack_emb.py @@ -0,0 +1,82 @@ +import torch +import triton +import triton.language as tl +from lightllm.models.qwen3_vl.infer_struct import Qwen3VLInferStateInfo +from lightllm.models.qwen3_vl.triton_kernel.deepstack_multimodal_emb import add_deepstack_embs + + +def test_deepstack_same_image_twice(): + device = "cuda" + + # 1. 构造 input_ids，包含两段相同的 image token 范围 [100, 101, 102] + input_ids = torch.tensor( + [1, 100, 101, 102, 2, 100, 101, 102, 3], + device=device, + dtype=torch.long, + ) + seq_len = input_ids.shape[0] + + hidden_size = 4 + token_len = 3 # 每张图 3 个 token + + # 2. 构造初始 embedding，全 0，方便看增量 + input_embeddings = torch.zeros(seq_len, hidden_size, device=device, dtype=torch.float32) + + # 3. 构造 deepstack_embs（这一层的 deepstack） + # 只有一张图片，所以 deepstack_embs 形状是 [token_len, hidden_size] + # 每一行是 [1,1,1,1], [2,2,2,2], [3,3,3,3] + deepstack_embs = torch.tensor( + [ + [1.0, 1.0, 1.0, 1.0], # 对应 token_id = 100 + [2.0, 2.0, 2.0, 2.0], # 对应 token_id = 101 + [3.0, 3.0, 3.0, 3.0], # 对应 token_id = 102 + ], + device=device, + dtype=torch.float32, + ) + + # 4. image 相关索引信息（与 multimodal_emb 一致的语义） + img_start_token_ids = torch.tensor([100], device=device, dtype=torch.long) # 只有一个 image handle，从 100 开始 + img_token_lens = torch.tensor([token_len], device=device, dtype=torch.long) + img_start_locs = torch.tensor([0], device=device, dtype=torch.long) # deepstack_embs 从第 0 行开始是这张图的 + + # 5. 保存一份原始 embedding，方便求差 + before = input_embeddings.clone() + + # 6. 调用 Triton 算子 + add_deepstack_embs( + out=input_embeddings, + input_ids=input_ids, + deepstack_embs=deepstack_embs, + img_token_lens=img_token_lens, + img_start_token_ids=img_start_token_ids, + img_start_locs_in_cache=img_start_locs, + ) + + # 7. 看看相同图片两段上的增量 + delta = input_embeddings - before + + print("input_ids:", input_ids) + print("delta:\n", delta) + + # 第一次 image：位置 1,2,3 + print("first image span delta:\n", delta[1:4]) + # 第二次 image：位置 5,6,7 + print("second image span delta:\n", delta[5:8]) + + # 8. 断言它们和预期一致 + expected = deepstack_embs # [3, 4] + + assert torch.allclose(delta[1:4], expected), "first image span does not match expected deepstack" + assert torch.allclose(delta[5:8], expected), "second image span does not match expected deepstack" + + # 其他位置应该仍然是 0 + assert torch.all(delta[0] == 0) + assert torch.all(delta[4] == 0) + assert torch.all(delta[8] == 0) + + print("OK: same image appears twice, both spans get deepstack added correctly.") + + +if __name__ == "__main__": + test_deepstack_same_image_twice() diff --git a/unit_tests/models/sensenova_loop/test_gdn_state.py b/unit_tests/models/sensenova_loop/test_gdn_state.py new file mode 100644 index 0000000000..c25576a19a --- /dev/null +++ b/unit_tests/models/sensenova_loop/test_gdn_state.py @@ -0,0 +1,151 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.attention_state import AttentionStateManager +from lightllm.common.attention_state.bounded_att import ConvSSMStateConfig, GDNState +from lightllm.common.basemodel.attention.base_att import AttControl +from lightllm.common.basemodel.attention.linear.triton import TritonLinearAttBackend +from lightllm.common.basemodel.infer_struct import InferStateInfo +from lightllm.models.sensenova_loop.config import get_loop_layer_indices, get_loop_layer_types +from lightllm.server.core.objs.start_args_type import StartArgs + + +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +@pytest.mark.parametrize("mtp_step", [0, 2]) +@pytest.mark.parametrize( + "start,end,count,logical_layer,storage_layer", + [(4, 28, 1, 4, 3), (4, 28, 3, 28, 21), (1, 6, 2, 13, 10)], + ids=["no-loop", "aligned-loop", "unaligned-loop"], +) +def test_gdn_decode_keeps_logical_layer_state( + monkeypatch, device, mtp_step, start, end, count, logical_layer, storage_layer +): + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA required") + torch.manual_seed(1585) + settings = dict(start=start, end=end, eval_count=count, use_gate=False) + config = dict( + num_hidden_layers=32, + full_attention_interval=4, + layer_types=["full_attention" if (i + 1) % 4 == 0 else "linear_attention" for i in range(32)], + linear_num_key_heads=1, + linear_num_value_heads=2, + linear_key_head_dim=32, + linear_value_head_dim=32, + linear_conv_kernel_dim=4, + hidden_act="silu", + ) + indices = get_loop_layer_indices(settings, 32) + kinds = get_loop_layer_types(settings, config) + layers = tuple(i for i, kind in enumerate(kinds) if kind == "linear_attention") + args = StartArgs(mtp_step=mtp_step, linear_att_ssm_data_type="float32") + monkeypatch.setattr("lightllm.common.basemodel.attention.linear.gdn.get_env_start_args", lambda: args) + monkeypatch.setattr("lightllm.common.basemodel.attention.base_att.get_env_start_args", lambda: args) + backend = object.__new__(TritonLinearAttBackend) + backend.__init__( + SimpleNamespace( + config=config, + tp_world_size_=1, + is_mtp_draft_model=False, + mtp_manager=SimpleNamespace(get_decode_draft_step=lambda _: mtp_step), + ) + ) + dtype = torch.bfloat16 if device == "cuda" else torch.float32 + layout = ConvSSMStateConfig(layers, 1, 2, 32, 32, 3, dtype, torch.float32) + resource = GDNState(layout, request_capacity=3, mtp_step=mtp_step, device=device) + resources = AttentionStateManager(request_capacity=3) + resources.register_state("main", resource) + resources.finish_state_registration() + conv_view, ssm_view = resource.get_layer_view(logical_layer_id=logical_layer) + assert conv_view.data_ptr() == resource.conv_state[storage_layer].data_ptr() + assert ssm_view.data_ptr() == resource.ssm_state[storage_layer].data_ptr() + physical_weights = [ + SimpleNamespace( + layer_num_=i, + linear_conv1d=SimpleNamespace( + mm_param=SimpleNamespace(weight=torch.randn(layout.conv_channels, 4, device=device, dtype=dtype) * 0.1), + bias=None, + ), + linear_A_log=SimpleNamespace(weight=torch.zeros(2, device=device)), + linear_dt_bias=SimpleNamespace(weight=torch.zeros(2, device=device)), + ) + for i in range(32) + ] + execution_weights = [physical_weights[i] for i in indices] + weight = execution_weights[logical_layer] + assert weight is physical_weights[indices[logical_layer]] + + if device == "cpu": + + def conv_pack(mixed, z, a, b, conv, *args): + conv[1].add_(1) + return (*backend._rearrange_mixed_qkv(mixed, decode=True), z, a, b) + + def conv_mtp(mixed, conv, *args, **kwargs): + conv[1].add_(1) + return mixed + + def recurrent(*, initial_state, v, **kwargs): + initial_state[mtp_step + 1 : 2 * (mtp_step + 1)].add_(1) + return v, None + + monkeypatch.setattr("lightllm.common.basemodel.attention.linear.gdn.conv_pack_gdn_decode_inputs", conv_pack) + monkeypatch.setattr( + "lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_mtp.causal_conv1d_update", conv_mtp + ) + monkeypatch.setattr( + "lightllm.common.basemodel.attention.linear.gdn.fused_recurrent_gated_delta_rule", recurrent + ) + monkeypatch.setattr( + "lightllm.common.basemodel.attention.linear.gdn.mtp_fused_recurrent_gated_delta_rule", + recurrent, + ) + + infer = InferStateInfo() + infer.att_state_manager, infer.att_state_role = resources, "main" + infer.b_req_idx = torch.tensor([1], device=device, dtype=torch.int32) + control = dict(layer_num=logical_layer, layer_weight=weight) + mixed_dim = backend.tp_qkvz_dim + backend.tp_ba_dim + if device == "cuda": + infer.b1_cu_q_seq_len = torch.tensor([0, 5], device=device, dtype=torch.int32) + infer.b_ready_cache_len = torch.zeros(1, device=device, dtype=torch.int32) + prefill = backend.create_att_prefill_state(infer) + prefill.init_state() + prefill.prefill_att( + None, + None, + None, + AttControl( + linear_att_prefill=True, + linear_att_prefill_dict=dict( + **control, mixed_qkvzba=torch.randn(5, mixed_dim, device=device, dtype=dtype) + ), + ), + ) + + infer.batch_size = mtp_step + 1 + infer.b_req_idx = infer.b_req_idx.repeat(infer.batch_size) + decode = backend.create_att_decode_state(infer) + untouched = [i for i in range(len(layers)) if i != storage_layer] + for _ in range(2): + before_conv, before_ssm = resource.conv_state.clone(), resource.ssm_state.clone() + decode.init_state() + output, _ = decode.decode_att( + None, + None, + None, + AttControl( + linear_att_decode=True, + linear_att_decode_dict=dict( + **control, mixed_qkvzba=torch.randn(infer.batch_size, mixed_dim, device=device, dtype=dtype) + ), + ), + ) + assert torch.isfinite(output).all() + assert not torch.equal(resource.conv_state[storage_layer], before_conv[storage_layer]) + assert not torch.equal(resource.ssm_state[storage_layer], before_ssm[storage_layer]) + assert torch.equal(resource.conv_state[untouched], before_conv[untouched]) + assert torch.equal(resource.ssm_state[untouched], before_ssm[untouched]) + assert weight.layer_num_ == indices[logical_layer] diff --git a/unit_tests/models/sensenova_loop/test_loop.py b/unit_tests/models/sensenova_loop/test_loop.py new file mode 100644 index 0000000000..7caaa09f3f --- /dev/null +++ b/unit_tests/models/sensenova_loop/test_loop.py @@ -0,0 +1,376 @@ +import dataclasses +import json +from copy import deepcopy +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.models.registry import ModelRegistry +from lightllm.models.sensenova_loop.config import get_loop_layer_indices, get_loop_layer_types +from lightllm.models.qwen3next.config import Qwen3NextResourceConfig +from lightllm.server.core.objs.start_args_type import StartArgs + + +def _config(layers=32, start=4, end=28, count=3): + return { + "model_type": "sensenova_loop", + "loop_settings": {"start": start, "end": end, "eval_count": count, "use_gate": False}, + "text_config": { + "num_hidden_layers": layers, + "full_attention_interval": 4, + "layer_types": ["full_attention" if (i + 1) % 4 == 0 else "linear_attention" for i in range(layers)], + "num_key_value_heads": 4, + "head_dim": 256, + "linear_num_key_heads": 16, + "linear_num_value_heads": 32, + "linear_key_head_dim": 128, + "linear_value_head_dim": 128, + "linear_conv_kernel_dim": 4, + }, + } + + +@pytest.mark.parametrize("start,end,count", [(4, 28, 3), (4, 28, 2), (1, 6, 2)]) +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.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() + 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) + full = linear = 0 + slots = {} + for execution, (physical, kind) in enumerate(zip(indices, kinds)): + if kind == "full_attention": + slot = full_layers.index(execution) + assert slot == full + full += 1 + else: + slot = linear_layers.index(execution) + assert slot == linear + linear += 1 + assert slot not in slots.setdefault(physical, set()) + slots[physical].add(slot) + 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): + assert (cache.all_layer_num, full, linear) == (56, 14, 42) + + +def test_physical_weights_are_created_and_loaded_once(monkeypatch): + from lightllm.models.sensenova_loop.model import SenseNovaLoopTpPartModel + from lightllm.models.qwen3_5.model import Qwen3_5TpPartModel + + model = object.__new__(SenseNovaLoopTpPartModel) + model.loop_layer_indices = get_loop_layer_indices(_config()["loop_settings"], 32) + model.num_physical_layers = 32 + model.config = {"n_layer": 80} + model.data_type, model.quant_cfg = torch.bfloat16, None + loaded = [] + + def create_weight(index, dtype, network_config, quant_cfg): + assert network_config is model.config + return object() + + def load_weights(self, weights): + loaded.extend(self.trans_layers_weight) + + model.pre_and_post_weight_class = lambda dtype, network_config: object() + model.transformer_weight_class = create_weight + monkeypatch.setattr(Qwen3_5TpPartModel, "load_weights", load_weights) + model._init_weights() + model.load_weights({}) + assert len(loaded) == len(set(loaded)) == 32 + assert len(model.trans_layers_weight) == 80 + for execution, physical in enumerate(model.loop_layer_indices): + assert model.trans_layers_weight[execution] is loaded[physical] + assert model.config["n_layer"] == 80 + + +@pytest.mark.parametrize("streams,count", [(2, 2), (4, 3)]) +def test_mhc_registry_and_embedded_mtp_reuse_existing_qwen_draft(tmp_path, streams, count): + from lightllm.models.draft_registry import get_draft_model_class + from lightllm.models.sensenova_loop.model import SenseNovaLoopTpPartModel + from lightllm.models.qwen3_5_mtp.model import Qwen3_5MTPModel + + config = _config(count=count) + config["mhc"] = dict(start=4, end=28, streams=streams, iterations=20) + config["text_config"].update(num_attention_heads=16, hidden_size=4096, vocab_size=128) + config["text_config"]["layer_types"] = ["linear_attention"] * 3 + ["full_attention"] + config["text_config"]["layer_types"] *= 8 + (tmp_path / "config.json").write_text(json.dumps(config)) + registration = ModelRegistry.get_model_config(config) + assert registration.model_class.endswith(":SenseNovaLoopTpPartModel") + assert registration.is_multimodal + assert get_draft_model_class(config, "eagle_with_att") is Qwen3_5MTPModel + + target = object.__new__(SenseNovaLoopTpPartModel) + target.weight_dir_, target.finetune_config, target.tp_world_size_ = str(tmp_path), None, 4 + target.args = SimpleNamespace(**dataclasses.asdict(StartArgs())) + target._init_config() + assert target.config["n_layer"] == 32 + 24 * (count - 1) + assert target.num_physical_layers == 32 + + draft = object.__new__(Qwen3_5MTPModel) + 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"] == ["full_attention"] + + +def test_loop_tokenizer_loads_without_hf_loop_autoconfig(tmp_path): + from tokenizers import Tokenizer, models + from transformers import PreTrainedTokenizerFast, Qwen2VLImageProcessorFast + from lightllm.server.tokenizer import get_tokenizer + + tokenizer = PreTrainedTokenizerFast( + tokenizer_object=Tokenizer(models.WordLevel({"hello": 0, "[UNK]": 1}, unk_token="[UNK]")), + unk_token="[UNK]", + ) + tokenizer.save_pretrained(tmp_path) + Qwen2VLImageProcessorFast().save_pretrained(tmp_path) + config = _config() + config.update( + architectures=["SenseNovaLoopForConditionalGeneration"], + vision_config={}, + vision_start_token_id=2, + vision_end_token_id=3, + image_token_id=4, + ) + (tmp_path / "config.json").write_text(json.dumps(config)) + wrapped = get_tokenizer(str(tmp_path)) + assert wrapped.encode("hello") == [0] + assert isinstance(wrapped.image_processor, Qwen2VLImageProcessorFast) + + +@pytest.mark.parametrize("custom_layout", [False, True]) +def test_model_and_worker_use_the_same_loop_cache_layout(monkeypatch, custom_layout): + from lightllm.models.sensenova_loop.model import SenseNovaLoopTpPartModel + + config = _config(start=1, end=6, count=2) + 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", 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, {})) + 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, + ) + (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 + + +def test_plain_loop_registry_does_not_enable_mhc(): + registration = ModelRegistry.get_model_config(_config()) + assert registration.model_class == "lightllm.models.sensenova_loop.model:SenseNovaLoopTpPartModel" + assert registration.is_multimodal + + +def test_awq_config_applies_to_main_model_and_excludes_float_mtp(): + from lightllm.models.sensenova_loop.model import SenseNovaLoopTpPartModel + from lightllm.models.qwen3_5_mtp.model import Qwen3_5MTPModel + + config = _config() + quantization = {"quant_method": "awq", "bits": 4, "group_size": 128, "modules_to_not_convert": ["mtp"]} + config["quantization_config"] = quantization + original = deepcopy(config) + assert SenseNovaLoopTpPartModel.normalize_model_config(config)["quantization_config"] == quantization + assert "quantization_config" not in Qwen3_5MTPModel.normalize_model_config(config) + assert config == original + + config["text_config"]["quantization_config"] = None + assert SenseNovaLoopTpPartModel.normalize_model_config(config)["quantization_config"] is None + + +@pytest.mark.parametrize("suffix,factor", [("qweight", 8), ("qzeros", 8), ("scales", 1)]) +def test_awq_fused_projections_split_on_output_axis(suffix, factor): + from lightllm.models.qwen3_5.layer_weights.transformer_layer_weight import Qwen35TransformerLayerWeight + + method = SimpleNamespace(method_name="awq_marlin", pack_factor=8) + layer = object.__new__(Qwen35TransformerLayerWeight) + layer.layer_num_ = 0 + layer.linear_num_k_heads = 2 + layer.linear_num_v_heads = 4 + layer.linear_k_head_dim = layer.linear_v_head_dim = 16 + layer.linear_in_proj = SimpleNamespace(quant_method=method) + prefix = "model.layers.0.linear_attn.in_proj_" + qkv = torch.arange(256 // factor).reshape(2, 128 // factor) + fused_name = f"{prefix}qkv.{suffix}" + weights = {fused_name: qkv} + layer._split_linear_in_proj_qkv(weights) + q_end, k_end = 32 // factor, 64 // factor + assert fused_name not in weights + assert torch.equal(weights[f"{prefix}q.{suffix}"], qkv[:, :q_end]) + assert torch.equal(weights[f"{prefix}k.{suffix}"], qkv[:, q_end:k_end]) + assert torch.equal(weights[f"{prefix}v.{suffix}"], qkv[:, k_end:]) + + layer.is_linear_attention_layer = False + layer.q_head_num_, layer.head_dim = 2, 16 + layer.qkvo_gate_proj = SimpleNamespace(quant_method=method) + layer._q_weight_name = "model.layers.0.self_attn.q_proj.weight" + layer._o_gate_weight_name = "model.layers.0.self_attn.o_gate_proj.weight" + q_name = f"model.layers.0.self_attn.q_proj.{suffix}" + gate_name = f"model.layers.0.self_attn.o_gate_proj.{suffix}" + head_width = layer.head_dim // factor + q_gate = torch.arange(8 * head_width).reshape(2, 4 * head_width) + q1, gate1, q2, gate2 = q_gate.split(head_width, dim=1) + weights = {q_name: q_gate} + layer._split_q_with_gate(weights) + assert torch.equal(weights[q_name], torch.cat((q1, q2), dim=1)) + assert torch.equal(weights[gate_name], torch.cat((gate1, gate2), dim=1)) + layer._split_q_with_gate(weights) + assert torch.equal(weights[q_name], torch.cat((q1, q2), dim=1)) + assert torch.equal(weights[gate_name], torch.cat((gate1, gate2), dim=1)) + + +def test_loop_binds_split_gdn_projection_only_for_weights_with_separate_gates(monkeypatch): + from lightllm.models.qwen3_5.model import Qwen3_5TpPartModel + from lightllm.models.qwen3_5.layer_infer.transformer_layer_infer import Qwen35TransformerLayerInfer + from lightllm.models.qwen3next.layer_infer.transformer_layer_infer import Qwen3NextTransformerLayerInfer + from lightllm.models.sensenova_loop.model import SenseNovaLoopTpPartModel + from lightllm.models.sensenova_loop.layer_infer.transformer_layer_infer import SenseNovaLoopTransformerLayerInfer + + qkvz = torch.arange(8).reshape(2, 4) + gates = torch.arange(8, 12).reshape(2, 2) + fused = torch.cat((qkvz, gates), dim=-1) + fused_weight = SimpleNamespace(linear_in_proj=SimpleNamespace(mm=lambda _: fused)) + split_weight = SimpleNamespace( + linear_in_proj=SimpleNamespace(mm=lambda _: qkvz), + linear_gate_proj=SimpleNamespace(mm=lambda _: gates), + ) + model = object.__new__(SenseNovaLoopTpPartModel) + model.trans_layers_weight = [fused_weight, split_weight, split_weight] + layers = [object.__new__(SenseNovaLoopTransformerLayerInfer) for _ in model.trans_layers_weight] + for layer in layers: + layer.embed_dim_ = 4 + monkeypatch.setattr(Qwen3_5TpPartModel, "_init_infer_layer", lambda self, _: setattr(self, "layers_infer", layers)) + + model._init_infer_layer() + assert Qwen35TransformerLayerInfer._linear_in_proj is Qwen3NextTransformerLayerInfer._linear_in_proj + assert layers[0]._linear_in_proj.__func__ is Qwen3NextTransformerLayerInfer._linear_in_proj + for layer in layers[1:]: + assert layer._linear_in_proj.__func__ is SenseNovaLoopTransformerLayerInfer._linear_in_proj_with_gates + for layer, weight in zip(layers, model.trans_layers_weight): + assert torch.equal(layer._linear_in_proj(torch.zeros(1, 2, 4), weight), fused) + + +def test_weight_load_failure_restores_execution_weights(monkeypatch): + from lightllm.models.sensenova_loop.model import SenseNovaLoopTpPartModel + from lightllm.models.qwen3_5.model import Qwen3_5TpPartModel + + model = object.__new__(SenseNovaLoopTpPartModel) + model.config = {"n_layer": 80} + model.physical_layer_weights = [object()] + model.trans_layers_weight = [object()] + execution_config, execution_weights = model.config, model.trans_layers_weight + + def fail(*args): + raise RuntimeError("weight load failed") + + monkeypatch.setattr(Qwen3_5TpPartModel, "load_weights", fail) + with pytest.raises(RuntimeError, match="weight load failed"): + model.load_weights({}) + assert model.config is execution_config + assert model.trans_layers_weight is execution_weights + + +@pytest.fixture +def layer_runtime(monkeypatch): + from lightllm.utils.envs_utils import get_env_start_args + + monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(dataclasses.asdict(StartArgs()))) + for key, value in { + "DP_WORLD_SIZE": 1, + "CURRENT_RANK_IN_DP": 0, + "CURRENT_DEVICE_ID": 0, + "CURRENT_RANK_IN_NODE": 0, + }.items(): + monkeypatch.setenv("LIGHTLLM_" + key, str(value)) + get_env_start_args.cache_clear() + yield + get_env_start_args.cache_clear() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("mhc_enabled", [False, True]) +def test_loop_layer_types_are_local_to_sensenova(layer_runtime, tmp_path, mhc_enabled): + from lightllm.common.quantization import Quantcfg + from lightllm.models.qwen3_5.layer_infer.transformer_layer_infer import Qwen35TransformerLayerInfer + from lightllm.models.qwen3_5.layer_weights.transformer_layer_weight import Qwen35TransformerLayerWeight + from lightllm.models.qwen3_5_mtp.model import Qwen3_5MTPModel + from lightllm.models.sensenova_loop.layer_infer.transformer_layer_infer import SenseNovaLoopTransformerLayerInfer + from lightllm.models.sensenova_loop.layer_weights.transformer_layer_weight import ( + SenseNovaLoopTransformerLayerWeight, + ) + + config = _config(layers=8, start=1, end=6, count=2) + text = config["text_config"] + text.update( + n_layer=8, + hidden_size=128, + intermediate_size=256, + num_attention_heads=1, + num_key_value_heads=1, + head_dim=128, + linear_num_key_heads=1, + linear_num_value_heads=1, + rms_norm_eps=1e-6, + vocab_size=128, + layer_types=["full_attention", "linear_attention"] * 4, + ) + if mhc_enabled: + text["mhc"] = dict(start=1, end=6, streams=2, iterations=20) + indices = get_loop_layer_indices(config["loop_settings"], 8) + execution_config = { + **text, + "loop_physical_indices": indices, + "layer_types": get_loop_layer_types(config["loop_settings"], text), + } + weights = [ + SenseNovaLoopTransformerLayerWeight(i, torch.bfloat16, execution_config, Quantcfg(execution_config)) + for i in range(8) + ] + for execution, physical in enumerate(indices): + layer = SenseNovaLoopTransformerLayerInfer(execution, execution_config) + expected_linear = text["layer_types"][physical] == "linear_attention" + assert layer.is_linear_attention_layer == weights[physical].is_linear_attention_layer == expected_linear + assert hasattr(weights[physical], "linear_in_proj") == expected_linear + assert hasattr(weights[physical], "qkvo_gate_proj") != expected_linear + + # The base model still uses its original periodic layout, even with layer_types present. + base_layer = Qwen35TransformerLayerInfer(0, text) + base_weight = Qwen35TransformerLayerWeight(0, torch.bfloat16, text, Quantcfg(text)) + assert base_layer.is_linear_attention_layer and base_weight.is_linear_attention_layer + + (tmp_path / "config.json").write_text(json.dumps(config)) + draft = object.__new__(Qwen3_5MTPModel) + draft.weight_dir_, draft.finetune_config, draft.tp_world_size_ = str(tmp_path), None, 1 + draft._init_config() + for global_index in (len(indices), len(indices) + 1): + assert not Qwen35TransformerLayerInfer(global_index, draft.config).is_linear_attention_layer diff --git a/unit_tests/models/sensenova_loop/test_loop_mhc.py b/unit_tests/models/sensenova_loop/test_loop_mhc.py new file mode 100644 index 0000000000..34d4fe9696 --- /dev/null +++ b/unit_tests/models/sensenova_loop/test_loop_mhc.py @@ -0,0 +1,107 @@ +import pytest +import torch +import torch.nn.functional as F + +from lightllm.models.sensenova_loop.triton_kernel.mhc import hc_contract, hc_expand, hc_post, hc_pre_norm +from lightllm.models.sensenova_loop.triton_kernel.mhc.pre_norm import has_deepgemm_mhc +from lightllm.common.basemodel.triton_kernel.norm.rmsnorm import rmsnorm_forward + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +def _route_reference(x, weight, bias, streams, iterations): + # Training normalizes the residual before the FP32 router projection. + flat = x.float() + logits = F.linear(flat * torch.rsqrt(flat.square().mean(-1, keepdim=True) + 1e-6), weight, bias) + pre = logits[:, :streams].sigmoid() + post = 2 * logits[:, streams : 2 * streams].sigmoid() + logits = logits[:, 2 * streams :].reshape(-1, streams, streams) + for _ in range(iterations): + logits = logits - torch.logsumexp(logits, dim=-1, keepdim=True) + logits = logits - torch.logsumexp(logits, dim=-2, keepdim=True) + mix = logits.exp() + branch = (pre.to(x.dtype).unsqueeze(1) @ x.view(-1, streams, x.shape[1] // streams)).squeeze(1) + return branch, mix, post + + +@pytest.mark.parametrize("tokens", [1, 19, 257]) +@pytest.mark.parametrize("streams", [2, 4]) +def test_loop_mhc_matches_training_routing_and_rounding(tokens, streams): + torch.manual_seed(35) + hidden = 4096 + x = torch.randn(tokens, streams * hidden, device="cuda", dtype=torch.bfloat16) + weight = torch.randn(2 * streams + streams * streams, streams * hidden, device="cuda") * 0.008 + bias = torch.randn(weight.shape[0], device="cuda") + norm = torch.randn(hidden, device="cuda", dtype=torch.bfloat16) + expected = _route_reference(x, weight, bias, streams, 20) + actual = hc_pre_norm( + x, + weight, + bias, + norm, + streams, + 1e-6, + 1e-6, + 20, + use_deepgemm=False, + ) + torch.testing.assert_close(actual[0], rmsnorm_forward(expected[0], norm, 1e-6), atol=0.02, rtol=0.02) + torch.testing.assert_close(actual[1], expected[1], atol=2e-6, rtol=2e-5) + torch.testing.assert_close(actual[2], expected[2], atol=2e-6, rtol=2e-5) + layer_out = torch.randn(tokens, hidden, device="cuda", dtype=torch.bfloat16) + # Qwen uses output,input order; both matmul and the post product round to bf16. + expected_out = expected[1].to(x.dtype) @ x.view(tokens, streams, hidden) + expected_out += expected[2].to(x.dtype).unsqueeze(-1) * layer_out.unsqueeze(1) + actual_out = hc_post(layer_out, x, actual[1].transpose(1, 2), actual[2], streams) + torch.testing.assert_close(actual_out.view_as(expected_out), expected_out, atol=0.02, rtol=0.01) + + +@pytest.mark.skipif(not has_deepgemm_mhc(), reason="DeepGEMM mHC projection unavailable") +def test_loop_deepgemm_projection_matches_torch(): + x = torch.randn(19, 16384, device="cuda", dtype=torch.bfloat16) + weight = torch.randn(24, 16384, device="cuda") * 0.005 + bias = torch.randn(24, device="cuda") + norm = torch.randn(4096, device="cuda", dtype=torch.bfloat16) + actual = hc_pre_norm(x, weight, bias, norm, 4, 1e-6, 1e-6, 20) + expected = hc_pre_norm(x, weight, bias, norm, 4, 1e-6, 1e-6, 20, use_deepgemm=False) + for a, b in zip(actual, expected): + torch.testing.assert_close(a, b, atol=0.02, rtol=0.02) + + +@pytest.mark.parametrize("streams", [2, 4]) +def test_loop_mhc_decode_can_capture_and_replay_cuda_graph(streams): + x = torch.randn(3, streams * 4096, device="cuda", dtype=torch.bfloat16) + weight = torch.randn(2 * streams + streams * streams, streams * 4096, device="cuda") * 0.005 + bias = torch.randn(weight.shape[0], device="cuda") + norm = torch.ones(4096, device="cuda", dtype=torch.bfloat16) + + def forward(): + branch, mix, post = hc_pre_norm( + x, + weight, + bias, + norm, + streams, + 1e-6, + 1e-6, + 20, + use_deepgemm=False, + ) + return hc_post(branch, x, mix.transpose(1, 2), post, streams) + + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + forward() + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = forward() + x.normal_() + expected = forward() + graph.replay() + torch.testing.assert_close(output, expected, atol=0, rtol=0) + collapsed = torch.randn(3, 4096, device="cuda", dtype=torch.bfloat16) + assert torch.equal(hc_contract(hc_expand(collapsed, streams), streams), collapsed) diff --git a/unit_tests/models/test_qwen3_dspark_model_output.py b/unit_tests/models/test_qwen3_dspark_model_output.py new file mode 100644 index 0000000000..8a1187a7f5 --- /dev/null +++ b/unit_tests/models/test_qwen3_dspark_model_output.py @@ -0,0 +1,517 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.basemodel import batch_objs +from lightllm.common.basemodel.batch_objs import ModelMtpOutputCollector, ModelOutput +from lightllm.models.qwen3_5_dflash.model import Qwen3_5DFlashModel +from lightllm.models.qwen3_5_dspark.model import Qwen3_5DSparkModel +from lightllm.models.qwen3_dflash.layer_infer import transformer_layer_infer as qwen3_dflash_layer_infer +from lightllm.models.qwen3_dflash.layer_infer.transformer_layer_infer import Qwen3DFlashTransformerLayerInfer +from lightllm.models.qwen3_dflash.model import Qwen3DFlashModel +from lightllm.models.qwen3_dspark.layer_weights import pre_and_post_layer_weight as dspark_pre_post_weight +from lightllm.models.qwen3_dspark.layer_infer.post_layer_infer import Qwen3DSparkPostLayerInfer +from lightllm.models.qwen3_dspark.model import Qwen3DSparkModel +from lightllm.models.llama.model import LlamaTpPartModel + + +@pytest.mark.parametrize("mtp_step", [1, 5, 7]) +@pytest.mark.parametrize("model_class", [Qwen3DFlashModel, Qwen3DSparkModel, Qwen3_5DSparkModel]) +def test_parallel_block_runtime_width_follows_mtp_step(monkeypatch, mtp_step, model_class): + monkeypatch.setattr(LlamaTpPartModel, "_verify_params", lambda self: None) + model = model_class.__new__(model_class) + model.args = SimpleNamespace(mtp_mode="dspark", mtp_step=mtp_step) + model.config = {"block_size": 7} + model.enable_tpsp_mix_mode = False + + model._verify_params() + + assert model.config["block_size"] == mtp_step + + +@pytest.mark.parametrize("model_class", [Qwen3DFlashModel, Qwen3DSparkModel, Qwen3_5DSparkModel]) +def test_parallel_block_rejects_mtp_step_above_checkpoint_capacity(monkeypatch, model_class): + monkeypatch.setattr(LlamaTpPartModel, "_verify_params", lambda self: None) + model = model_class.__new__(model_class) + model.args = SimpleNamespace(mtp_mode="dflash", mtp_step=8) + model.config = {"block_size": 7} + model.enable_tpsp_mix_mode = False + + with pytest.raises(AssertionError): + model._verify_params() + + +@pytest.mark.parametrize("model_class", [Qwen3DSparkModel, Qwen3_5DSparkModel]) +@pytest.mark.parametrize("mtp_step", [1, 6]) +def test_dspark_without_block_size_uses_mtp_step(model_class, mtp_step): + model = model_class.__new__(model_class) + model.args = SimpleNamespace(mtp_step=mtp_step) + model.config = {"num_attention_heads": 40, "num_key_value_heads": 8} + model.load_way = "HF" + model.tp_world_size_ = 4 + model.enable_tpsp_mix_mode = False + + model._verify_params() + + assert model.config["block_size"] == mtp_step + + +@pytest.mark.parametrize( + "model_class", + [Qwen3DFlashModel, Qwen3_5DFlashModel, Qwen3DSparkModel, Qwen3_5DSparkModel], +) +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_state_manager = object() + model.is_mtp_draft_model = True + model.mtp_previous_draft_models = [] + model.pre_post_weight = object() + + target_hiddens = torch.arange(6, dtype=torch.float32).view(2, 3) + mem_indexes = torch.tensor([7, 11]) + model._select_mem_indexes = lambda _: mem_indexes + observed_states = [] + + class PreInfer: + def context_forward(self, input_ids, infer_state, layer_weight): + assert input_ids is None + assert layer_weight is model.pre_post_weight + assert infer_state.mtp_draft_input_hiddens is target_hiddens + observed_states.append(infer_state) + return target_hiddens + 1 + + class TransformerLayer: + def __init__(self, increment): + self.increment = increment + + def context_forward(self, hidden, infer_state, layer_weight): + assert infer_state is observed_states[0] + assert layer_weight == self.increment + return hidden + self.increment + + model.pre_infer = PreInfer() + model.layers_infer = [TransformerLayer(2), TransformerLayer(3)] + model.trans_layers_weight = [2, 3] + model_input = SimpleNamespace( + batch_size=2, + b_seq_len=torch.tensor([3, 5]), + mtp_draft_input_hiddens=target_hiddens, + ) + + output = model._decode(model_input) + + assert isinstance(output, ModelOutput) + assert output.logits.shape == (2, 1) + infer_state = observed_states[0] + assert not hasattr(infer_state, "full_att_state") + 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]]) + assert torch.equal(infer_state.position_sin, model._sin_cached[[2, 4]]) + + +@pytest.mark.parametrize("model_class", [Qwen3DSparkModel, Qwen3_5DSparkModel]) +def test_dspark_decode_unpad_uses_common_output_and_slices_mtp_fields(model_class): + model = model_class.__new__(model_class) + output = ModelOutput( + logits=torch.arange(48).view(12, 4), + mtp_collector=ModelMtpOutputCollector( + spec_hidden=torch.arange(36).view(12, 3), + confidence_logits=torch.arange(12).view(3, 4), + draft_token_ids=torch.arange(12), + ), + ) + + unpadded = model._create_unpad_decode_model_output(output, origin_batch_size=8) + + assert isinstance(unpadded, ModelOutput) + assert unpadded.logits.shape == (8, 4) + assert unpadded.mtp_collector.spec_hidden.shape == (8, 3) + assert unpadded.mtp_collector.confidence_logits.shape == (2, 4) + assert unpadded.mtp_collector.draft_token_ids.shape == (8,) + assert output.logits.shape == (12, 4) + assert output.mtp_collector.confidence_logits.shape == (3, 4) + assert output.mtp_collector.draft_token_ids.shape == (12,) + + +def test_common_output_no_ref_conversion_includes_dspark_fields(monkeypatch): + monkeypatch.setattr(batch_objs, "tensor_to_no_ref_tensor", torch.clone) + output = ModelOutput( + logits=torch.ones((2, 4)), + mtp_collector=ModelMtpOutputCollector( + spec_hidden=torch.ones((2, 3)), + confidence_logits=torch.ones((1, 2)), + draft_token_ids=torch.ones((2,), dtype=torch.int64), + ), + ) + original_ptrs = ( + output.logits.data_ptr(), + output.mtp_collector.spec_hidden.data_ptr(), + output.mtp_collector.confidence_logits.data_ptr(), + output.mtp_collector.draft_token_ids.data_ptr(), + ) + + output.to_no_ref_tensor() + + converted_ptrs = ( + output.logits.data_ptr(), + output.mtp_collector.spec_hidden.data_ptr(), + output.mtp_collector.confidence_logits.data_ptr(), + output.mtp_collector.draft_token_ids.data_ptr(), + ) + assert all(converted != original for converted, original in zip(converted_ptrs, original_ptrs)) + + +def test_dspark_post_layer_publishes_head_results_through_collector(): + post_infer = Qwen3DSparkPostLayerInfer.__new__(Qwen3DSparkPostLayerInfer) + post_infer.block_size_ = 2 + post_infer.markov_rank_ = 2 + post_infer._slice_get_last_input = lambda input_embeddings, infer_state: ( + input_embeddings, + 4, + ) + post_infer._norm = lambda hidden, infer_state, layer_weight: hidden + post_infer._sample_markov = lambda *args, **kwargs: torch.tensor([[1, 2], [3, 4]]) + confidence_logits = torch.arange(4, dtype=torch.float32).view(2, 2) + post_infer.predict_confidence_logits = lambda *args, **kwargs: confidence_logits + + class RecordingCollector: + def add_mtp_outputs(self, **kwargs): + self.outputs = kwargs + + class LMHead: + def __call__(self, input, alloc_func): + return torch.empty((8, input.shape[1])) + + collector = RecordingCollector() + infer_state = SimpleNamespace( + is_prefill=False, + input_ids=torch.tensor([10, 0, 20, 0]), + hidden_collector=collector, + ) + layer_weight = SimpleNamespace(lm_head_weight_=LMHead()) + + output = post_infer.token_forward( + input_embdings=torch.randn(4, 3), + infer_state=infer_state, + layer_weight=layer_weight, + ) + + assert output.logits.shape == (4, 1) + assert output.logits_token_ids is None + assert torch.equal(collector.outputs["draft_token_ids"], torch.tensor([1, 2, 3, 4])) + assert collector.outputs["confidence_logits"] is confidence_logits + + +def test_dspark_non_markov_head_does_not_treat_candidate_columns_as_token_ids(): + post_infer = Qwen3DSparkPostLayerInfer.__new__(Qwen3DSparkPostLayerInfer) + post_infer.block_size_ = 2 + post_infer.markov_rank_ = 0 + post_infer._slice_get_last_input = lambda input_embeddings, infer_state: (input_embeddings, 4) + candidate_output = batch_objs.PostLayerOutput( + logits=torch.randn(4, 2), + logits_token_ids=torch.tensor([[10, 20], [30, 40], [50, 60], [70, 80]]), + ) + post_infer._draft_lm_head_and_gather = lambda *args: candidate_output + observed = {} + + def predict_confidence_logits(*args, sampled_tokens, **kwargs): + observed["sampled_tokens"] = sampled_tokens + return None + + post_infer.predict_confidence_logits = predict_confidence_logits + + class RecordingCollector: + def add_mtp_outputs(self, **kwargs): + self.outputs = kwargs + + collector = RecordingCollector() + infer_state = SimpleNamespace( + is_prefill=False, + input_ids=torch.tensor([10, 0, 20, 0]), + hidden_collector=collector, + ) + + output = post_infer.token_forward(torch.randn(4, 3), infer_state, object()) + + assert output is candidate_output + assert observed["sampled_tokens"] is None + assert collector.outputs["draft_token_ids"] is None + + +@pytest.mark.parametrize( + "head_config", + [ + {"enable_confidence_head": True}, + {"markov_rank": 4, "markov_head_type": "gated"}, + {"markov_rank": 4, "markov_head_type": "rnn"}, + ], +) +def test_biased_dspark_heads_do_not_inherit_model_quantization(monkeypatch, head_config): + captured_kwargs = [] + + def init_base_weight(self, data_type, network_config, quant_cfg): + self.data_type_ = data_type + self.quant_cfg = quant_cfg + + class RecordingROWMMWeight: + def __init__(self, **kwargs): + captured_kwargs.append(kwargs) + + class StubWeight: + def __init__(self, **kwargs): + pass + + monkeypatch.setattr( + dspark_pre_post_weight.Qwen3DFlashPreAndPostLayerWeight, + "__init__", + init_base_weight, + ) + monkeypatch.setattr(dspark_pre_post_weight, "EmbeddingWeight", StubWeight) + monkeypatch.setattr(dspark_pre_post_weight, "LMHeadWeight", StubWeight) + monkeypatch.setattr(dspark_pre_post_weight, "ROWMMWeight", RecordingROWMMWeight) + + quant_method = object() + quant_cfg = SimpleNamespace(get_quant_method=lambda *_: quant_method) + network_config = { + "hidden_size": 16, + "vocab_size": 32, + **head_config, + } + dspark_pre_post_weight.Qwen3DSparkPreAndPostLayerWeight( + data_type=torch.bfloat16, + network_config=network_config, + quant_cfg=quant_cfg, + ) + + assert captured_kwargs + assert all(kwargs["quant_method"] is None for kwargs in captured_kwargs) + + +def test_fixed_dspark_does_not_require_confidence_head(monkeypatch): + monkeypatch.setattr(Qwen3DFlashModel, "_verify_params", lambda self: None) + model = Qwen3DSparkModel.__new__(Qwen3DSparkModel) + model.args = SimpleNamespace(mtp_step=6) + model.config = {"enable_confidence_head": False} + + model._verify_params() + + +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): + rotary_factors = [] + monkeypatch.setattr(qwen3_dflash_layer_infer, "qk_rmsnorm_forward", lambda *args, **kwargs: None) + monkeypatch.setattr( + qwen3_dflash_layer_infer, + "rotary_emb_fwd", + lambda *args, **kwargs: rotary_factors.append(kwargs["partial_rotary_factor"]), + ) + + class Projection: + def __init__(self, width): + self.width = width + + def mm(self, input, **kwargs): + return input.new_zeros((input.shape[0], self.width)) + + class QKNorm: + k_weight = object() + + def __call__(self, *args, **kwargs): + return None + + layer = Qwen3DFlashTransformerLayerInfer.__new__(Qwen3DFlashTransformerLayerInfer) + layer.tp_q_head_num_ = 1 + layer.tp_k_head_num_ = 1 + layer.tp_v_head_num_ = 1 + layer.head_dim_ = 8 + layer.eps_ = 1e-6 + layer.partial_rotary_factor = 0.25 + layer._post_cache_kv = lambda *args, **kwargs: None + infer_state = SimpleNamespace(position_cos=object(), position_sin=object()) + layer_weight = SimpleNamespace( + q_proj=Projection(8), + kv_proj=Projection(16), + qk_norm_weight_=QKNorm(), + ) + inputs = torch.zeros((2, 4)) + + layer.context_forward(inputs, infer_state, layer_weight) + layer._get_qkv(inputs, infer_state, layer_weight) + + assert rotary_factors == [0.25, 0.25] + + +def _cuda_empty(*size, **kwargs): + kwargs.setdefault("device", "cuda") + return torch.empty(*size, **kwargs) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_vanilla_markov_local_sampling_matches_full_logits(): + post_infer = Qwen3DSparkPostLayerInfer.__new__(Qwen3DSparkPostLayerInfer) + post_infer.block_size_ = 3 + post_infer.markov_rank_ = 2 + post_infer.markov_head_type_ = "vanilla" + post_infer.tp_world_size_ = 1 + + vocab_size = 5 + request_count = 2 + local_logits = torch.randn(vocab_size, request_count * post_infer.block_size_, device="cuda") + + class MarkovEmbedding: + def __init__(self, weight): + self.weight = weight + + def __call__(self, input_ids, alloc_func): + return torch.nn.functional.embedding(input_ids, self.weight) + + class MarkovLMHead: + def __init__(self, weight): + self.weight = weight + self.tp_vocab_start_id = 0 + + def __call__(self, input, alloc_func): + return self.weight @ input + + markov_w1 = torch.randn(vocab_size, post_infer.markov_rank_, device="cuda") + markov_w2 = torch.randn(vocab_size, post_infer.markov_rank_, device="cuda") + layer_weight = SimpleNamespace( + markov_w1_weight_=MarkovEmbedding(markov_w1), + markov_w2_weight_=MarkovLMHead(markov_w2), + ) + anchor_token_ids = torch.tensor([1, 3], device="cuda") + block_hidden = torch.empty(request_count, post_infer.block_size_, 0, device="cuda") + post_infer.alloc_tensor = _cuda_empty + + sampled_tokens = post_infer._sample_markov( + local_logits=local_logits, + block_hidden=block_hidden, + infer_state=SimpleNamespace(), + anchor_token_ids=anchor_token_ids, + layer_weight=layer_weight, + ) + + base_logits = local_logits.T.reshape(request_count, post_infer.block_size_, vocab_size) + prev_token_ids = anchor_token_ids + expected_tokens = [] + for step_idx in range(post_infer.block_size_): + markov_bias = torch.nn.functional.linear(markov_w1[prev_token_ids], markov_w2) + prev_token_ids = torch.argmax(base_logits[:, step_idx] + markov_bias, dim=-1) + expected_tokens.append(prev_token_ids) + expected_tokens = torch.stack(expected_tokens, dim=1) + + torch.testing.assert_close(sampled_tokens, expected_tokens) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_vanilla_markov_tp4_sampling_matches_full_logits(monkeypatch): + torch.manual_seed(0) + post_infer = Qwen3DSparkPostLayerInfer.__new__(Qwen3DSparkPostLayerInfer) + post_infer.block_size_ = 3 + post_infer.markov_rank_ = 4 + post_infer.markov_head_type_ = "vanilla" + post_infer.tp_world_size_ = 4 + post_infer.alloc_tensor = _cuda_empty + + vocab_size = 11 + request_count = 2 + split_indexes = torch.linspace(0, vocab_size, post_infer.tp_world_size_ + 1, dtype=torch.int64) + tp_rank = 2 + local_start = int(split_indexes[tp_rank]) + local_end = int(split_indexes[tp_rank + 1]) + + base_logits = torch.randn(request_count, post_infer.block_size_, vocab_size, device="cuda") + markov_w1 = torch.randn(vocab_size, post_infer.markov_rank_, device="cuda") + markov_w2 = torch.randn(vocab_size, post_infer.markov_rank_, device="cuda") + anchor_token_ids = torch.tensor([1, 7], device="cuda") + + class MarkovEmbedding: + weight = markov_w1 + + class MarkovLMHead: + weight = markov_w2[local_start:local_end] + tp_vocab_start_id = local_start + + layer_weight = SimpleNamespace( + markov_w1_weight_=MarkovEmbedding(), + markov_w2_weight_=MarkovLMHead(), + ) + local_logits = base_logits[:, :, local_start:local_end].reshape(-1, local_end - local_start).T.contiguous() + + expected_tokens = [] + prev_token_ids = anchor_token_ids + for step_idx in range(post_infer.block_size_): + scores = base_logits[:, step_idx] + torch.nn.functional.linear(markov_w1[prev_token_ids], markov_w2) + prev_token_ids = torch.argmax(scores, dim=-1) + expected_tokens.append(prev_token_ids) + expected_tokens = torch.stack(expected_tokens, dim=1) + + step = 0 + + def gather_tp_winners(output, local_winners, group, async_op): + nonlocal step + prev_tokens = anchor_token_ids if step == 0 else expected_tokens[:, step - 1] + scores = base_logits[:, step] + torch.nn.functional.linear(markov_w1[prev_tokens], markov_w2) + winners = [] + for rank in range(post_infer.tp_world_size_): + start = int(split_indexes[rank]) + end = int(split_indexes[rank + 1]) + values, indexes = scores[:, start:end].max(dim=-1) + winners.append(torch.stack((values, (indexes + start).float()), dim=-1)) + torch.testing.assert_close(local_winners, winners[tp_rank]) + output.copy_(torch.stack(winners).reshape_as(output)) + step += 1 + + monkeypatch.setattr( + "lightllm.models.qwen3_dspark.layer_infer.post_layer_infer.all_gather_into_tensor", + gather_tp_winners, + ) + + sampled_tokens = post_infer._sample_markov( + local_logits=local_logits, + block_hidden=torch.empty(request_count, post_infer.block_size_, 0, device="cuda"), + infer_state=SimpleNamespace(dist_group=None), + anchor_token_ids=anchor_token_ids, + layer_weight=layer_weight, + ) + + torch.testing.assert_close(sampled_tokens, expected_tokens) diff --git a/unit_tests/models/test_qwen3_eagle_model_input.py b/unit_tests/models/test_qwen3_eagle_model_input.py new file mode 100644 index 0000000000..1395153351 --- /dev/null +++ b/unit_tests/models/test_qwen3_eagle_model_input.py @@ -0,0 +1,100 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaTransformerLayerInfer +from lightllm.models.llama.model import LlamaTpPartModel +from lightllm.models.qwen3_eagle.layer_infer.pre_layer_infer import Qwen3EaglePreLayerInfer +from lightllm.models.qwen3_eagle.layer_infer.transformer_layer_infer import Qwen3EagleTransformerLayerInfer +from lightllm.models.qwen3_eagle.model import Qwen3EagleModel + + +def test_qwen3_eagle_uses_configured_attention_head_dim(monkeypatch): + def init_llama_layer(self, layer_num, network_config): + self.head_dim_ = network_config["hidden_size"] // network_config["num_attention_heads"] + + monkeypatch.setattr(LlamaTransformerLayerInfer, "__init__", init_llama_layer) + layer = Qwen3EagleTransformerLayerInfer( + layer_num=0, + network_config={ + "hidden_size": 2560, + "num_attention_heads": 32, + "head_dim": 128, + }, + ) + + assert layer.head_dim_ == 128 + + +def test_qwen3_eagle_projects_concatenated_target_hiddens_before_inferstate(monkeypatch): + parent_calls = [] + monkeypatch.setattr( + LlamaTpPartModel, + "_create_inferstate", + lambda self, model_input, microbatch_index=0: parent_calls.append((model_input, microbatch_index)) + or SimpleNamespace(model_input=model_input), + ) + projected_hiddens = torch.empty((3, 4)) + projection_inputs = [] + model = Qwen3EagleModel.__new__(Qwen3EagleModel) + model.config = {"hidden_size": 4} + model.pre_post_weight = SimpleNamespace( + fc_weight_=SimpleNamespace( + mm=lambda hiddens: projection_inputs.append(hiddens) or projected_hiddens, + ) + ) + target_hiddens = torch.empty((3, 12)) + model_input = SimpleNamespace( + input_ids=torch.arange(3), + mtp_draft_input_hiddens=target_hiddens, + ) + + infer_state = model._create_inferstate(model_input=model_input, microbatch_index=1) + + assert projection_inputs == [target_hiddens] + assert len(parent_calls) == 1 + normalized_input, microbatch_index = parent_calls[0] + assert normalized_input is not model_input + assert normalized_input.mtp_draft_input_hiddens is projected_hiddens + assert model_input.mtp_draft_input_hiddens is target_hiddens + assert microbatch_index == 1 + assert infer_state.model_input is normalized_input + + +def test_qwen3_eagle_keeps_recursive_draft_hiddens_without_projection(monkeypatch): + monkeypatch.setattr( + LlamaTpPartModel, + "_create_inferstate", + lambda self, model_input, microbatch_index=0: model_input, + ) + model = Qwen3EagleModel.__new__(Qwen3EagleModel) + model.config = {"hidden_size": 4} + model.pre_post_weight = SimpleNamespace( + fc_weight_=SimpleNamespace( + mm=lambda _: pytest.fail("fixed-width recursive hidden must not be projected"), + ) + ) + model_input = SimpleNamespace( + input_ids=torch.arange(3), + mtp_draft_input_hiddens=torch.empty((3, 4)), + ) + + normalized_input = model._create_inferstate(model_input=model_input) + + assert normalized_input is model_input + + +def test_qwen3_eagle_pre_layer_rejects_non_normalized_hidden_width(): + pre_layer = Qwen3EaglePreLayerInfer.__new__(Qwen3EaglePreLayerInfer) + pre_layer.hidden_size_ = 4 + infer_state = SimpleNamespace(mtp_draft_input_hiddens=torch.empty((2, 12))) + + with pytest.raises(AssertionError): + pre_layer.prepare_spec_draft_hiddens(infer_state) + + normalized_hiddens = torch.empty((2, 4)) + infer_state.mtp_draft_input_hiddens = normalized_hiddens + pre_layer.prepare_spec_draft_hiddens(infer_state) + + assert infer_state.eagle_draft_hidden_states is normalized_hiddens diff --git a/unit_tests/models/test_yarn_rotary.py b/unit_tests/models/test_yarn_rotary.py new file mode 100644 index 0000000000..ef791c20d1 --- /dev/null +++ b/unit_tests/models/test_yarn_rotary.py @@ -0,0 +1,142 @@ +import json + +import pytest +import torch +from transformers.modeling_rope_utils import _compute_yarn_parameters +from transformers.models.qwen3_5.configuration_qwen3_5 import Qwen3_5TextConfig + +from lightllm.models.qwen2_vl.triton_kernel.mrope import mrope_triton_fused +from lightllm.models.qwen3_5.model import Qwen3_5TpPartModel + + +@pytest.fixture +def qwen35_model(tmp_path): + # Qwen3.5-27B with the documented 4x YaRN extension. + text_config = { + "hidden_size": 5120, + "head_dim": 256, + "num_attention_heads": 24, + "num_key_value_heads": 4, + "num_hidden_layers": 64, + "max_position_embeddings": 262144, + "rope_parameters": { + "mrope_interleaved": True, + "mrope_section": [11, 11, 10], + "rope_type": "yarn", + "rope_theta": 10000000, + "partial_rotary_factor": 0.25, + "factor": 4.0, + "original_max_position_embeddings": 262144, + }, + } + (tmp_path / "config.json").write_text(json.dumps({"text_config": text_config, "vision_config": {}})) + model = Qwen3_5TpPartModel.__new__(Qwen3_5TpPartModel) + model.weight_dir_ = str(tmp_path) + model.finetune_config = None + model.tp_world_size_ = 1 + model._init_config() + model.head_dim_ = model.config["head_dim"] + model.max_seq_length = 262148 + model.data_type = torch.float32 + return model + + +@pytest.mark.parametrize("type_key", ["rope_type", "type"]) +@pytest.mark.parametrize("rope_type,expected", [("default", "default"), ("mrope", "default"), ("yarn", "yarn")]) +def test_mrope_respects_explicit_scaling_type(qwen35_model, monkeypatch, type_key, rope_type, expected): + rope_scaling = qwen35_model.config["rope_scaling"] + rope_scaling.pop("rope_type") + rope_scaling[type_key] = rope_type + selected = [] + monkeypatch.setattr(qwen35_model, "_init_to_get_rotary", lambda: selected.append("default")) + monkeypatch.setattr(qwen35_model, "_init_to_get_yarn_rotary", lambda: selected.append("yarn")) + + qwen35_model._init_custom() + + assert selected == [expected] + + +def _reference_yarn_cache(model, position_ids): + config = Qwen3_5TextConfig( + hidden_size=model.config["hidden_size"], + head_dim=model.head_dim_, + num_attention_heads=model.config["num_attention_heads"], + max_position_embeddings=model.config["max_position_embeddings"], + rope_parameters=model.config["rope_parameters"].copy(), + partial_rotary_factor=model.config["partial_rotary_factor"], + ) + inv_freq, attention_factor = _compute_yarn_parameters(config, position_ids.device) + freqs = position_ids.float().unsqueeze(-1) * inv_freq + return (freqs.cos() * attention_factor).to(model.data_type), (freqs.sin() * attention_factor).to(model.data_type) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for rotary caches") +@pytest.mark.parametrize("data_type", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("partial_rotary_factor", [0.25, 1.0]) +@pytest.mark.parametrize("factor", [1.0, 4.0]) +def test_yarn_cache_matches_transformers(qwen35_model, data_type, partial_rotary_factor, factor): + model = qwen35_model + model.data_type = data_type + model.config["partial_rotary_factor"] = partial_rotary_factor + model.config["rope_scaling"]["partial_rotary_factor"] = partial_rotary_factor + model.config["rope_scaling"]["factor"] = factor + + model._init_custom() + + half_rotary_dim = int(model.head_dim_ * partial_rotary_factor) // 2 + assert model._cos_cached.shape == (model.max_seq_length, half_rotary_dim) + assert model._sin_cached.shape == model._cos_cached.shape + # Include both sides of the original context boundary. + positions = torch.tensor([0, 1, 127, 8191, 262143, 262144, 262147], device="cuda") + expected_cos, expected_sin = _reference_yarn_cache(model, positions) + torch.testing.assert_close(model._cos_cached[positions], expected_cos, rtol=1e-5, atol=1e-6) + torch.testing.assert_close(model._sin_cached[positions], expected_sin, rtol=1e-5, atol=1e-6) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for the MRoPE kernel") +@pytest.mark.parametrize("data_type", [torch.float32, torch.bfloat16]) +def test_qwen35_yarn_interleaved_mrope_matches_reference(qwen35_model, data_type): + model = qwen35_model + model.data_type = data_type + model._init_custom() + position_ids = torch.tensor( + [[0, 1, 127, 8191, 262143, 262147], [0, 7, 83, 8189, 262144, 262145], [0, 5, 61, 8190, 262145, 262146]], + device="cuda", + ) + cos, sin = _reference_yarn_cache(model, position_ids) + half_rotary_dim = cos.shape[-1] + rotary_dim = 2 * half_rotary_dim + channels = torch.arange(half_rotary_dim, device="cuda") + # [11, 11, 10] assigns the 32 frequencies to repeating T/H/W axes. + cos = cos[channels % 3, :, channels].T.unsqueeze(1).float() + sin = sin[channels % 3, :, channels].T.unsqueeze(1).float() + cos = torch.cat((cos, cos), dim=-1) + sin = torch.cat((sin, sin), dim=-1) + + def rotate_reference(x): + result = x.clone() + rotary = x[..., :rotary_dim].float() + rotated_half = torch.cat((-rotary[..., half_rotary_dim:], rotary[..., :half_rotary_dim]), dim=-1) + result[..., :rotary_dim] = (rotary * cos + rotated_half * sin).to(data_type) + return result + + torch.manual_seed(0) + q = torch.randn((position_ids.shape[1], 24, 256), dtype=data_type, device="cuda") + k = torch.randn((position_ids.shape[1], 4, 256), dtype=data_type, device="cuda") + expected_q, expected_k = rotate_reference(q), rotate_reference(k) + + mrope_triton_fused( + q, + k, + model._cos_cached[position_ids], + model._sin_cached[position_ids], + torch.tensor(model.config["rope_scaling"]["mrope_section"], dtype=torch.int32, device="cuda"), + is_interleaved=True, + partial_rotary_factor=model.config["partial_rotary_factor"], + ) + + tolerance = 1e-5 if data_type == torch.float32 else 2e-2 + torch.testing.assert_close(q, expected_q, rtol=tolerance, atol=tolerance) + torch.testing.assert_close(k, expected_k, rtol=tolerance, atol=tolerance) + torch.testing.assert_close(q[..., rotary_dim:], expected_q[..., rotary_dim:], rtol=0, atol=0) + torch.testing.assert_close(k[..., rotary_dim:], expected_k[..., rotary_dim:], rtol=0, atol=0) diff --git a/unit_tests/models/vit/test_flash_attention_forward.py b/unit_tests/models/vit/test_flash_attention_forward.py new file mode 100644 index 0000000000..f736ec4cd6 --- /dev/null +++ b/unit_tests/models/vit/test_flash_attention_forward.py @@ -0,0 +1,139 @@ +import torch +import math +import time +import pytest +from lightllm.models.vit.triton_kernel.flashattention_nopad import flash_attention_fwd +from lightllm.models.vit.triton_kernel import flashattention_nopad +from lightllm.common.basemodel.attention_vit.fa3.fp import Fa3VitAttBackend + + +_fa3_mark = pytest.mark.skipif( + not flashattention_nopad._flash_attn_v3_available + or not torch.cuda.is_available() + or torch.cuda.get_device_capability()[0] < 9, + reason="FA3 requires sgl_kernel and a Hopper GPU", +) + + +def test_fa3_failure_computes_fallback_on_same_invocation(monkeypatch): + calls = [] + q = torch.zeros(2, 1, 4) + out = torch.full_like(q, -123) + cu = torch.tensor([0, 2], dtype=torch.int32) + + def fail_fa3(q, k, v, o, cu_seqlens, max_seqlen): + calls.append("fa3") + o[0].fill_(99) # A backend may fail after partially writing output. + raise RuntimeError("shape-dependent FA3 failure") + + def fallback(q, k, v, o, cu_seqlens, max_seqlen): + calls.append("triton") + o.fill_(7) + + monkeypatch.setattr(flashattention_nopad, "_flash_attn_v3_available", True) + monkeypatch.setattr(flashattention_nopad, "is_hopper", lambda: True) + monkeypatch.setattr(flashattention_nopad, "flash_attention_v3_fwd", fail_fa3, raising=False) + monkeypatch.setattr(flashattention_nopad, "_flash_attention_triton_fwd", fallback) + flash_attention_fwd(q, q, q, out, cu, 2) + assert calls == ["fa3", "triton"] + assert torch.equal(out, torch.full_like(q, 7)) + assert not flashattention_nopad._flash_attn_v3_available + + # Later calls should use the disabled-backend path without retrying FA3. + out.fill_(-123) + flash_attention_fwd(q, q, q, out, cu, 2) + assert calls == ["fa3", "triton", "triton"] + assert torch.equal(out, torch.full_like(q, 7)) + + +def test_fa3_success_does_not_run_fallback(monkeypatch): + q = torch.zeros(2, 1, 4) + out = torch.full_like(q, -123) + cu = torch.tensor([0, 2], dtype=torch.int32) + + def fa3(q, k, v, o, cu_seqlens, max_seqlen): + o.fill_(9) + + monkeypatch.setattr(flashattention_nopad, "_flash_attn_v3_available", True) + monkeypatch.setattr(flashattention_nopad, "is_hopper", lambda: True) + monkeypatch.setattr(flashattention_nopad, "flash_attention_v3_fwd", fa3, raising=False) + monkeypatch.setattr( + flashattention_nopad, + "_flash_attention_triton_fwd", + lambda *args: pytest.fail("a successful FA3 call must not run fallback"), + ) + flash_attention_fwd(q, q, q, out, cu, 2) + assert torch.equal(out, torch.full_like(q, 9)) + assert flashattention_nopad._flash_attn_v3_available + + +def reference_attention_varlen(q, k, v, cu): + """ + q, k, v : (total_len, n_head, D) + cu_seqlen : prefix sums (batch+1,) + """ + total, n_head, d = q.shape + out = torch.empty_like(q) + scale = 1.0 / math.sqrt(d) + + for b in range(cu.numel() - 1): + s, e = cu[b].item(), cu[b + 1].item() + q_b, k_b, v_b = q[s:e], k[s:e], v[s:e] # (seq, head, D) + + q_hsd = q_b.permute(1, 0, 2) # (head, seq, D) + k_hds = k_b.permute(1, 2, 0) # (head, D, seq) + v_hsd = v_b.permute(1, 0, 2) # (head, seq, D) + + scores = torch.matmul(q_hsd, k_hds) * scale # (head, seq, seq) + probs = torch.softmax(scores.float(), dim=-1) + + out_hsd = torch.matmul(probs, v_hsd.float()) # (head, seq, D) + out[s:e] = out_hsd.permute(1, 0, 2).to(q.dtype) # back to (seq, head, D) + + return out + + +@pytest.mark.parametrize("dtype,atol", [(torch.float16, 1e-2), (torch.bfloat16, 2e-2)]) +@pytest.mark.parametrize( + "attention_fn", + [ + flash_attention_fwd, + pytest.param( + getattr(flashattention_nopad, "flash_attention_v3_fwd", None), + marks=_fa3_mark, + ), + pytest.param( + Fa3VitAttBackend._vit_att_fwd, + marks=_fa3_mark, + ), + ], +) +def test_varlen(dtype, atol, attention_fn, batch=4, heads=8, d=80, device="cuda:0"): + torch.manual_seed(0) + lengths = torch.randint(1, 257, (batch,)) + max_len = int(lengths.max().item()) + + cu = torch.zeros(batch + 1, dtype=torch.int32, device=device) + cu[1:] = torch.cumsum(lengths, 0) + tot = int(cu[-1]) + + q = torch.randn(tot, heads, d, dtype=dtype, device=device) + k = torch.randn_like(q) + v = torch.randn_like(q) + out_tri = torch.randn_like(q) + attention_fn(q, k, v, out_tri, cu, max_len) + a = time.time() + for _ in range(100): + attention_fn(q, k, v, out_tri, cu, max_len) + b = time.time() + print(f"flash_attention_fwd time: {(b - a) / 100 * 1000:.2f} ms") + out_ref = reference_attention_varlen(q, k, v, cu) + + max_err = (out_ref - out_tri).abs().max().item() + mean_err = (out_ref - out_tri).abs().mean().item() + print(f"{dtype}: max {max_err:.6f}, mean {mean_err:.6f}") + torch.testing.assert_close(out_tri, out_ref, atol=atol, rtol=0) + + +if __name__ == "__main__": + pytest.main() diff --git a/unit_tests/server/core/objs/test_out_token_circlequeue.py b/unit_tests/server/core/objs/test_out_token_circlequeue.py new file mode 100644 index 0000000000..80b91b7b5d --- /dev/null +++ b/unit_tests/server/core/objs/test_out_token_circlequeue.py @@ -0,0 +1,127 @@ +import pytest + +from lightllm.server.core.objs.out_token_circlequeue import CircularQueue, LIGHTLLM_OUT_TOKEN_QUEUE_SIZE + + +def test_queue_initialization(): + queue = CircularQueue() + assert queue.head == 0 + assert queue.tail == 0 + assert len(queue) == 0 + assert queue.is_empty() is True + assert queue.is_full() is False + + +def test_queue_push_and_pop(): + queue = CircularQueue() + + # Push an item + queue.push("token1", 0, False, 1) + assert len(queue) == 1 + assert queue.is_empty() is False + assert queue.is_full() is False + + # Pop the item + token, src_index, special, out_token_count = queue.pop() + assert token == "token1" + assert src_index == 0 + assert len(queue) == 0 + assert special is False + assert out_token_count == 1 + + +def test_queue_push_multiple_items(): + queue = CircularQueue() + + # Push multiple items + queue.push("token1", 0, False, 1) + queue.push("token2", 1, False, 2) + queue.push("token3", 2, False, 3) + + assert len(queue) == 3 + assert queue.is_empty() is False + if len(queue) == LIGHTLLM_OUT_TOKEN_QUEUE_SIZE - 1: + assert queue.is_full() is True + + # Pop items + token, src_index, special, out_token_count = queue.pop() + assert token == "token1" + assert src_index == 0 + assert special is False + assert out_token_count == 1 + + token, src_index, special, out_token_count = queue.pop() + assert token == "token2" + assert src_index == 1 + assert special is False + assert out_token_count == 2 + + token, src_index, special, out_token_count = queue.pop() + assert token == "token3" + assert src_index == 2 + assert special is False + assert out_token_count == 3 + + assert len(queue) == 0 + + +def test_queue_full_condition(): + queue = CircularQueue() + + # Fill the queue + for i in range(LIGHTLLM_OUT_TOKEN_QUEUE_SIZE - 1): + queue.push(f"token{i}", i, False, 1) + + assert queue.is_full() is True + + # Try to push another item (should raise an exception) + with pytest.raises(Exception, match="Queue is full"): + queue.push("token_overflow", LIGHTLLM_OUT_TOKEN_QUEUE_SIZE, False, 1) + + +def test_queue_pop_empty_condition(): + queue = CircularQueue() + + # Attempt to pop from an empty queue (should raise an exception) + with pytest.raises(Exception, match="Queue is empty"): + queue.pop() + + +def test_queue_wrap_around(): + queue = CircularQueue() + + # Push enough items to wrap around + for i in range(LIGHTLLM_OUT_TOKEN_QUEUE_SIZE - 1): + queue.push(f"token{i}", i, False, i) + + # Pop all items + for i in range(LIGHTLLM_OUT_TOKEN_QUEUE_SIZE - 1): + token, src_index, _, _ = queue.pop() + assert token == f"token{i}" + assert src_index == i + + # Now the queue should be empty + assert queue.is_empty() is True + + # Push again + queue.push("token1", 0, False, 1) + queue.push("token2", 1, False, 2) + + assert len(queue) == 2 + assert queue.is_empty() is False + + # Pop the items + token, src_index, _, _ = queue.pop() + assert token == "token1" + assert src_index == 0 + + token, src_index, _, _ = queue.pop() + assert token == "token2" + assert src_index == 1 + + assert len(queue) == 0 + + +# Run with pytest +if __name__ == "__main__": + pytest.main() diff --git a/unit_tests/server/core/objs/test_req.py b/unit_tests/server/core/objs/test_req.py new file mode 100644 index 0000000000..5966497ef1 --- /dev/null +++ b/unit_tests/server/core/objs/test_req.py @@ -0,0 +1,127 @@ +import pytest +import easydict +from types import SimpleNamespace +from lightllm.server.core.objs.req import Req, ChunkedPrefillReq, SamplingParams +from lightllm.server.core.objs.token_metadata import ReqFinalTokenMetadata +from lightllm.utils import shm_utils +from lightllm.utils.envs_utils import get_env_start_args, set_env_start_args + + +@pytest.fixture(scope="module", autouse=True) +def setup_module_env(): + with pytest.MonkeyPatch.context() as monkeypatch: + monkeypatch.setattr(shm_utils, "get_unique_server_name", lambda: "test_req_service_0") + monkeypatch.setenv("LIGHTLLM_START_ARGS", "{}") + set_env_start_args( + easydict.EasyDict( + { + "mtp_step": 0, + "llm_prefill_att_backend": ["None"], + "llm_decode_att_backend": ["None"], + "cpu_cache_token_page_size": 256, + "enable_cpu_cache": False, + "model_dir": "", + "page_size": 4, + } + ) + ) + get_env_start_args.cache_clear() + yield + get_env_start_args.cache_clear() + + +@pytest.fixture +def req(): + req_instance = Req() + req_instance.init(1, [1, 2, 3], {"max_new_tokens": 1}, None, chunked_prefill_size=128) + return req_instance + + +def test_req_init(req): + assert req.request_id == 1 + assert req.input_len == 3 + + +def test_req_init_validates_dictionary_params(): + with pytest.raises(ValueError, match="max_new_tokens must be at least 1"): + Req().init(1, [1], {"max_new_tokens": 0}, None) + + +def test_create_prompt_ids_shm_array(req): + assert hasattr(req, "shm_prompt_ids") + + +def test_detach_shm_arrays(req): + prompt_ids = req.shm_prompt_ids + logprobs = req.shm_logprobs + + req.detach_shm_arrays() + + assert prompt_ids.shm is None + assert logprobs.shm is None + assert not hasattr(req, "shm_prompt_ids") + assert not hasattr(req, "shm_logprobs") + + +def test_get_used_tokens(req): + req.shm_cur_kv_len = 5 + assert req.get_used_tokens() == 5 + + +def test_final_token_metadata_read_returns_actual_prompt_tokens(req): + req.sample_params.prompt_logprobs = 0 + req.shm_logprobs.arr["logprob"][1] = -0.5 + req.shm_logprobs.arr["logprob"][2] = -1.25 + req.shm_logprobs.arr["rank"][1] = 315 + req.shm_logprobs.arr["rank"][2] = 4 + + metadata = ReqFinalTokenMetadata(req).read() + + assert metadata["prompt_token_ids"] == [1, 2, 3] + assert metadata["prompt_logprobs"] == [ + None, + {2: {"logprob": -0.5, "rank": 315, "decoded_token": None}}, + {3: {"logprob": -1.25, "rank": 4, "decoded_token": None}}, + ] + + +def test_chunked_req_get_tuple_tokens_adds_page_and_async_reserve(): + req = SimpleNamespace( + input_len=10, + shm_cur_output_len=0, + shm_cur_kv_len=0, + sample_params=SimpleNamespace(ignore_eos=True, max_new_tokens=5), + ) + + assert ChunkedPrefillReq.get_tuple_tokens(req, False, 10) == (11, 26) + + +def test_finish_status(req): + req.finish_status.set_status(req.finish_status.FINISHED_STOP) + assert req.finish_status.is_finished() + assert not req.finish_status.is_error_finished() + assert req.finish_status.get_finish_reason() == "stop" + + req.finish_status.set_status(req.finish_status.FINISHED_LENGTH) + assert req.finish_status.is_finished() + assert not req.finish_status.is_error_finished() + + req.finish_status.set_status(req.finish_status.FINISHED_ABORTED) + assert req.finish_status.is_finished() + assert req.finish_status.is_error_finished() + + req.finish_status.set_status(req.finish_status.FINISHED_ERROR) + assert req.finish_status.is_finished() + assert req.finish_status.is_finished_error() + assert req.finish_status.is_error_finished() + assert req.finish_status.get_finish_reason() == "error" + + req.finish_status.set_status(req.finish_status.FINISHED_PD_DECODE_CAPACITY) + assert req.finish_status.is_finished() + assert req.finish_status.is_finished_pd_decode_capacity() + assert not req.finish_status.is_error_finished() + assert req.finish_status.get_finish_reason() == "length" + + +if __name__ == "__main__": + pytest.main() diff --git a/unit_tests/server/core/objs/test_shm_array.py b/unit_tests/server/core/objs/test_shm_array.py new file mode 100644 index 0000000000..af02c7de0b --- /dev/null +++ b/unit_tests/server/core/objs/test_shm_array.py @@ -0,0 +1,74 @@ +# test_shm_array.py + +import pytest +import numpy as np +from multiprocessing import shared_memory +from lightllm.server.core.objs.shm_array import ShmArray # Replace 'your_module' with the actual module name +from lightllm.utils import shm_utils + + +@pytest.fixture(scope="module", autouse=True) +def service_name(): + monkeypatch = pytest.MonkeyPatch() + monkeypatch.setattr(shm_utils, "get_unique_server_name", lambda: "test_shm_array_service_0") + yield + monkeypatch.undo() + + +@pytest.fixture(scope="module") +def shm_array(): + name = "test_shm_array" + shape = (10, 10) + dtype = np.float32 + shm = ShmArray(name, shape, dtype) + shm.create_shm() + yield shm + shm.close_shm() + + +def test_create_shm(shm_array): + assert shm_array.shm is not None + assert shm_array.arr.shape == shm_array.shape + assert shm_array.arr.dtype == shm_array.dtype + + +def test_link_shm(shm_array): + # Create a new instance to link the existing shm + shm_linked = ShmArray(shm_array.name, shm_array.shape, shm_array.dtype) + shm_linked.link_shm() + + assert shm_linked.shm is not None + assert shm_linked.arr.shape == shm_array.shape + assert shm_linked.arr.dtype == shm_array.dtype + + +def test_shm_size_mismatch(shm_array): + # Create a new shm with a different size + shape_mismatch = (5, 5) + shm_mismatch = ShmArray(shm_array.name, shape_mismatch, shm_array.dtype) + + with pytest.raises(Exception): + shm_mismatch.link_shm() + + +def test_close_shm(shm_array): + shm_array.close_shm() + assert shm_array.shm is None + + +def test_recreate_shm_after_close(shm_array): + shm_array.close_shm() + + # Attempting to create a new shm after close + shm_new = ShmArray(shm_array.name, shm_array.shape, shm_array.dtype) + shm_new.create_shm() + + assert shm_new.shm is not None + assert shm_new.arr.shape == shm_array.shape + assert shm_new.arr.dtype == shm_array.dtype + + shm_new.close_shm() + + +if __name__ == "__main__": + pytest.main() diff --git a/unit_tests/server/core/objs/test_token_ids_validation.py b/unit_tests/server/core/objs/test_token_ids_validation.py new file mode 100644 index 0000000000..d848137c6b --- /dev/null +++ b/unit_tests/server/core/objs/test_token_ids_validation.py @@ -0,0 +1,109 @@ +import pytest +from lightllm.server.core.objs.sampling_params import ( + StopSequence, + InvalidTokenIds, + SamplingParams, + _check_and_store_int_token_ids, + STOP_SEQUENCE_MAX_LENGTH, + INVALID_TOKEN_IDS_MAX_LENGTH, +) + + +def test_invalid_token_ids_accepts_valid_ints(): + invalid_ids = InvalidTokenIds() + invalid_ids.initialize([4, 5, 6]) + assert invalid_ids.size == 3 + assert invalid_ids.to_list() == [4, 5, 6] + + +@pytest.mark.parametrize("bad_ids", [[4, "5"], [4, 5.0]]) +def test_invalid_token_ids_rejects_non_int(bad_ids): + invalid_ids = InvalidTokenIds() + with pytest.raises(AssertionError): + invalid_ids.initialize(bad_ids) + + +def test_invalid_token_ids_rejects_too_many(): + invalid_ids = InvalidTokenIds() + with pytest.raises(AssertionError): + invalid_ids.initialize([1] * (INVALID_TOKEN_IDS_MAX_LENGTH + 1)) + + +def test_invalid_token_ids_rejects_negative(): + # A negative id survives the GPU-side vocab-size filter and is then used as a + # pointer offset into the logits tensor by the apply_invalid_token triton kernel. + invalid_ids = InvalidTokenIds() + with pytest.raises(AssertionError, match="invalid token ids"): + invalid_ids.initialize([-5]) + + +@pytest.mark.parametrize("bad_ids", [[2 ** 31], [4, 2 ** 40]]) +def test_invalid_token_ids_rejects_int32_overflow(bad_ids): + # An id >= 2**31 would wrap to a different value in the c_int buffer + # (2**31 -> -2147483648), silently banning the wrong token. + invalid_ids = InvalidTokenIds() + with pytest.raises(AssertionError, match="invalid token ids"): + invalid_ids.initialize(bad_ids) + + +def test_invalid_token_ids_accepts_int32_boundary_ids(): + invalid_ids = InvalidTokenIds() + invalid_ids.initialize([0, 2 ** 31 - 1]) + assert invalid_ids.size == 2 + assert invalid_ids.to_list() == [0, 2 ** 31 - 1] + + +def test_stop_sequence_rejects_non_int(): + seq = StopSequence() + with pytest.raises(AssertionError): + seq.initialize([1, "2"]) + + +@pytest.mark.parametrize("bad_ids", [[-7], [2 ** 31 + 5]]) +def test_stop_sequence_rejects_out_of_range_ids(bad_ids): + seq = StopSequence() + with pytest.raises(AssertionError, match="stop token ids"): + seq.initialize(bad_ids) + + +def test_stop_sequence_accepts_int32_boundary_ids(): + seq = StopSequence() + seq.initialize([0, 2 ** 31 - 1]) + assert seq.to_list() == [0, 2 ** 31 - 1] + + +def test_check_and_store_int_token_ids_returns_size_and_writes_buffer(): + import ctypes + + buf = (ctypes.c_int * 8)() + size = _check_and_store_int_token_ids(buf, [7, 8, 9], 8, "test ids") + assert size == 3 + assert list(buf[:size]) == [7, 8, 9] + + +def test_check_and_store_int_token_ids_rejects_overflow(): + import ctypes + + buf = (ctypes.c_int * 2)() + with pytest.raises(AssertionError): + _check_and_store_int_token_ids(buf, [1, 2, 3], 2, "test ids") + + +@pytest.mark.parametrize("bad_key", ["-3", str(2 ** 31)]) +def test_logit_bias_out_of_range_keys_rejected(bad_key): + # logit_bias keys are int()-converted and stored as invalid_token_ids by + # SamplingParams.init, so out-of-range keys must fail the same range check + # on the real production path, not only when initialize() is called directly. + params = SamplingParams() + with pytest.raises(AssertionError, match="invalid token ids"): + params.init(None, logit_bias={bad_key: 0.5}) + + +def test_logit_bias_boundary_keys_round_trip(): + params = SamplingParams() + params.init(None, logit_bias={"0": 1.0, str(2 ** 31 - 1): 0.5}) + assert params.invalid_token_ids.to_list() == [0, 2 ** 31 - 1] + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/unit_tests/server/grammar_helpers.py b/unit_tests/server/grammar_helpers.py new file mode 100644 index 0000000000..a27c025fb1 --- /dev/null +++ b/unit_tests/server/grammar_helpers.py @@ -0,0 +1,154 @@ +"""Shared byte-vocabulary fixtures and CPU references for grammar tests.""" + +import asyncio +import json +from itertools import count +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch + +from lightllm.server.core.objs import ConstraintOutParams, FinishStatus, StartArgs +from lightllm.server.router.model_infer.mode_backend.generic_post_process import ( + _get_post_sample_tensors, + g_pin_mem_manager, +) +from lightllm.utils.envs_utils import get_env_start_args + + +@pytest.fixture(autouse=True) +def reset_args_cache(): + get_env_start_args.cache_clear() + yield + get_env_start_args.cache_clear() + + +def configure(monkeypatch, **kwargs): + monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(vars(StartArgs(**kwargs)))) + get_env_start_args.cache_clear() + + +req_indices = count() + + +class SamplingRequest(SimpleNamespace): + def get_cur_total_len(self): + return self.shm_req.input_len + self.cur_output_len + + +def make_req(guided_reasoning_end=(), **constraints): + constraint_out_params = ConstraintOutParams() + constraint_out_params.initialize( + tokenizer=None, + regular_constraint=constraints.pop("regular_constraint", None) or "", + guided_grammar=constraints.pop("guided_grammar", None) or "", + guided_json=constraints.pop("guided_json", None) or "", + guided_reasoning_end=list(guided_reasoning_end), + ) + params = dict( + invalid_token_ids=[], + shm_param=SimpleNamespace( + constraint_out_params=constraint_out_params, + exponential_decay_length_penalty=SimpleNamespace(to_tuple=lambda: (0, 1.0)), + min_new_tokens=0, + temperature=1.0, + top_p=1.0, + top_k=1, + ), + ) + params.update(constraints) + return SamplingRequest( + req_idx=next(req_indices), + vocab_size=257, + generator=None, + sampling_param=SimpleNamespace(**params), + output_constraint=None, + finish_status=FinishStatus(), + cur_output_len=0, + shm_req=SimpleNamespace( + input_len=1, shm_prompt_ids=SimpleNamespace(arr=[ord("P")]), get_compiled_grammar=lambda: None + ), + ) + + +def init_request(compiler, req): + params = req.sampling_param.shm_param.constraint_out_params + for kind, value in ( + ("grammar", params.guided_grammar), + ("regex", params.regular_constraint), + ("json", params.guided_json), + ): + if value.length: + payload = asyncio.run(compiler.compile(kind, value.to_str())) + req.shm_req.get_compiled_grammar = lambda: payload + req.output_constraint = compiler.output_constraint_state_factory.create_state( + req.shm_req, req.sampling_param + ) + return + + +def make_mask_buffers(reqs, verify_width=1): + capacity = max((req.req_idx for req in reqs), default=0) + 1 + return SimpleNamespace( + vocab_size=257, + req_to_next_token_ids=torch.zeros((capacity, verify_width), dtype=torch.int64), + req_to_next_token_ids_cpu=torch.zeros((capacity, verify_width), dtype=torch.int64), + req_to_bitmask=torch.empty((capacity, verify_width, 9), dtype=torch.int32), + req_to_bitmask_enabled=torch.zeros(capacity, dtype=torch.bool), + ) + + +def prepare_sampling_tensors(reqs, manager, has_output=None): + # Exercise production sampling preparation with CPU parameter tensors; + # CUDA integration tests use the real pinned-memory manager and transfers. + def cpu_tensor(key, data, dtype): + tensor = torch.tensor(data, dtype=dtype) + return SimpleNamespace(cuda=lambda non_blocking: tensor) + + with patch.object(g_pin_mem_manager, "gen_from_list", side_effect=cpu_tensor): + return _get_post_sample_tensors(reqs, manager, has_output) + + +def build_mask(reqs, has_output=None, draft_input_ids=None): + manager = make_mask_buffers(reqs, len(reqs)) + if draft_input_ids is not None: + for index, req in enumerate(reqs): + if index == 0 or reqs[index - 1] is not req: + request_start = index + manager.req_to_next_token_ids_cpu[req.req_idx, index - request_start] = draft_input_ids[index] + return manager if prepare_sampling_tensors(reqs, manager, has_output)[-1] else None + + +def apply_masks(reqs, logits, manager, b_mtp_index=None): + if manager is None: + return + if b_mtp_index is None: + positions = [] + for index, req in enumerate(reqs): + if index == 0 or reqs[index - 1] is not req: + request_start = index + positions.append(index - request_start) + b_mtp_index = torch.tensor(positions, dtype=torch.int32, device=logits.device) + # Independent CPU reference; CUDA integration tests use the production kernel. + token_ids = torch.arange(manager.vocab_size) + for index, req in enumerate(reqs): + if manager.req_to_bitmask_enabled[req.req_idx]: + bitmask = manager.req_to_bitmask[req.req_idx, b_mtp_index[index]] + allowed = ((bitmask[token_ids // 32] >> (token_ids % 32)) & 1).bool() + logits[index, : manager.vocab_size].masked_fill_(~allowed, float("-inf")) + logits[:, manager.vocab_size :].fill_(float("-inf")) + + +def allowed(reqs, has_output=None): + mask = build_mask(reqs, has_output) + logits = torch.zeros(len(reqs), 257) + apply_masks(reqs, logits, mask) + return torch.isfinite(logits) + + +def commit_token(req, token): + req.shm_req.shm_prompt_ids.arr.append(token) + req.cur_output_len += 1 + req.output_constraint.commit(token) + return req.output_constraint.is_terminated() diff --git a/unit_tests/server/httpserver/test_pd_continuation.py b/unit_tests/server/httpserver/test_pd_continuation.py new file mode 100644 index 0000000000..bdf18c3fb3 --- /dev/null +++ b/unit_tests/server/httpserver/test_pd_continuation.py @@ -0,0 +1,93 @@ +import asyncio +import pickle +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import pytest +from tokenizers import Tokenizer, models, processors + +from lightllm.models.glm5_next.tokenizer import Glm5NextTokenizer +from lightllm.models.qwen3_5.model import QWen3_5Tokenizer +from lightllm.server.core.objs import SamplingParams +from lightllm.server.httpserver.manager import HttpServerManager +from lightllm.server.httpserver_for_pd_master.manager import HttpServerManagerForPDMaster +from lightllm.server.pd_io_struct import NodeRole + + +@pytest.mark.parametrize("multimodal", [False, True]) +@pytest.mark.parametrize("add_special_tokens", [False, True]) +def test_master_counts_prompt_tokens_and_prefill_preserves_output_ids(multimodal, add_special_tokens): + tokenizer = Tokenizer(models.BPE({"P": 0, "a": 1, "Pa": 2, "": 3}, [("P", "a")])) + tokenizer.post_processor = processors.TemplateProcessing(single=" $A", special_tokens=[("", 3)]) + master = object.__new__(HttpServerManagerForPDMaster) + master.tokenizer = SimpleNamespace( + encode=Mock(side_effect=lambda text, _, **kwargs: tokenizer.encode(text, **kwargs).ids) + ) + params = SamplingParams() + params.init(None, add_special_tokens=add_special_tokens) + resources = SimpleNamespace(images=[], audios=[], verify_resource_limits=lambda: None) + token_count = master.tokens("P", resources, params) + prompt_ids = ([3] if add_special_tokens else []) + [0] + assert token_count == len(prompt_ids) + master.tokenizer.encode.assert_called_once_with("P", None, add_special_tokens=add_special_tokens) + + manager = object.__new__(HttpServerManager) + manager.enable_multimodal = multimodal + manager.pd_mode = NodeRole.P + manager.vocab_size = 4 + manager._alloc_multimodal_resources = AsyncMock() + # P expands media placeholders in the supplied IDs without text tokenization. + manager.tokenizer = SimpleNamespace(encode=Mock(side_effect=lambda ids, *_args, **_kwargs: [2 ** 32] + ids)) + params.constraint_out_params.pd_previous_output_len = 1 + continuation_ids = pickle.loads(pickle.dumps(prompt_ids + [1])) + result = asyncio.run(manager._encode(continuation_ids, resources, params)) + + assert tokenizer.encode("Pa", add_special_tokens=False).ids == [2] + assert result == ([2 ** 32] if multimodal else []) + prompt_ids + [1] + if multimodal: + manager.tokenizer.encode.assert_called_once_with( + continuation_ids, resources, add_special_tokens=add_special_tokens + ) + manager._alloc_multimodal_resources.assert_awaited_once_with(resources, params) + else: + manager.tokenizer.encode.assert_not_called() + manager._alloc_multimodal_resources.assert_not_awaited() + + +@pytest.mark.parametrize("tokenizer_class", [QWen3_5Tokenizer, Glm5NextTokenizer], ids=["qwen3.5", "glm5.3-flash"]) +@pytest.mark.parametrize("image_count", [0, 1, 2]) +def test_multimodal_continuation_expands_images_without_retokenizing(tokenizer_class, image_count): + tokenizer = object.__new__(tokenizer_class) + tokenizer.tokenizer = SimpleNamespace(encode=Mock(return_value=[10, 12, 11] * image_count + [1])) + tokenizer.image_start_id, tokenizer.image_end_id, tokenizer.image_token_id = 10, 11, 12 + tokenizer.get_image_token_length = lambda image: image.token_num + images = [SimpleNamespace(token_id=1000, token_num=2) for _ in range(image_count)] + resources = SimpleNamespace(images=images, audios=[], verify_resource_limits=lambda: None) + params = SamplingParams() + params.init(None) + master = object.__new__(HttpServerManagerForPDMaster) + master.tokenizer = tokenizer + master.args = SimpleNamespace(max_image_token_count=16) + token_count = master.tokens("prompt", resources, params) + prompt_ids = [10, 11] * image_count + [1] + + manager = object.__new__(HttpServerManager) + manager.enable_multimodal = True + manager.pd_mode = NodeRole.P + manager.tokenizer = tokenizer + manager.args = master.args + manager._alloc_multimodal_resources = AsyncMock() + history = [2, 3] + params.constraint_out_params.pd_previous_output_len = len(history) + encoded = asyncio.run(manager._encode(prompt_ids + history, resources, params)) + assert encoded == [10, 1000, 1001, 11] * image_count + [1] + history + assert token_count + len(history) == len(encoded) + assert [image.start_idx for image in images] == [4 * index + 1 for index in range(image_count)] + + # D receives P's expanded IDs and must preserve them, including the history suffix. + manager.pd_mode = NodeRole.D + assert asyncio.run(manager._encode(encoded, resources, params)) == encoded + tokenizer.tokenizer.encode.assert_called_once_with("prompt", add_special_tokens=params.add_special_tokens) + manager._alloc_multimodal_resources.assert_awaited_once_with(resources, params) + assert master.tokens("prompt", resources, params) == token_count + assert manager.tokens("prompt", resources, params) == token_count diff --git a/unit_tests/server/httpserver/test_pd_generate_error.py b/unit_tests/server/httpserver/test_pd_generate_error.py new file mode 100644 index 0000000000..acb8abc714 --- /dev/null +++ b/unit_tests/server/httpserver/test_pd_generate_error.py @@ -0,0 +1,506 @@ +import asyncio +import pickle +from contextlib import suppress +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from lightllm.server.core.objs import FinishStatus, SamplingParams +from lightllm.server.httpserver.pd_loop import _pd_process_generate +from lightllm.server.httpserver_for_pd_master.manager import ( + HttpServerManagerForPDMaster, + ReqStatus, +) +from lightllm.server.multimodal_params import MultimodalParams +from lightllm.server.pd_io_struct import ObjType +from lightllm.utils.error_utils import InvalidRequestError, PDPrefillNodeStopGenToken, ServerBusyError + + +class _FailingManager: + args = SimpleNamespace(run_mode="prefill") + + async def generate(self, **_kwargs): + raise RuntimeError("prefill failed") + yield + + +class _CancelledManager: + args = SimpleNamespace(run_mode="prefill") + + async def generate(self, **_kwargs): + raise asyncio.CancelledError() + yield + + +class _SuccessfulManager: + args = SimpleNamespace(run_mode="prefill") + + async def generate(self, **_kwargs): + yield 123, "token", {}, FinishStatus(FinishStatus.FINISHED_STOP) + + +class _StopPrefillManager: + args = SimpleNamespace(run_mode="prefill") + + async def generate(self, **_kwargs): + raise PDPrefillNodeStopGenToken(group_request_id=123) + yield + + +class _FatalGenerateError(BaseException): + pass + + +class _FatalManager: + args = SimpleNamespace(run_mode="decode") + + async def generate(self, **_kwargs): + raise _FatalGenerateError("fatal generate failure") + yield + + +class _BusyManager: + args = SimpleNamespace(run_mode="decode") + + async def generate(self, **_kwargs): + raise ServerBusyError("decode node could not allocate a shm_req object within 20 seconds") + yield + + +def test_pd_node_reports_generate_error_to_master(): + async def run(): + sampling_params = SamplingParams() + sampling_params.group_request_id = 123 + websocket = AsyncMock() + + await _pd_process_generate( + manager=_FailingManager(), + prompt="prompt", + sampling_params=sampling_params, + multimodal_params=MagicMock(), + forwarding_queue=MagicMock(), + pd_upload_websocket=websocket, + pd_event=asyncio.Event(), + ) + + websocket.send.assert_awaited_once() + obj = pickle.loads(websocket.send.await_args.args[0]) + assert obj == ( + ObjType.PD_UPLOAD_GENERATE_ERROR, + 123, + "RuntimeError: prefill failed", + ) + + asyncio.run(run()) + + +def test_pd_node_reports_local_request_rejection_to_master(): + async def run(): + sampling_params = SamplingParams() + sampling_params.group_request_id = 123 + websocket = AsyncMock() + + await _pd_process_generate( + manager=_BusyManager(), + prompt="prompt", + sampling_params=sampling_params, + multimodal_params=MagicMock(), + forwarding_queue=MagicMock(), + pd_upload_websocket=websocket, + pd_event=asyncio.Event(), + ) + + websocket.send.assert_awaited_once() + obj = pickle.loads(websocket.send.await_args.args[0]) + assert obj == ( + ObjType.PD_UPLOAD_SERVER_BUSY, + 123, + "decode node could not allocate a shm_req object within 20 seconds", + ) + + asyncio.run(run()) + + +def test_pd_node_cancellation_finishes_without_reporting_generate_error(): + async def run(): + sampling_params = SamplingParams() + sampling_params.group_request_id = 123 + websocket = AsyncMock() + + await _pd_process_generate( + manager=_CancelledManager(), + prompt="prompt", + sampling_params=sampling_params, + multimodal_params=MagicMock(), + forwarding_queue=MagicMock(), + pd_upload_websocket=websocket, + pd_event=asyncio.Event(), + ) + + websocket.send.assert_not_awaited() + + asyncio.run(run()) + + +def test_pd_node_success_forwards_token_without_error_report(): + async def run(): + sampling_params = SamplingParams() + sampling_params.group_request_id = 123 + websocket = AsyncMock() + forwarding_queue = AsyncMock() + + await _pd_process_generate( + manager=_SuccessfulManager(), + prompt="prompt", + sampling_params=sampling_params, + multimodal_params=MagicMock(), + forwarding_queue=forwarding_queue, + pd_upload_websocket=websocket, + pd_event=asyncio.Event(), + ) + + forwarding_queue.put.assert_awaited_once() + forwarded = forwarding_queue.put.await_args.args[0] + assert forwarded[0] == 123 + assert forwarded[1] == "token" + assert forwarded[2]["node_mode"] == "prefill" + assert forwarded[3].is_finished() + websocket.send.assert_not_awaited() + + asyncio.run(run()) + + +def test_pd_prefill_full_cache_stop_is_not_reported_as_error(): + async def run(): + sampling_params = SamplingParams() + sampling_params.group_request_id = 123 + websocket = AsyncMock() + + await _pd_process_generate( + manager=_StopPrefillManager(), + prompt="prompt", + sampling_params=sampling_params, + multimodal_params=MagicMock(), + forwarding_queue=AsyncMock(), + pd_upload_websocket=websocket, + pd_event=asyncio.Event(), + ) + + websocket.send.assert_not_awaited() + + asyncio.run(run()) + + +def test_pd_node_reports_non_exception_base_exception(): + async def run(): + sampling_params = SamplingParams() + sampling_params.group_request_id = 123 + websocket = AsyncMock() + + await _pd_process_generate( + manager=_FatalManager(), + prompt="prompt", + sampling_params=sampling_params, + multimodal_params=MagicMock(), + forwarding_queue=AsyncMock(), + pd_upload_websocket=websocket, + pd_event=asyncio.Event(), + ) + + obj = pickle.loads(websocket.send.await_args.args[0]) + assert obj == ( + ObjType.PD_UPLOAD_GENERATE_ERROR, + 123, + "_FatalGenerateError: fatal generate failure", + ) + + asyncio.run(run()) + + +def test_pd_node_error_report_send_failure_is_contained(): + async def run(): + sampling_params = SamplingParams() + sampling_params.group_request_id = 123 + websocket = AsyncMock() + websocket.send.side_effect = ConnectionError("websocket closed") + + await _pd_process_generate( + manager=_FailingManager(), + prompt="prompt", + sampling_params=sampling_params, + multimodal_params=MagicMock(), + forwarding_queue=AsyncMock(), + pd_upload_websocket=websocket, + pd_event=asyncio.Event(), + ) + + websocket.send.assert_awaited_once() + + asyncio.run(run()) + + +def test_pd_master_generate_error_marks_request_and_wakes_all_waiters(): + async def run(): + manager = HttpServerManagerForPDMaster.__new__(HttpServerManagerForPDMaster) + manager.args = SimpleNamespace(config_server_host=None) + manager.pd_manager = MagicMock() + manager.timer_log = AsyncMock() + manager.infos_queues = None + + p_node = SimpleNamespace(websocket=SimpleNamespace(send_bytes=AsyncMock())) + d_node = SimpleNamespace(websocket=SimpleNamespace(send_bytes=AsyncMock())) + req_status = ReqStatus(123, p_node, d_node) + manager.req_id_to_out_inf = {123: req_status} + + handle_task = asyncio.create_task(manager.handle_loop()) + try: + while manager.infos_queues is None: + await asyncio.sleep(0) + await manager.put_to_handle_queue((ObjType.PD_UPLOAD_GENERATE_ERROR, 123, "RuntimeError: prefill failed")) + await asyncio.wait_for(req_status.event.wait(), timeout=1) + + assert req_status.error_info == "RuntimeError: prefill failed" + assert req_status.prefill_prompt_ids_event.is_set() + assert req_status.up_status_event.is_set() + assert manager.req_id_to_out_inf[123] is req_status + p_node.websocket.send_bytes.assert_not_awaited() + d_node.websocket.send_bytes.assert_not_awaited() + + with pytest.raises( + RuntimeError, + match="PD node generate failed: RuntimeError: prefill failed", + ): + req_status.raise_if_error() + finally: + handle_task.cancel() + with suppress(asyncio.CancelledError): + await handle_task + + asyncio.run(run()) + + +@pytest.mark.parametrize( + "terminal_status, output_index", + [ + (FinishStatus.FINISHED_ERROR, 1), + (FinishStatus.FINISHED_ABORTED, 1), + (FinishStatus.FINISHED_LENGTH, 2), + (FinishStatus.FINISHED_STOP, 2), + ], +) +def test_pd_master_first_token_dedup_preserves_terminal_status(terminal_status, output_index): + async def run(): + manager = HttpServerManagerForPDMaster.__new__(HttpServerManagerForPDMaster) + manager.args = SimpleNamespace(pd_node_id=1) + manager.req_id_to_out_inf = {} + p_node = SimpleNamespace(websocket=SimpleNamespace(send_bytes=AsyncMock())) + d_node = SimpleNamespace(websocket=SimpleNamespace(send_bytes=AsyncMock())) + sampling_params = SamplingParams() + sampling_params.group_request_id = 123 + sampling_params.max_new_tokens = 4 + + async def handshake(event, _request, **kwargs): + if kwargs["stage"] == "prefill": + event.prompt_ids = [1, 2] + else: + event.upkv_status = SimpleNamespace(pd_kv_trans_params=pickle.dumps(SimpleNamespace(ready_kv_len=0))) + status = manager.req_id_to_out_inf[123] + status.out_token_info_list.extend( + [ + ( + 123, + "first", + {"count_output_tokens": 1, "node_mode": "prefill"}, + FinishStatus(FinishStatus.FINISHED_LENGTH), + ), + ( + 123, + "duplicate", + {"count_output_tokens": 1, "node_mode": "decode"}, + FinishStatus(FinishStatus.NO_FINISH), + ), + ( + 123, + "terminal", + {"count_output_tokens": output_index, "node_mode": "decode"}, + FinishStatus(terminal_status), + ), + ] + ) + status.event.set() + + manager._wait_for_event_or_disconnect = AsyncMock(side_effect=handshake) + stream = manager.fetch_pd_stream( + p_node, + d_node, + "prompt", + sampling_params, + MultimodalParams(), + SimpleNamespace(is_disconnected=AsyncMock(return_value=False)), + ) + try: + first = await asyncio.wait_for(anext(stream), timeout=1) + assert first[1] == "first" + assert not first[3].is_finished() + terminal = await asyncio.wait_for(anext(stream), timeout=1) + assert terminal[1] == "terminal" + assert terminal[3].status == terminal_status + finally: + await stream.aclose() + + asyncio.run(run()) + + +def test_pd_master_request_rejection_becomes_server_busy_error(): + async def run(): + manager = HttpServerManagerForPDMaster.__new__(HttpServerManagerForPDMaster) + manager.args = SimpleNamespace(config_server_host=None) + manager.pd_manager = MagicMock() + manager.timer_log = AsyncMock() + manager.infos_queues = None + + req_status = ReqStatus(123, MagicMock(), MagicMock()) + manager.req_id_to_out_inf = {123: req_status} + + handle_task = asyncio.create_task(manager.handle_loop()) + try: + while manager.infos_queues is None: + await asyncio.sleep(0) + await manager.put_to_handle_queue( + ( + ObjType.PD_UPLOAD_SERVER_BUSY, + 123, + "prefill node could not allocate a shm_req object within 20 seconds", + ) + ) + await asyncio.wait_for(req_status.event.wait(), timeout=1) + + assert req_status.is_server_busy is True + with pytest.raises(ServerBusyError, match="prefill node could not allocate a shm_req object"): + req_status.raise_if_error() + finally: + handle_task.cancel() + with suppress(asyncio.CancelledError): + await handle_task + + asyncio.run(run()) + + +@pytest.mark.parametrize("event_name", ["prefill_prompt_ids_event", "up_status_event"]) +def test_pd_master_generate_error_wakes_resource_wait(event_name): + async def run(): + manager = HttpServerManagerForPDMaster.__new__(HttpServerManagerForPDMaster) + req_status = ReqStatus(123, MagicMock(), MagicMock()) + request = SimpleNamespace(is_disconnected=AsyncMock(return_value=False)) + wait_task = asyncio.create_task( + manager._wait_for_event_or_disconnect( + getattr(req_status, event_name), + request, + timeout=60, + group_request_id=123, + stage="test", + ) + ) + + await req_status.set_error("RuntimeError: node failed") + await asyncio.wait_for(wait_task, timeout=1) + + with pytest.raises( + RuntimeError, + match="PD node generate failed: RuntimeError: node failed", + ): + req_status.raise_if_error() + + asyncio.run(run()) + + +def test_pd_master_abort_removes_request_even_when_node_notifications_fail(): + async def run(): + manager = HttpServerManagerForPDMaster.__new__(HttpServerManagerForPDMaster) + manager.req_id_to_out_inf = {} + p_node = SimpleNamespace(websocket=SimpleNamespace(send_bytes=AsyncMock(side_effect=ConnectionError("p down")))) + d_node = SimpleNamespace(websocket=SimpleNamespace(send_bytes=AsyncMock(side_effect=ConnectionError("d down")))) + manager.req_id_to_out_inf[123] = ReqStatus(123, p_node, d_node) + + await manager.abort(123) + + assert 123 not in manager.req_id_to_out_inf + + asyncio.run(run()) + + +def test_pd_master_abort_uses_explicit_nodes_when_request_status_is_missing(): + async def run(): + manager = HttpServerManagerForPDMaster.__new__(HttpServerManagerForPDMaster) + manager.req_id_to_out_inf = {} + p_node = SimpleNamespace(websocket=SimpleNamespace(send_bytes=AsyncMock())) + d_node = SimpleNamespace(websocket=SimpleNamespace(send_bytes=AsyncMock())) + + await manager.abort(123, p_node=p_node, d_node=d_node) + + p_node.websocket.send_bytes.assert_awaited_once_with(pickle.dumps((ObjType.ABORT, 123))) + d_node.websocket.send_bytes.assert_awaited_once_with(pickle.dumps((ObjType.ABORT, 123))) + + asyncio.run(run()) + + +def test_pd_invalid_request_error_round_trip(): + class InvalidManager: + args = SimpleNamespace(run_mode="prefill") + + async def generate(self, **_kwargs): + raise InvalidRequestError("invalid image tag num: 0 images vs 1 tags") + yield + + async def run(): + sampling_params = SamplingParams() + sampling_params.group_request_id = 123 + websocket = AsyncMock() + await _pd_process_generate( + manager=InvalidManager(), + prompt="prompt", + sampling_params=sampling_params, + multimodal_params=MagicMock(), + forwarding_queue=AsyncMock(), + pd_upload_websocket=websocket, + pd_event=asyncio.Event(), + ) + kind, req_id, error_info = pickle.loads(websocket.send.await_args.args[0]) + assert kind == ObjType.PD_UPLOAD_GENERATE_ERROR + assert req_id == 123 + status = ReqStatus(req_id, MagicMock(), MagicMock()) + await status.set_error(error_info) + with pytest.raises(InvalidRequestError, match="invalid image tag num: 0 images vs 1 tags"): + status.raise_if_error() + + asyncio.run(run()) + + +def test_pd_stream_validation_returns_bad_request(monkeypatch): + from fastapi import FastAPI + from fastapi.testclient import TestClient + from lightllm.server import api_stream_obj + from lightllm.server.api_http import g_objs, invalid_request_exception_handler + from lightllm.server.api_openai import _safe_stream_wrapper + + monkeypatch.setattr( + api_stream_obj, "get_env_start_args", lambda: SimpleNamespace(disable_delay_response_start=False) + ) + monkeypatch.setattr(g_objs, "metric_client", MagicMock()) + app = FastAPI() + app.add_exception_handler(InvalidRequestError, invalid_request_exception_handler) + + async def generate(): + status = ReqStatus(123, MagicMock(), MagicMock()) + await status.set_error("InvalidRequestError: invalid image tag num: 0 images vs 1 tags") + status.raise_if_error() + yield + + @app.get("/") + async def stream(): + return api_stream_obj.CustomStreamingResponse(_safe_stream_wrapper(generate()), media_type="text/event-stream") + + response = TestClient(app).get("/") + assert response.status_code == 400 + assert "invalid image tag num" in response.text diff --git a/unit_tests/server/httpserver/test_pd_kv_status_channel.py b/unit_tests/server/httpserver/test_pd_kv_status_channel.py new file mode 100644 index 0000000000..9f00bbbf5f --- /dev/null +++ b/unit_tests/server/httpserver/test_pd_kv_status_channel.py @@ -0,0 +1,109 @@ +"""KV status routing through the master owners discovered by P/D nodes.""" + +import asyncio +import pickle +from contextlib import asynccontextmanager +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +from lightllm.server.pd_io_struct import PD_Master_Obj, PDUpKVStatus +from lightllm.server.router.model_infer.mode_backend.pd.decode_node_impl import up_status + + +def _manager(): + manager = up_status.UpStatusManager.__new__(up_status.UpStatusManager) + manager.args = SimpleNamespace(pd_node_id=20) + manager.id_to_handle_queue = {} + return manager + + +def _websocket(send): + return SimpleNamespace( + send=send, + recv=AsyncMock(), + transport=SimpleNamespace(get_extra_info=lambda _key: MagicMock()), + ) + + +def test_status_is_sent_directly_without_control_frames_or_reply(monkeypatch): + async def run(): + manager = _manager() + status = PDUpKVStatus(8, 100, b"decode metadata") + manager.id_to_handle_queue[100] = asyncio.Queue() + manager.id_to_handle_queue[100].put_nowait(status) + sent = [] + status_sent = asyncio.Event() + + async def send(data): + sent.append(pickle.loads(data)) + status_sent.set() + + websocket = _websocket(send) + + @asynccontextmanager + async def connect(uri): + assert uri == "ws://master:8100/kv_move_status" + yield websocket + + monkeypatch.setattr(up_status.websockets, "connect", connect) + task = asyncio.create_task(manager.up_kv_status_task(PD_Master_Obj(100, "master:8100"))) + try: + await asyncio.wait_for(status_sent.wait(), timeout=1) + assert sent == [status] + websocket.recv.assert_not_called() + finally: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + asyncio.run(run()) + + +def test_topology_refresh_connects_all_new_workers_and_retires_previous_ids(monkeypatch): + from lightllm.server.httpserver import pd_loop + + async def run(): + manager = _manager() + responses = asyncio.Queue() + started = asyncio.Queue() + retired = set() + children = [] + original_sleep = asyncio.sleep + + async def discover(_args): + return await responses.get() + + async def dispatch(): + children.append(asyncio.current_task()) + await asyncio.Event().wait() + + async def connect(master): + children.append(asyncio.current_task()) + started.put_nowait(master.node_id) + try: + await asyncio.Event().wait() + finally: + retired.add(master.node_id) + + async def skip_poll_interval(seconds): + assert seconds == 30, "Topology refresh unexpectedly failed" + await original_sleep(0) + + monkeypatch.setattr(pd_loop, "_get_pd_master_objs", discover) + monkeypatch.setattr(manager, "dispatch_task_loop", dispatch) + monkeypatch.setattr(manager, "up_kv_status_task", connect) + monkeypatch.setattr(up_status.asyncio, "sleep", skip_poll_interval) + task = asyncio.create_task(manager.task_loop()) + try: + responses.put_nowait({node_id: PD_Master_Obj(node_id, f"master:{node_id}") for node_id in (100, 101)}) + assert {await asyncio.wait_for(started.get(), timeout=1) for _ in range(2)} == {100, 101} + responses.put_nowait({node_id: PD_Master_Obj(node_id, f"master:{node_id}") for node_id in (200, 201)}) + assert {await asyncio.wait_for(started.get(), timeout=1) for _ in range(2)} == {200, 201} + assert retired == {100, 101} + assert set(manager.id_to_handle_queue) == {200, 201} + finally: + task.cancel() + for child in children: + child.cancel() + await asyncio.gather(task, *children, return_exceptions=True) + + asyncio.run(run()) diff --git a/unit_tests/server/httpserver/test_pd_master_cached_tokens.py b/unit_tests/server/httpserver/test_pd_master_cached_tokens.py new file mode 100644 index 0000000000..cb815fb469 --- /dev/null +++ b/unit_tests/server/httpserver/test_pd_master_cached_tokens.py @@ -0,0 +1,771 @@ +import asyncio +import copy +import pickle +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest +import numpy as np + +from lightllm.server.core.objs import StopStrMatchMode, FinishStatus, SamplingParams +from lightllm.server.router.model_infer.infer_batch import InferReq +from lightllm.server.httpserver_for_pd_master.manager import HttpServerManagerForPDMaster, PDManager, logger +from lightllm.server.httpserver_for_pd_master.pd_selector import PDSelectionExtraInfo +from lightllm.server.pd_io_struct import PD_Client_Obj + +pytestmark = pytest.mark.usefixtures("single_http_worker_start_args") + + +def _make_manager(monkeypatch): + monkeypatch.setattr( + "lightllm.server.httpserver.manager.HttpServerManager._check_and_repair_length", + classmethod(lambda cls, *a, **k: asyncio.sleep(0)), + ) + monkeypatch.setattr(SamplingParams, "from_buffer_copy", classmethod(lambda cls, other: copy.copy(other))) + mgr = object.__new__(HttpServerManagerForPDMaster) + mgr.active_req_ids = set() + mgr.req_id_to_out_inf = {} + mgr.args = SimpleNamespace( + disable_pd_master_decode_capacity_limit=True, + reasoning_parser=None, + enable_rl=False, + enable_return_routed_experts=False, + select_p_d_node_strategy="cache_aware", + pd_master_mode="elastic", + ) + mgr.enable_pd_node_self_request_limit = True + mgr.pd_node_resource_wait_timeout_seconds = -1 + mgr.pd_node_continuation_resource_wait_timeout_seconds = 60 + mgr.pd_node_busy_retry_timeout_seconds = 120 + mgr.pd_cache_high_priority_max_age_seconds = 60 + mgr.pd_cache_high_priority_min_prompt_tokens = 8192 + mgr.disable_pd_cache_high_priority = False + mgr.running_request_count = 0 + counter = [0] + + def gen_id(): + counter[0] += 1 + return counter[0] + + mgr.metric_client = SimpleNamespace(counter_inc=lambda *a, **k: None, histogram_observe=lambda *a, **k: None) + mgr.tokens = lambda prompt, *a, **k: 10 + len(prompt) - len("hello") + mgr.tokenizer = SimpleNamespace(encode=lambda text, _multimodal=None, **kwargs: [1, 2]) + mgr._log_req_header = lambda *a, **k: asyncio.sleep(0) + mgr.recorded_cache_hit_rates = [] + mgr.inserted_prompt_caches = [] + mgr.segment_inputs = [] + mgr.pd_manager = PDManager(mgr.args) + mgr.pd_node_stats = mgr.pd_manager.selector.pd_node_stats + mgr.id_gen = SimpleNamespace(generate_id=gen_id) + original_record_hit_rate = mgr.pd_manager.selector.record_prompt_cache_hit_rate + prompt_cache_tree = mgr.pd_manager.selector.policy.prompt_cache_tree + original_insert_cache = prompt_cache_tree.insert_sampled + + def record_hit_rate(hit_rate): + mgr.recorded_cache_hit_rates.append(hit_rate) + original_record_hit_rate(hit_rate) + + async def insert_cache(prompt_key, connection_key, sample_stride): + mgr.inserted_prompt_caches.append((prompt_key, connection_key)) + await original_insert_cache(prompt_key, connection_key, sample_stride) + + mgr.pd_manager.selector.record_prompt_cache_hit_rate = record_hit_rate + prompt_cache_tree.insert_sampled = insert_cache + p_node = PD_Client_Obj( + node_id=100, + client_ip_port="prefill:8000", + mode="prefill", + start_args={}, + websocket=SimpleNamespace(send_bytes=AsyncMock()), + ) + d_node = PD_Client_Obj( + node_id=200, + client_ip_port="decode:8000", + mode="decode", + start_args={}, + websocket=SimpleNamespace(send_bytes=AsyncMock()), + ) + mgr.pd_manager.prefill_nodes = [p_node] + mgr.pd_manager.decode_nodes = [d_node] + mgr.pd_manager.url_to_pd_nodes = {node.connection_key: node for node in [p_node, d_node]} + mgr.pd_manager.selector.update_nodes([p_node], [d_node]) + mgr.select_p_d_node = lambda *a, **k: asyncio.sleep(0, result=(p_node, d_node, PDSelectionExtraInfo())) + mgr.remove_req = lambda *a, **k: asyncio.sleep(0) + return mgr + + +def _collect( + mgr, + sampling_params, + monkeypatch, + segments, + first_node_mode="prefill", + needs_prefill_first_token=True, + first_token_text="", + prompt="hello", +): + segment_iter = iter(segments) + + async def fake_wait(p_node, d_node, start_time, prompt_ids, sp, multimodal_params, request): + mgr.segment_inputs.append( + pickle.loads(pickle.dumps((prompt_ids, sp.constraint_out_params.pd_previous_output_len))) + ) + sub_req_id = sp.group_request_id + token_count, final_status = next(segment_iter) + hit = sampling_params.max_new_tokens * 10 + for token_index in range(1, token_count + 1): + finish_status = FinishStatus() + if token_index == token_count: + finish_status = FinishStatus(final_status) + metadata = { + "id": ord("x"), + "prompt_tokens": 100, + "prompt_cache_len": hit if token_index == 1 else 0, + "count_output_tokens": token_index, + "node_mode": first_node_mode if token_index == 1 else "decode", + } + if token_index == 1 and needs_prefill_first_token is not None: + metadata["needs_prefill_first_token"] = needs_prefill_first_token + yield sub_req_id, first_token_text if token_index == 1 else "x", metadata, finish_status + + monkeypatch.setattr(mgr, "_wait_to_token_package", fake_wait) + + async def run(): + out = [] + async for sub_id, out_str, metadata, finish in mgr.generate( + prompt=prompt, + sampling_params=sampling_params, + multimodal_params=SimpleNamespace(images=[], audios=[], verify_and_preload=lambda req: asyncio.sleep(0)), + request=None, + ): + assert "needs_prefill_first_token" not in metadata + out.append(metadata.get("prompt_cache_len", -1)) + return out + + return asyncio.run(run()) + + +def test_single_block_prefill_hit_persists_past_decode_zeros(monkeypatch): + mgr = _make_manager(monkeypatch) + sp = SamplingParams() + sp.init(None) + sp.n = 1 + sp.max_new_tokens = 3 + sp.best_of = 1 + sp.group_request_id = 0 + cached = _collect(mgr, sp, monkeypatch, segments=[(3, FinishStatus.FINISHED_STOP)]) + assert cached and all(c == 30 for c in cached), cached + assert mgr.recorded_cache_hit_rates == [pytest.approx(0.3)] + assert len(mgr.inserted_prompt_caches) == 1 + assert mgr.inserted_prompt_caches[0][0] == "h" + + +@pytest.mark.parametrize("mode", [StopStrMatchMode.ALL_TEXT, StopStrMatchMode.CONTENT_ONLY]) +@pytest.mark.parametrize("hidden_prefix", [False, True]) +@pytest.mark.parametrize("token_chunks", [[[10], [11, 12]], [[10], [11], [12, 13]]]) +def test_pd_token_id_stop_matches_across_capacity_segments(monkeypatch, mode, hidden_prefix, token_chunks): + async def run(): + manager = _make_manager(monkeypatch) + params = SamplingParams() + stop_ids = list(range(10, 10 + len(token_chunks))) + params.init(None, stop_sequences=[stop_ids], stop_str_match_mode=mode, max_new_tokens=6) + segments = [] + + async def raw_segment(p_node, d_node, start_time, prompt, sp, multimodal, request): + index = len(segments) + segments.append(sp.group_request_id) + req = InferReq.__new__(InferReq) + req.stop_sequences = sp.stop_sequences.to_token_ids_list() + req.sampling_param = SimpleNamespace(shm_param=sp) + tokens = token_chunks[index] + req.shm_req = SimpleNamespace(input_len=1, shm_prompt_ids=SimpleNamespace(arr=np.array([1, *tokens]))) + for output_len, token_id in enumerate(tokens, 1): + req.finish_status = FinishStatus() + req.update_finish_status([], output_len) + if index == len(token_chunks) - 1 and output_len == len(tokens) and not req.finish_status.is_finished(): + req.finish_status = FinishStatus(FinishStatus.FINISHED_LENGTH) + text = "" if token_id == 10 and hidden_prefix else str(token_id) + yield sp.group_request_id, text, {"id": token_id, "prompt_tokens": 2}, req.finish_status + if req.finish_status.is_finished(): + return + yield sp.group_request_id, "", {"id": None}, FinishStatus(FinishStatus.FINISHED_PD_DECODE_CAPACITY) + + manager._wait_to_token_package = raw_segment + # 分段内仍由真实 InferReq 判断；跨分段匹配由对外生成流的包装层完成。 + from lightllm.server.httpserver.stop_seq_match import wrap_stop_match_stream + + raw_generator = manager._generate_one_attempt("prompt", params, None, None, 0, 80, 2) + results = [result async for result in wrap_stop_match_stream(manager, raw_generator, params, "prompt")] + assert [result[2]["id"] for result in results] == stop_ids + assert results[-1][3].is_stopped() + assert len(segments) == len(token_chunks) + + asyncio.run(asyncio.wait_for(run(), timeout=2)) + + +def test_dynamic_split_keeps_first_segment_hit(monkeypatch): + mgr = _make_manager(monkeypatch) + sp = SamplingParams() + sp.init(None) + sp.n = 1 + sp.max_new_tokens = 5 + sp.best_of = 1 + sp.group_request_id = 0 + cached = _collect( + mgr, + sp, + monkeypatch, + segments=[ + (3, FinishStatus.FINISHED_PD_DECODE_CAPACITY), + (2, FinishStatus.FINISHED_PD_DECODE_CAPACITY), + (1, FinishStatus.FINISHED_STOP), + ], + ) + assert cached == [50, 50, 50, 50], cached + assert mgr.recorded_cache_hit_rates == [pytest.approx(0.5)] + assert len(mgr.inserted_prompt_caches) == 1 + assert mgr.inserted_prompt_caches[0][0] == "h" + # 无约束请求不需要计算语法状态重放长度。 + assert mgr.segment_inputs == [ + ("hello", 0), + ("hellox", 0), + ("hellox", 0), + ] + + +@pytest.mark.parametrize("add_special_tokens", [False, True]) +@pytest.mark.parametrize("prompt,expected_lengths", [("P", (0, 1)), ("PP", (0, 0))]) +@pytest.mark.parametrize("constraint_enabled", [False, True]) +def test_continuation_uses_reencoded_length_difference( + monkeypatch, add_special_tokens, prompt, expected_lengths, constraint_enabled +): + from tokenizers import Tokenizer, models, processors + from unittest.mock import Mock, call + + tokenizer = Tokenizer(models.BPE({"P": 0, "x": 1, "Px": 2, "": 3, "PPx": 4}, [("P", "x"), ("P", "Px")])) + tokenizer.post_processor = processors.TemplateProcessing(single=" $A", special_tokens=[("", 3)]) + manager = _make_manager(monkeypatch) + del manager.tokens + manager.tokenizer = SimpleNamespace( + encode=Mock(side_effect=lambda text, _multimodal=None, **kwargs: tokenizer.encode(text, **kwargs).ids) + ) + params = SamplingParams() + params.init( + None, + max_new_tokens=4, + add_special_tokens=add_special_tokens, + regular_constraint=".*" if constraint_enabled else "", + ) + _collect( + manager, + params, + monkeypatch, + segments=[ + (2, FinishStatus.FINISHED_PD_DECODE_CAPACITY), + (2, FinishStatus.FINISHED_PD_DECODE_CAPACITY), + (2, FinishStatus.FINISHED_STOP), + ], + first_token_text="x", + prompt=prompt, + ) + # 约束输出按拼接文本的编码长度差重放；普通请求跳过后续编码。 + assert manager.segment_inputs == [ + (prompt, 0), + (prompt + "x", expected_lengths[0] if constraint_enabled else 0), + (prompt + "xx", expected_lengths[1] if constraint_enabled else 0), + ] + expected_encoded_prompts = (prompt, prompt + "x", prompt + "xx") if constraint_enabled else (prompt,) + assert manager.tokenizer.encode.call_args_list == [ + call(text, None, add_special_tokens=add_special_tokens) for text in expected_encoded_prompts + ] + + +def test_error_result_records_hit_rate_without_inserting_prompt_cache(monkeypatch): + mgr = _make_manager(monkeypatch) + sampling_params = SamplingParams() + sampling_params.init(None) + sampling_params.n = 1 + sampling_params.best_of = 1 + sampling_params.max_new_tokens = 1 + + async def failed_wait(_p_node, _d_node, _start_time, _prompt_ids, sp, *_args): + yield ( + sp.group_request_id, + "", + { + "id": None, + "prompt_tokens": 100, + "prompt_cache_len": 20, + "count_output_tokens": 0, + "needs_prefill_first_token": True, + }, + FinishStatus(FinishStatus.FINISHED_ERROR), + ) + + monkeypatch.setattr(mgr, "_wait_to_token_package", failed_wait) + + async def run(): + async for _ in mgr.generate( + prompt="hello", + sampling_params=sampling_params, + multimodal_params=SimpleNamespace( + images=[], + audios=[], + verify_and_preload=lambda req: asyncio.sleep(0), + ), + request=None, + ): + pass + + asyncio.run(run()) + + assert mgr.recorded_cache_hit_rates == [pytest.approx(0.2)] + assert mgr.inserted_prompt_caches == [] + + +def test_pd_raw_stop_does_not_use_configured_reasoning_parser(monkeypatch): + manager = _make_manager(monkeypatch) + manager.args.reasoning_parser = "qwen3" + manager.tokenizer = SimpleNamespace(encode=lambda text, _multimodal=None, **kwargs: [ord(c) for c in text]) + params = SamplingParams() + params.init( + manager.tokenizer, + stop_sequences="END", + max_new_tokens=10, + stop_str_match_mode=StopStrMatchMode.ALL_TEXT, + _initial_reasoning_state=1, + ) + + async def fake_wait(p_node, d_node, start_time, prompt, sp, multimodal_params, request): + assert sp.stop_str_match_mode.get_mode() == (StopStrMatchMode.ALL_TEXT) + metadata = {"prompt_tokens": 10, "prompt_cache_len": 0, "count_output_tokens": 1, "id": 10} + yield sp.group_request_id, "reasonENDextra", metadata, FinishStatus(FinishStatus.FINISHED_LENGTH) + + monkeypatch.setattr(manager, "_wait_to_token_package", fake_wait) + + async def run(): + return [ + result + async for result in manager.generate( + "prompt", + params, + SimpleNamespace(images=[], audios=[], verify_and_preload=lambda req: asyncio.sleep(0)), + None, + ) + ] + + results = asyncio.run(run()) + assert "".join(result[1] for result in results) == "reason" + assert results[-1][3].is_stopped() + + +@pytest.mark.parametrize("in_reasoning", [False, True]) +def test_pd_continuation_uses_local_reasoning_state(monkeypatch, in_reasoning): + + manager = _make_manager(monkeypatch) + manager.args.reasoning_parser = "qwen3" + manager.tokenizer = SimpleNamespace( + encode=lambda text, _multimodal=None, **kwargs: {"": [1000], "": [1001]}.get( + text, [ord(c) for c in text] + ) + ) + params = SamplingParams() + params.init( + manager.tokenizer, + stop_sequences="END", + max_new_tokens=3, + _initial_reasoning_state=1, + stop_str_match_mode=StopStrMatchMode.CONTENT_ONLY, + ) + segment_modes = [] + + async def fake_wait(p_node, d_node, start_time, prompt, sp, multimodal_params, request): + segment_modes.append(sp._initial_reasoning_state) + assert sp.stop_str_match_mode.get_mode() == StopStrMatchMode.CONTENT_ONLY + metadata = {"prompt_tokens": 10, "prompt_cache_len": 0, "count_output_tokens": 1} + if len(segment_modes) == 1: + yield sp.group_request_id, "reason", {**metadata, "id": 10}, FinishStatus() + if not in_reasoning: + yield sp.group_request_id, "", {**metadata, "id": 1001}, FinishStatus() + # A simulated delimiter at the capacity boundary must not change state. + yield ( + sp.group_request_id, + "", + {**metadata, "id": 1001}, + FinishStatus(FinishStatus.FINISHED_PD_DECODE_CAPACITY), + ) + else: + yield sp.group_request_id, "END", {**metadata, "id": 11}, FinishStatus(FinishStatus.FINISHED_STOP) + + monkeypatch.setattr(manager, "_wait_to_token_package", fake_wait) + + async def run(): + return [ + result + async for result in manager.generate( + "hello", + params, + SimpleNamespace(images=[], audios=[], verify_and_preload=lambda req: asyncio.sleep(0)), + None, + ) + ] + + results = asyncio.run(run()) + assert segment_modes == [1, 1] # reasoning 状态只保存在应用侧，不再回写节点参数。 + assert "".join(result[1] for result in results) == "reason" + ("END" if in_reasoning else "") + + +@pytest.mark.parametrize("node_mode", ["decode", None]) +def test_skipped_prefill_does_not_insert_prompt_cache(monkeypatch, node_mode): + mgr = _make_manager(monkeypatch) + sampling_params = SamplingParams() + sampling_params.init(None) + sampling_params.n = sampling_params.best_of = 1 + sampling_params.max_new_tokens = 3 + cached = _collect( + mgr, + sampling_params, + monkeypatch, + segments=[(3, FinishStatus.FINISHED_STOP)], + first_node_mode=node_mode, + needs_prefill_first_token=False, + ) + assert cached == [30, 30, 30] + assert mgr.recorded_cache_hit_rates == [pytest.approx(0.3)] + assert mgr.inserted_prompt_caches == [] + + +@pytest.mark.parametrize("needs_prefill_first_token", [True, None]) +def test_completed_prefill_inserts_cache_when_first_output_is_decode(monkeypatch, needs_prefill_first_token): + mgr = _make_manager(monkeypatch) + sampling_params = SamplingParams() + sampling_params.init(None) + sampling_params.n = sampling_params.best_of = 1 + sampling_params.max_new_tokens = 3 + cached = _collect( + mgr, + sampling_params, + monkeypatch, + segments=[(3, FinishStatus.FINISHED_STOP)], + first_node_mode="decode", + needs_prefill_first_token=needs_prefill_first_token, + ) + assert cached == [30, 30, 30] + assert len(mgr.inserted_prompt_caches) == 1 + assert mgr.inserted_prompt_caches[0][0] == "h" + + +@pytest.mark.parametrize("include", [False, True]) +@pytest.mark.parametrize("parser", [None, "qwen3"]) +@pytest.mark.parametrize( + "continuation,node_finish", + [ + ("NDextra", FinishStatus.NO_FINISH), + ("NDextra", FinishStatus.FINISHED_STOP), + ("NX", FinishStatus.FINISHED_LENGTH), + ("N", FinishStatus.FINISHED_LENGTH), + ], +) +def test_pd_stop_prefix_survives_capacity_split(monkeypatch, include, parser, continuation, node_finish): + manager = _make_manager(monkeypatch) + manager.args.reasoning_parser = parser + manager.tokenizer = SimpleNamespace( + encode=lambda text, _multimodal=None, **kwargs: {"": [1000], "": [1001]}.get( + text, [ord(c) for c in text] + ) + ) + manager.abort = AsyncMock() + stop_request = manager.stop_request + manager.stop_request = AsyncMock(side_effect=stop_request) + params = SamplingParams() + params.init( + manager.tokenizer, + stop_sequences="END", + include_stop_str_in_output=include, + _initial_reasoning_state=int(parser is not None), + stop_str_match_mode=StopStrMatchMode.ALL_TEXT if parser is None else StopStrMatchMode.CONTENT_ONLY, + max_new_tokens=10, + ) + segments = [] + + async def fake_wait(p_node, d_node, start_time, prompt, sp, multimodal_params, request): + segments.append(prompt) + metadata = {"prompt_tokens": 10, "prompt_cache_len": 0, "count_output_tokens": 1} + if len(segments) == 1: + if parser: + yield sp.group_request_id, "END", {**metadata, "id": 10}, FinishStatus() + yield sp.group_request_id, "", {**metadata, "id": 1001}, FinishStatus() + yield sp.group_request_id, "helloE", {**metadata, "id": 11}, FinishStatus() + yield sp.group_request_id, "", dict(metadata), FinishStatus(FinishStatus.FINISHED_PD_DECODE_CAPACITY) + else: + yield sp.group_request_id, continuation, {**metadata, "id": 12}, FinishStatus(node_finish) + if continuation == "NDextra" and node_finish == FinishStatus.NO_FINISH: + yield sp.group_request_id, "", {**metadata, "id": None}, FinishStatus(FinishStatus.FINISHED_ABORTED) + + monkeypatch.setattr(manager, "_wait_to_token_package", fake_wait) + + async def run(): + return [ + result + async for result in manager.generate( + "prompt", + params, + SimpleNamespace(images=[], audios=[], verify_and_preload=lambda req: asyncio.sleep(0)), + None, + ) + ] + + results = asyncio.run(run()) + assert len(segments) == 2 + assert segments == ["prompt", "prompt" + ("END" if parser else "") + "helloE"] + expected = "hello" + ("END" if include else "") if continuation == "NDextra" else "helloE" + continuation + assert "".join(result[1] for result in results) == ("END" if parser else "") + expected + assert [result[2]["id"] for result in results] == ([10, 1001] if parser else []) + [11, 12] + assert results[-1][3].get_finish_reason() == ("stop" if continuation == "NDextra" else "length") + assert all(not result[3].is_finished() for result in results[:-1]) + if continuation == "NDextra" and node_finish == FinishStatus.NO_FINISH: + manager.stop_request.assert_awaited_once() + else: + manager.stop_request.assert_not_awaited() + + +def test_pd_stop_request_marks_only_matching_choice(monkeypatch): + manager = _make_manager(monkeypatch) + manager.active_req_ids.update((80, 81)) + + async def run(): + await manager.stop_request(80) + assert manager.active_req_ids == {81} + await manager.stop_request(80) + await manager.stop_request(999) + assert manager.active_req_ids == {81} + + asyncio.run(run()) + + +def test_pd_abort_removes_output_channel_by_default(monkeypatch): + manager = _make_manager(monkeypatch) + p_node = SimpleNamespace(websocket=SimpleNamespace(send_bytes=AsyncMock())) + d_node = SimpleNamespace(websocket=SimpleNamespace(send_bytes=AsyncMock())) + manager.req_id_to_out_inf = {88: SimpleNamespace(p_node=p_node, d_node=d_node), 96: object()} + + asyncio.run(manager.abort(88)) + + assert set(manager.req_id_to_out_inf) == {96} + for node in (p_node, d_node): + node.websocket.send_bytes.assert_awaited_once() + + +def test_application_stop_discards_late_pd_output_and_closes_segment(monkeypatch): + manager = _make_manager(monkeypatch) + manager.tokenizer = SimpleNamespace(encode=lambda text, _multimodal=None, **kwargs: [ord(c) for c in text]) + params = SamplingParams() + params.init(manager.tokenizer, stop_sequences="END", max_new_tokens=10) + + async def run(): + segments = [] + closed = asyncio.Event() + stopped = asyncio.Event() + manager.abort = AsyncMock() + stop_request = manager.stop_request + + async def stop_and_notify(request_id): + await stop_request(request_id) + stopped.set() + + manager.stop_request = stop_and_notify + + async def raw_segment(p_node, d_node, start_time, prompt, sp, multimodal_params, request): + segments.append(prompt) + metadata = {"prompt_tokens": 10, "prompt_cache_len": 0, "count_output_tokens": 1, "id": 10} + try: + yield sp.group_request_id, "helloENDextra", metadata, FinishStatus() + await stopped.wait() + yield sp.group_request_id, "late", {**metadata, "id": 11}, FinishStatus() + finally: + closed.set() + + manager._wait_to_token_package = raw_segment + results = [ + result + async for result in manager.generate( + "prompt", + params, + SimpleNamespace(images=[], audios=[], verify_and_preload=lambda req: asyncio.sleep(0)), + None, + ) + ] + assert segments == ["prompt"] + assert results[-1][1] == "hello" + assert results[-1][3].is_stopped() + assert not manager.active_req_ids + await asyncio.wait_for(closed.wait(), timeout=1) + assert closed.is_set() + manager.abort.assert_awaited_once() + + asyncio.run(asyncio.wait_for(run(), timeout=2)) + + +@pytest.mark.parametrize( + "exit_mode, token_count", [("finish", 2), ("close", 1), ("error", 1), ("empty", 0), ("empty_prompt", 2)] +) +def test_pd_segment_summary_is_recorded_on_exit(monkeypatch, caplog, exit_mode, token_count): + import math + import time + from unittest.mock import MagicMock + + manager = _make_manager(monkeypatch) + manager.first_time_costs = MagicMock() + manager.per_token_costs = MagicMock() + manager.metric_client = MagicMock() + params = SamplingParams() + params.init(SimpleNamespace(encode=lambda *_args, **_kwargs: []), max_new_tokens=4, group_request_id=80) + request = SimpleNamespace(headers={"X-Request-Id": "summary-test"}, is_disconnected=lambda: asyncio.sleep(0, False)) + caplog.set_level("INFO", logger="lightllm.server.httpserver_for_pd_master.manager") + monkeypatch.setattr(logger, "handlers", [caplog.handler]) + + async def raw_stream(*_args): + if exit_mode == "empty": + return + metadata = { + "prompt_tokens": 0 if exit_mode == "empty_prompt" else 10, + "prompt_cache_len": 0 if exit_mode == "empty_prompt" else 3, + } + yield 80, "first", dict(metadata), FinishStatus() + if exit_mode == "error": + raise RuntimeError("decode failed") + yield 80, "last", dict(metadata), FinishStatus(FinishStatus.FINISHED_LENGTH) + + manager.fetch_pd_stream = raw_stream + + async def run(): + generator = manager._wait_to_token_package(None, None, time.time(), "prompt", params, None, request) + if exit_mode == "close": + await generator.__anext__() + await generator.aclose() + elif exit_mode == "error": + with pytest.raises(RuntimeError, match="decode failed"): + async for _ in generator: + pass + else: + assert len([result async for result in generator]) == token_count + + asyncio.run(run()) + summaries = [record.getMessage() for record in caplog.records if "out_token_counter:" in record.getMessage()] + if token_count: + assert len(summaries) == 1 + assert "X-Request-Id:summary-test" in summaries[0] + assert f"out_token_counter:{token_count}" in summaries[0] + manager.first_time_costs.add.assert_called_once() + manager.per_token_costs.add.assert_called_once() + assert math.isfinite(manager.per_token_costs.add.call_args.args[0]) + generated = [ + args + for args, _ in manager.metric_client.histogram_observe.call_args_list + if args[0] == "lightllm_request_generated_tokens" + ] + assert generated == [("lightllm_request_generated_tokens", token_count)] + else: + assert not summaries + manager.per_token_costs.add.assert_not_called() + manager.metric_client.histogram_observe.assert_not_called() + + +@pytest.mark.parametrize("late_status", [FinishStatus.NO_FINISH, FinishStatus.FINISHED_PD_DECODE_CAPACITY]) +def test_pd_stopped_choice_records_segment_summary(monkeypatch, caplog, late_status): + import time + from unittest.mock import AsyncMock, MagicMock + + manager = _make_manager(monkeypatch) + manager.first_time_costs = MagicMock() + manager.per_token_costs = MagicMock() + manager.metric_client = MagicMock() + manager.abort = AsyncMock() + params = SamplingParams() + params.init(SimpleNamespace(encode=lambda *_args, **_kwargs: []), max_new_tokens=4) + request = SimpleNamespace(headers={}, is_disconnected=lambda: asyncio.sleep(0, False)) + caplog.set_level("INFO", logger="lightllm.server.httpserver_for_pd_master.manager") + monkeypatch.setattr(logger, "handlers", [caplog.handler]) + + async def raw_stream(_p_node, _d_node, _prompt, sampling_params, *_args): + metadata = {"id": 1, "prompt_tokens": 10, "prompt_cache_len": 3} + yield sampling_params.group_request_id, "first", dict(metadata), FinishStatus() + yield sampling_params.group_request_id, "late", dict(metadata), FinishStatus(late_status) + + manager.fetch_pd_stream = raw_stream + + async def run(): + generator = manager._generate_one("prompt", params, None, request, time.time(), 80, 10) + assert (await generator.__anext__())[1] == "first" + await manager.stop_request(80) + # stop_request 只做标记，恢复生成时才停止当前 P/D 分段。 + manager.abort.assert_not_awaited() + with pytest.raises(StopAsyncIteration): + await generator.__anext__() + assert not manager.active_req_ids + manager.abort.assert_awaited_once() + summaries = [record.getMessage() for record in caplog.records if "out_token_counter:" in record.getMessage()] + assert len(summaries) == 1 + assert "lightllm_req_id:1" in summaries[0] + manager.metric_client.histogram_observe.assert_any_call("lightllm_request_generated_tokens", 2) + + asyncio.run(run()) + + +def test_application_pd_stop_leaves_sibling_choice_running(monkeypatch): + manager = _make_manager(monkeypatch) + manager.tokenizer = SimpleNamespace(encode=lambda text, _multimodal=None, **kwargs: [ord(c) for c in text]) + params = SamplingParams() + params.init(manager.tokenizer, stop_sequences="END", n=2, best_of=2, max_new_tokens=10) + + async def run(): + stopped = asyncio.Event() + cleaned = asyncio.Event() + segments = [] + manager.abort = AsyncMock(side_effect=lambda *args, **kwargs: stopped.set()) + original_remove_req = manager.remove_req + + async def remove_req(*args, **kwargs): + await original_remove_req(*args, **kwargs) + if not manager.active_req_ids: + cleaned.set() + + manager.remove_req = remove_req + + async def raw_segment(p_node, d_node, start_time, prompt, sp, multimodal_params, request): + index = len(segments) + segments.append(sp.group_request_id) + metadata = {"prompt_tokens": 10, "prompt_cache_len": 0, "count_output_tokens": 1, "id": 10} + if index == 0: + yield sp.group_request_id, "helloENDextra", dict(metadata), FinishStatus() + await asyncio.sleep(0.01) + yield sp.group_request_id, "late", dict(metadata), FinishStatus() + else: + yield sp.group_request_id, "otherE", dict(metadata), FinishStatus() + await stopped.wait() + # 保持 sibling 未结束，验证另一个 choice 停止后仍能继续生成。 + await asyncio.sleep(0.01) + yield sp.group_request_id, "X", {**metadata, "id": 11}, FinishStatus(FinishStatus.FINISHED_LENGTH) + + manager._wait_to_token_package = raw_segment + results = [ + result + async for result in manager.generate( + "prompt", + params, + SimpleNamespace(images=[], audios=[], verify_and_preload=lambda req: asyncio.sleep(0)), + None, + ) + ] + ids = sorted({result[0] for result in results}) + assert len(ids) == 2 + assert "".join(result[1] for result in results if result[0] == ids[0]) == "hello" + assert "".join(result[1] for result in results if result[0] == ids[1]) == "otherEX" + assert {result[0]: result[3].get_finish_reason() for result in results if result[3].is_finished()} == { + ids[0]: "stop", + ids[1]: "length", + } + await asyncio.wait_for(cleaned.wait(), timeout=1) + assert not manager.active_req_ids + + asyncio.run(asyncio.wait_for(run(), timeout=2)) diff --git a/unit_tests/server/httpserver/test_pd_master_discovery.py b/unit_tests/server/httpserver/test_pd_master_discovery.py new file mode 100644 index 0000000000..25c4a49163 --- /dev/null +++ b/unit_tests/server/httpserver/test_pd_master_discovery.py @@ -0,0 +1,371 @@ +import asyncio +import base64 +import pickle +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import httpx +import pytest + +from lightllm.server.httpserver import pd_loop +from lightllm.server.pd_io_struct import PD_Master_Obj + + +def _registry_payload(masters): + return {"data": base64.b64encode(pickle.dumps(masters)).decode("ascii")} + + +def _install_http(monkeypatch, handler): + original_client = httpx.AsyncClient + options = [] + + def client(**kwargs): + options.append(kwargs) + return original_client(transport=httpx.MockTransport(handler), **kwargs) + + monkeypatch.setattr(pd_loop.httpx, "AsyncClient", client) + monkeypatch.setattr( + pd_loop, "get_shm_port_args", lambda: SimpleNamespace(pd_master_port=8000, config_server_port=9000) + ) + return options + + +@pytest.mark.parametrize("config_host,config_port", [(None, None), ("config", None), (None, 9000)]) +@pytest.mark.parametrize("worker_ports", [[8000], [8100, 8101]]) +def test_without_config_server_reads_worker_ports_before_connecting( + monkeypatch, config_host, config_port, worker_ports +): + requests = [] + + def handler(request): + requests.append(str(request.url)) + if request.url.path == "/get_server_info": + worker_port = request.url.port + return httpx.Response( + 200, + json={ + "run_mode": "pd_master", + "pd_master_multi_ports": worker_ports, + "pd_node_id": 100 + max(0, worker_port - 8100), + "host": "internal-master", + "port": worker_port, + "httpserver_workers": len(worker_ports), + }, + ) + raise AssertionError(f"unexpected request: {request.url}") + + options = _install_http(monkeypatch, handler) + args = SimpleNamespace(config_server_host=config_host, config_server_port=config_port, pd_master_ip="seed") + masters = asyncio.run(pd_loop._get_pd_master_objs(args)) + if len(worker_ports) > 1: + assert masters == { + 8100: PD_Master_Obj(8100, "seed:8100"), + 8101: PD_Master_Obj(8101, "seed:8101"), + } + assert requests == ["http://seed:8000/get_server_info"] + else: + assert masters == {100: PD_Master_Obj(100, "seed:8000")} + assert requests == ["http://seed:8000/get_server_info"] + assert options == [{"trust_env": False}] + + +@pytest.mark.parametrize("status", [404, 405, 500, 503]) +def test_mode_probe_http_errors_do_not_guess_a_single_master(monkeypatch, status): + requests = [] + + def handler(request): + requests.append(str(request.url)) + return httpx.Response(status) + + _install_http(monkeypatch, handler) + sleep = AsyncMock() + monkeypatch.setattr(pd_loop.asyncio, "sleep", sleep) + args = SimpleNamespace(config_server_host=None, config_server_port=None, pd_master_ip="seed") + + assert asyncio.run(pd_loop._get_pd_master_objs(args)) is None + assert requests == ["http://seed:8000/get_server_info"] + sleep.assert_awaited_once_with(10) + + +@pytest.mark.parametrize( + "info", + [ + {"run_mode": "prefill", "pd_master_multi_ports": [8000], "pd_node_id": 0}, + {"run_mode": "pd_master", "pd_node_id": 0}, + {"run_mode": "pd_master", "pd_master_multi_ports": [], "pd_node_id": 0}, + ], +) +def test_invalid_mode_responses_do_not_guess_a_single_master(monkeypatch, info): + requests = [] + + def handler(request): + requests.append(str(request.url)) + return httpx.Response(200, json=info) + + _install_http(monkeypatch, handler) + monkeypatch.setattr(pd_loop.asyncio, "sleep", AsyncMock()) + args = SimpleNamespace(config_server_host=None, config_server_port=None, pd_master_ip="seed") + + assert asyncio.run(pd_loop._get_pd_master_objs(args)) is None + assert requests == ["http://seed:8000/get_server_info"] + + +def test_discovered_multiprocess_master_does_not_query_worker_ports(monkeypatch): + requests = [] + + def handler(request): + requests.append(str(request.url)) + return httpx.Response( + 200, + json={ + "run_mode": "pd_master", + "httpserver_workers": 2, + "pd_master_multi_ports": [8100, 8101], + }, + ) + + _install_http(monkeypatch, handler) + args = SimpleNamespace(config_server_host=None, config_server_port=None, pd_master_ip="seed") + + assert asyncio.run(pd_loop._get_pd_master_objs(args)) == { + 8100: PD_Master_Obj(8100, "seed:8100"), + 8101: PD_Master_Obj(8101, "seed:8101"), + } + assert requests == ["http://seed:8000/get_server_info"] + + +def test_single_worker_discovery_keeps_the_reachable_external_endpoint(monkeypatch): + requests = [] + + def handler(request): + requests.append(str(request.url)) + return httpx.Response( + 200, + json={ + "run_mode": "pd_master", + "httpserver_workers": 1, + "pd_master_multi_ports": [8080], + "pd_node_id": 100, + "port": 8080, + }, + ) + + _install_http(monkeypatch, handler) + monkeypatch.setattr( + pd_loop, "get_shm_port_args", lambda: SimpleNamespace(pd_master_port=28080, config_server_port=9000) + ) + args = SimpleNamespace(config_server_host=None, config_server_port=None, pd_master_ip="master") + + assert asyncio.run(pd_loop._get_pd_master_objs(args)) == {100: PD_Master_Obj(100, "master:28080")} + assert requests == ["http://master:28080/get_server_info"] + + +def test_config_server_discovery_keeps_the_original_registry_protocol_and_uses_http(monkeypatch): + requests = [] + registered = {100: PD_Master_Obj(100, "master:8100"), 101: PD_Master_Obj(101, "master:8101")} + + def handler(request): + requests.append(str(request.url)) + return httpx.Response(200, json=_registry_payload(registered)) + + options = _install_http(monkeypatch, handler) + args = SimpleNamespace(config_server_host="config", config_server_port=9000) + + assert asyncio.run(pd_loop._get_pd_master_objs(args)) == registered + assert requests == ["http://config:9000/registered_objects"] + assert options == [{"trust_env": False}] + + +@pytest.mark.parametrize("status", [404, 405, 400, 401, 500, 503]) +def test_config_server_http_errors_do_not_fall_back_to_a_direct_master(monkeypatch, status): + _install_http(monkeypatch, lambda request: httpx.Response(status)) + args = SimpleNamespace(config_server_host="config", config_server_port=9000, pd_master_ip="seed") + + assert asyncio.run(pd_loop._get_pd_master_objs(args)) is None + + +@pytest.mark.parametrize("failure", ["network", "invalid_pickle", "base_exception"]) +def test_discovery_retries_only_regular_exceptions(monkeypatch, failure): + class DiscoveryFailure(BaseException): + pass + + def handler(request): + if failure == "network": + raise httpx.ConnectError("not ready", request=request) + if failure == "base_exception": + raise DiscoveryFailure("discovery failed") + return httpx.Response(200, json={"data": base64.b64encode(b"invalid-pickle").decode()}) + + _install_http(monkeypatch, handler) + sleep = AsyncMock() + monkeypatch.setattr(pd_loop.asyncio, "sleep", sleep) + args = SimpleNamespace(config_server_host="config", config_server_port=9000, pd_master_ip="seed") + + if failure == "base_exception": + with pytest.raises(DiscoveryFailure): + asyncio.run(pd_loop._get_pd_master_objs(args)) + sleep.assert_not_awaited() + else: + assert asyncio.run(pd_loop._get_pd_master_objs(args)) is None + sleep.assert_awaited_once_with(10) + + +def test_discovery_follows_config_registry_registration_and_removal(monkeypatch): + registered = {} + _install_http(monkeypatch, lambda request: httpx.Response(200, json=_registry_payload(registered))) + args = SimpleNamespace(config_server_host="config", config_server_port=9000) + + assert asyncio.run(pd_loop._get_pd_master_objs(args)) == {} + registered[100] = PD_Master_Obj(100, "master:8100") + registered[101] = PD_Master_Obj(101, "[2001:db8::1]:8101") + assert asyncio.run(pd_loop._get_pd_master_objs(args)) == registered + del registered[100] + assert asyncio.run(pd_loop._get_pd_master_objs(args)) == registered + registered.clear() + assert asyncio.run(pd_loop._get_pd_master_objs(args)) == {} + + +@pytest.mark.parametrize("host", ["::1", "[::1]", "2001:db8::1234"]) +def test_ipv6_direct_master_address_is_bracketed(monkeypatch, host): + _install_http( + monkeypatch, + lambda request: httpx.Response( + 200, + json={ + "run_mode": "pd_master", + "httpserver_workers": 1, + "pd_master_multi_ports": [8000], + "pd_node_id": 0, + "port": 8000, + }, + ), + ) + args = SimpleNamespace(config_server_host=None, config_server_port=None, pd_master_ip=host) + masters = asyncio.run(pd_loop._get_pd_master_objs(args)) + expected_host = host if host.startswith("[") else f"[{host}]" + + assert masters == {0: PD_Master_Obj(0, f"{expected_host}:8000")} + + +def test_ipv6_config_server_address_is_bracketed_for_http(monkeypatch): + requests = [] + + def handler(request): + requests.append(str(request.url)) + return httpx.Response(200, json=_registry_payload({})) + + _install_http(monkeypatch, handler) + args = SimpleNamespace(config_server_host="2001:db8::1", config_server_port=9000) + + assert asyncio.run(pd_loop._get_pd_master_objs(args)) == {} + assert requests == ["http://[2001:db8::1]:9000/registered_objects"] + + +def test_discovery_cancellation_closes_http_client_without_retry(monkeypatch): + async def run(): + started = asyncio.Event() + closed = asyncio.Event() + + class Client: + async def __aenter__(self): + return self + + async def __aexit__(self, *_exc): + closed.set() + + async def get(self, _url): + started.set() + await asyncio.Event().wait() + + monkeypatch.setattr(pd_loop.httpx, "AsyncClient", lambda **_kwargs: Client()) + monkeypatch.setattr(pd_loop, "get_shm_port_args", lambda: SimpleNamespace(config_server_port=9000)) + sleep = AsyncMock() + monkeypatch.setattr(pd_loop.asyncio, "sleep", sleep) + args = SimpleNamespace(config_server_host="config", config_server_port=9000) + task = asyncio.create_task(pd_loop._get_pd_master_objs(args)) + await started.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + sleep.assert_not_awaited() + assert closed.is_set() + + asyncio.run(run()) + + +def test_registry_changes_connect_all_owners_and_retire_removed_ids(monkeypatch): + async def run(): + manager = SimpleNamespace(args=SimpleNamespace(host="prefill")) + responses = asyncio.Queue() + discoveries = asyncio.Queue() + started = asyncio.Queue() + children = [] + live = set() + original_sleep = asyncio.sleep + + async def skip_poll_interval(seconds): + assert seconds == 30, "Topology refresh unexpectedly failed" + await original_sleep(0) + + async def discover(_args): + value = await responses.get() + discoveries.put_nowait(value) + return value + + async def connect(_manager, master, _rank): + children.append(asyncio.current_task()) + live.add(master.node_id) + started.put_nowait(master.node_id) + try: + await asyncio.Event().wait() + finally: + live.remove(master.node_id) + + monkeypatch.setattr(pd_loop.asyncio, "sleep", skip_poll_interval) + monkeypatch.setattr(pd_loop, "_get_pd_master_objs", discover) + monkeypatch.setattr(pd_loop, "_pd_handle_task", connect) + task = asyncio.create_task(pd_loop.pd_handle_loop(manager)) + try: + responses.put_nowait({node_id: PD_Master_Obj(node_id, f"master:{node_id}") for node_id in (100, 101)}) + assert {await asyncio.wait_for(started.get(), timeout=1) for _ in range(2)} == {100, 101} + await discoveries.get() + responses.put_nowait(None) + assert await discoveries.get() is None + await original_sleep(0) + assert live == {100, 101} + responses.put_nowait({node_id: PD_Master_Obj(node_id, f"master:{node_id}") for node_id in (101, 200)}) + assert await asyncio.wait_for(started.get(), timeout=1) == 200 + assert live == {101, 200} + await discoveries.get() + responses.put_nowait({}) + assert await discoveries.get() == {} + await original_sleep(0) + assert not live + finally: + task.cancel() + for child in children: + child.cancel() + await asyncio.gather(task, *children, return_exceptions=True) + + asyncio.run(run()) + + +@pytest.mark.parametrize("dp_index,load", [(None, 0.5), (0, 0.2), (1, 0.8)]) +def test_load_report_keeps_its_original_fields_and_formats_ipv6_keys(monkeypatch, dp_index, load): + from lightllm.server import api_http + + monkeypatch.setattr( + api_http, + "g_objs", + SimpleNamespace( + args=SimpleNamespace(dp=2, nnodes=1), + httpserver_manager=SimpleNamespace(host_ip="2001:db8::1"), + shared_token_load=SimpleNamespace(get_dynamic_max_load=lambda rank: [0.2, 0.8][rank]), + ), + ) + monkeypatch.setattr(pd_loop, "get_shm_port_args", lambda: SimpleNamespace(port=8100)) + + assert pd_loop._get_load_info(dp_index) == { + "total_token_usage_rate": load, + "connection_key": f"[2001:db8::1]:8100/dp{dp_index}", + } diff --git a/unit_tests/server/httpserver/test_pd_split_connections.py b/unit_tests/server/httpserver/test_pd_split_connections.py new file mode 100644 index 0000000000..e4fe76eafb --- /dev/null +++ b/unit_tests/server/httpserver/test_pd_split_connections.py @@ -0,0 +1,294 @@ +import asyncio +import json +import pickle +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from lightllm.server.api_cli import make_argument_parser +from lightllm.server.core.objs import FinishStatus, SamplingParams, StartArgs +from lightllm.server.httpserver import pd_loop +from lightllm.server.httpserver.manager import HttpServerManager +from lightllm.server.httpserver_for_pd_master.manager import PDManager +from lightllm.server.pd_io_struct import ObjType, PD_Master_Obj + +pytestmark = pytest.mark.usefixtures("single_http_worker_start_args") + + +def _registration(args, dp_index=None, mode="prefill"): + info = { + "node_id": 1, + "client_ip_port": "p:8000" if mode == "prefill" else "d:8000", + "mode": mode, + "start_args": dict(vars(args), dp=8, nnodes=2), + } + if dp_index is not None: + info["dp_index"] = dp_index + return info + + +def test_split_connections_register_select_update_and_reconnect_independently(): + args = StartArgs(select_p_d_node_strategy="cache_aware", pd_master_mode="1p1d") + manager = PDManager(args) + sockets = [object() for _ in range(4)] + for rank, websocket in enumerate(sockets): + manager.register_pd(_registration(args, rank), websocket) + manager.register_pd(_registration(args, mode="decode"), object()) + assert len(manager.prefill_nodes) == 4 + assert manager.is_pd_nodes_ready() + assert {node.to_llm_url() for node in manager.prefill_nodes} == {"http://p:8000/pd_generate_stream"} + + prompt = "x" * 1025 + cached_node = manager.prefill_nodes[3] + asyncio.run(manager.selector.insert_prompt_cache(prompt, cached_node)) + selected, _, extra = asyncio.run(manager.selector.select_p_d_node(prompt, None, None)) + assert selected is cached_node + assert extra.estimated_cache_hit_rate == 1.0 + for rank in range(4): + manager.update_node_load_info({"connection_key": f"p:8000/dp{rank}", "total_token_usage_rate": rank / 10}) + assert [node.run_status.total_token_usage_rate for node in manager.prefill_nodes] == [0, 0.1, 0.2, 0.3] + replacement_socket = object() + replacement_info = dict(_registration(args, 3), connection_id="replacement") + manager.register_pd(replacement_info, replacement_socket) + replacement = manager.url_to_pd_nodes["p:8000/dp3"] + assert replacement.connection_key == cached_node.connection_key + manager.remove_pd(_registration(args, 3)) + assert manager.url_to_pd_nodes["p:8000/dp3"] is replacement + assert replacement in manager.prefill_nodes + assert replacement in manager.selector.prefill_nodes + manager.remove_pd(replacement_info) + assert len(manager.prefill_nodes) == 3 + assert "p:8000/dp3" not in manager.url_to_pd_nodes + + +@pytest.mark.parametrize("mode", ["prefill", "decode"]) +def test_unsplit_reconnection_ignores_old_connection_id_cleanup(mode): + args = StartArgs() + manager = PDManager(args) + info = dict(_registration(args, mode=mode), connection_id="old") + new_info = dict(info, connection_id="new") + old_socket, new_socket = object(), object() + manager.register_pd(info, old_socket) + manager.register_pd(new_info, new_socket) + key = f"{info['client_ip_port']}/dpNone" + replacement = manager.url_to_pd_nodes[key] + manager.remove_pd(info) + assert manager.url_to_pd_nodes[key] is replacement + assert replacement.websocket is new_socket + manager.remove_pd(new_info) + assert key not in manager.url_to_pd_nodes + + +@pytest.mark.parametrize( + "split, dp, nnodes, expected", + [(False, 8, 2, [None]), (True, 8, 2, [0, 1, 2, 3]), (True, 1, 1, [0]), (True, 1, 2, [0])], +) +def test_prefill_opens_local_dp_connections(monkeypatch, split, dp, nnodes, expected): + async def run(): + args = StartArgs(host="p", run_mode="prefill", dp=dp, nnodes=nnodes, use_dp_split_mode_connect_pd_master=split) + manager = HttpServerManager.__new__(HttpServerManager) + manager.args = args + manager.pd_mode = pd_loop.NodeRole.P + manager.is_multinode_tp_slave = False + manager.recycle_resource_loop = AsyncMock() + started = [] + + async def connect(_manager, dp_index=None): + started.append(dp_index) + + async def receive(): + while len(started) < len(expected): + await asyncio.sleep(0) + raise asyncio.CancelledError() + + manager.zmq_recv_socket = SimpleNamespace(recv_pyobj=receive) + timer = AsyncMock() + monkeypatch.setattr(pd_loop, "timer_log", timer) + monkeypatch.setattr(pd_loop, "pd_handle_loop", connect) + with pytest.raises(asyncio.CancelledError): + await manager.handle_loop() + assert started == expected + timer.assert_awaited_once_with(manager) + + asyncio.run(run()) + + +@pytest.mark.parametrize("dp_index", [None, 0, 3]) +def test_pd_handle_loop_passes_its_fixed_rank_to_each_master(monkeypatch, dp_index): + async def run(): + manager = SimpleNamespace(args=StartArgs(host="p", run_mode="prefill")) + started = [] + tasks = [] + masters = {rank: PD_Master_Obj(rank, f"master{rank}:8000") for rank in (1, 2)} + + async def connect(_manager, master, rank): + started.append((master.node_id, rank)) + tasks.append(asyncio.current_task()) + await asyncio.Event().wait() + + async def stop_after_connections(_seconds): + while len(started) < len(masters): + await original_sleep(0) + raise asyncio.CancelledError() + + original_sleep = asyncio.sleep + monkeypatch.setattr(pd_loop, "_get_pd_master_objs", AsyncMock(return_value=masters)) + monkeypatch.setattr(pd_loop, "_pd_handle_task", connect) + monkeypatch.setattr(pd_loop.asyncio, "sleep", stop_after_connections) + try: + with pytest.raises(asyncio.CancelledError): + await pd_loop.pd_handle_loop(manager, dp_index) + assert started == [(1, dp_index), (2, dp_index)] + finally: + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + + asyncio.run(run()) + + +@pytest.mark.parametrize("dp_index", [None, 0, 3]) +def test_connection_registers_rank_binds_requests_and_reports_rank_load(monkeypatch, dp_index): + from lightllm.server import api_http + + async def run(): + args = StartArgs(run_mode="prefill", dp=8, nnodes=2) + manager = SimpleNamespace(args=args, host_ip="p", pd_mode=SimpleNamespace(value="prefill")) + monkeypatch.setattr( + api_http, + "g_objs", + SimpleNamespace( + args=args, + httpserver_manager=manager, + shared_token_load=SimpleNamespace(get_dynamic_max_load=lambda rank: rank / 10), + ), + ) + monkeypatch.setattr(pd_loop, "get_shm_port_args", lambda: SimpleNamespace(port=8000)) + params = SamplingParams() + params.init(None, max_new_tokens=1) + params.group_request_id = 123 + params.suggested_dp_index = -1 + dispatched = asyncio.Event() + payloads = [] + + async def generate(**kwargs): + req_params = kwargs["sampling_params"] + assert req_params.suggested_dp_index == (-1 if dp_index is None else dp_index) + assert kwargs["pd_upload_websocket"] is websocket + await kwargs["forwarding_queue"].put((123, "P", {}, FinishStatus())) + dispatched.set() + + requests = 0 + + async def receive(): + nonlocal requests + requests += 1 + if requests == 1: + return pickle.dumps((ObjType.REQ, ("prompt", params, None))) + await dispatched.wait() + while not any(isinstance(payload, tuple) and payload[0] == ObjType.TOKEN_PACKS for payload in payloads): + await asyncio.sleep(0) + raise asyncio.CancelledError() + + async def send(payload): + payloads.append(json.loads(payload) if isinstance(payload, str) else pickle.loads(payload)) + + websocket = SimpleNamespace( + send=send, recv=receive, transport=SimpleNamespace(get_extra_info=lambda _key: MagicMock()) + ) + + class Connection: + async def __aenter__(self): + return websocket + + async def __aexit__(self, *_args): + return False + + monkeypatch.setattr(pd_loop.websockets, "connect", lambda *_args, **_kwargs: Connection()) + monkeypatch.setattr(pd_loop, "_pd_process_generate", generate) + await asyncio.wait_for(pd_loop._pd_handle_task(manager, PD_Master_Obj(1, "master:8000"), dp_index), 2) + assert payloads[0].get("dp_index") == dp_index + assert "dp_index" in payloads[0] + for payload in payloads[1:]: + if payload[0] != ObjType.TOKEN_PACKS: + continue + load = payload[-1] + assert "dp_index" not in load + assert load["connection_key"] == f"p:8000/dp{dp_index}" + assert load["total_token_usage_rate"] == pytest.approx(0.15 if dp_index is None else dp_index / 10) + assert args.host != "p" # Registration must not mutate shared startup arguments. + + asyncio.run(run()) + + +@pytest.mark.parametrize("dp_index", [None, 0, 3]) +@pytest.mark.parametrize("with_source", [False, True], ids=["legacy", "connection-source"]) +def test_reported_load_key_matches_registered_connection_without_rewriting(monkeypatch, dp_index, with_source): + from fastapi import WebSocketDisconnect + from lightllm.server import api_http, api_http_pd + + args = StartArgs(dp=8, nnodes=2, pd_node_id=1) + manager = PDManager(args) + info = _registration(args, dp_index) + if with_source: + info["connection_id"] = f"load-source-{dp_index}" + monkeypatch.setattr( + api_http, + "g_objs", + SimpleNamespace( + args=args, + httpserver_manager=SimpleNamespace(host_ip="p"), + shared_token_load=SimpleNamespace(get_dynamic_max_load=lambda _rank: 0.7), + ), + ) + monkeypatch.setattr(pd_loop, "get_shm_port_args", lambda: SimpleNamespace(port=8000)) + load_info = pd_loop._get_load_info(dp_index) + message = (ObjType.TOKEN_PACKS, [], load_info) + messages = iter([pickle.dumps(message)]) + updated_nodes = [] + + async def receive(): + try: + return next(messages) + except StopIteration: + raise WebSocketDisconnect() + + async def handle(obj): + assert obj[2] == load_info + manager.update_node_load_info(obj[2]) + updated_nodes.append(manager.url_to_pd_nodes[obj[2]["connection_key"]]) + + websocket = SimpleNamespace( + accept=AsyncMock(), + client=("p", 8000), + receive_text=AsyncMock(return_value=json.dumps(info)), + receive_bytes=receive, + ) + monkeypatch.setattr( + api_http, + "g_objs", + SimpleNamespace( + httpserver_manager=SimpleNamespace( + register_pd=AsyncMock(side_effect=manager.register_pd), + remove_pd=AsyncMock(side_effect=manager.remove_pd), + put_to_handle_queue=handle, + ) + ), + ) + asyncio.run(api_http_pd.register_and_keep_alive(websocket)) + assert len(updated_nodes) == 1 + assert updated_nodes[0].dp_index == dp_index + assert updated_nodes[0].websocket is websocket + assert updated_nodes[0].run_status.total_token_usage_rate == 0.7 + + +def test_split_option_defaults_off_and_is_prefill_only(): + from lightllm.server.api_server import launch_server + + parser = make_argument_parser() + assert parser.parse_args([]).use_dp_split_mode_connect_pd_master is False + assert parser.parse_args(["--use_dp_split_mode_connect_pd_master"]).use_dp_split_mode_connect_pd_master is True + assert StartArgs().use_dp_split_mode_connect_pd_master is False + with pytest.raises(ValueError, match="only supports prefill"): + launch_server(StartArgs(run_mode="decode", use_dp_split_mode_connect_pd_master=True)) diff --git a/unit_tests/server/httpserver/test_pd_websocket_lifecycle.py b/unit_tests/server/httpserver/test_pd_websocket_lifecycle.py new file mode 100644 index 0000000000..0c261b5dda --- /dev/null +++ b/unit_tests/server/httpserver/test_pd_websocket_lifecycle.py @@ -0,0 +1,81 @@ +"""Existing PD WebSocket registration and message-forwarding behavior.""" + +import asyncio +import json +import pickle +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from lightllm.server import api_http, api_http_pd + + +def _websocket(**overrides): + values = dict( + accept=AsyncMock(), + close=AsyncMock(), + client=("pd", 18000), + receive_text=AsyncMock(return_value=json.dumps({"node_id": 1, "mode": "prefill"})), + receive_bytes=AsyncMock(), + send_bytes=AsyncMock(), + ) + values.update(overrides) + return SimpleNamespace(**values) + + +def _manager(monkeypatch, **overrides): + values = dict( + register_pd=AsyncMock(), + remove_pd=AsyncMock(), + put_to_handle_queue=AsyncMock(), + args=SimpleNamespace(pd_node_id=100), + update_req_status=AsyncMock(), + ) + values.update(overrides) + manager = SimpleNamespace(**values) + monkeypatch.setattr(api_http, "g_objs", SimpleNamespace(httpserver_manager=manager)) + return manager + + +def test_rejected_registration_does_not_remove_an_existing_connection(monkeypatch): + manager = _manager(monkeypatch, register_pd=AsyncMock(side_effect=ValueError("wrong model parameters"))) + websocket = _websocket() + with pytest.raises(ValueError, match="wrong model parameters"): + asyncio.run(api_http_pd.register_and_keep_alive(websocket)) + manager.remove_pd.assert_not_awaited() + + +def test_heartbeat_wait_is_bounded(monkeypatch): + async def run(): + original_wait_for = asyncio.wait_for + timeouts = [] + + async def wait_forever(): + await asyncio.Event().wait() + + async def short_heartbeat_wait(awaitable, timeout): + timeouts.append(timeout) + return await original_wait_for(awaitable, timeout=0.01) + + manager = _manager(monkeypatch) + monkeypatch.setattr(api_http_pd.asyncio, "wait_for", short_heartbeat_wait) + websocket = _websocket(receive_bytes=wait_forever) + await original_wait_for(api_http_pd.register_and_keep_alive(websocket), timeout=1) + assert timeouts == [30] + assert websocket.close.await_args.kwargs["code"] == 1011 + manager.remove_pd.assert_awaited_once() + + asyncio.run(run()) + + +def test_kv_status_is_forwarded_without_registration_or_reply(monkeypatch): + from fastapi import WebSocketDisconnect + from lightllm.server.pd_io_struct import PDUpKVStatus + + status = PDUpKVStatus(8, 100, b"metadata") + manager = _manager(monkeypatch) + websocket = _websocket(receive_bytes=AsyncMock(side_effect=[pickle.dumps(status), WebSocketDisconnect()])) + asyncio.run(api_http_pd.kv_move_status(websocket)) + manager.update_req_status.assert_awaited_once_with(status) + websocket.send_bytes.assert_not_awaited() diff --git a/unit_tests/server/httpserver/test_running_request_lifecycle.py b/unit_tests/server/httpserver/test_running_request_lifecycle.py new file mode 100644 index 0000000000..0487e45281 --- /dev/null +++ b/unit_tests/server/httpserver/test_running_request_lifecycle.py @@ -0,0 +1,565 @@ +import asyncio +import pickle +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from lightllm.server.core.objs import StopStrMatchMode, FinishStatus, SamplingParams +from lightllm.server.httpserver.manager import HttpServerManager, ReqStatus +from lightllm.server.pd_io_struct import NodeRole, ObjType +from lightllm.utils.error_utils import ClientDisconnected, PDPrefillNodeStopGenToken, ServerBusyError + + +class _ValueMark: + def __init__(self, value=0): + self.value = value + + def get_value(self): + return self.value + + def set_value(self, value): + self.value = value + + +def _make_manager(mode: NodeRole): + manager = HttpServerManager.__new__(HttpServerManager) + manager.args = SimpleNamespace( + run_mode=mode.value, + running_max_req_size=2, + reasoning_parser=None, + enable_rl=False, + enable_return_routed_experts=False, + ) + manager.pd_mode = mode + manager.is_multinode_tp_slave = False + manager.alloc_req_id = MagicMock(return_value=123) + manager.is_multinode_tp_master = False + manager.rl_controller = None + manager.output_grammar_compiler = None + manager.req_id_to_out_inf = {} + manager._log_stage_timing = MagicMock() + manager._log_req_header = AsyncMock() + manager._encode = AsyncMock(return_value=[10, 11, 12]) + manager._check_and_repair_length = AsyncMock(side_effect=lambda prompt_ids, _params: prompt_ids) + manager._release_multimodal_resources = AsyncMock() + manager.abort = AsyncMock() + manager._register_running_request = AsyncMock() + manager._unregister_running_request = AsyncMock() + manager.metric_client = MagicMock() + manager._run_reqs_count_lock = asyncio.Lock() + manager.run_reqs_count_mark = _ValueMark() + manager.shm_req_manager = SimpleNamespace( + async_alloc_req_indexes=AsyncMock(side_effect=RuntimeError("alloc failed")) + ) + return manager + + +def _sampling_params(): + sampling_params = SamplingParams() + sampling_params.group_request_id = 123 + sampling_params.n = 1 + sampling_params.max_new_tokens = 1 + sampling_params.pd_node_resource_wait_timeout_seconds = 20 + return sampling_params + + +def _multimodal_params(): + return SimpleNamespace(audios=[], images=[], verify_and_preload=AsyncMock()) + + +@pytest.mark.parametrize("mode", [NodeRole.NORMAL, NodeRole.P, NodeRole.D]) +@pytest.mark.parametrize("warm_cache", [False, True]) +def test_http_publishes_grammar_for_all_choices_before_forwarding(compiler, mode, warm_cache): + import xgrammar as xgr + + async def run(): + manager = _make_manager(mode) + manager.output_grammar_compiler = compiler + manager.tokenizer = None + manager.enable_multimodal = False + manager.args.chunked_prefill_size = 1 + manager._alloc_shm_req_indexes = AsyncMock(return_value=[0, 1]) + reqs = [MagicMock(), MagicMock()] + manager.shm_req_manager.async_get_req_obj_by_index = AsyncMock(side_effect=reqs) + + async def forward(*args): + payloads = [] + for req in reqs: + req.init.assert_called_once() + req.set_compiled_grammar.assert_called_once() + payload = req.set_compiled_grammar.call_args.args[0] + payloads.append(payload) + grammar = compiler.output_constraint_state_factory.get_grammar(payload) + matcher = xgr.GrammarMatcher(grammar) + assert matcher.accept_token(ord("a")) and matcher.accept_token(ord("b")) + assert matcher.accept_token(256) and matcher.is_terminated() + assert payloads[0] is payloads[1] + + manager.transfer_to_next_module_or_node = AsyncMock(side_effect=forward) + + async def results(*args): + for req_id in (123, 124): + yield req_id, "", {"id": 256}, FinishStatus(FinishStatus.FINISHED_STOP) + + manager._wait_to_token_package = results + params = SamplingParams() + params.init(None, regular_constraint="ab", n=2, best_of=2, max_new_tokens=1, group_request_id=123) + with patch.object(compiler, "_compile", wraps=compiler._compile) as compile_grammar: + if warm_cache: + await params.verify_cpu_async(compiler) + params = pickle.loads(pickle.dumps(SamplingParams.from_buffer_copy(params))) + outputs = [result async for result in manager.generate("prompt", params, _multimodal_params(), None)] + assert len(outputs) == 2 + compile_grammar.assert_called_once_with("regex", "ab") + manager.transfer_to_next_module_or_node.assert_awaited_once() + assert not hasattr(params, "compiled_grammar") + + asyncio.run(run()) + + +@pytest.mark.parametrize("mode", [NodeRole.NORMAL, NodeRole.P, NodeRole.D]) +def test_grammar_get_failure_precedes_request_admission_and_pd_handshake(mode): + async def run(): + manager = _make_manager(mode) + manager.output_grammar_compiler = SimpleNamespace(compile=AsyncMock(side_effect=ValueError("compile failed"))) + manager._alloc_shm_req_indexes = AsyncMock() + params = SamplingParams() + params.init(None, regular_constraint="ab", max_new_tokens=1, group_request_id=123) + multimodal_params = _multimodal_params() + websocket = AsyncMock() if mode == NodeRole.P else None + pd_event = asyncio.Event() if mode == NodeRole.P else None + with pytest.raises(ValueError, match="compile failed"): + await _drain_generate(manager, params, multimodal_params, websocket, pd_event) + manager.alloc_req_id.assert_not_called() + manager._encode.assert_not_awaited() + multimodal_params.verify_and_preload.assert_not_awaited() + manager._alloc_shm_req_indexes.assert_not_awaited() + manager._register_running_request.assert_not_awaited() + manager._unregister_running_request.assert_not_awaited() + manager.abort.assert_not_awaited() + if websocket is not None: + websocket.send.assert_not_awaited() + + asyncio.run(run()) + + +@pytest.mark.parametrize("mode", [NodeRole.NORMAL, NodeRole.D, NodeRole.P]) +@pytest.mark.parametrize("include", [False, True]) +@pytest.mark.parametrize("parser", [None, "qwen3"]) +@pytest.mark.parametrize("stop_mode", [None, StopStrMatchMode.ALL_TEXT, StopStrMatchMode.CONTENT_ONLY]) +def test_http_forwards_packaged_tokens_and_preserves_stop_policy(mode, include, parser, stop_mode): + async def run(): + manager = _make_manager(mode) + manager.tokenizer = SimpleNamespace(encode=lambda text, **kwargs: [ord(c) for c in text]) + manager.enable_multimodal = False + manager.args.chunked_prefill_size = 1 + manager.args.reasoning_parser = parser + manager._alloc_shm_req_indexes = AsyncMock(return_value=[0]) + manager.shm_req_manager.async_get_req_obj_by_index = AsyncMock(return_value=MagicMock()) + manager.transfer_to_next_module_or_node = AsyncMock() + + async def results(*args): + yield 123, "hello", {"id": 1}, FinishStatus() + yield 123, "END" if include else "", {"id": 2}, FinishStatus(FinishStatus.FINISHED_STOP) + + manager._wait_to_token_package = results + params = _sampling_params() + params.init( + manager.tokenizer, + stop_sequences="END", + include_stop_str_in_output=include, + max_new_tokens=1, + group_request_id=123, + pd_node_resource_wait_timeout_seconds=20, + ) + # P/D nodes keep the matching policy supplied by Master. + if stop_mode is not None: + params.stop_str_match_mode.set_mode(stop_mode) + websocket = AsyncMock() if mode == NodeRole.P else None + pd_event = None + if mode == NodeRole.P: + pd_event = asyncio.Event() + pd_event.decode_node_info = SimpleNamespace(ready_kv_len=1) + pd_event.set() + + results = [ + result + async for result in manager.generate("prompt", params, _multimodal_params(), None, websocket, pd_event) + ] + expected_mode = StopStrMatchMode.ALL_TEXT if stop_mode is None else stop_mode + assert params.stop_str_match_mode.get_mode() == expected_mode + expected = "helloEND" if include else "hello" + assert "".join(result[1] for result in results) == expected + assert [result[2]["id"] for result in results] == [1, 2] + assert results[0][2]["input_usage"]["input_text_tokens"] == 3 + assert results[-1][3].get_finish_reason() == "stop" + manager.shm_req_manager.async_get_req_obj_by_index.return_value.set_compiled_grammar.assert_not_called() + manager._unregister_running_request.assert_awaited_once() + + asyncio.run(run()) + + +def _req_status(reqs): + return ReqStatus(123, None, reqs, 0) + + +async def _drain_generate(manager, sampling_params, multimodal_params, websocket=None, pd_event=None): + async for _ in manager.generate( + prompt="prompt", + sampling_params=sampling_params, + multimodal_params=multimodal_params, + request=None, + pd_upload_websocket=websocket, + pd_event=pd_event, + ): + pass + + +@pytest.mark.parametrize("mode", [NodeRole.NORMAL, NodeRole.D]) +def test_non_prefill_request_is_counted_before_early_failure_and_always_unregistered( + mode, +): + async def run(): + manager = _make_manager(mode) + manager._encode.side_effect = RuntimeError("encode failed") + + with pytest.raises(RuntimeError, match="encode failed"): + await _drain_generate(manager, _sampling_params(), _multimodal_params()) + + manager._register_running_request.assert_awaited_once() + manager._unregister_running_request.assert_awaited_once() + + asyncio.run(run()) + + +def test_prefill_encode_failure_is_outside_running_request_count(): + async def run(): + manager = _make_manager(NodeRole.P) + manager._encode.side_effect = RuntimeError("encode failed") + + with pytest.raises(RuntimeError, match="encode failed"): + await _drain_generate( + manager, + _sampling_params(), + _multimodal_params(), + AsyncMock(), + asyncio.Event(), + ) + + manager._register_running_request.assert_not_awaited() + manager._unregister_running_request.assert_not_awaited() + + asyncio.run(run()) + + +def test_prefill_full_cache_hit_never_enters_running_request_count(): + async def run(): + manager = _make_manager(NodeRole.P) + websocket = AsyncMock() + pd_event = asyncio.Event() + pd_event.decode_node_info = SimpleNamespace(ready_kv_len=2) + pd_event.set() + + with pytest.raises(PDPrefillNodeStopGenToken): + await _drain_generate(manager, _sampling_params(), _multimodal_params(), websocket, pd_event) + + manager._register_running_request.assert_not_awaited() + manager._unregister_running_request.assert_not_awaited() + websocket.send.assert_awaited_once() + + asyncio.run(run()) + + +def test_prefill_upload_preserves_int64_multimodal_prompt_ids(): + async def run(): + manager = _make_manager(NodeRole.P) + multimodal_token_id = 2 ** 32 + manager._encode.return_value = [10, multimodal_token_id, 12] + websocket = AsyncMock() + pd_event = asyncio.Event() + pd_event.decode_node_info = SimpleNamespace(ready_kv_len=2) + pd_event.set() + + with pytest.raises(PDPrefillNodeStopGenToken): + await _drain_generate(manager, _sampling_params(), _multimodal_params(), websocket, pd_event) + + obj_type, group_request_id, prompt_ids = pickle.loads(websocket.send.await_args.args[0]) + assert obj_type == ObjType.PD_UPLOAD_PREFILL_PROMPT_IDS + assert group_request_id == 123 + assert prompt_ids.itemsize == 8 + assert prompt_ids.tolist() == [10, multimodal_token_id, 12] + + asyncio.run(run()) + + +def test_prefill_waiting_for_decode_assignment_can_be_cancelled_without_touching_count(): + async def run(): + manager = _make_manager(NodeRole.P) + websocket = AsyncMock() + pd_event = asyncio.Event() + task = asyncio.create_task( + _drain_generate(manager, _sampling_params(), _multimodal_params(), websocket, pd_event) + ) + + for _ in range(100): + if websocket.send.await_count: + break + await asyncio.sleep(0) + assert websocket.send.await_count == 1 + assert not task.done() + manager._register_running_request.assert_not_awaited() + + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + manager._register_running_request.assert_not_awaited() + manager._unregister_running_request.assert_not_awaited() + + asyncio.run(run()) + + +def test_prefill_is_counted_after_decode_assignment_and_unregistered_on_following_failure(): + async def run(): + manager = _make_manager(NodeRole.P) + websocket = AsyncMock() + pd_event = asyncio.Event() + pd_event.decode_node_info = SimpleNamespace(ready_kv_len=1) + pd_event.set() + + with pytest.raises(RuntimeError, match="alloc failed"): + await _drain_generate(manager, _sampling_params(), _multimodal_params(), websocket, pd_event) + + manager._register_running_request.assert_awaited_once() + manager._unregister_running_request.assert_awaited_once() + + asyncio.run(run()) + + +@pytest.mark.parametrize("mode", [NodeRole.P, NodeRole.D]) +def test_pd_node_returns_busy_when_shm_req_allocation_times_out(mode): + async def run(): + manager = _make_manager(mode) + manager.shm_req_manager = SimpleNamespace( + async_alloc_req_indexes=AsyncMock(return_value=None), + async_release_req_index=AsyncMock(), + ) + + websocket = AsyncMock() if mode == NodeRole.P else None + pd_event = None + if mode == NodeRole.P: + pd_event = asyncio.Event() + pd_event.decode_node_info = SimpleNamespace(ready_kv_len=1) + pd_event.set() + + with patch("lightllm.server.httpserver.manager.time.monotonic", side_effect=[0, 21]): + with pytest.raises(ServerBusyError, match=f"PD {mode.value} node is busy"): + await _drain_generate( + manager, + _sampling_params(), + _multimodal_params(), + websocket, + pd_event, + ) + + manager.shm_req_manager.async_alloc_req_indexes.assert_awaited_once_with(1) + manager.shm_req_manager.async_release_req_index.assert_not_awaited() + manager._register_running_request.assert_awaited_once() + manager._unregister_running_request.assert_awaited_once() + + asyncio.run(run()) + + +@pytest.mark.parametrize("mode", [NodeRole.P, NodeRole.D]) +def test_pd_node_returns_busy_while_first_token_request_waits_in_router(mode): + async def run(): + manager = _make_manager(mode) + req = SimpleNamespace( + request_id=123, + is_aborted=False, + router_arrival_time=0, + infer_start_time=0, + sample_params=SimpleNamespace( + high_priority_request=False, + pd_node_resource_wait_timeout_seconds=20, + ), + ) + req_status = _req_status([req]) + req_status.event.set() + sampling_params = _sampling_params() + + with patch("lightllm.server.httpserver.manager.time.monotonic", return_value=21): + req.router_arrival_time = 1.0 + output_generator = manager._wait_to_token_package( + start_time=0, + prompt_ids=[], + group_request_id=123, + sampling_params=sampling_params, + req_status=req_status, + request=None, + ) + with pytest.raises(ServerBusyError, match="request did not enter inference"): + await output_generator.__anext__() + + asyncio.run(run()) + + +def test_multinode_tp_slave_does_not_apply_router_wait_timeout(): + async def run(): + manager = _make_manager(NodeRole.D) + manager.is_multinode_tp_slave = True + req = SimpleNamespace( + router_arrival_time=1.0, + infer_start_time=0.0, + sample_params=SimpleNamespace(pd_node_resource_wait_timeout_seconds=0), + ) + req_status = _req_status([req]) + req_status.event.set() + request = SimpleNamespace(is_disconnected=AsyncMock(return_value=True)) + + with patch("lightllm.server.httpserver.manager.time.monotonic", return_value=2): + # ReqStatus 只判断时间条件；是否允许 slave 执行超时策略由 manager 在调用处决定。 + assert req_status.has_timed_out_waiting_for_inference() is True + output_generator = manager._wait_to_token_package( + start_time=0, + prompt_ids=[], + group_request_id=123, + sampling_params=_sampling_params(), + req_status=req_status, + request=request, + ) + with pytest.raises(ClientDisconnected): + await output_generator.__anext__() + + asyncio.run(run()) + + +def test_httpserver_keeps_started_requests_and_requests_with_remaining_master_timeout(): + waiting_req = SimpleNamespace( + router_arrival_time=1.0, + infer_start_time=0.0, + sample_params=SimpleNamespace( + high_priority_request=False, + pd_node_resource_wait_timeout_seconds=60, + ), + ) + started_req = SimpleNamespace( + router_arrival_time=1.0, + infer_start_time=2.0, + sample_params=SimpleNamespace(high_priority_request=False), + ) + req_status = _req_status([waiting_req, started_req]) + + with patch("lightllm.server.httpserver.manager.time.monotonic", return_value=30): + assert req_status.has_timed_out_waiting_for_inference() is False + + started_req.infer_start_time = 0.0 + assert req_status.has_timed_out_waiting_for_inference() is False + + +def test_pd_node_allocation_failure_does_not_hold_partial_slots(): + async def run(): + manager = _make_manager(NodeRole.D) + manager.shm_req_manager = SimpleNamespace( + async_alloc_req_indexes=AsyncMock(side_effect=RuntimeError("alloc failed")), + async_release_req_index=AsyncMock(), + ) + sampling_params = _sampling_params() + sampling_params.n = 2 + + with pytest.raises(RuntimeError, match="alloc failed"): + await _drain_generate(manager, sampling_params, _multimodal_params()) + + manager.shm_req_manager.async_alloc_req_indexes.assert_awaited_once_with(2) + manager.shm_req_manager.async_release_req_index.assert_not_awaited() + + asyncio.run(run()) + + +@pytest.mark.parametrize("high_priority", [False, True]) +def test_batch_allocation_retries_with_priority_backoff(high_priority): + async def run(): + manager = _make_manager(NodeRole.D) + manager.shm_req_manager.async_alloc_req_indexes = AsyncMock(side_effect=[None, None, [0, 1]]) + with patch("lightllm.server.httpserver.manager.asyncio.sleep", new_callable=AsyncMock) as sleep: + assert await manager._alloc_shm_req_indexes(2, high_priority_request=high_priority) == [0, 1] + factor = 0.2 if high_priority else 1 + assert [call.args[0] for call in sleep.await_args_list] == pytest.approx([0.1 * factor, 0.11 * factor]) + assert [call.args for call in manager.shm_req_manager.async_alloc_req_indexes.await_args_list] == [(2,)] * 3 + + asyncio.run(run()) + + +def test_batch_allocation_rejects_requests_larger_than_capacity(): + async def run(): + manager = _make_manager(NodeRole.NORMAL) + with pytest.raises(ValueError, match="must not exceed running_max_req_size"): + await manager._alloc_shm_req_indexes(3) + manager.shm_req_manager.async_alloc_req_indexes.assert_not_awaited() + + asyncio.run(run()) + + +def test_multinode_tp_slave_does_not_apply_allocation_wait_timeout(): + async def run(): + manager = _make_manager(NodeRole.D) + manager.is_multinode_tp_slave = True + manager.shm_req_manager.async_alloc_req_indexes = AsyncMock(side_effect=[None, [0, 1]]) + with patch("lightllm.server.httpserver.manager.asyncio.sleep", new_callable=AsyncMock): + with patch("lightllm.server.httpserver.manager.time.monotonic") as monotonic: + assert await manager._alloc_shm_req_indexes(2, pd_node_resource_wait_timeout_seconds=0) == [0, 1] + monotonic.assert_not_called() + + asyncio.run(run()) + + +def test_cancelled_batch_allocation_does_not_hold_slots(): + async def run(): + manager = _make_manager(NodeRole.NORMAL) + manager.shm_req_manager = SimpleNamespace( + async_alloc_req_indexes=AsyncMock(return_value=None), + async_release_req_index=AsyncMock(), + ) + task = asyncio.create_task(manager._alloc_shm_req_indexes(2)) + await asyncio.sleep(0) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + manager.shm_req_manager.async_alloc_req_indexes.assert_awaited_once_with(2) + manager.shm_req_manager.async_release_req_index.assert_not_awaited() + + asyncio.run(run()) + + +def test_running_request_helpers_are_atomic_and_refresh_timestamp_only_on_idle_to_running_transition(): + async def run(): + manager = HttpServerManager.__new__(HttpServerManager) + manager._run_reqs_count_lock = asyncio.Lock() + manager.run_reqs_count_mark = _ValueMark() + manager.latest_success_infer_time_mark = _ValueMark(7) + + with patch("lightllm.server.httpserver.manager.time.time", side_effect=[1000, 2000]): + await asyncio.gather(*(manager._register_running_request() for _ in range(100))) + assert manager.run_reqs_count_mark.get_value() == 100 + assert manager.latest_success_infer_time_mark.get_value() == 1000 + + await asyncio.gather(*(manager._unregister_running_request() for _ in range(100))) + assert manager.run_reqs_count_mark.get_value() == 0 + + await manager._register_running_request() + assert manager.run_reqs_count_mark.get_value() == 1 + assert manager.latest_success_infer_time_mark.get_value() == 2000 + await manager._unregister_running_request() + assert manager.run_reqs_count_mark.get_value() == 0 + + asyncio.run(run()) + + +def test_real_supported_max_req_total_len_reserves_mtp_overlap_tokens(): + manager = HttpServerManager.__new__(HttpServerManager) + manager.shm_max_total_token_num = _ValueMark(1000) + manager.max_req_total_len = 2000 + manager.args = SimpleNamespace(mtp_step=2, page_size=4) + + assert manager.get_real_supported_max_req_total_len() == 951 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 new file mode 100644 index 0000000000..1b2675737a --- /dev/null +++ b/unit_tests/server/multi_level_kv_cache/test_cache_placement_controller.py @@ -0,0 +1,367 @@ +from types import SimpleNamespace +from concurrent.futures import ThreadPoolExecutor +import inspect +import pytest + +from lightllm.server.api_cli import make_argument_parser +from lightllm.server.multi_level_kv_cache import ( + AdaptiveCachePlacementController, + CacheCapacityConfig, + CachePlacementController, + CacheTier, + GpuOnlyCachePlacementController, + LegacyCachePlacementController, + create_cache_placement_controller, +) + + +def make_controller( + gpu_tokens: int, + cpu_tokens: int, + disk_tokens: int, + history_size: int = AdaptiveCachePlacementController.MAX_HISTORY_SIZE, + initial_history_size: int = None, + enable_cpu_cache: bool = None, + enable_disk_cache: bool = None, +) -> CachePlacementController: + if enable_cpu_cache is None: + enable_cpu_cache = cpu_tokens > 0 + if enable_disk_cache is None: + enable_disk_cache = disk_tokens > 0 + if initial_history_size is None: + initial_history_size = min(AdaptiveCachePlacementController.INITIAL_HISTORY_SIZE, history_size) + return AdaptiveCachePlacementController( + capacity=CacheCapacityConfig( + gpu_tokens=gpu_tokens, + cpu_tokens=cpu_tokens, + disk_tokens=disk_tokens, + ), + args=SimpleNamespace( + enable_cpu_cache=enable_cpu_cache, + enable_disk_cache=enable_disk_cache, + ), + max_history_size=history_size, + initial_history_size=initial_history_size, + ) + + +def make_req(request_id: int, input_len: int, group_req_id: int = None, cache_tiers=(CacheTier.GPU,)): + if group_req_id is None: + group_req_id = request_id + return SimpleNamespace( + cache_tiers=cache_tiers, + shm_req=SimpleNamespace( + request_id=request_id, + group_req_id=group_req_id, + input_len=input_len, + ), + ) + + +def test_controller_exposes_one_public_operation(): + controller = make_controller(1, 1, 1) + + public_methods = [ + name for name, method in inspect.getmembers(controller, predicate=inspect.ismethod) if not name.startswith("_") + ] + + assert public_methods == ["set_req_cache_way"] + + +def test_controller_is_an_interface(): + assert inspect.isabstract(CachePlacementController) + + +@pytest.mark.parametrize( + "cache_tiers", + [ + (CacheTier.DISK,), + (CacheTier.GPU, CacheTier.DISK), + ], +) +def test_controller_rejects_illegal_existing_cache_tiers(cache_tiers): + req = make_req(0, 100, cache_tiers=cache_tiers) + + with pytest.raises(AssertionError): + make_controller(1, 1, 1).set_req_cache_way([req]) + + +def test_gpu_only_controller_always_selects_gpu(): + reqs = [make_req(0, 100), make_req(1, 1000)] + + GpuOnlyCachePlacementController().set_req_cache_way(reqs) + + assert tuple(req.cache_tiers for req in reqs) == ((CacheTier.GPU,), (CacheTier.GPU,)) + + +def test_legacy_controller_selects_all_enabled_cache_levels(): + req = make_req(0, 100) + + LegacyCachePlacementController(enable_cpu_cache=True, enable_disk_cache=True).set_req_cache_way([req]) + + assert req.cache_tiers == (CacheTier.GPU, CacheTier.CPU, CacheTier.DISK) + + +def test_legacy_controller_does_not_select_disabled_disk_cache(): + req = make_req(0, 100) + + LegacyCachePlacementController(enable_cpu_cache=True, enable_disk_cache=False).set_req_cache_way([req]) + + assert req.cache_tiers == (CacheTier.GPU, CacheTier.CPU) + + +def test_legacy_controller_keeps_diverse_slave_request_on_gpu(): + req = make_req(request_id=2, group_req_id=1, input_len=100) + + LegacyCachePlacementController(enable_cpu_cache=True, enable_disk_cache=True).set_req_cache_way([req]) + + assert req.cache_tiers == (CacheTier.GPU,) + + +def test_cache_placement_strategy_cli_defaults_to_adaptive_and_accepts_legacy(): + parser = make_argument_parser() + + assert parser.parse_args([]).cache_placement_strategy == "adaptive" + assert parser.parse_args(["--cache_placement_strategy", "legacy"]).cache_placement_strategy == "legacy" + + +def test_factory_creates_gpu_only_controller_when_cpu_cache_is_disabled(): + args = SimpleNamespace(enable_cpu_cache=False) + + controller = create_cache_placement_controller(args=args, radix_cache=None) + + assert isinstance(controller, GpuOnlyCachePlacementController) + + +def test_factory_creates_legacy_controller_for_legacy_strategy(): + args = SimpleNamespace( + enable_cpu_cache=True, + enable_disk_cache=True, + cache_placement_strategy="legacy", + ) + + controller = create_cache_placement_controller(args=args, radix_cache=None) + + assert isinstance(controller, LegacyCachePlacementController) + + +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, 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( + enable_cpu_cache=True, + enable_disk_cache=True, + cache_placement_strategy="adaptive", + cpu_cache_token_page_size=16, + disk_cache_storage_size=1, + dp=4, + nnodes=2, + ) + + controller = create_cache_placement_controller( + args=args, + radix_cache=SimpleNamespace(total_token_num=80), + ) + + assert isinstance(controller, AdaptiveCachePlacementController) + assert controller._capacity == CacheCapacityConfig( + gpu_tokens=80, + cpu_tokens=160, + disk_tokens=int(1024 ** 3) // 64 * 16, + ) + + +@pytest.mark.parametrize( + ("configured_ratio", "expected_ratio"), + [ + (None, 0.8), + ("0.5", 0.5), + ("1", 1.0), + ], +) +def test_cache_placement_gpu_capacity_ratio_from_environment(monkeypatch, configured_ratio, expected_ratio): + from lightllm.utils.envs_utils import get_cache_placement_gpu_capacity_ratio + + if configured_ratio is None: + monkeypatch.delenv("LIGHTLLM_CACHE_PLACEMENT_GPU_CAPACITY_RATIO", raising=False) + else: + monkeypatch.setenv("LIGHTLLM_CACHE_PLACEMENT_GPU_CAPACITY_RATIO", configured_ratio) + + assert get_cache_placement_gpu_capacity_ratio() == expected_ratio + + +@pytest.mark.parametrize("configured_ratio", ["0", "-0.1", "1.1", "nan"]) +def test_cache_placement_gpu_capacity_ratio_rejects_out_of_range_values(monkeypatch, configured_ratio): + from lightllm.utils.envs_utils import get_cache_placement_gpu_capacity_ratio + + monkeypatch.setenv("LIGHTLLM_CACHE_PLACEMENT_GPU_CAPACITY_RATIO", configured_ratio) + + with pytest.raises(AssertionError): + get_cache_placement_gpu_capacity_ratio() + + +def test_set_req_cache_way_uses_recent_input_lengths_and_capacity_ratio(): + controller = make_controller(gpu_tokens=1, cpu_tokens=1, disk_tokens=2, history_size=4) + warmup_reqs = [ + make_req(0, 400), + make_req(1, 100), + make_req(2, 300), + make_req(3, 200), + ] + controller.set_req_cache_way(warmup_reqs) + reqs = [ + make_req(4, 400), + make_req(5, 100), + make_req(6, 300), + make_req(7, 200), + ] + + result = controller.set_req_cache_way(reqs) + + assert result is None + assert tuple(req.cache_tiers for req in warmup_reqs) == ((CacheTier.GPU, CacheTier.CPU, CacheTier.DISK),) * 4 + assert tuple(req.cache_tiers for req in reqs) == ( + (CacheTier.CPU, CacheTier.DISK), + (CacheTier.GPU,), + (CacheTier.GPU,), + (CacheTier.GPU,), + ) + + +def test_adaptive_controller_uses_enable_flags_instead_of_capacity_to_select_disk_path(): + controller = make_controller( + gpu_tokens=1, + cpu_tokens=1, + disk_tokens=0, + enable_cpu_cache=True, + enable_disk_cache=True, + ) + warmup_req = make_req(0, 100) + + controller.set_req_cache_way([warmup_req]) + + assert warmup_req.cache_tiers == (CacheTier.GPU, CacheTier.CPU, CacheTier.DISK) + + +def test_length_boundary_uses_max_cpu_and_disk_capacity(): + controller = make_controller(gpu_tokens=1, cpu_tokens=1, disk_tokens=3, history_size=4) + controller.set_req_cache_way([make_req(0, 100), make_req(1, 200), make_req(2, 300), make_req(3, 400)]) + assert controller._gpu_max_input_len == 200 + req = make_req(4, 150) + + controller.set_req_cache_way([req]) + + assert req.cache_tiers == (CacheTier.GPU,) + + +def test_zero_gpu_capacity_does_not_place_requests_in_gpu_after_warmup(): + controller = make_controller(gpu_tokens=0, cpu_tokens=1, disk_tokens=0, history_size=2) + controller.set_req_cache_way([make_req(0, 100), make_req(1, 200)]) + req = make_req(2, 100) + + controller.set_req_cache_way([req]) + + assert controller._gpu_max_input_len == 0 + assert req.cache_tiers == (CacheTier.CPU,) + + +def test_set_req_cache_way_does_not_use_disabled_tiers(): + controller = make_controller(gpu_tokens=1, cpu_tokens=3, disk_tokens=0, history_size=4) + warmup_reqs = [make_req(index, length) for index, length in enumerate((100, 200, 300, 400))] + controller.set_req_cache_way(warmup_reqs) + reqs = [make_req(index + 4, length) for index, length in enumerate((100, 200, 300, 400))] + + controller.set_req_cache_way(reqs) + + assert tuple(req.cache_tiers for req in warmup_reqs) == ((CacheTier.GPU, CacheTier.CPU),) * 4 + assert tuple(req.cache_tiers for req in reqs) == ( + (CacheTier.GPU,), + (CacheTier.GPU,), + (CacheTier.CPU,), + (CacheTier.CPU,), + ) + + +def test_length_boundary_keeps_history_after_initial_update(): + controller = make_controller(1, 1, 1, history_size=3) + + controller.set_req_cache_way([make_req(0, 100), make_req(1, 200)]) + assert controller._gpu_max_input_len is None + assert list(controller._recent_input_lengths) == [100, 200] + + controller.set_req_cache_way([make_req(2, 300)]) + assert controller._gpu_max_input_len == 300 + assert list(controller._recent_input_lengths) == [100, 200, 300] + + latest_req = make_req(3, 1000) + controller.set_req_cache_way([latest_req]) + + assert controller._gpu_max_input_len == 300 + assert list(controller._recent_input_lengths) == [200, 300, 1000] + assert latest_req.cache_tiers == (CacheTier.CPU, CacheTier.DISK) + + +def test_length_boundary_updates_every_36_steps_after_initial_update(): + controller = make_controller(1, 1, 1, history_size=4, initial_history_size=2) + + controller.set_req_cache_way([make_req(0, 100), make_req(1, 200)]) + assert controller._gpu_max_input_len == 200 + assert controller._steps_since_last_update == 0 + assert list(controller._recent_input_lengths) == [100, 200] + + for step in range(35): + controller.set_req_cache_way([make_req(step + 2, 1000)]) + assert controller._gpu_max_input_len == 200 + assert controller._steps_since_last_update == 35 + assert list(controller._recent_input_lengths) == [1000, 1000, 1000, 1000] + + controller.set_req_cache_way([make_req(37, 1000)]) + assert controller._gpu_max_input_len == 1000 + assert controller._steps_since_last_update == 0 + assert list(controller._recent_input_lengths) == [1000, 1000, 1000, 1000] + + +def test_history_window_keeps_latest_requests_when_single_extend_exceeds_limit(): + controller = make_controller(1, 1, 1, history_size=3) + reqs = [make_req(index, length) for index, length in enumerate((100, 200, 300, 400, 500))] + + controller.set_req_cache_way(reqs) + + assert controller._gpu_max_input_len == 400 + assert list(controller._recent_input_lengths) == [300, 400, 500] + + +def test_repeated_assignment_is_rejected(): + controller = make_controller(1, 1, 1) + req = make_req(0, 100) + controller.set_req_cache_way([req]) + + req.shm_req.input_len = 10000 + with pytest.raises(AssertionError): + controller.set_req_cache_way([req]) + + assert list(controller._recent_input_lengths) == [100] + + +def test_concurrent_reassignment_is_rejected(): + controller = make_controller(1, 1, 1) + req = make_req(0, 100) + + with pytest.raises(AssertionError): + with ThreadPoolExecutor(max_workers=2) as executor: + list(executor.map(lambda _: controller.set_req_cache_way([req]), range(2))) + + assert list(controller._recent_input_lengths) == [100] + + +def test_diverse_slave_request_stays_on_gpu_and_is_not_recorded(): + controller = make_controller(0, 1, 1) + slave_req = make_req(request_id=2, group_req_id=1, input_len=1000) + + controller.set_req_cache_way([slave_req]) + + assert slave_req.cache_tiers == (CacheTier.GPU,) + assert list(controller._recent_input_lengths) == [] diff --git a/unit_tests/server/pd_master_fault_launcher.py b/unit_tests/server/pd_master_fault_launcher.py new file mode 100644 index 0000000000..1b804bedf6 --- /dev/null +++ b/unit_tests/server/pd_master_fault_launcher.py @@ -0,0 +1,94 @@ +"""Test-only launcher that injects failures at real startup boundaries.""" + +import argparse +import json +from pathlib import Path +import tempfile + +import psutil + +from lightllm.server.api_cli import add_cli_args +from lightllm.server.api_server import launch_server +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.server.httpserver_for_pd_master import multiprocess +from lightllm.utils.envs_utils import get_unique_server_name + + +def observe_cleanup(observation_path): + original_stop = multiprocess.process_manager.terminate_all_processes + + def stop(): + prefix = "lightllm-pd-master-" + get_unique_server_name() + "-*" + directories = list(Path(tempfile.gettempdir()).glob(prefix)) + before = [directory.exists() for directory in directories] + processes = list(multiprocess.process_manager.processes) + original_stop() + + def process_stopped(process): + try: + return not process.is_running() + except psutil.NoSuchProcess: + return True + + observation_path.write_text( + json.dumps( + { + "runtime_paths": [str(directory) for directory in directories], + "runtime_before_stop": before, + "runtime_after_stop": [directory.exists() for directory in directories], + "backends_stopped": all(process_stopped(process) for process in processes), + } + ), + encoding="utf-8", + ) + + multiprocess.process_manager.terminate_all_processes = stop + + +def inject_failure(stage, haproxy_binary): + if stage == "backend_start": + original_get_context = multiprocess.multiprocessing.get_context + + def get_context(method): + context = original_get_context(method) + process_count = 0 + + class ProcessContext: + def Process(self, *args, **kwargs): + nonlocal process_count + process_count += 1 + if process_count == 2: + process = context.Process(*args, **kwargs) + + def fail_to_start(): + raise RuntimeError("injected backend startup failure") + + process.start = fail_to_start + return process + return context.Process(*args, **kwargs) + + return ProcessContext() + + multiprocess.multiprocessing.get_context = get_context + elif stage == "proxy_start": + original_popen = multiprocess.subprocess.Popen + + def popen(command, *args, **kwargs): + if command[0] == haproxy_binary and kwargs.get("start_new_session"): + raise OSError("injected HAProxy startup failure") + return original_popen(command, *args, **kwargs) + + multiprocess.subprocess.Popen = popen + + +if __name__ == "__main__": + parser = add_cli_args(argparse.ArgumentParser()) + parser.add_argument("--failure-stage", required=True) + parser.add_argument("--cleanup-observations", type=Path, required=True) + options = vars(parser.parse_args()) + failure_stage = options.pop("failure_stage") + observe_cleanup(options.pop("cleanup_observations")) + args = StartArgs(**options) + haproxy_binary = multiprocess.validate_multiprocess_pd_master_args(args) + inject_failure(failure_stage, haproxy_binary) + launch_server(args) diff --git a/unit_tests/server/pd_master_test_utils.py b/unit_tests/server/pd_master_test_utils.py new file mode 100644 index 0000000000..075cd2765e --- /dev/null +++ b/unit_tests/server/pd_master_test_utils.py @@ -0,0 +1,188 @@ +"""Helpers for real CPU-only PD master integration tests.""" + +from contextlib import contextmanager +from dataclasses import dataclass +import json +import os +from pathlib import Path +import re +import signal +import socket +import subprocess +import sys +import time +from urllib.error import HTTPError, URLError +from urllib.request import ProxyHandler, build_opener + +import psutil +import pytest + +from lightllm.utils.start_utils import kill_recursive + + +def get_test_haproxy_binary(): + try: + from ray_haproxy import get_haproxy_binary + except ImportError: + pytest.skip("Install ray-haproxy to run native proxy integration tests") + return get_haproxy_binary() + + +def create_tiny_tokenizer(model_dir: Path) -> Path: + from tokenizers import Tokenizer + from tokenizers.models import WordLevel + from tokenizers.pre_tokenizers import Whitespace + from transformers import PreTrainedTokenizerFast + + model_dir.mkdir(parents=True, exist_ok=True) + tokenizer = Tokenizer(WordLevel({"[UNK]": 0, "[EOS]": 1, "hello": 2, "world": 3}, unk_token="[UNK]")) + tokenizer.pre_tokenizer = Whitespace() + fast_tokenizer = PreTrainedTokenizerFast(tokenizer_object=tokenizer, unk_token="[UNK]", eos_token="[EOS]") + fast_tokenizer.save_pretrained(model_dir) + (model_dir / "config.json").write_text( + json.dumps( + { + "model_type": "llama", + "architectures": ["LlamaForCausalLM"], + "max_position_embeddings": 128, + "vocab_size": 4, + "eos_token_id": 1, + } + ), + encoding="utf-8", + ) + (model_dir / "generation_config.json").write_text('{"eos_token_id":1}', encoding="utf-8") + return model_dir + + +@dataclass +class RunningPDMaster: + process: subprocess.Popen + public_url: str + workers: list[dict] + opener: object + log_path: Path + + def get_json(self, path, base_url=None): + with self.opener.open((base_url or self.public_url) + path, timeout=3) as response: + return json.load(response) + + def descendants(self): + try: + return psutil.Process(self.process.pid).children(recursive=True) + except psutil.NoSuchProcess: + return [] + + def haproxy_runtime_dir(self): + match = re.search(r"PD master HAProxy configuration file (.+)/haproxy.cfg", self.log_path.read_text()) + assert match is not None, self.log_path.read_text() + return Path(match.group(1)) + + +def assert_processes_exit(processes, timeout=8): + def is_alive(process): + try: + return process.is_running() and process.status() != psutil.STATUS_ZOMBIE + except psutil.NoSuchProcess: + return False + + deadline = time.monotonic() + timeout + while True: + alive = [process for process in processes if is_alive(process)] + if not alive: + return + assert time.monotonic() < deadline, f"Processes survived shutdown: {[process.pid for process in alive]}" + time.sleep(0.1) + + +def discover_test_workers(opener, public_url, timeout=10): + """Build local worker identities from the ports published by the launcher.""" + deadline = time.monotonic() + timeout + while True: + try: + with opener.open(public_url + "/get_server_info", timeout=3) as response: + public_info = json.load(response) + break + except (HTTPError, OSError, URLError): + if time.monotonic() >= deadline: + raise + time.sleep(0.1) + host = public_url.split("://", 1)[1].rsplit(":", 1)[0] + single_worker = public_info["httpserver_workers"] == 1 + workers = [] + for worker_port in public_info["pd_master_multi_ports"]: + host_ip_port = f"{host}:{worker_port}" + workers.append( + { + "node_id": public_info["pd_node_id"] if single_worker else worker_port, + "host_ip_port": host_ip_port, + "port": worker_port, + } + ) + return workers + + +@contextmanager +def running_pd_master(runtime_dir: Path, *, extra_args=(), wait_for_startup=True, launcher_script=None): + """Run real metric/scheduler/API workers and HAProxy without model weights.""" + get_test_haproxy_binary() + runtime_dir.mkdir(parents=True, exist_ok=True) + model_dir = create_tiny_tokenizer(runtime_dir / "tiny-tokenizer") + with socket.socket() as listener: + listener.bind(("", 0)) + public_port = listener.getsockname()[1] + command = [sys.executable, "-m", "lightllm.server.api_server"] + if launcher_script is not None: + command = [sys.executable, str(launcher_script)] + command.extend( + [ + "--run_mode", + "pd_master", + "--host", + "0.0.0.0", + "--port", + str(public_port), + "--httpserver_workers", + "2", + "--model_dir", + str(model_dir), + "--max_req_total_len", + "128", + ] + ) + command.extend(extra_args) + env = os.environ.copy() + env.update(CUDA_VISIBLE_DEVICES="", HF_HUB_OFFLINE="1", TRANSFORMERS_OFFLINE="1", TOKENIZERS_PARALLELISM="false") + env["PYTHONPATH"] = str(Path(__file__).resolve().parents[2]) + os.pathsep + env.get("PYTHONPATH", "") + opener = build_opener(ProxyHandler({})) + public_url = f"http://127.0.0.1:{public_port}" + log_path = runtime_dir / "launcher.log" + process = None + with log_path.open("wb") as log: + try: + process = subprocess.Popen(command, env=env, stdout=log, stderr=subprocess.STDOUT) + if not wait_for_startup: + yield RunningPDMaster(process, public_url, [], opener, log_path) + return + deadline = time.monotonic() + 90 + while True: + assert process.poll() is None, log_path.read_text() + try: + with opener.open(public_url + "/liveness", timeout=0.5) as response: + assert response.status == 200 + break + except (OSError, URLError): + if time.monotonic() >= deadline: + pytest.fail("PD master launcher did not become live:\n" + log_path.read_text()) + time.sleep(0.1) + workers = discover_test_workers(opener, public_url) + yield RunningPDMaster(process, public_url, workers, opener, log_path) + finally: + if process is not None: + if process.poll() is None: + process.send_signal(signal.SIGTERM) + try: + process.wait(timeout=15) + except subprocess.TimeoutExpired: + kill_recursive(process) + process.wait(timeout=5) diff --git a/unit_tests/server/router/dynamic_prompt/test_radix_cache.py b/unit_tests/server/router/dynamic_prompt/test_radix_cache.py new file mode 100644 index 0000000000..c2d5a67c19 --- /dev/null +++ b/unit_tests/server/router/dynamic_prompt/test_radix_cache.py @@ -0,0 +1,722 @@ +from types import SimpleNamespace + +from sortedcontainers import SortedDict + +import pytest +import torch + +from lightllm.server.router.dynamic_prompt.radix_cache import RadixCache, TreeNode +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(): + monkeypatch = pytest.MonkeyPatch() + monkeypatch.setattr(shm_utils, "get_unique_server_name", lambda: "test_radix_cache_service_0") + yield + monkeypatch.undo() + + +def test_case1(): + tree = RadixCache(100, 0) + ans, _ = tree.insert(torch.tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], dtype=torch.int64, device="cpu")) + assert ans == 0 + tree.print_self() + ans, _ = tree.insert(torch.tensor([0, 1, 2, 3, 4, 7, 8, 9], dtype=torch.int64, device="cpu")) + assert ans == 5 + tree.print_self() + ans, _ = tree.insert(torch.tensor([0, 1, 2, 3, 4, 7, 8, 9], dtype=torch.int64, device="cpu")) + assert ans == 8 + tree.print_self() + + assert tree.get_refed_tokens_num() == 0 + assert tree.get_tree_total_tokens_num() == 13 + + # print("evict") + tree.evict(9, lambda x: x) + tree.print_self() + assert tree.get_refed_tokens_num() == 0 and tree.get_tree_total_tokens_num() == 0 + + +def test_case2(): + tree = RadixCache(100, 1) + 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")) + tree.print_self() + + tree_node, size, values = tree.match_prefix( + torch.tensor([0, 1, 2, 3, 4], dtype=torch.int64, device="cpu"), update_refs=False + ) + assert tree_node.node_prefix_total_len == 5 and size == 5 and len(values) == 5 + tree_node, size, values = tree.match_prefix( + torch.tensor([0, 1, 2, 3, 4, 9], dtype=torch.int64, device="cpu"), update_refs=False + ) + assert tree_node.node_prefix_total_len == 5 and size == 5 and len(values) == 5 + tree_node, size, values = tree.match_prefix( + torch.tensor([0, 1, 2, 3, 4, 7, 8], dtype=torch.int64, device="cpu"), update_refs=False + ) + assert tree_node.node_prefix_total_len == 7 and size == 7 and len(values) == 7 + tree_node, size, values = tree.match_prefix( + torch.tensor([0, 1, 2, 3, 4, 7, 9], dtype=torch.int64, device="cpu"), update_refs=False + ) + assert tree_node.node_prefix_total_len == 6 and size == 6 and len(values) == 6 + print(ans) + return + + +def test_case3(): + 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")) + tree.print_self() + + tree_node, size, values = tree.match_prefix( + torch.tensor([0, 1, 2, 3, 4], dtype=torch.int64, device="cpu"), update_refs=True + ) + assert tree_node.node_prefix_total_len == 5 and size == 5 and len(values) == 5 + assert tree.get_refed_tokens_num() == 5 and tree.get_tree_total_tokens_num() == 13 + + tree_node, size, values = tree.match_prefix( + torch.tensor([0, 1, 2, 3, 4, 7, 9], dtype=torch.int64, device="cpu"), update_refs=True + ) + assert tree_node.node_prefix_total_len == 6 and size == 6 and len(values) == 6 + assert tree.get_refed_tokens_num() == 6 and tree.get_tree_total_tokens_num() == 13 + + tree.print_self() + tree.evict(2, lambda x: x) + assert tree.get_refed_tokens_num() == 6 and tree.get_tree_total_tokens_num() == 8 + tree.print_self() + + tree.dec_node_ref_counter(tree_node) + tree.print_self() + print(ans) + return + + +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")) + tree.print_self() + + tree.clear_tree_nodes() + print(ans) + return + + +def test_case5(): + """ + 测试场景：一个简单的父子节点链 (A -> B)，在 ref_counter 都为 0 时，应该成功合并。 + """ + print("\nTest Case 5: Merging simple parent-child nodes when ref_counter is 0\n") + tree = RadixCache(100, 0) + + _, node_a = tree.insert(torch.tensor([1, 2, 3], dtype=torch.int64)) + _, node_b = tree.insert(torch.tensor([1, 2, 3, 4, 5], dtype=torch.int64)) + tree.print_self() + + # 验证初始状态：A -> B 结构，且 ref_counter 均为 0 + assert node_b.parent == node_a + assert torch.equal(node_a.token_id_key, torch.tensor([1, 2, 3], dtype=torch.int64)) + assert len(node_a.children) == 1 + assert node_a.ref_counter == 0 + assert node_b.ref_counter == 0 + assert tree.get_tree_total_tokens_num() == 5 + + # 执行合并 + tree.merge_unreferenced_nodes() + tree.print_self() + + assert torch.equal(node_b.token_id_key, torch.tensor([1, 2, 3, 4, 5], dtype=torch.int64)) + assert node_b.is_leaf() + assert tree.get_tree_total_tokens_num() == 5 + assert tree.root_node.children[1] is node_b + + +def test_case6(): + """ + 测试场景：一个长的节点链 (A -> B -> C)，在 ref_counter 都为 0 时，应该级联合并成一个节点。 + """ + print("\nTest Case 6: Merging long nodes when ref_counter is 0\n") + tree = RadixCache(100, 0) + _, node_a = tree.insert(torch.tensor([1], dtype=torch.int64)) + _, node_b = tree.insert(torch.tensor([1, 2], dtype=torch.int64)) + _, node_c = tree.insert(torch.tensor([1, 2, 3, 4], dtype=torch.int64)) + tree.print_self() + + assert node_c.parent == node_b + assert node_b.parent == node_a + assert tree.get_tree_total_tokens_num() == 4 + tree.merge_unreferenced_nodes() + tree.print_self() + + assert len(tree.root_node.children) == 1 + # 节点 C 的 key 应该是完整的 [1, 2, 3, 4] + assert torch.equal(node_c.token_id_key, torch.tensor([1, 2, 3, 4], dtype=torch.int64)) + assert node_c.is_leaf() + assert tree.get_tree_total_tokens_num() == 4 + + +def test_case7(): + """ + 测试场景：由于父节点或子节点的 ref_counter > 0，合并不应该发生。 + """ + print("\nTest Case 7: Merging when parent or child ref_counter > 0\n") + tree = RadixCache(100, 0) + + _, node_a = tree.insert(torch.tensor([1, 2, 3], dtype=torch.int64)) + _, node_b = tree.insert(torch.tensor([1, 2, 3, 4, 5], dtype=torch.int64)) + tree.print_self() + + matched_node, _, _ = tree.match_prefix(torch.tensor([1, 2, 3], dtype=torch.int64), update_refs=True) + assert matched_node is node_a + assert node_a.ref_counter == 1 + assert node_b.ref_counter == 0 + + tree.merge_unreferenced_nodes() + tree.print_self() + + assert torch.equal(node_a.token_id_key, torch.tensor([1, 2, 3], dtype=torch.int64)) + assert not node_a.is_leaf() + assert node_b.parent is node_a + + +def test_case8(): + """ + 测试场景：由于父节点有多个子节点，合并不应该发生。 + """ + print("\nTest Case 8: Merging when parent has multiple children\n") + tree = RadixCache(100, 0) + + _, node_a = tree.insert(torch.tensor([1, 2], dtype=torch.int64)) + _, node_b = tree.insert(torch.tensor([1, 2, 3], dtype=torch.int64)) + _, node_c = tree.insert(torch.tensor([1, 2, 4], dtype=torch.int64)) + tree.print_self() + + assert len(node_a.children) == 2 + assert node_a.ref_counter == 0 + assert node_b.ref_counter == 0 + assert node_c.ref_counter == 0 + + tree.merge_unreferenced_nodes() + tree.print_self() + + assert len(node_a.children) == 2 + assert torch.equal(node_a.token_id_key, torch.tensor([1, 2], dtype=torch.int64)) + assert tree.root_node.children[1].children[3] is node_b + assert tree.root_node.children[1].children[4] is node_c + + +def test_case9(): + """ + 测试场景：在一个复杂的树中，只有满足条件的分支被合并。 + """ + print("\nTest Case 9: Merging in a complex tree with mixed conditions\n") + tree = RadixCache(100, 0) + + # 分支1: 可合并的链 A -> B + _, node_a = tree.insert(torch.tensor([1, 2], dtype=torch.int64)) + _, node_b = tree.insert(torch.tensor([1, 2, 3], dtype=torch.int64)) + + # 分支2: 不可合并的链 C -> D (因为 C 被引用) + _, node_c = tree.insert(torch.tensor([4, 5], dtype=torch.int64)) + _, node_d = tree.insert(torch.tensor([4, 5, 6], dtype=torch.int64)) + + # 增加 C 的引用计数 + tree.match_prefix(torch.tensor([4, 5], dtype=torch.int64), update_refs=True) + assert node_c.ref_counter == 1 + tree.print_self() + + tree.merge_unreferenced_nodes() + tree.print_self() + + merged_node_b = tree.root_node.children[1] + assert torch.equal(merged_node_b.token_id_key, torch.tensor([1, 2, 3], dtype=torch.int64)) + assert merged_node_b.is_leaf() + + unmerged_node_c = tree.root_node.children[4] + assert torch.equal(unmerged_node_c.token_id_key, torch.tensor([4, 5], dtype=torch.int64)) + assert not unmerged_node_c.is_leaf() + assert len(unmerged_node_c.children) == 1 + + unmerged_node_d = unmerged_node_c.children[6] + assert torch.equal(unmerged_node_d.token_id_key, torch.tensor([6], dtype=torch.int64)) + + +def test_flush_cache_requires_released_references_and_returns_allocator_pages(): + from lightllm.common.attention_state.allocators import FullAttKVAllocator + + allocator = FullAttKVAllocator(100, pin_memory=False) + 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]) + 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() + 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 + + +@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 + + monkeypatch.setattr( + "lightllm.common.attention_state.bounded_att.utils.get_env_start_args", + lambda: SimpleNamespace(enable_dp_prompt_cache_fetch=False), + ) + resources = composition(device="cuda") + 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_runtime_to_pages( + req_indices=[0], + state_token_lens=[8], + page_kind="big", + page_slot_ids=[big], + ) + resources.save_bounded_att_runtime_to_pages( + req_indices=[1], + state_token_lens=[12], + page_kind="small", + page_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) + + prefix_len, _ = tree.insert(torch.tensor([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]), values) + assert prefix_len == 0 + assert tree.get_tree_total_tokens_num() == 8 + + # A mismatch inside a page cannot produce a partial-page cache hit. + node, matched_len, matched_values = tree.match_prefix(torch.tensor([1, 2, 0, 4, 5, 6, 7, 8])) + assert node is None + assert matched_len == 0 + assert matched_values is None + + node, matched_len, matched_values = tree.match_prefix(torch.tensor([1, 2, 3, 4, 5, 6, 7, 8, 11])) + assert node is not None + assert matched_len == 8 + assert matched_values.tolist() == list(range(100, 108)) + + # The second sequence shares exactly one page and then branches by its + # complete second-page key. + prefix_len, _ = tree.insert( + torch.tensor([1, 2, 3, 4, 5, 6, 0, 8]), + torch.arange(200, 208, dtype=torch.int64), + ) + assert prefix_len == 4 + assert tree.get_tree_total_tokens_num() == 12 + + +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_state_manager=SimpleNamespace(full_att_allocator=mem_manager), page_size=4) + + +@pytest.mark.parametrize("update_refs", [False, True]) +@pytest.mark.parametrize( + "spans,expected", + [ + ([], 24), + ([(10, 21), (22, 30)], 0), + ([(0, 10), (11, 30)], 0), + ([(8, 24)], 24), + ([(24, 30)], 24), + ], +) +def test_mm_gpu_match_uses_existing_ancestors_and_preserves_other_refs(update_refs, spans, expected): + tree = RadixCache(64, 103, page_size=4) + key = torch.arange(32) + values = key + 100 + tree.insert(key, values) + held_node, _, _ = tree.match_prefix(key[:24], update_refs=True) + + node, matched_len, matched_values = tree.match_prefix(key[:27], update_refs=update_refs, mm_input_spans=spans) + assert matched_len == expected + if expected: + assert node.node_prefix_total_len == expected + assert torch.equal(matched_values, values[:expected]) + else: + assert node is None and matched_values is None + assert tree.get_refed_tokens_num() == 24 + + tree.dec_node_ref_counter(held_node) + assert tree.get_refed_tokens_num() == (expected if update_refs else 0) + if update_refs: + tree.dec_node_ref_counter(node) + assert tree.get_refed_tokens_num() == 0 + assert tree.root_node.ref_counter == 1 + freed = [] + tree.evict(32, lambda indexes: freed.extend(indexes.tolist())) + assert sorted(freed) == values.tolist() + + +@pytest.mark.parametrize("update_refs", [False, True]) +@pytest.mark.parametrize("block_start,expected", [(1, 0), (5, 4), (9, 4), (13, 12)]) +def test_mm_gpu_fallback_walks_parents_without_revisiting(monkeypatch, update_refs, block_start, expected): + tree = RadixCache(64, 105, page_size=4) + tree.insert(torch.arange(40, 44), torch.arange(200, 204)) + key = torch.arange(24) + values = key + 100 + cached_nodes = {} + for length in (4, 12, 24): + _, cached_nodes[length] = tree.insert(key[:length], values[:length]) + + # Retreat across existing nodes, even when the image starts inside an ancestor. + visited = [] + update_time = TreeNode.update_time + + def record_visit(node): + visited.append(node) + update_time(node) + + monkeypatch.setattr(TreeNode, "update_time", record_visit) + node, matched_len, matched_values = tree.match_prefix( + key, update_refs=update_refs, mm_input_spans=[(block_start, 30)] + ) + assert matched_len == expected + assert node is cached_nodes.get(expected) + assert [len(n.token_id_key) for n in cached_nodes.values()] == [4, 8, 12] + assert len(visited) == len(set(visited)) == 4 + assert tree.get_refed_tokens_num() == (expected if update_refs else 0) + if expected: + assert node.node_prefix_total_len == expected + assert torch.equal(matched_values, values[:expected]) + else: + assert node is None and matched_values is None + + # The untouched branch is still oldest; changed tail refs must preserve eviction ordering. + freed = [] + tree.evict(4, lambda indexes: freed.extend(indexes.tolist())) + assert freed == list(range(200, 204)) + if update_refs: + tree.dec_node_ref_counter(node) + assert tree.get_refed_tokens_num() == 0 and tree.root_node.ref_counter == 1 + tree.evict(24, lambda indexes: freed.extend(indexes.tolist())) + assert sorted(freed) == sorted(values.tolist() + list(range(200, 204))) + assert not tree.root_node.children + + +def test_page_key_bytes_does_not_share_tensor_memory(): + token_ids = torch.tensor([1, 2, 3, 4], dtype=torch.int64) + expected_key = token_ids.clone().numpy().tobytes() + key = TreeNode(page_size=4).get_child_key(token_ids) + + token_ids[0] = 100 + + assert isinstance(key, bytes) + assert key == expected_key + 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 = [] + tree = HybridAttPagedRadixCache( + total_token_num=100, + rank_in_node=101, + hash_page_size=4, + big_page_num=2, + 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) + tree.insert( + key, + values, + block_hashs=[11, 22, 33], + block_state_idxs=[None, None, 7], + len_to_big_page_id=SortedDict({8: 5}), + ) + + big_page = tree.root_node.children[22] + tail_page = big_page.children[33] + node, matched_len, _ = tree.match_prefix(key, block_hashs=[11, 22, 33], update_refs=True) + assert node is tail_page and matched_len == 12 + assert tree.get_refed_tokens_num() == 12 + tree.dec_node_ref_counter(node) + assert tree.get_refed_tokens_num() == 0 + + tree.free_one_small_page_buffer() + assert freed_state_ids == [7] + node, matched_len, matched_values = tree.match_prefix(key, block_hashs=[11, 22, 33], update_refs=True) + assert node is big_page and matched_len == 8 + assert matched_values.tolist() == list(range(100, 108)) + assert big_page.ref_counter == 1 and tail_page.ref_counter == 0 + assert tree.get_refed_tokens_num() == 8 + tree.dec_node_ref_counter(node) + assert big_page.ref_counter == tail_page.ref_counter == tree.get_refed_tokens_num() == 0 + assert tree.root_node.ref_counter == 1 + + +@pytest.fixture +def hybrid_tree(): + freed_big_states = [] + tree = HybridAttPagedRadixCache( + total_token_num=100, + rank_in_node=102, + hash_page_size=4, + big_page_num=2, + att_state_manager=hybrid_resources(freed_big=freed_big_states), + ) + return tree, freed_big_states + + +def test_hybrid_match_traverses_pages_inside_image(hybrid_tree): + tree, freed = hybrid_tree + key = torch.arange(20) + tree.insert( + key, + key + 100, + block_hashs=[1, 2, 3, 4, 5], + block_state_idxs=[None, None, None, None, 7], + len_to_big_page_id=SortedDict({8: 5, 16: 6}), + ) + first_page = tree.root_node.children[2] + second_page = first_page.children[4] + assert tree.get_big_page_ids_by_node(second_page) == [5, 6] + + for length, expected in [(8, 0), (12, 0), (16, 16), (20, 20), (24, 20)]: + node, matched, values = tree.match_prefix( + torch.arange(length), + block_hashs=list(range(1, length // 4 + 1)), + update_refs=True, + mm_input_spans=[(4, 12)], + ) + assert matched == expected + if expected: + assert not 4 < node.node_prefix_total_len < 12 + assert values.tolist() == list(range(100, 100 + expected)) + assert tree.get_refed_tokens_num() == expected + tree.dec_node_ref_counter(node) + assert tree.get_refed_tokens_num() == 0 + + # Duplicate requests release their slots; no per-page validity metadata is needed. + tree.insert(key[:8], key[:8] + 100, [1, 2], [None, None], SortedDict({8: 8})) + tree.insert(key[:16], key[:16] + 100, [1, 2, 3, 4], [None] * 4, SortedDict({8: 9, 16: 10})) + assert first_page.big_page_buffer_idx == 5 + assert second_page.big_page_buffer_idx == 6 + assert freed == [8, 9, 10] + tree._evict(20, lambda *args: None) + assert sorted(freed) == [5, 6, 8, 9, 10] + assert not tree.root_node.children + + +@pytest.mark.parametrize( + "small_state_valid,mm_spans,match_len,expected", + [ + (True, (), 640, 512), + (True, [(448, 640)], 384, 384), + (True, [(448, 640)], 640, 256), + (False, [(448, 640)], 640, 256), + (True, [(448, 640)], 768, 768), + (True, [(448, 896)], 896, 256), + (False, [(448, 896)], 896, 256), + (True, [(448, 896)], 1024, 1024), + (True, [(700, 800)], 896, 512), + ], +) +def test_hybrid_match_falls_back_along_matched_branch(monkeypatch, small_state_valid, mm_spans, match_len, expected): + freed_big_states = [] + tree = HybridAttPagedRadixCache( + total_token_num=2048, + rank_in_node=104, + hash_page_size=64, + big_page_num=4, + 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) + tree.insert( + torch.arange(384), + old_values, + block_hashs=list(range(1, 7)), + block_state_idxs=[None] * 5 + [7], + len_to_big_page_id=SortedDict({256: 5}), + ) + # A longer request keeps the old small branch and owns a separate big-page tail. + new_values = torch.cat([old_values[:256], torch.arange(2000, 2768)]) + big_slots = {512: 6, 768: 8, 1024: 9} + prefix_len, _ = tree.insert( + torch.arange(1024), + new_values, + block_hashs=list(range(1, 17)), + block_state_idxs=[None] * 16, + len_to_big_page_id=SortedDict(big_slots), + ) + assert prefix_len == 256 + if not small_state_valid: + tree.free_one_small_page_buffer() + + common = tree.root_node.children[4] + small_320 = common.children[5] + small_384 = small_320.children[6] + big_512 = common.children[8] + big_768 = big_512.children[12] + big_1024 = big_768.children[16] + visited = [] + update_time = type(common).update_time + + def record_visit(node): + visited.append(node) + update_time(node) + + monkeypatch.setattr(type(common), "update_time", record_visit) + node, matched_len, values = tree.match_prefix( + torch.arange(match_len), + block_hashs=list(range(1, match_len // 64 + 1)), + update_refs=True, + mm_input_spans=mm_spans, + ) + assert matched_len == expected + # Retain small pages only when already on the matched path; never rematch sibling branches. + assert len(visited) == len(set(visited)) + assert not any(start < matched_len < end for start, end in mm_spans) + expected_values = old_values if expected <= 384 else new_values + assert torch.equal(values, expected_values[:expected]) + assert tree.get_refed_tokens_num() == expected + assert tree.root_node.ref_counter == 2 and common.ref_counter == 1 + assert small_320.ref_counter == small_384.ref_counter == int(expected == 384) + for big in (big_512, big_768, big_1024): + assert big.ref_counter == int(big.node_prefix_total_len <= expected) + + tree.dec_node_ref_counter(node) + assert tree.get_refed_tokens_num() == 0 and tree.root_node.ref_counter == 1 + assert all(n.ref_counter == 0 for n in (common, small_320, small_384, big_512, big_768, big_1024)) + freed_indexes = [] + tree._evict(tree.get_tree_total_tokens_num(), lambda indexes, _: freed_indexes.extend(indexes.tolist())) + assert sorted(freed_indexes) == sorted(old_values.tolist() + new_values[256:].tolist()) + assert sorted(freed_big_states) == [5, 6, 8, 9] + assert not tree.root_node.children + + +@pytest.mark.parametrize("first_outside_image", [False, True]) +@pytest.mark.parametrize("require_state", [False, True]) +def test_hybrid_fallback_distinguishes_kv_ownership_from_state_restore(hybrid_tree, first_outside_image, require_state): + tree, freed = hybrid_tree + tree.insert( + torch.arange(20), + block_hashs=[1, 2, 3, 4, 5], + block_state_idxs=[None] * 4 + [7], + len_to_big_page_id=SortedDict({8: 6, 16: 8}), + ) + spans = [(12, 18)] if first_outside_image else [(4, 18)] + node, _, _ = tree.match_prefix( + torch.arange(20), block_hashs=[1, 2, 3, 4, 5], update_refs=True, mm_input_spans=spans + ) + node = tree.deref_to_first_big_page_node(node, mm_input_spans=spans if require_state else ()) + expected = 16 if not require_state else (8 if first_outside_image else 0) + assert (0 if node is None else node.node_prefix_total_len) == expected + assert tree.get_refed_tokens_num() == expected + tree.dec_node_ref_counter(node) + assert tree.get_refed_tokens_num() == 0 + tree._evict(20, lambda *args: None) + assert sorted(freed) == [6, 8] # Invalid state slots have the same ownership and lifetime. + + +@pytest.mark.parametrize( + "computed,cache_len,retained,expected", + [(12, 20, 0, 0), (18, 20, 16, 16), (20, 20, 20, 20), (26, 20, 20, 20), (26, 24, 24, 16)], +) +def test_hybrid_request_release_frees_or_transfers_image_page_slots( + monkeypatch, hybrid_tree, computed, cache_len, retained, expected +): + from lightllm.server.router.model_infer import infer_batch + from lightllm.server.router.model_infer.infer_batch import InferReq, InferenceContext + + tree, freed = hybrid_tree + context = InferenceContext() + context.is_hybrid_att_model = True + context.args = SimpleNamespace(linear_att_hash_page_size=4, linear_att_page_block_num=2) + 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_state_manager = tree.att_state_manager + table = torch.arange(100, 132).view(1, -1) + context.req_manager = SimpleNamespace() + req = InferReq.__new__(InferReq) + req.req_idx = req.cur_output_len = 0 + req.token_index = SimpleNamespace( + get_cpu_token_indexes=lambda start, end: table[0, start:end].clone(), + get_free_token_indexes=lambda start, end: table[0, start:end].clone(), + reset=lambda: None, + ) + req.cur_kv_len, req.hold_kv_len, req.hybrid_cache_len = computed, max(24, computed), cache_len + req.shared_kv_node = None + req.hybrid_len_to_big_page_id = SortedDict({end: end // 8 + 4 for end in range(8, min(computed, cache_len) + 1, 8)}) + big_state_ids = list(req.hybrid_len_to_big_page_id.values()) + req.mm_input_spans = [(4, 12)] + ([(20, cache_len + 1)] if cache_len >= 24 else []) + req.tail_small_page_buffer_id = 7 if computed >= cache_len and cache_len % 8 else None + block_hashs = list(range(1, cache_len // 4 + 1)) + req.shm_req = SimpleNamespace( + input_len=cache_len + 1, + shm_prompt_ids=SimpleNamespace(arr=list(range(cache_len + 1))), + hybrid_token_hash_list=SimpleNamespace(get_all=lambda: block_hashs), + ) + released = [] + context._hybrid_att_free_req(released, req) + assert [index for block in released for index in block.tolist()] == list( + range(100 + retained, 100 + max(24, computed)) + ) + assert not req.hybrid_len_to_big_page_id + assert req.tail_small_page_buffer_id is None + assert tree.get_tree_total_tokens_num() == retained + assert freed == ([] if retained else big_state_ids) + node, length, _ = tree.match_prefix( + torch.arange(cache_len), block_hashs=block_hashs, update_refs=True, mm_input_spans=req.mm_input_spans + ) + assert length == expected + assert not any(start < length < end for start, end in req.mm_input_spans) + tree.dec_node_ref_counter(node) + tree._evict(retained, lambda *args: None) + assert sorted(freed) == big_state_ids + + +if __name__ == "__main__": + pytest.main() 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..4bf7aea081 --- /dev/null +++ b/unit_tests/server/router/model_infer/mode_backend/test_dp_big_pages.py @@ -0,0 +1,286 @@ +"""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.attention_state.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, + _copy_req_token_indexes_to_gpu=lambda reqs: 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), (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) + 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_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: + 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), + full_att_allocator=None, + 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( + is_hybrid_att_model=hybrid, + is_master_in_dp=True, + node_nccl_group="node", + model=SimpleNamespace(att_state_manager=manager), + dp_att_state_managers={}, + dp_req_to_token_indices={}, + _copy_req_token_indexes_to_gpu=lambda reqs: None, + ) + 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, + 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, + req_manager=SimpleNamespace(get_need_alloc_token_num=allocation_need, alloc_token_indexes=allocate_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, + ) + + 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_state_manager=SimpleNamespace( + bounded_att_big_page_allocator=allocator, full_att_allocator=SimpleNamespace(page_size=4) + ) + ), + dp_req_to_token_indices={1: torch.arange(64, dtype=torch.int32).reshape(2, 32)}, + ) + 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]]])) + requests = SimpleNamespace( + get_need_alloc_token_num=lambda req, end, page_size=None: max(0, ((end + 3) // 4) * 4 - req.hold_kv_len), + ) + req.kv_cache_alloc_need = lambda target: requests.get_need_alloc_token_num(req, target) + req.alloc_kv_mem = lambda alloc_token_num: allocate(req, alloc_token_num) + monkeypatch.setattr( + dp_shared_kv_trans, + "g_infer_context", + SimpleNamespace(get_can_alloc_token_num=lambda: 100, req_manager=requests, radix_cache=None), + ) + 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) + requests.get_need_alloc_token_num = lambda req, end, page_size=None: max(0, ((end + 3) // 4) * 4 - req.hold_kv_len) + + def allocate_tokens(req, end, allocator, page_size, radix_cache=None): + count = requests.get_need_alloc_token_num(req, end) + radix_cache.free_radix_cache_to_get_enough_token(count) + indexes = full.alloc(count) + start = req.hold_kv_len + requests.req_to_token_indexs[req.req_idx, start : start + count] = indexes + req.hold_kv_len += count + return requests.req_to_token_indexs[req.req_idx, req.cur_kv_len : end].clone() + + req.kv_cache_alloc_need = lambda target: requests.get_need_alloc_token_num(req, target) + req.alloc_kv_mem = lambda alloc_token_num: allocate_tokens(req, req.cur_kv_len + alloc_token_num, None, 4, radix) + 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 diff --git a/unit_tests/server/router/model_infer/mode_backend/test_dp_overlap_spec_engine.py b/unit_tests/server/router/model_infer/mode_backend/test_dp_overlap_spec_engine.py new file mode 100644 index 0000000000..fdc29fb87a --- /dev/null +++ b/unit_tests/server/router/model_infer/mode_backend/test_dp_overlap_spec_engine.py @@ -0,0 +1,491 @@ +from contextlib import nullcontext +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.server.router.model_infer.mode_backend.dp_backend import ( + impl as dp_backend_impl, +) +from lightllm.server.router.model_infer.mode_backend.dp_backend.impl import ( + DPChunkedPrefillBackend, +) +from lightllm.server.router.model_infer.mode_backend.base_backend import ModeBackend +from lightllm.server.router.model_infer.mode_backend.chunked_prefill.impl import ( + ChunkedPrefillBackend, +) +from lightllm.server.router.model_infer.mtp_speculative.dp_overlap_engine import ( + DPOverlapSpecEngine, +) +from lightllm.server.router.model_infer.mtp_speculative.engine import SpecEngine +from lightllm.server.router.model_infer.mtp_speculative.planner import ( + lightspec as lightspec_planner_impl, +) +from lightllm.server.router.model_infer.mtp_speculative.planner import ( + FixedSpecPlanner, + LightSpecPlanner, + SpecDecodePlan, +) +from lightllm.server.router.model_infer.mtp_speculative.proposers.base import ( + SpecProposal, +) + + +def test_dp_backend_reuses_common_engine_outside_overlap(): + args = SimpleNamespace(mtp_mode="eagle3", mtp_dynamic_verify=False, dp=1) + backend = ChunkedPrefillBackend.__new__(ChunkedPrefillBackend) + backend.args = args + backend.max_draft_step = 2 + backend.init_spec_engine() + dp_backend = DPChunkedPrefillBackend.__new__(DPChunkedPrefillBackend) + dp_backend.args = args + dp_backend.max_draft_step = 2 + dp_backend.dp_size = 1 + dp_backend.enable_prefill_microbatch_overlap = False + dp_backend.enable_decode_microbatch_overlap = True + dp_backend.init_spec_engine() + + assert "spec_engine_class" not in ModeBackend.__dict__ + assert type(backend.spec_engine) is SpecEngine + assert type(dp_backend.spec_engine) is SpecEngine + assert not dp_backend.spec_engine.proposer.enable_dynmaic_mtp + assert type(dp_backend.spec_engine.planner) is FixedSpecPlanner + assert type(dp_backend.dp_overlap_spec_engine) is DPOverlapSpecEngine + assert dp_backend.dp_overlap_spec_engine.common_engine is dp_backend.spec_engine + assert dp_backend.prefill_draft_engine is dp_backend.spec_engine + assert dp_backend.decode_draft_engine is dp_backend.dp_overlap_spec_engine + assert not issubclass(DPOverlapSpecEngine, SpecEngine) + + +def test_dp_backend_dspark_uses_common_engine_without_overlap(): + backend = DPChunkedPrefillBackend.__new__(DPChunkedPrefillBackend) + backend.args = SimpleNamespace(mtp_mode="dspark", mtp_dynamic_verify=False, dp=2) + backend.max_draft_step = 5 + backend.dp_size = 2 + backend.draft_models = [SimpleNamespace(block_size=5, graph=None)] + backend.enable_prefill_microbatch_overlap = False + backend.enable_decode_microbatch_overlap = False + + backend.init_spec_engine() + + assert type(backend.spec_engine) is SpecEngine + assert backend.dp_overlap_spec_engine is None + assert backend.prefill_draft_engine is backend.spec_engine + assert backend.decode_draft_engine is backend.spec_engine + + +def test_lightspec_planner_reduces_draft_step_with_max(monkeypatch): + group = object() + planner = LightSpecPlanner.__new__(LightSpecPlanner) + planner._draft_step_group = group + planner._draft_step_tensor = torch.zeros((1,), dtype=torch.int32) + planner._draft_step_stream = object() + planner.draft_steps = (1, 2, 3) + planner.pre_draft_step = 1 + planner.backend = SimpleNamespace(dp_size=2) + + def all_reduce(tensor, op, group, async_op): + assert op == torch.distributed.ReduceOp.MAX + assert group is planner._draft_step_group + assert not async_op + tensor.fill_(3) + + monkeypatch.setattr(lightspec_planner_impl.dist, "all_reduce", all_reduce) + monkeypatch.setattr(lightspec_planner_impl.torch.cuda, "stream", lambda stream: nullcontext()) + + plan = planner.plan(decode_reqs=[], origin_batch_size=2) + + assert plan.dynamic_batch_size == 2 + assert plan.draft_step == 3 + assert planner.pre_draft_step == 3 + + +def test_dp_backend_builds_dedicated_global_nccl_group(monkeypatch): + created_group = object() + created_stream = object() + new_group_args = {} + real_torch_zeros = torch.zeros + + def new_group(*, ranks, backend): + new_group_args.update(ranks=ranks, backend=backend) + return created_group + + def zeros_without_cuda(*args, **kwargs): + kwargs.pop("device", None) + return real_torch_zeros(*args, **kwargs) + + monkeypatch.setattr(lightspec_planner_impl, "get_global_world_size", lambda: 4) + monkeypatch.setattr(lightspec_planner_impl.dist, "new_group", new_group) + monkeypatch.setattr(lightspec_planner_impl.torch, "zeros", zeros_without_cuda) + monkeypatch.setattr(lightspec_planner_impl.torch.cuda, "Stream", lambda: created_stream) + + backend = DPChunkedPrefillBackend.__new__(DPChunkedPrefillBackend) + backend.args = SimpleNamespace(mtp_mode="eagle3", mtp_dynamic_verify=True, dp=2) + backend.max_draft_step = 3 + backend.dp_size = 2 + backend.model = SimpleNamespace(graph=None) + backend.draft_models = [SimpleNamespace(graph=None)] + backend.enable_prefill_microbatch_overlap = False + backend.enable_decode_microbatch_overlap = False + backend.init_spec_engine() + + assert new_group_args == {"ranks": [0, 1, 2, 3], "backend": "nccl"} + planner = backend.spec_engine.planner + assert planner._draft_step_group is created_group + assert planner._draft_step_tensor.shape == (1,) + assert planner._draft_step_stream is created_stream + + +def test_dp_backend_does_not_build_group_for_single_draft_step_mode(monkeypatch): + monkeypatch.setattr( + lightspec_planner_impl.dist, + "new_group", + lambda **kwargs: pytest.fail(f"unexpected NCCL group: {kwargs}"), + ) + + backend = DPChunkedPrefillBackend.__new__(DPChunkedPrefillBackend) + backend.args = SimpleNamespace(mtp_mode="vanilla_with_att", mtp_dynamic_verify=True, dp=2) + backend.max_draft_step = 3 + backend.dp_size = 2 + backend.model = SimpleNamespace(graph=None) + backend.draft_models = [SimpleNamespace(graph=None, block_size=4)] + backend.enable_prefill_microbatch_overlap = False + backend.enable_decode_microbatch_overlap = False + + backend.init_spec_engine() + + assert backend.spec_engine.planner._draft_step_group is None + assert backend.spec_engine.planner._draft_step_tensor is None + assert backend.spec_engine.planner._draft_step_stream is None + + +def test_dp_backend_does_not_build_group_for_fixed_planner(monkeypatch): + monkeypatch.setattr( + lightspec_planner_impl.dist, + "new_group", + lambda **kwargs: pytest.fail(f"unexpected NCCL group: {kwargs}"), + ) + + backend = DPChunkedPrefillBackend.__new__(DPChunkedPrefillBackend) + backend.args = SimpleNamespace(mtp_mode="eagle3", mtp_dynamic_verify=False, dp=2) + backend.max_draft_step = 3 + backend.dp_size = 2 + backend.enable_prefill_microbatch_overlap = False + backend.enable_decode_microbatch_overlap = False + + backend.init_spec_engine() + + assert not hasattr(backend.spec_engine.planner, "_draft_step_group") + + +def test_dp_backend_overlap_decode_reuses_common_lightspec_group(monkeypatch): + created_group = object() + created_stream = object() + real_torch_zeros = torch.zeros + + monkeypatch.setattr(lightspec_planner_impl, "get_global_world_size", lambda: 2) + monkeypatch.setattr(lightspec_planner_impl.dist, "new_group", lambda **kwargs: created_group) + monkeypatch.setattr( + lightspec_planner_impl.torch, + "zeros", + lambda *args, **kwargs: real_torch_zeros( + *args, **{key: value for key, value in kwargs.items() if key != "device"} + ), + ) + monkeypatch.setattr(lightspec_planner_impl.torch.cuda, "Stream", lambda: created_stream) + + backend = DPChunkedPrefillBackend.__new__(DPChunkedPrefillBackend) + backend.args = SimpleNamespace(mtp_mode="eagle3", mtp_dynamic_verify=True, dp=2) + backend.max_draft_step = 3 + backend.dp_size = 2 + backend.model = SimpleNamespace(graph=None) + backend.draft_models = [SimpleNamespace(graph=None)] + backend.enable_prefill_microbatch_overlap = False + backend.enable_decode_microbatch_overlap = True + + backend.init_spec_engine() + + assert backend.spec_engine.planner._draft_step_group is created_group + assert backend.spec_engine.planner._draft_step_stream is created_stream + assert backend.decode_draft_engine is backend.dp_overlap_spec_engine + assert backend.decode_draft_engine.common_engine.planner is backend.spec_engine.planner + + +def test_dp_backend_keeps_dynamic_planning_before_global_reduction(): + backend = DPChunkedPrefillBackend.__new__(DPChunkedPrefillBackend) + backend.args = SimpleNamespace(mtp_mode="eagle3", mtp_dynamic_verify=True, dp=1) + backend.max_draft_step = 3 + backend.dp_size = 1 + backend.model = SimpleNamespace(graph=None) + backend.draft_models = [SimpleNamespace(graph=None)] + backend.enable_prefill_microbatch_overlap = False + backend.enable_decode_microbatch_overlap = False + + backend.init_spec_engine() + + assert isinstance(backend.spec_engine.planner, LightSpecPlanner) + assert backend.spec_engine.proposer.enable_dynmaic_mtp + assert backend.spec_engine.planner._draft_step_group is None + + +def test_dp_prefill_and_decode_select_overlap_engine_independently(): + args = SimpleNamespace(mtp_mode="eagle3", mtp_dynamic_verify=False, dp=1) + + for prefill_overlap in (False, True): + for decode_overlap in (False, True): + backend = DPChunkedPrefillBackend.__new__(DPChunkedPrefillBackend) + backend.args = args + backend.max_draft_step = 2 + backend.dp_size = 1 + backend.enable_prefill_microbatch_overlap = prefill_overlap + backend.enable_decode_microbatch_overlap = decode_overlap + backend.init_spec_engine() + + expected_prefill_engine = backend.dp_overlap_spec_engine if prefill_overlap else backend.spec_engine + expected_decode_engine = backend.dp_overlap_spec_engine if decode_overlap else backend.spec_engine + assert backend.prefill_draft_engine is expected_prefill_engine + assert backend.decode_draft_engine is expected_decode_engine + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_dp_decode_mtp_runs_common_engine_for_empty_batch(monkeypatch): + device = "cuda" + empty_i32 = torch.empty((0,), dtype=torch.int32, device=device) + model_input = SimpleNamespace( + batch_size=0, + b_req_idx=empty_i32, + b_mtp_index=empty_i32, + ) + model_output = SimpleNamespace(logits=torch.empty((0, 8), device=device)) + calls = [] + + class _CommonEngine: + def plan_decode(self, **kwargs): + calls.append("plan") + return SpecDecodePlan(0, 0, 2, 2) + + def prepare_decode_model_input(self, **kwargs): + calls.append("prepare") + return kwargs["model_input"], None + + def propose_next(self, **kwargs): + calls.append("propose") + assert kwargs["target_next_token_ids"].shape == (0,) + assert kwargs["b_req_mtp_start_loc"].shape == (0,) + assert kwargs["accept_len"].shape == (0,) + return SpecProposal(token_ids=torch.empty((0, 2), dtype=torch.int64, device=device)) + + backend = DPChunkedPrefillBackend.__new__(DPChunkedPrefillBackend) + backend.spec_engine = _CommonEngine() + backend.args = SimpleNamespace(penalty_counter_mode="gpu_counter") + backend.model = SimpleNamespace(forward=lambda _: model_output) + event_pack = SimpleNamespace( + notify_post_handle_and_wait_pre_post_handle=lambda: calls.append("post_wait"), + notify_forward_and_wait_post_handle=lambda: calls.append("forward_wait"), + notify_pre_post_handle=lambda: calls.append("pre_post"), + ) + monkeypatch.setattr( + dp_backend_impl, + "prepare_decode_inputs", + lambda req_objs: (model_input, []), + ) + monkeypatch.setattr( + dp_backend_impl, + "g_infer_context", + SimpleNamespace(get_overlap_stream=lambda: torch.cuda.current_stream()), + ) + backend.decode_mtp(event_pack=event_pack, decode_reqs=[]) + + assert calls == [ + "plan", + "prepare", + "post_wait", + "propose", + "forward_wait", + "pre_post", + ] + + +def test_dp_overlap_engine_delegates_raw_verify_layout_to_proposer(): + calls = {} + + class _Proposer: + def propose_next_overlap(self, **kwargs): + calls.update(kwargs) + return SpecProposal(token_ids=kwargs["target_next_token_ids0"].new_empty((3, 7))) + + engine = DPOverlapSpecEngine.__new__(DPOverlapSpecEngine) + engine.proposer = _Proposer() + model_input0 = SimpleNamespace(batch_size=8) + model_input1 = SimpleNamespace(batch_size=16) + model_output0 = SimpleNamespace() + model_output1 = SimpleNamespace() + target_next_token_ids0 = torch.arange(8, dtype=torch.int64) + target_next_token_ids1 = torch.arange(8, 24, dtype=torch.int64) + accept_len0 = torch.tensor([2], dtype=torch.int32) + accept_len1 = torch.tensor([3, 4], dtype=torch.int32) + + proposal = engine.propose_next_overlap( + target_model_input0=model_input0, + target_model_output0=model_output0, + target_next_token_ids0=target_next_token_ids0, + accept_len0=accept_len0, + target_model_input1=model_input1, + target_model_output1=model_output1, + target_next_token_ids1=target_next_token_ids1, + accept_len1=accept_len1, + draft_step=7, + ) + + assert calls["target_model_input0"] is model_input0 + assert calls["target_model_input1"] is model_input1 + assert calls["target_model_output0"] is model_output0 + assert calls["target_model_output1"] is model_output1 + assert calls["target_next_token_ids0"] is target_next_token_ids0 + assert calls["target_next_token_ids1"] is target_next_token_ids1 + assert calls["accept_len0"] is accept_len0 + assert calls["accept_len1"] is accept_len1 + assert calls["draft_step"] == 7 + assert proposal.token_ids.shape == (3, 7) + + +def test_dp_overlap_engine_distributes_dynamic_verify_budget_by_request_count(): + prepare_calls = [] + + class _CommonEngine: + def prepare_decode_model_input(self, **kwargs): + prepare_calls.append(kwargs) + model_input = kwargs["model_input"] + model_input.batch_size = kwargs["plan"].dynamic_batch_size + return model_input, f"mask{len(prepare_calls)}" + + engine = DPOverlapSpecEngine.__new__(DPOverlapSpecEngine) + engine.common_engine = _CommonEngine() + model_input0 = SimpleNamespace(batch_size=12) + model_input1 = SimpleNamespace(batch_size=8) + plan = SpecDecodePlan( + origin_batch_size=20, + dynamic_batch_size=11, + draft_step=2, + pre_draft_step=3, + ) + + compacted_input0, mask0, compacted_input1, mask1 = engine.prepare_decode_model_inputs( + model_input0=model_input0, + req_num0=3, + model_input1=model_input1, + req_num1=2, + plan=plan, + ) + + assert compacted_input0.batch_size == 7 + assert compacted_input1.batch_size == 4 + assert mask0 == "mask1" + assert mask1 == "mask2" + assert prepare_calls[0]["req_num"] == 3 + assert prepare_calls[0]["plan"] == SpecDecodePlan(12, 7, 2, 3) + assert prepare_calls[1]["req_num"] == 2 + assert prepare_calls[1]["plan"] == SpecDecodePlan(8, 4, 2, 3) + + +@pytest.mark.parametrize( + ("batch_size0", "batch_size1", "dynamic_batch_size", "expected_batch_sizes"), + ( + (4, 8, 9, (4, 5)), + (8, 4, 10, (6, 4)), + ), +) +def test_dp_overlap_engine_moves_verify_rows_to_the_side_with_capacity( + batch_size0, + batch_size1, + dynamic_batch_size, + expected_batch_sizes, +): + dynamic_batch_sizes = [] + + class _CommonEngine: + def prepare_decode_model_input(self, **kwargs): + dynamic_batch_sizes.append(kwargs["plan"].dynamic_batch_size) + return kwargs["model_input"], None + + engine = DPOverlapSpecEngine.__new__(DPOverlapSpecEngine) + engine.common_engine = _CommonEngine() + engine.prepare_decode_model_inputs( + model_input0=SimpleNamespace(batch_size=batch_size0), + req_num0=2, + model_input1=SimpleNamespace(batch_size=batch_size1), + req_num1=2, + plan=SpecDecodePlan( + origin_batch_size=batch_size0 + batch_size1, + dynamic_batch_size=dynamic_batch_size, + draft_step=2, + pre_draft_step=3, + ), + ) + + assert tuple(dynamic_batch_sizes) == expected_batch_sizes + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_dp_overlap_decode_delegates_empty_layout(monkeypatch): + device = "cuda" + empty_i32 = torch.empty((0,), dtype=torch.int32, device=device) + model_input0 = SimpleNamespace(batch_size=0, b_req_idx=empty_i32, b_mtp_index=empty_i32) + model_input1 = SimpleNamespace(batch_size=0, b_req_idx=empty_i32, b_mtp_index=empty_i32) + model_output0 = SimpleNamespace(logits=torch.empty((0, 8), device=device)) + model_output1 = SimpleNamespace(logits=torch.empty((0, 8), device=device)) + calls = [] + + class _OverlapEngine: + def plan_decode(self, **kwargs): + calls.append("plan") + return SpecDecodePlan(0, 0, 2, 2) + + def prepare_decode_model_inputs(self, **kwargs): + calls.append("prepare") + return kwargs["model_input0"], None, kwargs["model_input1"], None + + def propose_next_overlap(self, **kwargs): + calls.append("propose") + assert kwargs["target_next_token_ids0"].shape == (0,) + assert kwargs["target_next_token_ids0"].dtype == torch.int64 + assert kwargs["target_next_token_ids1"].shape == (0,) + assert kwargs["target_next_token_ids1"].dtype == torch.int64 + assert kwargs["accept_len0"].shape == (0,) + assert kwargs["accept_len0"].dtype == torch.int32 + assert kwargs["accept_len1"].shape == (0,) + assert kwargs["accept_len1"].dtype == torch.int32 + assert kwargs["draft_step"] == 2 + return SpecProposal(token_ids=torch.empty((0, 2), dtype=torch.int64, device=device)) + + backend = DPChunkedPrefillBackend.__new__(DPChunkedPrefillBackend) + backend.decode_draft_engine = _OverlapEngine() + backend.args = SimpleNamespace(penalty_counter_mode="gpu_counter") + backend.model = SimpleNamespace( + microbatch_overlap_decode=lambda input0, input1: (model_output0, model_output1), + ) + event_pack = SimpleNamespace( + notify_post_handle_and_wait_pre_post_handle=lambda: calls.append("post_wait"), + notify_forward_and_wait_post_handle=lambda: calls.append("forward_wait"), + notify_pre_post_handle=lambda: calls.append("pre_post"), + ) + monkeypatch.setattr( + dp_backend_impl, + "overlap_prepare_decode_inputs", + lambda req_objs: (model_input0, [], [], model_input1, [], []), + ) + monkeypatch.setattr( + dp_backend_impl, + "g_infer_context", + SimpleNamespace(get_overlap_stream=lambda: torch.cuda.current_stream()), + ) + backend.decode_overlap_mtp(event_pack=event_pack, decode_reqs=[]) + + assert calls == [ + "plan", + "prepare", + "post_wait", + "propose", + "forward_wait", + "pre_post", + ] diff --git a/unit_tests/server/router/model_infer/mode_backend/test_generic_pre_process.py b/unit_tests/server/router/model_infer/mode_backend/test_generic_pre_process.py new file mode 100644 index 0000000000..42519a7415 --- /dev/null +++ b/unit_tests/server/router/model_infer/mode_backend/test_generic_pre_process.py @@ -0,0 +1,176 @@ +from types import SimpleNamespace + +import torch +from lightllm.server.router.model_infer.mode_backend import generic_pre_process + + +def _patch_empty_input_context(monkeypatch): + mem_manager = SimpleNamespace( + HOLD_TOKEN_MEMINDEXES=(-1,), + alloc=lambda size: torch.empty((size,), dtype=torch.int32), + ) + infer_context = SimpleNamespace( + args=SimpleNamespace(page_size=1), + req_manager=SimpleNamespace( + HOLD_REQUEST_ID=-1, + mem_manager=mem_manager, + req_to_token_indexs=torch.zeros((16, 32), dtype=torch.int32), + ), + radix_cache=None, + ) + monkeypatch.setattr(generic_pre_process, "g_infer_context", infer_context) + return infer_context + + +def _patch_overlap_input_context(monkeypatch): + return _patch_empty_input_context(monkeypatch) + + +def _make_prefill_req(req_idx: int, token_num: int): + input_token_ids = [req_idx] * token_num + return SimpleNamespace( + req_idx=req_idx, + cur_kv_len=0, + hold_kv_len=token_num, + multimodal_params={"images": [], "audios": []}, + get_chuncked_input_token_ids=lambda: input_token_ids, + get_input_token_ids=lambda: input_token_ids, + get_cur_total_len=lambda: token_num, + ) + + +def _make_decode_req(req_idx: int): + return SimpleNamespace( + req_idx=req_idx, + cur_kv_len=3, + hold_kv_len=4, + mtp_step=0, + multimodal_params={"images": [], "audios": []}, + shared_kv_node=None, + get_cur_total_len=lambda: 4, + get_radix_cache_shared_len=lambda: 0, + ) + + +def test_prepare_prefill_inputs_allows_empty_batch(monkeypatch): + _patch_empty_input_context(monkeypatch) + + model_input, run_reqs = generic_pre_process.prepare_prefill_inputs([], is_chuncked_mode=True) + + assert run_reqs == [] + assert model_input.batch_size == 0 + assert model_input.input_ids.shape == (0,) + assert model_input.b_req_idx.shape == (0,) + assert model_input.b_prefill_start_loc.shape == (0,) + assert model_input.b_prefill_has_output_cpu == [] + assert model_input.max_q_seq_len == 0 + assert model_input.max_kv_seq_len == 0 + + +def test_prepare_decode_inputs_allows_empty_batch(monkeypatch): + _patch_empty_input_context(monkeypatch) + + model_input, run_reqs = generic_pre_process.prepare_decode_inputs([]) + + assert run_reqs == [] + assert model_input.batch_size == 0 + assert model_input.input_ids is None + assert model_input.b_req_idx.shape == (0,) + assert model_input.b_position_delta.shape == (0,) + assert model_input.b_shared_seq_len.shape == (0,) + assert model_input.b_shared_radix_node_id.shape == (0,) + assert model_input.max_q_seq_len == 1 + assert model_input.max_kv_seq_len == 0 + + +def test_overlap_prefill_balances_request_token_load_without_padding(monkeypatch): + _patch_overlap_input_context(monkeypatch) + reqs = [ + _make_prefill_req(req_idx=0, token_num=8), + _make_prefill_req(req_idx=1, token_num=7), + _make_prefill_req(req_idx=2, token_num=6), + _make_prefill_req(req_idx=3, token_num=5), + _make_prefill_req(req_idx=4, token_num=1), + ] + + ( + model_input0, + run_reqs0, + model_input1, + run_reqs1, + ) = generic_pre_process.overlap_prepare_prefill_inputs(reqs) + + assert [req.req_idx for req in run_reqs0] == [0, 3, 4] + assert [req.req_idx for req in run_reqs1] == [1, 2] + assert model_input0.b_req_idx.tolist() == [0, 3, 4] + assert model_input1.b_req_idx.tolist() == [1, 2] + assert model_input0.input_ids.shape == (14,) + assert model_input1.input_ids.shape == (13,) + assert model_input0.batch_size == 3 + assert model_input1.batch_size == 2 + + +def test_overlap_prefill_balances_single_token_request_normally(monkeypatch): + _patch_overlap_input_context(monkeypatch) + req = _make_prefill_req(req_idx=7, token_num=1) + + ( + model_input0, + run_reqs0, + model_input1, + run_reqs1, + ) = generic_pre_process.overlap_prepare_prefill_inputs([req]) + + assert run_reqs0 == [req] + assert model_input0.batch_size == 1 + assert model_input0.input_ids.tolist() == [7] + assert model_input0.b_req_idx.tolist() == [7] + assert run_reqs1 == [] + assert model_input1.batch_size == 0 + assert model_input1.input_ids.shape == (0,) + assert model_input1.b_req_idx.shape == (0,) + + +def test_overlap_decode_builds_two_unpadded_inputs(monkeypatch): + _patch_overlap_input_context(monkeypatch) + reqs = [_make_decode_req(req_idx=index) for index in range(3)] + + ( + model_input0, + run_reqs0, + decode_reqs0, + model_input1, + run_reqs1, + decode_reqs1, + ) = generic_pre_process.overlap_prepare_decode_inputs(reqs) + + assert decode_reqs0 == reqs[:2] + assert decode_reqs1 == reqs[2:] + assert run_reqs0 == reqs[:2] + assert run_reqs1 == reqs[2:] + assert model_input0.batch_size == 2 + assert model_input1.batch_size == 1 + assert model_input0.b_req_idx.tolist() == [0, 1] + assert model_input1.b_req_idx.tolist() == [2] + + +def test_overlap_decode_preserves_empty_microbatch(monkeypatch): + _patch_overlap_input_context(monkeypatch) + req = _make_decode_req(req_idx=7) + + ( + model_input0, + run_reqs0, + decode_reqs0, + model_input1, + run_reqs1, + decode_reqs1, + ) = generic_pre_process.overlap_prepare_decode_inputs([req]) + + assert decode_reqs0 == [req] + assert decode_reqs1 == [] + assert run_reqs0 == [req] + assert model_input0.batch_size == 1 + assert run_reqs1 == [] + assert model_input1.batch_size == 0 + assert model_input1.b_req_idx.shape == (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 new file mode 100644 index 0000000000..ece87927e0 --- /dev/null +++ b/unit_tests/server/router/model_infer/mode_backend/test_multi_level_kv_cache.py @@ -0,0 +1,507 @@ +from types import SimpleNamespace +from collections import deque + +import pytest +import torch + +from lightllm.server.multi_level_kv_cache import ( + AdaptiveCachePlacementController, + CacheCapacityConfig, + CacheTier, +) +from lightllm.server.router.model_infer.mode_backend.multi_level_kv_cache import MultiLevelKvCacheModule +from lightllm.server.router.model_infer.mode_backend import multi_level_kv_cache as multi_level_kv_cache_impl +from lightllm.server.router.model_infer.infer_batch import InferenceContext + + +def test_running_cpu_cache_tasks_are_not_ready(monkeypatch): + from lightllm.server.router.model_infer.infer_batch import InferReq + from lightllm.server.router.model_infer.mode_backend.base_backend import ModeBackend + from lightllm.server.router.model_infer.mode_backend import base_backend as base_backend_impl + + status = InferReq._CpuCacheTaskStatus + + def req(load_status, offload_status): + return SimpleNamespace( + cpu_cache_load_task_status=load_status, + cpu_cache_offload_task_status=offload_status, + ) + + running_load = req(status.RUNNING, status.NOT_STARTED) + running_offload = req(status.FINISHED, status.RUNNING) + ready = req(status.NOT_STARTED, status.NOT_STARTED) + backend = ModeBackend.__new__(ModeBackend) + backend.args = SimpleNamespace(enable_cpu_cache=True) + monkeypatch.setattr( + base_backend_impl, + "g_infer_context", + SimpleNamespace(requests_mapping={1: running_load, 2: running_offload, 3: ready}), + ) + + assert backend._filter_not_ready_reqs([1, 2, 3]) == [ready] + backend.args.enable_cpu_cache = False + assert backend._filter_not_ready_reqs([1, 2, 3]) == [running_load, running_offload, ready] + + +def test_cache_tiers_reassignment_is_rejected(): + context = InferenceContext() + context.cache_placement_controller = AdaptiveCachePlacementController( + capacity=CacheCapacityConfig(1, 1, 1), + args=SimpleNamespace(enable_cpu_cache=True, enable_disk_cache=True), + ) + reqs = [ + SimpleNamespace( + cache_tiers=(CacheTier.GPU,), + shm_req=SimpleNamespace(group_req_id=index, request_id=index, input_len=length), + ) + for index, length in enumerate((100, 200, 300)) + ] + + context.cache_placement_controller.set_req_cache_way(reqs) + context.cache_placement_controller = AdaptiveCachePlacementController( + capacity=CacheCapacityConfig(0, 1, 0), + args=SimpleNamespace(enable_cpu_cache=True, enable_disk_cache=False), + ) + + with pytest.raises(AssertionError): + context.cache_placement_controller.set_req_cache_way(reqs) + + +def test_non_gpu_cache_tiers_release_owned_tokens_without_radix_insert(): + released_refs = [] + context = InferenceContext() + context.is_hybrid_att_model = False + context.req_manager = SimpleNamespace() + context.radix_cache = SimpleNamespace(dec_node_ref_counter=released_refs.append) + shared_node = SimpleNamespace(node_prefix_total_len=2) + req = SimpleNamespace( + req_idx=0, + cur_kv_len=5, + hold_kv_len=7, + shared_kv_node=shared_node, + token_index=SimpleNamespace( + get_free_token_indexes=lambda start, end: torch.arange(10 + start, 10 + end), + get_cpu_token_indexes=lambda start, end: torch.arange(10 + start, 10 + end), + reset=lambda: None, + ), + ) + free_token_index_blocks = [] + + context._free_req_mem_without_radix_insert(free_token_index_blocks, req) + + assert free_token_index_blocks[0].tolist() == [12, 13, 14, 15, 16] + assert released_refs == [shared_node] + assert req.shared_kv_node is None + + +def test_legacy_cache_tiers_still_insert_gpu_radix_cache(): + context = InferenceContext() + context.radix_cache = object() + context.is_hybrid_att_model = False + inserted_reqs = [] + context._full_att_free_req = lambda free_token_index, req: inserted_reqs.append(req) + req = SimpleNamespace( + cache_tiers=(CacheTier.GPU, CacheTier.CPU, CacheTier.DISK), + cur_kv_len=3, + shm_req=SimpleNamespace(shm_cur_kv_len=3), + token_index=SimpleNamespace(reset=lambda: None), + ) + + context.free_a_req_mem([], req) + + assert inserted_reqs == [req] + + +def test_finished_batch_routes_cpu_and_disk_offloads_separately(monkeypatch): + class NotStartedStatus: + @staticmethod + def is_finished(): + return False + + @staticmethod + def is_running(): + return False + + @staticmethod + def is_not_started(): + return True + + module = MultiLevelKvCacheModule.__new__(MultiLevelKvCacheModule) + module.args = SimpleNamespace(cpu_cache_token_page_size=64, linear_att_hash_page_size=64) + module.backend = SimpleNamespace(radix_cache=object()) + module.offload_tasks = deque() + module.need_sync_compute_stream = lambda: False + offload_calls = [] + + def start_offload(req): + offload_calls.append((req.shm_req.request_id, CacheTier.DISK in req.cache_tiers)) + return SimpleNamespace(req=req) + + module.write_req_to_cpu = start_offload + current = torch.cuda.current_stream() + offload_stream = torch.cuda.Stream() + waited = [] + offload_stream.wait_stream = lambda stream: waited.append(stream) + module.get_cpu_offload_stream = lambda: offload_stream + monkeypatch.setattr(multi_level_kv_cache_impl.g_infer_context, "is_hybrid_att_model", False) + cache_tiers = ( + (CacheTier.GPU,), + (CacheTier.CPU,), + (CacheTier.CPU, CacheTier.DISK), + ) + reqs = [ + SimpleNamespace( + cache_tiers=req_cache_tiers, + cur_kv_len=length, + cpu_cache_offload_task_status=NotStartedStatus(), + shm_req=SimpleNamespace( + group_req_id=index, + request_id=index, + input_len=length, + ), + ) + for index, (length, req_cache_tiers) in enumerate(zip((100, 200, 300), cache_tiers)) + ] + offload_reqs = [req for req in reqs if CacheTier.CPU in req.cache_tiers or CacheTier.DISK in req.cache_tiers] + + true_finished_reqs = module.offload_finished_reqs_to_cpu_cache(offload_reqs) + + assert waited == [current] + assert true_finished_reqs == [] + assert offload_calls == [(1, False), (2, True)] + assert len(module.offload_tasks) == 2 + + overlap = torch.cuda.Stream() + waited_by_overlap = [] + overlap.wait_stream = lambda stream: waited_by_overlap.append(stream) + monkeypatch.setattr(multi_level_kv_cache_impl.g_infer_context, "get_overlap_stream", lambda: overlap) + module.need_sync_compute_stream = lambda: True + waited.clear() + assert module.offload_finished_reqs_to_cpu_cache([]) == [] + assert waited == [] and waited_by_overlap == [] + short_req = SimpleNamespace( + cache_tiers=(CacheTier.CPU,), + cur_kv_len=1, + cpu_cache_offload_task_status=NotStartedStatus(), + shm_req=SimpleNamespace(group_req_id=0, request_id=0, input_len=1), + ) + assert module.offload_finished_reqs_to_cpu_cache([short_req]) == [short_req] + assert waited == [current, overlap] + assert waited_by_overlap == [offload_stream] + + +def test_non_gpu_linear_cache_tiers_release_pending_state_pages(): + freed_small_pages = [] + freed_big_pages = [] + context = InferenceContext() + context.is_hybrid_att_model = True + context.req_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), + ) + req = SimpleNamespace( + req_idx=0, + cur_kv_len=3, + hold_kv_len=3, + token_index=SimpleNamespace( + get_free_token_indexes=lambda start, end: torch.arange(10 + start, 10 + end), + reset=lambda: None, + ), + shared_kv_node=None, + tail_small_page_buffer_id=7, + hybrid_len_to_big_page_id={128: 8, 256: 9}, + ) + free_token_index_blocks = [] + + context._free_req_mem_without_radix_insert(free_token_index_blocks, req) + + assert free_token_index_blocks[0].tolist() == [10, 11, 12] + assert freed_small_pages == [7] + assert freed_big_pages == [8, 9] + assert req.tail_small_page_buffer_id is None + assert req.hybrid_len_to_big_page_id == {} + + +def test_cpu_cache_load_uses_exact_aligned_size(monkeypatch): + table = torch.full((1, 136), -1, dtype=torch.int32) + table[0, :4] = torch.arange(4, dtype=torch.int32) + next_index = 4 + alloc_sizes = [] + loaded_indexes = [] + + def alloc(need_size): + nonlocal next_index + alloc_sizes.append(need_size) + mem_indexes = torch.arange(next_index, next_index + need_size, dtype=torch.int32) + next_index += need_size + return mem_indexes + + def alloc_token_indexes(req, end, allocator, page_size, radix_cache=None): + alloc_token_num = end - req.hold_kv_len + mem_indexes = alloc(alloc_token_num) + new_hold_kv_len = req.hold_kv_len + alloc_token_num + table[req.req_idx, req.hold_kv_len : new_hold_kv_len] = mem_indexes + req.hold_kv_len = new_hold_kv_len + return mem_indexes + + resources = SimpleNamespace( + 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() + ), + ), + full_att_allocator=None, + ) + backend = SimpleNamespace( + is_master_in_dp=False, + dp_world_size=1, + is_hybrid_att_model=False, + radix_cache=None, + model=SimpleNamespace( + att_state_manager=resources, + req_manager=SimpleNamespace( + req_to_token_indexs=table, alloc_token_indexes=alloc_token_indexes, wait_token_indexes=lambda: None + ), + ), + ) + module = MultiLevelKvCacheModule.__new__(MultiLevelKvCacheModule) + module.load_tasks = deque() + module.args = SimpleNamespace(page_size=4, cpu_cache_token_page_size=132) + module.backend = backend + module.init_sync_group = object() + module.need_sync_compute_stream = lambda: False + module.cpu_cache_client = SimpleNamespace(cpu_kv_cache_tensor=None) + context = SimpleNamespace( + req_manager=SimpleNamespace(req_to_token_indexs=table, alloc_token_indexes=alloc_token_indexes), + radix_cache=None, + get_can_alloc_token_num=lambda: 200, + ) + req = SimpleNamespace( + req_idx=0, + cur_kv_len=4, + hold_kv_len=4, + mm_input_spans=[], + shm_req=SimpleNamespace( + input_len=256, + disk_prompt_cache_len=0, + cpu_cache_match_page_indexes=SimpleNamespace(get_all=lambda: [7]), + token_hash_page_len_list=SimpleNamespace(get_all=lambda: [132]), + ), + sampling_param=SimpleNamespace(shm_param=SimpleNamespace(prompt_logprobs=-1)), + args=SimpleNamespace(page_size=4), + ) + req.token_index = SimpleNamespace(get_gpu_token_indexes=lambda start, end: table[req.req_idx, start:end]) + req.kv_cache_alloc_need = lambda target: max(0, (target + 3) // 4 * 4 - req.hold_kv_len) + req.alloc_kv_mem = lambda alloc_token_num: alloc_token_indexes(req, req.hold_kv_len + alloc_token_num, None, 4) + monkeypatch.setattr(multi_level_kv_cache_impl, "g_infer_context", context) + monkeypatch.setattr(multi_level_kv_cache_impl.dist, "barrier", lambda group: None) + monkeypatch.setattr(torch.Tensor, "cuda", lambda self, **kwargs: self) + + module.load_cpu_cache_to_reqs([req]) + + assert alloc_sizes == [128] + assert req.cur_kv_len == 132 + assert req.hold_kv_len == 132 + assert loaded_indexes == list(range(132)) + assert table[0, :132].tolist() == list(range(132)) + + +def test_standard_load_releases_unloaded_pages_immediately(monkeypatch): + module = MultiLevelKvCacheModule.__new__(MultiLevelKvCacheModule) + module.backend = SimpleNamespace(is_master_in_dp=True, dp_world_size=1) + module.init_sync_group = object() + module.need_sync_compute_stream = lambda: False + events = [] + module.cpu_cache_client = SimpleNamespace( + lock=SimpleNamespace( + acquire_sleep1ms=lambda: events.append("lock"), + release=lambda: events.append("unlock"), + ), + deref_pages=lambda page_list: events.append(("deref", list(page_list))), + ) + reqs = [ + SimpleNamespace( + cur_kv_len=0, + mm_input_spans=[], + sampling_param=SimpleNamespace(shm_param=SimpleNamespace(prompt_logprobs=logprobs)), + shm_req=SimpleNamespace( + input_len=64, + disk_prompt_cache_len=0, + cpu_cache_match_page_indexes=SimpleNamespace(get_all=lambda page=page: [page]), + token_hash_page_len_list=SimpleNamespace(get_all=lambda: [64]), + ), + ) + for page, logprobs in ((4, 0), (7, -1)) + ] + monkeypatch.setattr(multi_level_kv_cache_impl.g_infer_context, "get_can_alloc_token_num", lambda: 128) + + module.load_cpu_cache_to_reqs(reqs) + + assert events == ["lock", ("deref", [4]), ("deref", [7]), "unlock"] + assert not getattr(module, "load_tasks", None) + assert all(req.cpu_cache_load_task_status.is_finished() for req in reqs) + + +def _load_module(monkeypatch, *, sync, master=True, world=1): + events = [] + + def alloc_token_indexes(req, end, allocator, page_size, radix_cache=None): + req.hold_kv_len = end + return torch.arange(req.cur_kv_len, end, dtype=torch.int32) + + module = MultiLevelKvCacheModule.__new__(MultiLevelKvCacheModule) + module.args = SimpleNamespace(page_size=4, cpu_cache_token_page_size=128) + module.backend = SimpleNamespace( + is_master_in_dp=master, + dp_world_size=world, + is_hybrid_att_model=False, + model=SimpleNamespace( + req_manager=SimpleNamespace( + req_to_token_indexs=torch.arange(512, dtype=torch.int32).view(1, -1), + alloc_token_indexes=alloc_token_indexes, + wait_token_indexes=lambda: None, + ), + att_state_manager=SimpleNamespace( + get_cpu_cache_operator=lambda: SimpleNamespace( + get_page_layouts=lambda: {}, + read_req=lambda **kwargs: events.append("copy"), + ), + full_att_allocator=None, + ), + ), + ) + module.need_sync_compute_stream = lambda: sync + module.init_sync_group = module.filter_group = object() + module.load_sync_tensor = object() + module.load_tasks = deque() + module.cpu_cache_client = SimpleNamespace( + cpu_kv_cache_tensor=None, + lock=SimpleNamespace(acquire_sleep1ms=lambda: events.append("lock"), release=lambda: events.append("unlock")), + deref_pages=lambda page_list: events.append(("deref", list(page_list))), + ) + req = SimpleNamespace( + req_idx=0, + cur_kv_len=0, + hold_kv_len=0, + mm_input_spans=[], + sampling_param=SimpleNamespace(shm_param=SimpleNamespace(prompt_logprobs=-1)), + shm_req=SimpleNamespace( + input_len=256, + disk_prompt_cache_len=0, + cpu_prompt_cache_len=0, + shm_cur_kv_len=0, + cpu_cache_match_page_indexes=SimpleNamespace(get_all=lambda: [3, 4]), + token_hash_page_len_list=SimpleNamespace(get_all=lambda: [128, 256]), + ), + args=SimpleNamespace(page_size=4), + ) + token_table = module.backend.model.req_manager.req_to_token_indexs + req.token_index = SimpleNamespace(get_gpu_token_indexes=lambda start, end: token_table[req.req_idx, start:end]) + req.kv_cache_alloc_need = lambda target: max(0, (target + 3) // 4 * 4 - req.hold_kv_len) + req.alloc_kv_mem = lambda alloc_token_num: alloc_token_indexes(req, req.hold_kv_len + alloc_token_num, None, 4) + monkeypatch.setattr(multi_level_kv_cache_impl.g_infer_context, "get_can_alloc_token_num", lambda: 512) + monkeypatch.setattr(multi_level_kv_cache_impl.g_infer_context, "get_overlap_stream", torch.cuda.Stream) + monkeypatch.setattr(multi_level_kv_cache_impl.dist, "barrier", lambda group: events.append("barrier")) + monkeypatch.setattr(torch.Tensor, "cuda", lambda self, **kwargs: self) + return module, req, events + + +def test_async_load_defers_deref_until_the_event_is_notified(monkeypatch): + module, req, events = _load_module(monkeypatch, sync=False) + current = torch.cuda.current_stream() + overlap = torch.cuda.Stream() + waited_by_load = [] + overlap.wait_stream = lambda stream: events.append("wait") + load_stream = module.get_cpu_load_stream() + load_stream.wait_stream = lambda stream: waited_by_load.append(stream) + monkeypatch.setattr(multi_level_kv_cache_impl.g_infer_context, "get_overlap_stream", lambda: overlap) + recorded = [] + + class Event: + def __init__(self): + self.ready = False + + def record(self): + recorded.append(self) + + def query(self): + return self.ready + + monkeypatch.setattr(multi_level_kv_cache_impl.torch.cuda, "Event", Event) + module.load_cpu_cache_to_reqs([req]) + assert waited_by_load == [current] + assert events == ["copy"] and len(recorded) == 1 and len(module.load_tasks) == 1 + module.update_cpu_cache_load_task_states() + assert "deref" not in [item[0] for item in events if isinstance(item, tuple)] + recorded[0].ready = True + module.update_cpu_cache_load_task_states() + assert events[-3:] == ["lock", ("deref", [3, 4]), "unlock"] and not module.load_tasks + + +def test_synchronous_backend_waits_overlap_for_the_load_stream(monkeypatch): + module, req, events = _load_module(monkeypatch, sync=True) + current = torch.cuda.current_stream() + overlap = torch.cuda.Stream() + waited_by_overlap = [] + waited_by_load = [] + overlap.wait_stream = lambda stream: waited_by_overlap.append(stream) + load_stream = module.get_cpu_load_stream() + load_stream.wait_stream = lambda stream: waited_by_load.append(stream) + monkeypatch.setattr(multi_level_kv_cache_impl.g_infer_context, "get_overlap_stream", lambda: overlap) + module.load_cpu_cache_to_reqs([req]) + assert waited_by_load == [current, overlap] + assert waited_by_overlap == [load_stream] + assert events == ["copy"] and len(module.load_tasks) == 1 + module.load_tasks[0].event.synchronize() + module.update_cpu_cache_load_task_states() + assert events[-3:] == ["lock", ("deref", [3, 4]), "unlock"] and not module.load_tasks + + +def test_async_load_notifies_every_rank_before_deref(monkeypatch): + module, req, events = _load_module(monkeypatch, sync=False, world=2) + collectives = [] + + class Event: + def __init__(self): + self.ready = True + + def record(self): + return None + + def query(self): + return self.ready + + monkeypatch.setattr(multi_level_kv_cache_impl.torch.cuda, "Event", Event) + monkeypatch.setattr( + multi_level_kv_cache_impl.dist, + "all_reduce", + lambda tensor, group: collectives.append(("all_reduce", group)), + ) + monkeypatch.setattr( + multi_level_kv_cache_impl.dist, + "broadcast_object_list", + lambda data, group, group_src: collectives.append(("broadcast", data[0], group)), + ) + module.load_cpu_cache_to_reqs([req]) + assert collectives == [("all_reduce", module.init_sync_group)] and not any( + isinstance(item, tuple) and item[0] == "deref" for item in events + ) + module.update_cpu_cache_load_task_states() + assert collectives[1][0] == "broadcast" and collectives[1][1] == 1 + assert ("deref", [3, 4]) in events + + +def test_fa3_cpu_cache_sync_defaults_off(): + from lightllm.common.basemodel.attention.base_att import BaseAttBackend + from lightllm.common.basemodel.attention.fa3.fp import Fa3AttBackend + from lightllm.common.basemodel.attention.fa3.fp8 import Fp8Fa3AttBackend + from lightllm.common.basemodel.attention.fa3.mla import MlaFa3AttBackend + from lightllm.common.basemodel.attention.fa3.neo import NeoFa3AttBackend + + assert BaseAttBackend.requires_cpu_cache_sync is False + assert all( + backend.requires_cpu_cache_sync is False + for backend in (Fa3AttBackend, Fp8Fa3AttBackend, MlaFa3AttBackend, NeoFa3AttBackend) + ) 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..17d564987d --- /dev/null +++ b/unit_tests/server/router/model_infer/mode_backend/test_nccl_control_channel.py @@ -0,0 +1,151 @@ +from __future__ import annotations + +import concurrent.futures +import socket +import threading +from types import SimpleNamespace +from unittest.mock import Mock + +import rpyc +import pytest + +from lightllm.server.router.model_infer.mode_backend.pd import nccl_kv_transporter as module + + +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.mode == "send" + assert handle.finish_cuda_event is event + event.query.return_value = True + assert handle.check_status() == "DONE" + + 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 + 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): + 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 new file mode 100644 index 0000000000..623f9564b0 --- /dev/null +++ b/unit_tests/server/router/model_infer/mode_backend/test_paged_side_allocations.py @@ -0,0 +1,253 @@ +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +# 先加载通用预处理模块，沿用生产初始化顺序，避免 basemodel 与 req_manager +# 在独立运行本测试文件时形成模块级循环导入。 +from lightllm.server.router.model_infer.mode_backend import generic_pre_process # noqa: F401 +from lightllm.server.router.model_infer.infer_batch import InferReq +from lightllm.server.router.model_infer.mode_backend.dp_backend import dp_shared_kv_trans +from lightllm.server.router.model_infer.mode_backend.pd.decode_node_impl import decode_impl +from lightllm.server.router.model_infer.mode_backend.pd.prefill_node_impl import prefill_impl + + +@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) + table[0, :4] = torch.arange(4, dtype=torch.int32) + alloc_sizes = [] + next_index = [4] + tasks = [] + backend = decode_impl.PDDecodeNode.__new__(decode_impl.PDDecodeNode) + backend.args = SimpleNamespace(pd_kv_page_size=4, page_size=4) + req_manager = _make_req_manager(table, alloc_sizes, next_index) + backend.model = SimpleNamespace( + req_manager=req_manager, + att_state_manager=SimpleNamespace(full_att_allocator=None), + ) + backend.is_master_in_dp = False + backend.radix_cache = None + + def create_task(**kwargs): + tasks.append( + ( + kwargs["kv_start_index"], + kwargs["kv_end_index"], + kwargs["mem_indexes"], + kwargs.get("page_kind", "kv"), + ) + ) + kwargs["group"].task_list.append(SimpleNamespace()) + + backend._create_pd_trans_task = create_task + monkeypatch.setattr(decode_impl.g_infer_context, "is_hybrid_att_model", is_hybrid_att_model) + req = _bind_alloc_need( + SimpleNamespace( + req_idx=0, + cur_kv_len=4, + hold_kv_len=4, + pd_trans_kv_start_index=0, + shm_req=SimpleNamespace(input_len=10), + ), + page_size=4, + ) + req.kv_cache_alloc_need = lambda target: req_manager.get_need_alloc_token_num(req, target) + req.alloc_kv_mem = lambda alloc_token_num: req_manager.alloc_token_indexes( + req, req.cur_kv_len + alloc_token_num, None, req.args.page_size + ) + + backend._decode_node_gen_trans_tasks(req) + + assert alloc_sizes == [8] + assert req.cur_kv_len == 10 + assert req.hold_kv_len == 12 + expected_tasks = [ + (4, 8, [4, 5, 6, 7], "kv"), + (8, 10, [8, 9], "kv"), + ] + if is_hybrid_att_model: + expected_tasks.append((10, 10, [], "att_state")) + assert tasks == expected_tasks + assert table[0, :12].tolist() == list(range(12)) + + +def test_pd_decode_rejects_unaligned_transfer_page_size(): + backend = decode_impl.PDDecodeNode.__new__(decode_impl.PDDecodeNode) + backend.args = SimpleNamespace(pd_kv_page_size=3, page_size=4) + 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._decode_node_gen_trans_tasks(req) + + +def test_pd_prefill_rejects_unaligned_transfer_page_size(): + backend = prefill_impl.PDChunkedPrefillForPrefillNode.__new__(prefill_impl.PDChunkedPrefillForPrefillNode) + backend.args = SimpleNamespace(pd_kv_page_size=3, page_size=4) + 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_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)]) +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 = [] + next_index = [4] + req = _bind_alloc_need( + SimpleNamespace( + req_idx=0, + cur_kv_len=4, + hold_kv_len=4, + get_cur_total_len=lambda: 20, + ), + page_size=4, + ) + source_table = torch.arange(32, dtype=torch.int32).reshape(2, 16) + backend = SimpleNamespace( + model=SimpleNamespace( + mem_manager=SimpleNamespace(), + att_state_manager=SimpleNamespace(full_att_allocator=SimpleNamespace(page_size=4)), + ), + dp_world_size=tp_size, + rank_in_dp=tp_rank, + node_nccl_group=None, + dp_req_to_token_indices={0: source_table, 1: source_table}, + is_hybrid_att_model=False, + ) + + module = dp_shared_kv_trans.DPKVSharedMoudle.__new__(dp_shared_kv_trans.DPKVSharedMoudle) + module.backend = backend + module.dp_rank_in_node = 0 + module.shared_req_infos = SimpleNamespace(arr=np.array([[[4, 0], [10, 0]]], dtype=np.int64)) + req_manager = _make_req_manager(table, alloc_sizes, next_index) + req.kv_cache_alloc_need = lambda target: req_manager.get_need_alloc_token_num(req, target) + req.alloc_kv_mem = lambda alloc_token_num: req_manager.alloc_token_indexes( + req, req.cur_kv_len + alloc_token_num, None, req.args.page_size + ) + context = SimpleNamespace( + req_manager=req_manager, + radix_cache=None, + get_can_alloc_token_num=lambda: 100, + ) + monkeypatch.setattr(dp_shared_kv_trans, "g_infer_context", context) + monkeypatch.setattr(dp_shared_kv_trans.dist, "barrier", lambda group: None) + + tasks = module.build_shared_kv_trans_tasks([req], req_dp_ranks=[0]) + + assert alloc_sizes == [8] + assert req.cur_kv_len == 4 + assert req.hold_kv_len == 12 + assert len(tasks) == 1 + assert tasks[0].src_dp_rank == 1 + assert tasks[0].dst_token_indices.tolist() == list(range(4, 10)) + assert len(tasks[0].src_token_indices) == 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( + get_dp_shared_operator=lambda rank=rank: SimpleNamespace( + register_page_pools_pin_memory=lambda: 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_state_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_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_state_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) + 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 + + +def _make_req_manager(table, alloc_sizes, next_index): + def need(req, end, page_size=None): + page_size = req.args.page_size if page_size is None else page_size + return max(0, (end + page_size - 1) // page_size * page_size - req.hold_kv_len) + + def allocate(req, end, allocator, page_size, radix_cache=None): + _reserve_into_table(table, req, alloc_sizes, next_index, need(req, end)) + return table[req.req_idx, req.cur_kv_len : end].clone() + + return SimpleNamespace(req_to_token_indexs=table, alloc_token_indexes=allocate, get_need_alloc_token_num=need) diff --git a/unit_tests/server/router/model_infer/mode_backend/test_pd_dynamic_split.py b/unit_tests/server/router/model_infer/mode_backend/test_pd_dynamic_split.py new file mode 100644 index 0000000000..38e60dd255 --- /dev/null +++ b/unit_tests/server/router/model_infer/mode_backend/test_pd_dynamic_split.py @@ -0,0 +1,244 @@ +from types import MethodType, SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +from lightllm.common.req_manager import ReqManager +from lightllm.server.core.objs import FinishStatus +from lightllm.server.core.objs import req as req_module +from lightllm.server.router.model_infer.mode_backend import base_backend +from lightllm.server.router.model_infer.mode_backend.pd.decode_node_impl import ( + decode_impl as pd_decode_impl, +) +from lightllm.server.router.req_queue import _get_req_queue_class +from lightllm.server.router.req_queue.chunked_prefill.impl_for_pd_decode import PDDecodeQueue +from lightllm.server.router.req_queue.chunked_prefill.impl_for_pd_prefill import PDPrefillQueue + + +def _make_infer_req(cur_output_len: int, shm_output_len: int): + return SimpleNamespace( + req_id=1, + filter_mark=False, + wait_pause=False, + paused=False, + infer_aborted=False, + finish_status=FinishStatus(), + cpu_cache_offload_task_status=SimpleNamespace(is_not_started=MagicMock(return_value=True)), + cur_kv_len=10, + hold_kv_len=8, + cur_output_len=cur_output_len, + shm_req=SimpleNamespace(shm_cur_output_len=shm_output_len), + sampling_param=SimpleNamespace( + shm_param=SimpleNamespace(max_new_tokens=65535), + ), + get_cur_total_len=MagicMock(return_value=11), + args=SimpleNamespace(page_size=4), + decode_need_token_num=MagicMock(return_value=(1, 4)), + kv_cache_alloc_need=lambda target: max(0, (target + 3) // 4 * 4 - 8), + alloc_kv_mem=lambda alloc_token_num: None, + ) + + +def _classify_without_token_capacity(monkeypatch, req, support_overlap=True): + reqs = req if isinstance(req, list) else [req] + backend = pd_decode_impl.PDDecodeNode.__new__(pd_decode_impl.PDDecodeNode) + backend.args = SimpleNamespace( + enable_cpu_cache=False, + enable_prefill_decode_mixed=False, + run_mode="decode", + page_size=4, + ) + backend.support_overlap = support_overlap + backend.is_master_in_dp = True + logger = MagicMock() + backend.logger = logger + backend._timer_merge_radix_tree = MagicMock() + backend._filter_not_ready_reqs = MagicMock(return_value=reqs) + backend.prefill_queue_strategy = SimpleNamespace(reorder=MagicMock(side_effect=lambda reqs, **kwargs: reqs)) + + backend.batch_max_tokens = 16 + infer_context = base_backend.g_infer_context + manager = ReqManager.__new__(ReqManager) + manager.page_size = 4 + monkeypatch.setattr(infer_context, "req_manager", manager) + monkeypatch.setattr( + infer_context, + "att_state_manager", + SimpleNamespace(full_att_allocator=SimpleNamespace(page_size=4, can_use_mem_size=0)), + ) + monkeypatch.setattr(infer_context, "get_can_alloc_token_num", MagicMock(return_value=0)) + monkeypatch.setattr( + infer_context, + "cache_placement_controller", + SimpleNamespace(set_req_cache_way=MagicMock()), + ) + filter_reqs = MagicMock() + monkeypatch.setattr(infer_context, "filter_reqs", filter_reqs) + monkeypatch.setattr(infer_context, "pause_reqs", MagicMock()) + + backend._get_classed_reqs(req_ids=[req.req_id for req in reqs]) + return filter_reqs, logger + + +def test_pd_decode_capacity_shortage_is_delayed_for_overlap(monkeypatch): + req = _make_infer_req(cur_output_len=5, shm_output_len=4) + + filter_reqs, logger = _classify_without_token_capacity(monkeypatch, req, support_overlap=True) + + assert req.filter_mark + assert req.finished_by_pd_decode_capacity + assert not req.finish_status.is_finished() + assert req.sampling_param.shm_param.max_new_tokens == 65535 + assert not req.wait_pause + filter_reqs.assert_called_once_with(finished_reqs=[]) + assert logger.info.call_args.args[0] == ( + "force early finish for PD decode req_id=1 because token capacity is insufficient" + ) + + filter_reqs, _ = _classify_without_token_capacity(monkeypatch, req, support_overlap=True) + filter_reqs.assert_called_once_with(finished_reqs=[req]) + + +def test_pd_decode_capacity_shortage_is_filtered_without_overlap(monkeypatch): + req = _make_infer_req(cur_output_len=5, shm_output_len=4) + + filter_reqs, logger = _classify_without_token_capacity(monkeypatch, req, support_overlap=False) + + assert not req.filter_mark + assert req.finished_by_pd_decode_capacity + assert not req.finish_status.is_finished() + assert req.sampling_param.shm_param.max_new_tokens == 65535 + assert not req.wait_pause + filter_reqs.assert_called_once_with(finished_reqs=[req]) + assert logger.info.call_args.args[0] == ( + "force early finish for PD decode req_id=1 because token capacity is insufficient" + ) + + +def test_pd_decode_capacity_shortage_only_handles_two_requests_per_iteration(monkeypatch): + reqs = [_make_infer_req(cur_output_len=5, shm_output_len=4) for _ in range(3)] + for req_id, req in enumerate(reqs, start=1): + req.req_id = req_id + + filter_reqs, _ = _classify_without_token_capacity(monkeypatch, reqs, support_overlap=True) + + assert [req.filter_mark for req in reqs] == [True, True, False] + filter_reqs.assert_called_once_with(finished_reqs=[]) + + +def test_pd_decode_capacity_finish_status_is_written_to_shm(): + shm_req = SimpleNamespace(mark_simulated_finished=MagicMock()) + req = SimpleNamespace( + finish_status=FinishStatus(), + finished_by_pd_decode_capacity=True, + shm_req=shm_req, + cur_output_len=3, + ) + + pd_decode_impl.InferReq.mark_shm_aborted_finished(req) + + shm_req.mark_simulated_finished.assert_called_once_with( + FinishStatus.FINISHED_PD_DECODE_CAPACITY, + output_len=3, + ) + + +def test_request_finish_status_takes_priority_over_pd_decode_capacity_marker(): + shm_req = SimpleNamespace(mark_simulated_finished=MagicMock()) + req = SimpleNamespace( + finish_status=FinishStatus(FinishStatus.FINISHED_STOP), + finished_by_pd_decode_capacity=True, + shm_req=shm_req, + cur_output_len=3, + ) + + pd_decode_impl.InferReq.mark_shm_aborted_finished(req) + + shm_req.mark_simulated_finished.assert_not_called() + + +def test_pd_decode_capacity_limit_never_extends_original_length(monkeypatch): + req = _make_infer_req(cur_output_len=5, shm_output_len=1) + req.sampling_param.shm_param.max_new_tokens = 3 + + _classify_without_token_capacity(monkeypatch, req) + assert req.sampling_param.shm_param.max_new_tokens == 3 + + +@pytest.mark.parametrize( + "overrides", + [{}, {"diverse_mode": True}, {"output_constraint_mode": "xgrammar"}], +) +def test_pd_nodes_use_pd_queue(overrides): + base_args = { + "diverse_mode": False, + "token_healing_mode": False, + "output_constraint_mode": "none", + "disable_chunked_prefill": False, + } + + base_args.update(overrides) + prefill_args = SimpleNamespace(**base_args, run_mode="prefill") + decode_args = SimpleNamespace(**base_args, run_mode="decode") + + assert _get_req_queue_class(prefill_args, router=None, dp_size_in_node=1) is PDPrefillQueue + assert _get_req_queue_class(decode_args, router=None, dp_size_in_node=1) is PDDecodeQueue + + +@pytest.mark.parametrize("page_size", [1, 16]) +def test_pd_decode_queue_uses_ema_for_prefill_stage_output_length(page_size, monkeypatch): + monkeypatch.setattr(req_module, "get_env_start_args", lambda: SimpleNamespace(page_size=page_size, mtp_step=0)) + queue = PDDecodeQueue.__new__(PDDecodeQueue) + queue.args = SimpleNamespace(run_mode="decode", page_size=page_size) + queue.dp_index = 0 + queue.max_total_tokens = 4096 + queue.running_max_req_size = 8 + queue.router = SimpleNamespace( + router_statics=SimpleNamespace(ema_req_out_len=128), + shared_token_load=SimpleNamespace( + set_estimated_peak_token_count=MagicMock(), + set_dynamic_max_load=MagicMock(), + ), + ) + queue.is_busy = MagicMock(return_value=False) + + req = SimpleNamespace( + input_len=10, + shm_cur_output_len=0, + shm_cur_kv_len=0, + sample_params=SimpleNamespace(suggested_dp_index=0, max_new_tokens=1024, ignore_eos=False), + is_infer_decode=MagicMock(return_value=False), + ) + req.get_tuple_tokens = MagicMock(wraps=MethodType(req_module.ChunkedPrefillReq.get_tuple_tokens, req)) + batch = SimpleNamespace(reqs=[req]) + + assert queue._caclu_batch_estimated_peak_token_num(batch) == 156 + page_size + req.get_tuple_tokens.assert_called_once_with(False, 128) + req.get_tuple_tokens.reset_mock() + assert queue._can_add_new_req(req, estimated_peak_token_num=0, batch_req_num=0) == (True, 156 + page_size, 1) + req.get_tuple_tokens.assert_called_once_with(False, 128) + + # 接近上下文上限的请求可能只剩很少输出额度，预计输出长度仍受 max_new_tokens 约束， + # 并在此基础上计入分页、MTP 和异步退出余量，避免因 EMA 偏大而永久滞留在 Decode 等待队列。 + req.sample_params.max_new_tokens = 20 + assert queue._caclu_batch_estimated_peak_token_num(batch) == 48 + page_size + assert queue._can_add_new_req(req, estimated_peak_token_num=0, batch_req_num=0) == (True, 48 + page_size, 1) + + req.sample_params.max_new_tokens = 1024 + req.get_tuple_tokens.reset_mock() + prefill_queue = PDPrefillQueue.__new__(PDPrefillQueue) + prefill_queue.args = SimpleNamespace(run_mode="prefill", page_size=page_size) + prefill_queue.dp_index = queue.dp_index + prefill_queue.max_total_tokens = queue.max_total_tokens + prefill_queue.running_max_req_size = queue.running_max_req_size + prefill_queue.router = queue.router + prefill_queue.is_busy = MagicMock(return_value=False) + assert prefill_queue._caclu_batch_estimated_peak_token_num(batch) == 156 + page_size + req.get_tuple_tokens.assert_called_once_with(False, 128) + req.get_tuple_tokens.reset_mock() + assert prefill_queue._can_add_new_req(req, estimated_peak_token_num=0, batch_req_num=0) == ( + True, + 156 + page_size, + 1, + ) + req.get_tuple_tokens.assert_called_once_with(False, 128) diff --git a/unit_tests/server/router/model_infer/mode_backend/test_pd_manager_startup.py b/unit_tests/server/router/model_infer/mode_backend/test_pd_manager_startup.py new file mode 100644 index 0000000000..9dbf0c8025 --- /dev/null +++ b/unit_tests/server/router/model_infer/mode_backend/test_pd_manager_startup.py @@ -0,0 +1,174 @@ +import ctypes +import importlib +import multiprocessing as mp +import os +import signal +import sys +import time +import traceback +from types import ModuleType, SimpleNamespace + +import psutil +import pytest + +from lightllm.utils import process_check + + +pytestmark = pytest.mark.skipif(sys.platform != "linux", reason="requires Linux subreaper and fork") + + +def _kill_and_reap_children(): + for child in psutil.Process().children(recursive=True): + try: + os.kill(child.pid, signal.SIGKILL) + except ProcessLookupError: + pass + while True: + try: + os.waitpid(-1, 0) + except ChildProcessError: + break + + +def _run_startup_failure(mode, late_monitor, result_connection): + """Own orphaned workers in an isolated subreaper so the test can always reap them.""" + result = None + try: + libc = ctypes.CDLL(None, use_errno=True) + if libc.prctl(36, 1, 0, 0, 0) != 0: # PR_SET_CHILD_SUBREAPER + raise OSError(ctypes.get_errno(), "could not become a child subreaper") + + ctx = mp.get_context("fork") + events, event_writer = ctx.Pipe(duplex=False) + start_monitor = ctx.Event() + if not late_monitor: + start_monitor.set() + + # Speed up the real monitor without changing the shared time module's sleep. + process_check.time = SimpleNamespace(sleep=lambda _: time.sleep(0.01)) + + def worker_main(started): + if late_monitor: + started.set() + assert start_monitor.wait(timeout=10) + event_writer.send(("monitor_parent", os.getpid(), os.getppid())) + process_check.start_parent_check_thread() + started.set() + while True: + time.sleep(100) + + def start_worker(*args, **kwargs): + started = ctx.Event() + worker = ctx.Process(target=worker_main, args=(started,)) + worker.start() + event_writer.send(("worker", worker.pid)) + assert not worker.daemon + assert started.wait(timeout=5) + + module_name = "lightllm.server.router.model_infer.mode_backend.pd." f"{mode}_node_impl.{mode}_kv_move_manager" + manager_module = importlib.import_module(module_name) + manager_module.get_unique_server_name = lambda: "startup-regression" + original_logger = manager_module.logger + + def log_failure(message): + original_logger.exception(message) + event_writer.send(("init_error",)) + + manager_module.logger = SimpleNamespace(exception=log_failure) + + def fail_initialization(*args, **kwargs): + start_worker() + raise RuntimeError("injected transport initialization failure") + + manager_class = "PrefillKVMoveManager" if mode == "prefill" else "DecodeKVMoveManager" + setattr(manager_module, manager_class, fail_initialization) + + # Exercise the real entry point while avoiding CUDA and network initialization. + package_name = module_name.rsplit(".", 1)[0] + transport = ModuleType(f"{package_name}.{mode}_trans_process") + setattr(transport, f"start_{mode}_trans_process", lambda *args, **kwargs: None) + sys.modules[transport.__name__] = transport + sys.modules["lightllm.utils.rpyc_fix_utils"] = ModuleType("lightllm.utils.rpyc_fix_utils") + if mode == "decode": + up_status = ModuleType(f"{package_name}.up_status") + up_status.start_up_kv_status_process = start_worker + sys.modules[up_status.__name__] = up_status + + ready = ctx.Event() + manager = ctx.Process(target=manager_module._init_env, args=(SimpleNamespace(), None, ready)) + manager.start() + manager.join(timeout=5) + assert not manager.is_alive(), "manager hung while joining its non-daemon children" + assert manager.exitcode == 1 + assert not ready.is_set() + + # In the late case workers are already reparented when they start the monitor. + start_monitor.set() + worker_count = 2 if mode == "decode" else 1 + workers = set() + monitor_parents = {} + logged_failure = False + deadline = time.monotonic() + 5 + while len(workers) < worker_count or len(monitor_parents) < worker_count or not logged_failure: + assert time.monotonic() < deadline, "missing worker or initialization failure report" + if not events.poll(0.05): + continue + event = events.recv() + if event[0] == "worker": + workers.add(event[1]) + elif event[0] == "monitor_parent": + monitor_parents[event[1]] = event[2] + elif event[0] == "init_error": + logged_failure = True + + expected_parent = os.getpid() if late_monitor else manager.pid + assert set(monitor_parents) == workers + assert set(monitor_parents.values()) == {expected_parent} + + for pid in workers: + while True: + reaped_pid, status = os.waitpid(pid, os.WNOHANG) + if reaped_pid: + assert os.WIFSIGNALED(status) + assert os.WTERMSIG(status) == signal.SIGKILL + break + assert time.monotonic() < deadline, "worker did not detect its dead parent" + time.sleep(0.01) + result = "ok" + except BaseException: + result = traceback.format_exc() + finally: + _kill_and_reap_children() + result_connection.send(result) + + +@pytest.mark.parametrize("mode", ["prefill", "decode"]) +@pytest.mark.parametrize("late_monitor", [False, True], ids=["monitor-running", "monitor-after-parent-exit"]) +def test_failed_manager_exits_and_workers_detect_parent_death(mode, late_monitor): + # Import before forking so costly framework imports happen only once per test run. + importlib.import_module( + f"lightllm.server.router.model_infer.mode_backend.pd.{mode}_node_impl.{mode}_kv_move_manager" + ) + ctx = mp.get_context("fork") + result_reader, result_writer = ctx.Pipe(duplex=False) + controller = ctx.Process(target=_run_startup_failure, args=(mode, late_monitor, result_writer)) + controller.start() + try: + controller.join(timeout=15) + assert not controller.is_alive(), "startup regression controller timed out" + assert controller.exitcode == 0 + assert result_reader.poll(timeout=1), "startup regression controller returned no result" + assert result_reader.recv() == "ok" + finally: + if controller.is_alive(): + # The controller normally owns cleanup; retain a bounded test-runner fallback. + for child in psutil.Process(controller.pid).children(recursive=True): + try: + os.kill(child.pid, signal.SIGKILL) + except ProcessLookupError: + pass + controller.kill() + controller.join(timeout=5) + controller.close() + result_reader.close() + result_writer.close() 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 new file mode 100644 index 0000000000..8f99c6add3 --- /dev/null +++ b/unit_tests/server/router/model_infer/mode_backend/test_pd_notify.py @@ -0,0 +1,223 @@ +import pickle +import subprocess +import sys +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +from lightllm.server.pd_io_struct import PDChunckedTransTask +from lightllm.server.router.model_infer.mode_backend.pd.decode_node_impl.decode_trans_process import ( + _SharedTimeOutMark, +) +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, +) + + +def make_task(page_kind): + task = PDChunckedTransTask( + request_id=123, + start_kv_index=4, + end_kv_index=5 if page_kind == "kv" else 4, + time_out_secs=10, + pd_master_node_id=2, + prefill_dp_index=0, + decode_dp_index=1, + src_device_id=0, + dst_device_id=1, + mem_indexes=[7] if page_kind == "kv" else [], + prefill_agent_name="prefill", + prefill_agent_metadata=b"prefill metadata", + prefill_num_pages=8, + prefill_page_reg_desc=b"prefill pages", + decode_agent_name="decode", + decode_agent_metadata=b"decode metadata", + decode_num_pages=9, + decode_page_reg_desc=b"decode pages", + first_gen_token_id=42, + first_gen_token_logprob=-0.5, + page_kind=page_kind, + req_idx=17, + src_page_index=3, + dst_page_index=6, + xfer_handle=object(), + ) + return task + + +@pytest.mark.parametrize("page_kind", ["kv", "att_state"]) +@pytest.mark.parametrize("with_mark", [False, True]) +def test_copy_for_notify_preserves_local_task(page_kind, with_mark): + task = make_task(page_kind) + if with_mark: + task.shared_time_out_mark = _SharedTimeOutMark() + time_mark = task.shared_time_out_mark.time_mark + original = vars(task).copy() + mem_indexes = list(task.mem_indexes) + + notify = task.copy_for_notify() + + assert notify is not task + assert not hasattr(notify, "shared_time_out_mark") + assert notify.mem_indexes is None + assert notify.xfer_handle is None + assert vars(notify) == { + **{key: value for key, value in original.items() if key != "shared_time_out_mark"}, + "mem_indexes": None, + "xfer_handle": None, + } + assert vars(task) == original + assert all(vars(task)[key] is value for key, value in original.items()) + assert task.mem_indexes == mem_indexes + if with_mark: + assert task.shared_time_out_mark.time_mark == time_mark + + +_UNPICKLE_WITHOUT_DECODE = """ +import importlib.abc +import pickle +import sys +from lightllm.server.pd_io_struct import PDChunckedTransTask + +class BlockDecode(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path=None, target=None): + if 'decode_node_impl' in fullname: + raise ImportError('decode modules are unavailable to the receiver') + +assert not any('decode_node_impl' in name for name in sys.modules) +sys.meta_path.insert(0, BlockDecode()) +task = pickle.loads(sys.stdin.buffer.read()) +assert isinstance(task, PDChunckedTransTask) +assert not hasattr(task, 'shared_time_out_mark') +assert task.mem_indexes is None +assert task.xfer_handle is None +assert not any('decode_node_impl' in name for name in sys.modules) +""" + + +@pytest.mark.parametrize("backend", [NcclKVTransporter, NixlKVTransporter]) +@pytest.mark.parametrize("page_kind", ["kv", "att_state"]) +@pytest.mark.parametrize( + "method,stage,destination", + [ + ("send_write_done_task_to_decode_node", "done", "decode"), + ("send_write_request_task_to_decode_node", "request", "decode"), + ("send_write_ready_task_to_prefill_node", "ready", "prefill"), + ("send_error_info_to_prefill_node", "error", "prefill"), + ("send_error_info_to_decode_node", "error", "decode"), + ], +) +def test_notification_payload(backend, page_kind, method, stage, destination): + task = make_task(page_kind) + task.error_info = "transfer failed" if stage == "error" else None + task.shared_time_out_mark = _SharedTimeOutMark() + original = vars(task).copy() + time_mark = task.shared_time_out_mark.time_mark + transporter = backend.__new__(backend) + transporter.node_id, transporter.tp_idx = 10, 0 + transporter.num_pages = 12 + transporter.remote_agents = {"prefill": object(), "decode": object()} + 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="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: + transporter.page_reg_desc = object() + transporter.nixl_agent = SimpleNamespace( + get_agent_metadata=lambda: b"local metadata", + get_serialized_descs=lambda _: b"local pages", + send_notif=Mock(), + ) + + getattr(transporter, method)(task) + + if backend is NcclKVTransporter: + transporter.control_channel.send_notif.assert_called_once() + ( + peer, + host, + port, + payload, + ) = transporter.control_channel.send_notif.call_args.args + assert (peer, host, port) == (destination, "127.0.0.2", 4321) + else: + transporter.nixl_agent.send_notif.assert_called_once() + kwargs = transporter.nixl_agent.send_notif.call_args.kwargs + assert kwargs["remote_agent_name"] == destination + payload = kwargs["notif_msg"] + result = subprocess.run( + [sys.executable, "-c", _UNPICKLE_WITHOUT_DECODE], + input=payload, + capture_output=True, + timeout=30, + ) + assert result.returncode == 0, result.stderr.decode() + notify = pickle.loads(payload) + sender = "prefill" if destination == "decode" else "decode" + expected = {key: value for key, value in original.items() if key != "shared_time_out_mark"} + expected.update(mem_indexes=None, xfer_handle=None, write_stage=stage) + expected.update( + { + f"{sender}_agent_name": transporter.agent_name, + f"{sender}_agent_metadata": transporter.agent_metadata, + f"{sender}_num_pages": transporter.num_pages, + f"{sender}_page_reg_desc": transporter.local_page_mem_desc, + } + ) + if backend is NixlKVTransporter and stage == "done": + expected.update(decode_agent_metadata=None, decode_page_reg_desc=None) + assert vars(notify) == expected + assert vars(task) == original + assert all(vars(task)[key] is value for key, value in original.items()) + assert task.shared_time_out_mark.time_mark == time_mark + + +def test_receiver_rejects_payload_with_decode_local_mark(): + task = make_task("kv") + task.shared_time_out_mark = _SharedTimeOutMark() + result = subprocess.run( + [sys.executable, "-c", _UNPICKLE_WITHOUT_DECODE], + input=pickle.dumps(task), + capture_output=True, + timeout=30, + ) + 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_bytes = 128 + transporter.nixl_agent = SimpleNamespace(get_xfer_descs=Mock(), prep_xfer_dlist=Mock()) + 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_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)) + 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_bytes": 256}), + ) + with pytest.raises(AssertionError, match="same page byte capacity"): + transporter.connect_add_remote_agent(peer) + assert not transporter.remote_agents + peer.page_reg_desc = transporter.local_page_mem_desc + 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_pd_page_io.py b/unit_tests/server/router/model_infer/mode_backend/test_pd_page_io.py new file mode 100644 index 0000000000..d2adf388f4 --- /dev/null +++ b/unit_tests/server/router/model_infer/mode_backend/test_pd_page_io.py @@ -0,0 +1,53 @@ +from types import SimpleNamespace + +import torch +import pytest + + +@pytest.mark.parametrize("role", ["prefill", "decode"]) +def test_hybrid_pd_uses_the_shared_att_state_task_without_model_flags(monkeypatch, role): + from lightllm.server.router.model_infer.infer_batch import g_infer_context + from lightllm.server.router.model_infer.mode_backend.pd.prefill_node_impl.prefill_impl import ( + PDChunkedPrefillForPrefillNode, + ) + from lightllm.server.router.model_infer.mode_backend.pd.decode_node_impl import decode_impl + from lightllm.server.router.model_infer.mode_backend.pd.decode_node_impl.decode_impl import PDDecodeNode + + monkeypatch.setattr(g_infer_context, "is_hybrid_att_model", True) + cls = PDChunkedPrefillForPrefillNode if role == "prefill" else PDDecodeNode + backend = cls.__new__(cls) + backend.args = SimpleNamespace(pd_kv_page_size=256, page_size=256) + backend.is_master_in_dp = role == "decode" + groups, tasks = [], [] + backend.info_queue = SimpleNamespace(put=groups.append) + req = SimpleNamespace( + shm_req=SimpleNamespace(input_len=512), + cur_kv_len=512 if role == "prefill" else 0, + hold_kv_len=0, + pd_trans_kv_start_index=0, + ) + + def create(**kwargs): + task = SimpleNamespace( + page_kind=kwargs.get("page_kind", "kv"), first_gen_token_id=None, first_gen_token_logprob=None + ) + tasks.append(task) + if "group" in kwargs: + kwargs["group"].task_list.append(task) + return task + + backend._create_pd_trans_task = create + if role == "prefill": + backend._prefill_chuncked_handle_func(req, next_token_id=17, next_token_prob=-0.2, output_len=1) + assert tasks[-1].first_gen_token_id == 17 + else: + backend.radix_cache = None + backend.model = SimpleNamespace( + att_state_manager=SimpleNamespace(full_att_allocator=SimpleNamespace(page_size=backend.args.page_size)), + ) + req.args = SimpleNamespace(page_size=backend.args.page_size) + req.kv_cache_alloc_need = lambda target: target + req.alloc_kv_mem = lambda alloc_token_num: torch.arange(alloc_token_num) + backend._decode_node_gen_trans_tasks(req) + assert len(groups) == 1 + assert [task.page_kind for task in tasks] == ["kv", "kv", "att_state"] diff --git a/unit_tests/server/router/model_infer/mode_backend/test_pd_prefill_abort.py b/unit_tests/server/router/model_infer/mode_backend/test_pd_prefill_abort.py new file mode 100644 index 0000000000..5fff4ddc2e --- /dev/null +++ b/unit_tests/server/router/model_infer/mode_backend/test_pd_prefill_abort.py @@ -0,0 +1,199 @@ +import queue +import threading +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +from lightllm.server.pd_io_struct import PDAbortReq +from lightllm.server.router.model_infer.mode_backend.pd.decode_node_impl import decode_impl +from lightllm.server.router.model_infer.mode_backend.pd.decode_node_impl.decode_impl import PDDecodeNode +from lightllm.server.router.model_infer.mode_backend.pd.prefill_node_impl import prefill_impl, prefill_trans_process +from lightllm.server.router.model_infer.mode_backend.pd.prefill_node_impl.prefill_impl import ( + PDChunkedPrefillForPrefillNode, +) +from lightllm.server.router.model_infer.mode_backend.pd.prefill_node_impl.prefill_kv_move_manager import ( + PrefillKVMoveManager, +) +from lightllm.server.router.model_infer.mode_backend.pd.prefill_node_impl.prefill_trans_process import ( + _PrefillTransModule, +) + + +class _StopLoop(Exception): + pass + + +class _SequenceQueue: + def __init__(self, *values): + self.values = list(values) + + def get(self): + if self.values: + return self.values.pop(0) + raise _StopLoop() + + +def test_prefill_infer_forwards_abort_to_kv_move_manager(monkeypatch): + request_id = 123 + req = SimpleNamespace( + req_id=request_id, + infer_aborted=True, + pd_task_num=2, + pd_task_failed_num=0, + pd_task_success_num=0, + pd_trans_device_id=1, + cur_kv_len=4, + shm_req=SimpleNamespace(input_len=8), + ) + monkeypatch.setattr(prefill_impl.g_infer_context, "requests_mapping", {request_id: req}) + + backend = PDChunkedPrefillForPrefillNode.__new__(PDChunkedPrefillForPrefillNode) + backend.is_master_in_dp = True + backend.args = SimpleNamespace(enable_cpu_cache=False) + backend.info_queue = queue.Queue() + + ready_reqs = backend._filter_not_ready_reqs([request_id]) + + assert ready_reqs == [] + abort_req = backend.info_queue.get_nowait() + assert abort_req == PDAbortReq(request_id=request_id, device_id=1) + assert req.pd_abort_req_send_count == 1 + + +def test_prefill_infer_sends_abort_at_most_six_times(monkeypatch): + request_id = 123 + req = SimpleNamespace( + req_id=request_id, + infer_aborted=True, + pd_task_num=2, + pd_task_failed_num=0, + pd_task_success_num=0, + pd_trans_device_id=1, + cur_kv_len=4, + shm_req=SimpleNamespace(input_len=8), + ) + monkeypatch.setattr(prefill_impl.g_infer_context, "requests_mapping", {request_id: req}) + + backend = PDChunkedPrefillForPrefillNode.__new__(PDChunkedPrefillForPrefillNode) + backend.is_master_in_dp = True + backend.args = SimpleNamespace(enable_cpu_cache=False) + backend.info_queue = queue.Queue() + + for _ in range(8): + assert backend._filter_not_ready_reqs([request_id]) == [] + + abort_reqs = [backend.info_queue.get_nowait() for _ in range(6)] + assert all(abort_req == PDAbortReq(request_id=request_id, device_id=1) for abort_req in abort_reqs) + assert backend.info_queue.empty() + assert req.pd_abort_req_send_count == 6 + + +def test_decode_infer_sends_abort_at_most_six_times(monkeypatch): + request_id = 123 + req = SimpleNamespace( + req_id=request_id, + infer_aborted=True, + pd_task_num=2, + pd_task_failed_num=0, + pd_task_success_num=0, + pd_trans_device_id=1, + ) + monkeypatch.setattr(decode_impl.g_infer_context, "requests_mapping", {request_id: req}) + + backend = PDDecodeNode.__new__(PDDecodeNode) + backend.is_master_in_dp = True + backend.info_queue = queue.Queue() + + for _ in range(8): + assert backend._filter_not_ready_reqs([request_id]) == [] + + abort_reqs = [backend.info_queue.get_nowait() for _ in range(6)] + assert all(abort_req == PDAbortReq(request_id=request_id, device_id=1) for abort_req in abort_reqs) + assert backend.info_queue.empty() + assert req.pd_abort_req_send_count == 6 + + +def test_prefill_filter_skips_request_while_cpu_cache_load_is_running(monkeypatch): + from lightllm.server.router.model_infer.infer_batch import InferReq + + request_id = 123 + status = InferReq._CpuCacheTaskStatus + req = SimpleNamespace( + req_id=request_id, + infer_aborted=False, + pd_task_num=0, + pd_task_failed_num=0, + pd_task_success_num=0, + cur_kv_len=8192, + shm_req=SimpleNamespace(input_len=8192), + cpu_cache_load_task_status=status.RUNNING, + cpu_cache_offload_task_status=status.NOT_STARTED, + ) + monkeypatch.setattr(prefill_impl.g_infer_context, "requests_mapping", {request_id: req}) + + backend = PDChunkedPrefillForPrefillNode.__new__(PDChunkedPrefillForPrefillNode) + backend.is_master_in_dp = True + backend.args = SimpleNamespace(enable_cpu_cache=True) + backend.info_queue = queue.Queue() + + assert backend._filter_not_ready_reqs([request_id]) == [] + + req.cpu_cache_load_task_status = status.FINISHED + assert backend._filter_not_ready_reqs([request_id]) == [req] + + +def test_prefill_kv_move_manager_routes_abort_to_target_device(): + abort_req = PDAbortReq(request_id=123, device_id=1) + target_queue = queue.Queue() + manager = PrefillKVMoveManager.__new__(PrefillKVMoveManager) + manager.info_queue = _SequenceQueue(abort_req) + manager.kv_trans_processes = [ + SimpleNamespace(task_in_queue=queue.Queue()), + SimpleNamespace(task_in_queue=target_queue), + ] + + with pytest.raises(_StopLoop): + manager.task_dispatcher_loop() + + assert target_queue.get_nowait() is abort_req + assert manager.kv_trans_processes[0].task_in_queue.empty() + + +def test_prefill_trans_process_handles_abort_without_allocating_page(monkeypatch): + abort_req = PDAbortReq(request_id=123, device_id=0) + module = _PrefillTransModule.__new__(_PrefillTransModule) + module.device_id = 0 + module.task_in_queue = _SequenceQueue(abort_req) + module.page_index_queue = MagicMock() + module._abort = MagicMock() + monkeypatch.setattr(prefill_trans_process.torch.cuda, "set_device", lambda _device_id: None) + + with pytest.raises(_StopLoop): + module.recv_task_loop() + + module._abort.assert_called_once_with(request_id=123) + module.page_index_queue.get.assert_not_called() + + +def test_prefill_abort_only_fails_tasks_not_submitted_to_transporter(): + pending_task = SimpleNamespace(request_id=123, xfer_handle=None, error_info=None) + active_task = SimpleNamespace(request_id=123, xfer_handle=object(), error_info=None) + unrelated_task = SimpleNamespace(request_id=456, xfer_handle=None, error_info=None) + module = _PrefillTransModule.__new__(_PrefillTransModule) + module.waiting_dict_lock = threading.Lock() + module.waiting_dict = { + "pending": pending_task, + "active": active_task, + "unrelated": unrelated_task, + } + module.failed_queue = queue.Queue() + + module._abort(request_id=123) + + assert module.failed_queue.get_nowait() is pending_task + assert pending_task.error_info == "aborted req" + assert module.waiting_dict == { + "active": active_task, + "unrelated": unrelated_task, + } 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..726df8dbee --- /dev/null +++ b/unit_tests/server/router/model_infer/mode_backend/test_pd_token_indices.py @@ -0,0 +1,126 @@ +"""PD 索引由普通 CPU Tensor 上传，使用实际拷贝流并跳过运行态任务。""" + +from types import SimpleNamespace +from unittest.mock import Mock +from queue import Queue + +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_operator=lambda: SimpleNamespace(get_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, + xfer_handle=None, + ) + 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_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) + + +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() diff --git a/unit_tests/server/router/model_infer/mode_backend/test_pd_transfer_progress.py b/unit_tests/server/router/model_infer/mode_backend/test_pd_transfer_progress.py new file mode 100644 index 0000000000..b3ca2f9b64 --- /dev/null +++ b/unit_tests/server/router/model_infer/mode_backend/test_pd_transfer_progress.py @@ -0,0 +1,176 @@ +import pickle +import queue +import threading +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +from lightllm.server.pd_io_struct import PDChunckedTransTask, PDChunckedTransTaskGroup +from lightllm.server.router.model_infer.mode_backend.pd.decode_node_impl import decode_trans_process + + +class _StopLoop(Exception): + pass + + +class _SequenceQueue: + def __init__(self, *values): + self.values = list(values) + + def get(self): + if self.values: + return self.values.pop(0) + raise _StopLoop() + + +def make_task(request_id, page, empty=False): + return PDChunckedTransTask( + request_id=request_id, + start_kv_index=page, + end_kv_index=page if empty else page + 1, + time_out_secs=10, + pd_master_node_id=1, + prefill_dp_index=0, + decode_dp_index=0, + src_device_id=0, + dst_device_id=0, + mem_indexes=[] if empty else [page], + prefill_agent_name="prefill", + prefill_agent_metadata=b"prefill metadata", + prefill_num_pages=1, + prefill_page_reg_desc=b"prefill page", + decode_agent_name=None, + decode_agent_metadata=None, + decode_num_pages=None, + decode_page_reg_desc=None, + first_gen_token_id=42, + first_gen_token_logprob=-0.5, + ) + + +def make_module(*tasks): + module = decode_trans_process._DecodeTransModule.__new__(decode_trans_process._DecodeTransModule) + module.device_id = 0 + module.args = SimpleNamespace(pd_node_id=2, pd_detail_log=False) + module.waiting_dict_lock = threading.Lock() + module.waiting_dict = {task.get_key(): task for task in tasks} + module.transporter = SimpleNamespace( + agent_name="decode", + agent_metadata=b"decode metadata", + num_pages=1, + local_page_mem_desc=b"decode page", + send_error_info_to_prefill_node=Mock(), + ) + for name in ( + "request_page_task_queue", + "ready_page_task_queue", + "success_queue", + "failed_queue", + "up_status_in_queue", + ): + setattr(module, name, queue.Queue()) + return module + + +def test_dispatch_shares_timeout_mark_with_all_pages_in_each_group(monkeypatch): + monkeypatch.setattr(decode_trans_process.time, "time", lambda: 100.0) + first, second, empty, unrelated = make_task(1, 0), make_task(1, 1), make_task(1, 2, empty=True), make_task(2, 0) + module = make_module() + module.recv_task_group_queue = _SequenceQueue( + PDChunckedTransTaskGroup([first, second, empty]), PDChunckedTransTaskGroup([unrelated]) + ) + + with pytest.raises(_StopLoop): + module.dispatch_task_loop() + + assert first.shared_time_out_mark is second.shared_time_out_mark is empty.shared_time_out_mark + assert unrelated.shared_time_out_mark is not first.shared_time_out_mark + assert first.shared_time_out_mark.time_mark == 100.0 + assert module.waiting_dict == {task.get_key(): task for task in (first, second, unrelated)} + assert module.success_queue.get_nowait() == (None, None, empty) + assert empty.start_trans_time == 100.0 + restored = pickle.loads(pickle.dumps(PDChunckedTransTaskGroup([first, second, empty]))).task_list + assert restored[0].shared_time_out_mark is restored[1].shared_time_out_mark is restored[2].shared_time_out_mark + + +def test_waiting_pages_timeout_from_shared_progress_and_initial_deadline(monkeypatch): + now = [100.0] + monkeypatch.setattr(decode_trans_process.time, "time", lambda: now[0]) + first, second, unrelated = make_task(1, 0), make_task(1, 1), make_task(2, 0) + first.shared_time_out_mark = second.shared_time_out_mark = decode_trans_process._SharedTimeOutMark() + unrelated.shared_time_out_mark = decode_trans_process._SharedTimeOutMark() + for task in (first, second, unrelated): + task.create_time = 0.0 + module = make_module(first, second, unrelated) + + now[0] = 109.0 + module._check_tasks_time_out() + assert len(module.waiting_dict) == 3 + first.shared_time_out_mark.update() + now[0] = 111.0 + module._check_tasks_time_out() + assert module.failed_queue.get_nowait() is unrelated + assert len(module.waiting_dict) == 2 + now[0] = 120.0 + module._check_tasks_time_out() + assert [module.failed_queue.get_nowait(), module.failed_queue.get_nowait()] == [first, second] + assert not module.waiting_dict + + +@pytest.mark.parametrize("stage", ["request", "done"]) +def test_matched_notifications_refresh_local_siblings_only(monkeypatch, stage): + now = [100.0] + monkeypatch.setattr(decode_trans_process.time, "time", lambda: now[0]) + monkeypatch.setattr(decode_trans_process.torch.cuda, "set_device", lambda _device: None) + first, sibling, unrelated = make_task(1, 0), make_task(1, 1), make_task(2, 0) + first.shared_time_out_mark = sibling.shared_time_out_mark = decode_trans_process._SharedTimeOutMark() + unrelated.shared_time_out_mark = decode_trans_process._SharedTimeOutMark() + module = make_module(first, sibling, unrelated) + remote = make_task(1, 0) + remote.write_stage = stage + remote.shared_time_out_mark = decode_trans_process._SharedTimeOutMark() + remote.shared_time_out_mark.time_mark = -100.0 + module.transporter.get_new_notifs = lambda: {"prefill": [pickle.dumps(remote)]} + + now[0] = 109.0 + with monkeypatch.context() as patch: + patch.setattr(module, "_check_tasks_time_out", Mock(side_effect=_StopLoop)) + with pytest.raises(_StopLoop): + module.accept_peer_task_loop() + + assert sibling.shared_time_out_mark.time_mark == 109.0 + assert unrelated.shared_time_out_mark.time_mark == 100.0 + target_queue = module.request_page_task_queue if stage == "request" else module.ready_page_task_queue + assert target_queue.get_nowait() is first + + # The repeated notify cannot keep siblings alive once its waiting task was removed. + now[0] = 110.0 + with monkeypatch.context() as patch: + patch.setattr(module, "_check_tasks_time_out", Mock(side_effect=_StopLoop)) + with pytest.raises(_StopLoop): + module.accept_peer_task_loop() + assert sibling.shared_time_out_mark.time_mark == 109.0 + module.transporter.send_error_info_to_prefill_node.assert_called_once() + + +def test_started_transfer_keeps_real_per_page_timeout(monkeypatch): + now = [100.0] + monkeypatch.setattr(decode_trans_process.time, "time", lambda: now[0]) + task = make_task(1, 0) + task.start_trans_time = 50.0 + task.shared_time_out_mark = decode_trans_process._SharedTimeOutMark() + module = make_module(task) + + now[0] = 111.0 + module._check_tasks_time_out() + assert module.waiting_dict + now[0] = 116.0 + task.shared_time_out_mark.update() + module._check_tasks_time_out() + assert module.waiting_dict + now[0] = 117.0 + module._check_tasks_time_out() + assert module.failed_queue.get_nowait() is task + assert task.error_info == "time out in accept_peer_task_loop" + assert not module.waiting_dict diff --git a/unit_tests/server/router/model_infer/mode_backend/test_pd_transfer_readiness.py b/unit_tests/server/router/model_infer/mode_backend/test_pd_transfer_readiness.py new file mode 100644 index 0000000000..1ad7a981dc --- /dev/null +++ b/unit_tests/server/router/model_infer/mode_backend/test_pd_transfer_readiness.py @@ -0,0 +1,38 @@ +import queue +from unittest.mock import Mock + +import pytest + +from lightllm.server.router.model_infer.mode_backend.pd.trans_process_obj import KVTransProcess + + +@pytest.mark.parametrize("status, expected", [("module_ready", True), ("init_failed", False)]) +def test_transfer_process_requires_module_ready(status, expected): + worker = KVTransProcess(process=Mock(), task_out_queue=queue.Queue(), device_id=0) + worker.task_out_queue.put(status) + assert worker.wait_until_ready() is expected + + +def test_transfer_process_waits_through_slow_initialization(): + output = Mock() + output.get.side_effect = [queue.Empty(), queue.Empty(), "module_ready"] + worker = KVTransProcess(process=Mock(), task_out_queue=output, device_id=0) + worker.process.is_alive.return_value = True + assert worker.wait_until_ready() + assert output.get.call_count == 3 + + +def test_transfer_process_detects_exit_before_ready(): + worker = KVTransProcess(process=Mock(), task_out_queue=Mock(), device_id=0) + worker.task_out_queue.get.side_effect = queue.Empty() + worker.process.is_alive.return_value = False + assert not worker.wait_until_ready() + assert worker.task_out_queue.get.call_count == 1 + + +def test_transfer_process_initialization_has_bounded_wait(): + worker = KVTransProcess(process=Mock(), task_out_queue=Mock(), device_id=0) + worker.task_out_queue.get.side_effect = queue.Empty() + worker.process.is_alive.return_value = True + assert not worker.wait_until_ready() + assert worker.task_out_queue.get.call_count == 600 diff --git a/unit_tests/server/router/model_infer/mode_backend/test_prefill_queue_strategy.py b/unit_tests/server/router/model_infer/mode_backend/test_prefill_queue_strategy.py new file mode 100644 index 0000000000..606e616aa3 --- /dev/null +++ b/unit_tests/server/router/model_infer/mode_backend/test_prefill_queue_strategy.py @@ -0,0 +1,401 @@ +from types import SimpleNamespace + +import pytest + +from lightllm.server.api_cli import make_argument_parser +from lightllm.server.core.objs import FinishStatus +from lightllm.server.router.model_infer.mode_backend.base_backend import ModeBackend +from lightllm.server.router.model_infer.infer_batch import InferReq +from lightllm.server.router.model_infer.mode_backend.prefill_queue_strategy import ( + PREFILL_QUEUE_STRATEGIES, + PrefillQueueStrategy, + create_prefill_queue_strategy, +) + + +def req( + group_id, + input_len, + kv_len=0, + priority=0, + router_priority=False, + total_len=None, + finish_status=FinishStatus.NO_FINISH, + infer_aborted=False, + filter_mark=False, + wait_pause=False, + paused=False, +): + request = InferReq.__new__(InferReq) + request.req_id = group_id + request.cur_kv_len = kv_len + request.finish_status = FinishStatus(finish_status) + request.infer_aborted = infer_aborted + request.filter_mark = filter_mark + request.wait_pause = wait_pause + request.paused = paused + request.shm_req = SimpleNamespace( + input_len=input_len, + group_req_id=group_id, + sample_params=SimpleNamespace( + high_priority_request=router_priority, + infer_high_priority=priority, + ), + ) + request.get_cur_total_len = lambda: input_len if total_len is None else total_len + return request + + +def strategy(name="default", processed_prefill_tokens=0): + backend = ModeBackend.__new__(ModeBackend) + backend.args = SimpleNamespace(prefill_queue_strategy=name) + backend.processed_prefill_tokens = processed_prefill_tokens + return create_prefill_queue_strategy(backend) + + +@pytest.mark.parametrize("name", PREFILL_QUEUE_STRATEGIES) +def test_empty_and_singleton(name): + policy = strategy(name) + assert isinstance(policy, PrefillQueueStrategy) + assert policy.reorder([]) == [] + request = req(1, 10) + assert policy.reorder([request]) == [request] + + +@pytest.mark.parametrize( + "name, expected", + [ + ("default", [2, 4, 1, 3]), + ("promote_shortest", [2, 4, 3, 1]), + ("hrrn", [2, 4, 1, 3]), + ], +) +def test_ordering_and_input_preservation(name, expected): + requests = [req(1, 20), req(2, 50, priority=-2), req(3, 102, 100), req(4, 5, priority=-1)] + original = list(requests) + result = strategy(name).reorder(requests) + assert [r.shm_req.group_req_id for r in result] == expected + assert all(a is b for a, b in zip(requests, original)) + assert sorted(map(id, result)) == sorted(map(id, original)) + + +def test_inference_priority_is_independent_from_router_priority(): + router_high_priority = req(1, 10, priority=0, router_priority=True) + infer_high_priority = req(2, 10, priority=-1, router_priority=False) + + assert strategy().reorder([router_high_priority, infer_high_priority]) == [ + infer_high_priority, + router_high_priority, + ] + + +def test_promote_shortest_clamps_remaining_tokens_and_preserves_ties(): + requests = [req(9, 5, 6), req(8, 5, 10), req(7, 5, 5), req(6, 10, 8), req(1, 2)] + assert strategy("promote_shortest").reorder(requests) == requests + + +def test_promote_shortest_ignores_finished_request_when_selecting_shortest(): + finished = req(1, 5, kv_len=5, finish_status=FinishStatus.FINISHED_LENGTH) + long_request = req(2, 20) + short_request = req(3, 2) + + assert strategy("promote_shortest").reorder([finished, long_request, short_request], no_decode=True) == [ + short_request, + long_request, + finished, + ] + + +def test_promote_shortest_preserves_arrival_order_for_equal_positive_remaining_tokens(): + long_request = req(1, 20) + first_short = req(3, 10, kv_len=8) + second_short = req(2, 2) + + assert strategy("promote_shortest").reorder([long_request, first_short, second_short]) == [ + first_short, + long_request, + second_short, + ] + + +@pytest.mark.parametrize("name", PREFILL_QUEUE_STRATEGIES) +@pytest.mark.parametrize("mode", [{}, {"no_decode": True}, {"strict_prefill": True}]) +def test_finished_aborted_and_cleanup_requests_stay_last(name, mode): + finished_normal = req(9, 5, kv_len=5, finish_status=FinishStatus.FINISHED_LENGTH) + finished_high = req(8, 2, priority=-10, finish_status=FinishStatus.FINISHED_STOP) + finished_decode = req(7, 5, kv_len=4, finish_status=FinishStatus.FINISHED_ERROR) + aborted = req(6, 2, priority=-10, infer_aborted=True) + cleanup = req(5, 5, kv_len=4, filter_mark=True) + long_request = req(1, 20) + short_request = req(2, 2) + high_priority = req(3, 100, priority=-1) + decode = req(4, 5, kv_len=4) + active = [long_request, short_request, high_priority, decode] + finished = [finished_normal, finished_high, finished_decode, aborted, cleanup] + requests = [ + finished_normal, + long_request, + finished_high, + short_request, + finished_decode, + high_priority, + aborted, + decode, + cleanup, + ] + original = list(requests) + policy = strategy(name) + + expected_active = policy.reorder(active, **mode) + result = policy.reorder(requests, **mode) + + assert result == expected_active + finished + assert all(a is b for a, b in zip(requests, original)) + assert sorted(map(id, result)) == sorted(map(id, original)) + for request in finished: + assert not hasattr(request, "arrival_processed_prefill_tokens") + assert not hasattr(request, "uncached_prefill_tokens") + + +@pytest.mark.parametrize("name", PREFILL_QUEUE_STRATEGIES) +def test_all_finished_requests_preserve_arrival_order(name): + requests = [ + req(6, 2, finish_status=FinishStatus.FINISHED_ABORTED), + req(5, 5, kv_len=5, finish_status=FinishStatus.FINISHED_LENGTH), + req(4, 10, priority=-10, finish_status=FinishStatus.FINISHED_STOP), + req(3, 5, kv_len=4, infer_aborted=True), + req(2, 2, filter_mark=True), + req(1, 5, kv_len=4, finish_status=FinishStatus.FINISHED_PD_DECODE_CAPACITY), + ] + + assert strategy(name).reorder(requests) == requests + + +@pytest.mark.parametrize("pause_state", ["wait_pause", "paused"]) +def test_promote_shortest_ignores_paused_request_when_selecting_shortest(pause_state): + suspended = req(1, 2, **{pause_state: True}) + long_request = req(2, 20) + short_request = req(3, 5) + + assert strategy("promote_shortest").reorder([suspended, long_request, short_request]) == [ + short_request, + long_request, + suspended, + ] + assert getattr(suspended, pause_state) is True + + +@pytest.mark.parametrize("name", PREFILL_QUEUE_STRATEGIES) +@pytest.mark.parametrize("mode", [{}, {"no_decode": True}, {"strict_prefill": True}]) +def test_paused_requests_keep_arrival_order_at_tail(name, mode): + paused_high = req(5, 2, priority=-10, paused=True) + waiting = req(6, 5, kv_len=4, wait_pause=True) + paused_normal = req(4, 2, paused=True) + finished = req(7, 5, finish_status=FinishStatus.FINISHED_STOP) + active = [req(1, 20), req(2, 5), req(3, 100, priority=-1)] + requests = [paused_high, active[0], waiting, active[1], finished, paused_normal, active[2]] + original = list(requests) + policy = strategy(name) + + assert policy.reorder(requests, **mode) == policy.reorder(active, **mode) + [ + paused_high, + waiting, + finished, + paused_normal, + ] + assert all(a is b for a, b in zip(requests, original)) + + +@pytest.mark.parametrize("pause_state", ["wait_pause", "paused"]) +def test_hrrn_initializes_paused_request_only_after_recovery(pause_state): + suspended = req(1, 100, **{pause_state: True}) + active = req(2, 20) + policy = strategy("hrrn", processed_prefill_tokens=1000) + + assert policy.reorder([suspended, active]) == [active, suspended] + assert not hasattr(suspended, "arrival_processed_prefill_tokens") + assert not hasattr(suspended, "uncached_prefill_tokens") + + setattr(suspended, pause_state, False) + suspended.cur_kv_len = 90 + policy.backend.processed_prefill_tokens = 2000 + assert policy.reorder([suspended, active]) == [active, suspended] + assert suspended.arrival_processed_prefill_tokens == 2000 + assert suspended.uncached_prefill_tokens == 10 + + +def test_hrrn_preserves_existing_aging_snapshot_across_pause_and_recovery(): + request = req(1, 100) + policy = strategy("hrrn", processed_prefill_tokens=1000) + policy.reorder([request]) + + request.paused = True + request.cur_kv_len = 0 + policy.backend.processed_prefill_tokens = 2000 + assert policy.reorder([request]) == [request] + assert request.paused is True + + request.paused = False + request.cur_kv_len = 90 + policy.backend.processed_prefill_tokens = 3000 + assert policy.reorder([request]) == [request] + assert request.arrival_processed_prefill_tokens == 1000 + assert request.uncached_prefill_tokens == 100 + + +def test_promote_shortest_keeps_high_priority_requests_first(): + long_high_priority = req(1, 1000, priority=-1) + short_normal_priority = req(2, 2, priority=0) + + assert strategy("promote_shortest").reorder([short_normal_priority, long_high_priority]) == [ + long_high_priority, + short_normal_priority, + ] + + +def test_promote_shortest_preserves_high_priority_arrival_order(): + first_high_priority = req(1, 100, priority=-1) + second_high_priority = req(2, 100, priority=-2) + + assert strategy("promote_shortest").reorder([first_high_priority, second_high_priority]) == [ + first_high_priority, + second_high_priority, + ] + + +def test_promote_shortest_moves_only_one_normal_request_and_preserves_other_order(): + first_normal = req(1, 20) + shortest_normal = req(2, 2) + third_normal = req(3, 10) + fourth_normal = req(4, 5) + + assert strategy("promote_shortest").reorder([first_normal, shortest_normal, third_normal, fourth_normal]) == [ + shortest_normal, + first_normal, + third_normal, + fourth_normal, + ] + + +def test_hrrn_prioritizes_short_work_for_equally_aged_requests(): + requests = [req(1, 100), req(2, 10), req(3, 2)] + policy = strategy("hrrn") + policy.reorder(requests) + policy.backend.processed_prefill_tokens = 1000 + + assert policy.reorder(requests) == [ + requests[2], + requests[1], + requests[0], + ] + + +def test_hrrn_aging_allows_long_request_to_overtake_new_short_request(): + long_old = req(1, 100) + short_new = req(2, 2) + policy = strategy("hrrn") + policy.reorder([long_old]) + policy.backend.processed_prefill_tokens = 100_000 + + assert policy.reorder([short_new, long_old]) == [ + long_old, + short_new, + ] + + +def test_hrrn_aging_starts_at_first_sort_and_survives_later_sorts(): + long_request = req(1, 100) + short_request = req(2, 2) + requests = [short_request, long_request] + policy = strategy("hrrn", processed_prefill_tokens=1000) + + # 首轮两个请求的等待量均为零，使用请求 ID 打破平局。 + assert policy.reorder(requests) == [long_request, short_request] + policy.backend.processed_prefill_tokens = 2000 + # 后续排序保留首次快照，等龄请求中短请求的响应比更高。 + assert policy.reorder(requests) == [short_request, long_request] + + +def test_hrrn_binds_uncached_tokens_once_and_uses_request_id_tiebreaker(): + less_cached = req(2, 100, kv_len=0) + equally_cached_later_id = req(3, 100, kv_len=90) + more_cached = req(1, 100, kv_len=90) + policy = strategy("hrrn") + + requests = [less_cached, equally_cached_later_id, more_cached] + policy.reorder(requests) + policy.backend.processed_prefill_tokens = 1000 + assert policy.reorder(requests) == [more_cached, equally_cached_later_id, less_cached] + assert [request.uncached_prefill_tokens for request in requests] == [100, 10, 10] + + less_cached.cur_kv_len = 95 + more_cached.cur_kv_len = 0 + policy.backend.processed_prefill_tokens = 2000 + assert policy.reorder(requests) == [more_cached, equally_cached_later_id, less_cached] + + +def test_hrrn_keeps_internal_high_priority_requests_first(): + shortest_normal = req(1, 2, priority=0) + long_high_priority = req(2, 1000, priority=-1) + + assert strategy("hrrn", processed_prefill_tokens=1000).reorder([shortest_normal, long_high_priority]) == [ + long_high_priority, + shortest_normal, + ] + + +@pytest.mark.parametrize("name", PREFILL_QUEUE_STRATEGIES) +def test_decode_requests_stay_left_and_are_not_reordered(name): + decode_normal = req(1, 5, kv_len=4, priority=0) + prefill_high = req(2, 100, priority=-2) + decode_high = req(3, 5, kv_len=4, priority=-10) + prefill_normal = req(4, 2, priority=0) + + assert strategy(name).reorder([decode_normal, prefill_normal, decode_high, prefill_high]) == [ + decode_normal, + decode_high, + prefill_high, + prefill_normal, + ] + + +def test_no_decode_and_strict_prefill_match_backend_classification_rules(): + boundary_req = req(1, 5, kv_len=4, priority=0) + high_priority_prefill = req(2, 10, priority=-1) + + policy = strategy() + assert policy.reorder([boundary_req, high_priority_prefill]) == [boundary_req, high_priority_prefill] + assert policy.reorder([boundary_req, high_priority_prefill], no_decode=True) == [ + high_priority_prefill, + boundary_req, + ] + assert policy.reorder([boundary_req, high_priority_prefill], strict_prefill=True) == [ + high_priority_prefill, + boundary_req, + ] + + +@pytest.mark.parametrize("name", PREFILL_QUEUE_STRATEGIES) +def test_cli_strategy_selection(name): + args = make_argument_parser().parse_args(["--prefill_queue_strategy", name]) + backend = SimpleNamespace(args=args) + policy = create_prefill_queue_strategy(backend) + assert isinstance(policy, PREFILL_QUEUE_STRATEGIES[name]) + assert policy.backend is backend + + +def test_cli_default_and_invalid_strategy(): + parser = make_argument_parser() + assert parser.parse_args([]).prefill_queue_strategy == "auto" + assert parser.parse_args(["--prefill_queue_strategy", "auto"]).prefill_queue_strategy == "auto" + assert set(PREFILL_QUEUE_STRATEGIES) == {"default", "promote_shortest", "hrrn"} + for name in ("unknown", "none"): + with pytest.raises(SystemExit): + parser.parse_args(["--prefill_queue_strategy", name]) + with pytest.raises(ValueError, match="Unknown prefill queue strategy"): + strategy(name) + + +def test_strategy_base_is_abstract(): + with pytest.raises(TypeError): + PrefillQueueStrategy(SimpleNamespace()) 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 new file mode 100644 index 0000000000..784848cb13 --- /dev/null +++ b/unit_tests/server/router/model_infer/mtp_speculative/test_dflash.py @@ -0,0 +1,180 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.server.router.model_infer.mtp_speculative.proposers.dflash import DFlashProposer +from lightllm.server.router.model_infer.mtp_speculative.proposers.proposal_type import DFlashSpecProposal +from lightllm.server.router.model_infer.pin_mem_manager import g_pin_mem_manager + + +def test_dflash_prefill_uses_a_shallow_copy_for_target_hidden(): + forwarded_inputs = [] + draft_model = SimpleNamespace(forward=forwarded_inputs.append) + proposer = DFlashProposer( + backend=SimpleNamespace(draft_models=[draft_model]), + enable_dynmaic_mtp=False, + ) + model_input = SimpleNamespace( + is_prefill=True, + b_position_delta=None, + b_req_idx=torch.tensor([3, 5], dtype=torch.int32), + input_ids=torch.arange(5, dtype=torch.int64), + mtp_draft_input_hiddens=None, + ) + target_hidden = torch.empty((5, 8)) + + proposer.fill_draft_model_kv_state( + target_model_input=model_input, + target_model_output=SimpleNamespace( + mtp_collector=SimpleNamespace(spec_hidden=target_hidden), + ), + target_next_token_ids=torch.tensor([11, 13], dtype=torch.int64), + ) + + assert len(forwarded_inputs) == 1 + assert forwarded_inputs[0] is not model_input + assert forwarded_inputs[0].mtp_draft_input_hiddens is target_hidden + assert model_input.mtp_draft_input_hiddens is None + + +def test_dflash_commits_verify_kv_and_builds_parallel_block(monkeypatch): + block_size = 3 + flat_draft_token_ids = torch.tensor([30, 31, 32, 40, 41, 42], dtype=torch.int64) + flat_draft_token_probs = torch.tensor([0.9, 0.8, 0.7, 0.6, 0.5, 0.4], dtype=torch.float16) + forwarded_inputs = [] + + def forward(model_input): + forwarded_inputs.append(model_input) + return SimpleNamespace() + + draft_model = SimpleNamespace( + block_size=block_size, + mask_token_id=99, + forward=forward, + ) + proposer = DFlashProposer( + backend=SimpleNamespace( + draft_models=[draft_model], + _gen_argmax_token_ids_and_prob=lambda _: (flat_draft_token_ids, flat_draft_token_probs), + ), + enable_dynmaic_mtp=True, + ) + monkeypatch.setattr(torch.Tensor, "cuda", lambda self, non_blocking=False: self) + monkeypatch.setattr( + g_pin_mem_manager, + "get_const_gpu_tensor", + lambda *, shape, fill_value, dtype, **_: torch.full(shape, fill_value, dtype=dtype), + ) + + target_hidden = torch.empty((5, 8)) + model_input = SimpleNamespace( + is_prefill=False, + batch_size=5, + total_token_num=40, + max_q_seq_len=1, + max_kv_seq_len=9, + input_ids=torch.tensor([10, 11, 12, 20, 21], dtype=torch.int64), + b_req_idx=torch.tensor([7, 7, 7, 9, 9], dtype=torch.int32), + b_mtp_index=torch.arange(5, dtype=torch.int32), + b_seq_len=torch.tensor([4, 5, 6, 8, 9], dtype=torch.int32), + b_position_delta=torch.tensor([1, 1, 1, 2, 2], dtype=torch.int32), + b_shared_seq_len=torch.tensor([3, 3, 3, 6, 6], dtype=torch.int32), + b_shared_radix_node_id=torch.tensor([17, 17, 17, 19, 19], dtype=torch.int64), + multimodal_params=[{"images": [], "audios": []} for _ in range(5)], + mtp_draft_input_hiddens=None, + ) + + proposal = proposer.propose_next( + target_model_input=model_input, + target_model_output=SimpleNamespace( + mtp_collector=SimpleNamespace(spec_hidden=target_hidden), + ), + target_next_token_ids=model_input.input_ids, + b_req_mtp_start_loc=torch.tensor([0, 3], dtype=torch.int32), + draft_step=2, + accept_len=torch.tensor([2, 2], dtype=torch.int32), + ) + + assert isinstance(proposal, DFlashSpecProposal) + assert len(forwarded_inputs) == 2 + verify_draft_input, block_draft_input = forwarded_inputs + assert verify_draft_input is not model_input + assert verify_draft_input.mtp_draft_input_hiddens is target_hidden + assert model_input.mtp_draft_input_hiddens is None + + assert torch.equal(block_draft_input.input_ids, torch.tensor([11, 99, 99, 21, 99, 99])) + assert torch.equal(block_draft_input.b_req_idx, torch.tensor([7, 7, 7, 9, 9, 9], dtype=torch.int32)) + assert torch.equal(block_draft_input.b_mtp_index, torch.zeros(6, dtype=torch.int32)) + assert torch.equal(block_draft_input.b_seq_len, torch.tensor([6, 7, 8, 10, 11, 12], dtype=torch.int32)) + assert torch.equal(block_draft_input.b_position_delta, torch.tensor([1, 1, 1, 2, 2, 2], dtype=torch.int32)) + assert block_draft_input.mtp_draft_input_hiddens is None + + assert torch.equal(proposal.token_ids, torch.tensor([[30, 31], [40, 41]], dtype=torch.int64)) + torch.testing.assert_close( + 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() diff --git a/unit_tests/server/router/model_infer/mtp_speculative/test_dflash2.py b/unit_tests/server/router/model_infer/mtp_speculative/test_dflash2.py new file mode 100644 index 0000000000..021b6d47fa --- /dev/null +++ b/unit_tests/server/router/model_infer/mtp_speculative/test_dflash2.py @@ -0,0 +1,117 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.basemodel.batch_objs import ModelMtpOutputCollector, ModelOutput, PostLayerOutput +from lightllm.server.router.model_infer.mtp_speculative.proposers.proposal_type import DFlashSpecProposal +from lightllm.server.router.model_infer.mtp_speculative.proposers.dflash2 import DFlash2Proposer + + +@pytest.mark.parametrize("dynamic_verify", [False, True]) +@pytest.mark.parametrize("draft_step", [1, 3]) +def test_dflash2_proposal_returns_greedy_tokens(dynamic_verify, draft_step): + tokens = torch.arange(6).view(2, 3) + physical_tokens = torch.tensor([100, 0, 1, 2, 200, 3, 4, 5]) + proposer = DFlash2Proposer(backend=SimpleNamespace(), enable_dynmaic_mtp=dynamic_verify) + confidence = torch.tensor([[0.0, 1.0, -1.0], [float("inf"), -100.0, 2.0]]) + output = SimpleNamespace( + mtp_collector=ModelMtpOutputCollector(draft_token_ids=physical_tokens, confidence_logits=confidence) + ) + proposal = proposer._build_proposal(output, 2, 4, draft_step) + assert type(proposal) is DFlashSpecProposal + torch.testing.assert_close(proposal.token_ids, tokens[:, :draft_step]) + if dynamic_verify: + torch.testing.assert_close(proposal.schedule_scores, confidence[:, :draft_step].sigmoid().clamp(0.01, 0.99)) + else: + assert proposal.schedule_scores is None + + +@pytest.mark.parametrize("dynamic_verify", [False, True]) +@pytest.mark.parametrize("req_num", [0, 1, 2]) +def test_dflash2_collector_unpads_physical_token_rows(dynamic_verify, req_num): + tokens = torch.arange(6).view(2, 3) + physical_tokens = torch.tensor([100, 0, 1, 2, 200, 3, 4, 5]) + confidence = tokens.float() if dynamic_verify else None + collector = ModelMtpOutputCollector(draft_token_ids=physical_tokens, confidence_logits=confidence) + unpadded = collector.unpad_decode(padded_batch_size=8, origin_batch_size=req_num * 4) + torch.testing.assert_close(unpadded.draft_token_ids, physical_tokens[: req_num * 4]) + if dynamic_verify: + torch.testing.assert_close(unpadded.confidence_logits, confidence[:req_num]) + else: + assert unpadded.confidence_logits is None + assert collector.draft_token_ids.shape == (8,) + + proposer = DFlash2Proposer(backend=SimpleNamespace(), enable_dynmaic_mtp=dynamic_verify) + proposal = proposer._build_proposal(SimpleNamespace(mtp_collector=unpadded), req_num, 4, 3) + torch.testing.assert_close(proposal.token_ids, tokens[:req_num]) + if dynamic_verify: + torch.testing.assert_close(proposal.schedule_scores, confidence[:req_num].sigmoid().clamp(0.01, 0.99)) + + +def test_dflash2_dynamic_verification_requires_confidence(): + proposer = DFlash2Proposer(backend=SimpleNamespace(), enable_dynmaic_mtp=True) + output = ModelOutput( + logits=torch.empty(4, 1), + mtp_collector=ModelMtpOutputCollector(draft_token_ids=torch.ones(4, dtype=torch.int64)), + ) + with pytest.raises(RuntimeError, match="selector confidence"): + proposer._build_proposal(output, 1, 4, 3) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("width", [1, 3, 7]) +@pytest.mark.parametrize("dynamic_verify", [False, True]) +def test_greedy_selector_conditions_on_selected_predecessor(width, dynamic_verify): + from lightllm.models.qwen3_dflash2.triton_kernel.selector_walk import greedy_selector_walk + + scores = torch.randn(3, width, 16, 16, generator=torch.Generator().manual_seed(42)) + scores[0].zero_() # Ties select the first candidate. + candidates = torch.arange(3 * width * 16).view(3, width, 16) + expected = torch.empty(3, width, dtype=torch.int64) + expected_probs = torch.empty(3, width) + for req in range(3): + previous = 0 + for slot in range(width): + row = scores[req, slot, previous] + previous = int(row.argmax()) + expected_probs[req, slot] = row.softmax(-1)[previous] + expected[req, slot] = candidates[req, slot, previous] + scores, candidates = scores.cuda(), candidates.cuda() + rng_before = torch.cuda.get_rng_state() + confidence = torch.empty(3, width, device="cuda") if dynamic_verify else None + actual = greedy_selector_walk(scores, candidates, confidence_logits=confidence) + if dynamic_verify: + torch.testing.assert_close(confidence.sigmoid().cpu(), expected_probs) + torch.testing.assert_close(actual.cpu(), expected) + assert torch.equal(rng_before, torch.cuda.get_rng_state()) + + +@pytest.mark.parametrize("is_prefill", [True, False]) +def test_dflash2_post_layer_satisfies_base_model_output_contract(is_prefill): + from lightllm.models.qwen3_dflash2.layer_infer.post_layer_infer import Qwen3DFlash2PostLayerInfer + + layer = Qwen3DFlash2PostLayerInfer.__new__(Qwen3DFlash2PostLayerInfer) + layer.block_size_ = 4 + layer.selector_top_k_ = 2 + hidden = torch.randn(4, 8) + ids = torch.arange(6).view(3, 2) + unary = torch.randn(3, 2) + selected = ids[:, 0].view(1, 3) + layer._slice_get_last_input = lambda *_: (hidden, 4) + layer._norm = lambda x, *_: x + layer._compute_candidates = lambda **_: (ids, unary) + confidence = torch.randn(1, 3) + layer._select_path = lambda **_: (selected, confidence) + collected = {} + state = SimpleNamespace( + is_prefill=is_prefill, + input_ids=torch.tensor([99, 0, 0, 0]), + hidden_collector=SimpleNamespace(add_mtp_outputs=lambda **kwargs: collected.update(kwargs)), + ) + output = layer.token_forward(hidden, state, None) + assert isinstance(output, PostLayerOutput) + assert output.logits.shape == ((0,) if is_prefill else (4, 1)) + if not is_prefill: + torch.testing.assert_close(collected["draft_token_ids"], torch.tensor([99, 0, 2, 4])) + torch.testing.assert_close(collected["confidence_logits"], confidence) diff --git a/unit_tests/server/router/model_infer/mtp_speculative/test_eagle_overlap.py b/unit_tests/server/router/model_infer/mtp_speculative/test_eagle_overlap.py new file mode 100644 index 0000000000..d5dae90b4b --- /dev/null +++ b/unit_tests/server/router/model_infer/mtp_speculative/test_eagle_overlap.py @@ -0,0 +1,363 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.basemodel.batch_objs import ModelMtpOutputCollector, ModelOutput +from lightllm.server.router.model_infer.mtp_speculative.dp_overlap_proposers import eagle_with_att +from lightllm.server.router.model_infer.mtp_speculative.dp_overlap_proposers.utils import ( + get_dp_overlap_req_start_rows, +) +from lightllm.server.router.model_infer.mtp_speculative.dp_overlap_proposers.eagle3 import ( + DpOverlapEagle3Proposer, +) +from lightllm.server.router.model_infer.mtp_speculative.dp_overlap_proposers.eagle_no_att import ( + DpOverlapEagleNoAttProposer, +) +from lightllm.server.router.model_infer.mtp_speculative.dp_overlap_proposers.eagle_with_att import ( + DpOverlapEagleWithAttProposer, +) +from lightllm.server.router.model_infer.mtp_speculative.proposers.eagle3 import ( + Eagle3Proposer, +) + + +class _DraftModel: + def __init__(self): + self.extend_batch_sizes = None + self.extend_inputs = None + self.decode_batch_sizes = [] + self.decode_inputs = [] + + def _microbatch_overlap_prefill_cuda(self, input0, input1): + self.extend_batch_sizes = (input0.batch_size, input1.batch_size) + self.extend_inputs = (input0, input1) + return tuple( + ModelOutput( + logits=torch.arange(model_input.batch_size, dtype=torch.float32).view(-1, 1), + mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.ones((model_input.batch_size, 2))), + ) + for model_input in (input0, input1) + ) + + def _microbatch_overlap_decode_cuda(self, input0, input1): + self.decode_batch_sizes.append((input0.batch_size, input1.batch_size)) + self.decode_inputs.append((input0, input1)) + return tuple( + ModelOutput( + logits=torch.arange( + model_input.batch_size, + dtype=torch.float32, + device=model_input.input_ids.device, + ).view(-1, 1), + mtp_collector=ModelMtpOutputCollector( + spec_hidden=torch.ones( + (model_input.batch_size, 2), + device=model_input.input_ids.device, + ) + ), + ) + for model_input in (input0, input1) + ) + + def map_draft_vocab_to_main_vocab(self, token_ids): + return token_ids + + +def _target_input(batch_size, b_mtp_index=None, device="cpu"): + if b_mtp_index is None: + b_mtp_index = torch.arange(batch_size, dtype=torch.int32, device=device) % 3 + return SimpleNamespace( + batch_size=batch_size, + total_token_num=batch_size, + input_ids=torch.arange(batch_size, dtype=torch.int64, device=device), + b_seq_len=torch.arange(batch_size, dtype=torch.int32, device=device) + 4, + b_req_idx=torch.arange(batch_size, dtype=torch.int32, device=device), + b_mtp_index=b_mtp_index, + b_position_delta=torch.zeros(batch_size, dtype=torch.int32, device=device), + b_shared_seq_len=torch.zeros(batch_size, dtype=torch.int32, device=device), + b_shared_radix_node_id=torch.arange(batch_size, dtype=torch.int64, device=device), + max_kv_seq_len=16, + max_cache_len=16, + is_prefill=False, + multimodal_params=[{"images": [], "audios": []}] * batch_size, + ) + + +def _patch_cpu_req_start_rows(monkeypatch): + def get_cpu_req_start_rows(b_mtp_index, req_num): + req_start_rows = torch.nonzero(b_mtp_index == 0, as_tuple=False).flatten().to(dtype=torch.int32) + assert req_start_rows.shape == (req_num,) + return req_start_rows + + monkeypatch.setattr(eagle_with_att, "get_dp_overlap_req_start_rows", get_cpu_req_start_rows) + + +def test_dp_overlap_req_start_rows_rejects_nonempty_cpu_input(): + with pytest.raises(AssertionError, match="must be a CUDA tensor"): + get_dp_overlap_req_start_rows( + b_mtp_index=torch.tensor([0, 1], dtype=torch.int32), + req_num=1, + ) + + +def test_overlap_eagle_supports_variable_verify_layout(monkeypatch): + _patch_cpu_req_start_rows(monkeypatch) + draft_model = _DraftModel() + backend = SimpleNamespace( + max_draft_step=2, + draft_models=[draft_model], + model=SimpleNamespace( + req_manager=SimpleNamespace( + mem_manager=SimpleNamespace(HOLD_TOKEN_MEMINDEXES=(99,), page_size=1), + ) + ), + _gen_argmax_token_ids=lambda output: output.logits[:, 0].to(torch.int64), + ) + proposer = DpOverlapEagleWithAttProposer(backend=backend, enable_dynmaic_mtp=False) + model_input0 = _target_input(batch_size=3) + model_input1 = _target_input(batch_size=6) + + proposal = proposer.propose_next_overlap( + target_model_input0=model_input0, + target_model_output0=ModelOutput( + logits=torch.empty((3, 1)), + mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.ones((3, 2))), + ), + target_next_token_ids0=torch.arange(3, dtype=torch.int64), + accept_len0=torch.tensor([2], dtype=torch.int32), + target_model_input1=model_input1, + target_model_output1=ModelOutput( + logits=torch.empty((6, 1)), + mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.ones((6, 2))), + ), + target_next_token_ids1=torch.arange(10, 16, dtype=torch.int64), + accept_len1=torch.tensor([1, 3], dtype=torch.int32), + draft_step=2, + ) + + assert draft_model.extend_batch_sizes is None + assert draft_model.decode_batch_sizes == [(3, 6), (1, 2)] + assert proposal.token_ids.shape == (3, 2) + assert torch.equal(proposal.token_ids, torch.tensor([[1, 0], [0, 0], [5, 1]])) + + +def test_overlap_eagle_supports_empty_verify_rows(): + draft_model = _DraftModel() + backend = SimpleNamespace( + max_draft_step=2, + draft_models=[draft_model], + model=SimpleNamespace( + req_manager=SimpleNamespace( + mem_manager=SimpleNamespace(HOLD_TOKEN_MEMINDEXES=(99,), page_size=1), + ) + ), + _gen_argmax_token_ids=lambda output: output.logits[:, 0].to(torch.int64), + ) + proposer = DpOverlapEagleWithAttProposer(backend=backend, enable_dynmaic_mtp=False) + + proposal = proposer.propose_next_overlap( + target_model_input0=_target_input(batch_size=0), + target_model_output0=ModelOutput( + logits=torch.empty((0, 1)), + mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.ones((0, 2))), + ), + target_next_token_ids0=torch.empty((0,), dtype=torch.int64), + accept_len0=torch.empty((0,), dtype=torch.int32), + target_model_input1=_target_input(batch_size=0), + target_model_output1=ModelOutput( + logits=torch.empty((0, 1)), + mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.ones((0, 2))), + ), + target_next_token_ids1=torch.empty((0,), dtype=torch.int64), + accept_len1=torch.empty((0,), dtype=torch.int32), + draft_step=2, + ) + + assert proposal.token_ids.shape == (0, 2) + assert draft_model.decode_batch_sizes == [(0, 0), (0, 0)] + + +def test_overlap_eagle_returns_dynamic_schedule_scores(monkeypatch): + _patch_cpu_req_start_rows(monkeypatch) + draft_model = _DraftModel() + backend = SimpleNamespace( + max_draft_step=2, + draft_models=[draft_model], + model=SimpleNamespace( + req_manager=SimpleNamespace( + mem_manager=SimpleNamespace(HOLD_TOKEN_MEMINDEXES=(99,), page_size=1), + ) + ), + _gen_argmax_token_ids_and_prob=lambda output: ( + output.logits[:, 0].to(torch.int64), + output.logits[:, 0] / 10 + 0.5, + ), + ) + proposer = DpOverlapEagleWithAttProposer(backend=backend, enable_dynmaic_mtp=True) + + proposal = proposer.propose_next_overlap( + target_model_input0=_target_input( + batch_size=2, + b_mtp_index=torch.tensor([0, 1], dtype=torch.int32), + ), + target_model_output0=ModelOutput( + logits=torch.empty((2, 1)), + mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.ones((2, 2))), + ), + target_next_token_ids0=torch.arange(2, dtype=torch.int64), + accept_len0=torch.tensor([2], dtype=torch.int32), + target_model_input1=_target_input( + batch_size=3, + b_mtp_index=torch.tensor([0, 1, 0], dtype=torch.int32), + ), + target_model_output1=ModelOutput( + logits=torch.empty((3, 1)), + mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.ones((3, 2))), + ), + target_next_token_ids1=torch.arange(2, 5, dtype=torch.int64), + accept_len1=torch.tensor([1, 1], dtype=torch.int32), + draft_step=2, + ) + + assert torch.equal(proposal.token_ids, torch.tensor([[1, 0], [0, 0], [2, 1]])) + assert torch.allclose( + proposal.schedule_scores, + torch.tensor([[0.6, 0.5], [0.5, 0.5], [0.7, 0.6]]), + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_overlap_eagle_no_att_supports_dynamic_draft_step(): + device = "cuda" + draft_model = _DraftModel() + backend = SimpleNamespace( + max_draft_step=3, + draft_models=[draft_model], + _gen_argmax_token_ids_and_prob=lambda output: ( + output.logits[:, 0].to(torch.int64), + output.logits[:, 0] / 10 + 0.5, + ), + ) + proposer = DpOverlapEagleNoAttProposer(backend=backend, enable_dynmaic_mtp=True) + model_input0 = _target_input( + batch_size=2, + b_mtp_index=torch.tensor([0, 1], dtype=torch.int32, device=device), + device=device, + ) + model_input1 = _target_input( + batch_size=3, + b_mtp_index=torch.tensor([0, 1, 0], dtype=torch.int32, device=device), + device=device, + ) + + proposal = proposer.propose_next_overlap( + target_model_input0=model_input0, + target_model_output0=ModelOutput( + logits=torch.empty((2, 1), device=device), + mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.ones((2, 2), device=device)), + ), + target_next_token_ids0=torch.arange(2, dtype=torch.int64, device=device), + accept_len0=torch.tensor([2], dtype=torch.int32, device=device), + target_model_input1=model_input1, + target_model_output1=ModelOutput( + logits=torch.empty((3, 1), device=device), + mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.ones((3, 2), device=device)), + ), + target_next_token_ids1=torch.arange(2, 5, dtype=torch.int64, device=device), + accept_len1=torch.tensor([1, 1], dtype=torch.int32, device=device), + draft_step=2, + ) + + assert proposal.token_ids.tolist() == [[0, 0], [0, 0], [1, 1]] + assert torch.allclose( + proposal.schedule_scores, + torch.tensor( + [[0.5, 0.5], [0.5, 0.5], [0.6, 0.6]], + device=device, + ), + ) + assert draft_model.decode_batch_sizes == [(1, 2), (1, 2)] + + +def test_autoregressive_eagle_reuses_overlap_inputs(monkeypatch): + _patch_cpu_req_start_rows(monkeypatch) + draft_model = _DraftModel() + backend = SimpleNamespace( + max_draft_step=2, + draft_models=[draft_model], + model=SimpleNamespace( + req_manager=SimpleNamespace( + mem_manager=SimpleNamespace(HOLD_TOKEN_MEMINDEXES=(99,), page_size=1), + ) + ), + _gen_argmax_token_ids=lambda output: output.logits[:, 0].to(torch.int64), + ) + proposer = DpOverlapEagle3Proposer(backend=backend, enable_dynmaic_mtp=False) + model_input0 = _target_input(batch_size=3) + model_input1 = _target_input(batch_size=6) + + proposal = proposer.propose_next_overlap( + target_model_input0=model_input0, + target_model_output0=ModelOutput( + logits=torch.empty((3, 1)), + mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.ones((3, 2))), + ), + target_next_token_ids0=torch.arange(3, dtype=torch.int64), + accept_len0=torch.tensor([2], dtype=torch.int32), + target_model_input1=model_input1, + target_model_output1=ModelOutput( + logits=torch.empty((6, 1)), + mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.ones((6, 2))), + ), + target_next_token_ids1=torch.arange(10, 16, dtype=torch.int64), + accept_len1=torch.tensor([1, 3], dtype=torch.int32), + draft_step=2, + ) + + assert draft_model.extend_inputs is None + assert len(draft_model.decode_inputs) == 2 + assert draft_model.decode_inputs[0][0] is not model_input0 + assert draft_model.decode_inputs[0][1] is not model_input1 + assert draft_model.extend_batch_sizes is None + assert draft_model.decode_batch_sizes == [(3, 6), (1, 2)] + assert proposal.token_ids.shape == (3, 2) + assert torch.equal(proposal.token_ids, torch.tensor([[1, 0], [0, 0], [5, 1]])) + + +def test_eagle3_maps_draft_token_ids_in_proposer(): + proposer = Eagle3Proposer.__new__(Eagle3Proposer) + proposer.backend = SimpleNamespace( + draft_models=[SimpleNamespace(map_draft_vocab_to_main_vocab=lambda token_ids: token_ids + 100)], + _gen_argmax_token_ids=lambda _: torch.tensor([1, 2]), + _gen_argmax_token_ids_and_prob=lambda _: ( + torch.tensor([3, 4]), + torch.tensor([0.8, 0.7]), + ), + ) + + token_ids = proposer._gen_argmax_token_ids(ModelOutput(logits=torch.empty(0))) + token_ids_with_prob, probs = proposer._gen_argmax_token_ids_and_prob(ModelOutput(logits=torch.empty(0))) + + assert torch.equal(token_ids, torch.tensor([101, 102])) + assert torch.equal(token_ids_with_prob, torch.tensor([103, 104])) + assert torch.equal(probs, torch.tensor([0.8, 0.7])) + + +def test_dp_overlap_eagle3_maps_draft_token_ids_in_proposer(): + proposer = DpOverlapEagle3Proposer.__new__(DpOverlapEagle3Proposer) + proposer.backend = SimpleNamespace( + draft_models=[SimpleNamespace(map_draft_vocab_to_main_vocab=lambda token_ids: token_ids + 100)], + _gen_argmax_token_ids=lambda _: torch.tensor([1, 2]), + _gen_argmax_token_ids_and_prob=lambda _: ( + torch.tensor([3, 4]), + torch.tensor([0.8, 0.7]), + ), + ) + + token_ids = proposer._gen_argmax_token_ids(ModelOutput(logits=torch.empty(0))) + token_ids_with_prob, probs = proposer._gen_argmax_token_ids_and_prob(ModelOutput(logits=torch.empty(0))) + + assert torch.equal(token_ids, torch.tensor([101, 102])) + assert torch.equal(token_ids_with_prob, torch.tensor([103, 104])) + assert torch.equal(probs, torch.tensor([0.8, 0.7])) diff --git a/unit_tests/server/router/model_infer/mtp_speculative/test_eagle_with_att.py b/unit_tests/server/router/model_infer/mtp_speculative/test_eagle_with_att.py new file mode 100644 index 0000000000..365acb30b2 --- /dev/null +++ b/unit_tests/server/router/model_infer/mtp_speculative/test_eagle_with_att.py @@ -0,0 +1,213 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.basemodel.batch_objs import ModelMtpOutputCollector, ModelOutput +from lightllm.server.router.model_infer.mtp_speculative.proposers.eagle3 import Eagle3Proposer +from lightllm.server.router.model_infer.mtp_speculative.proposers.eagle_with_att import EagleWithAttProposer +from lightllm.server.router.model_infer.mtp_speculative.proposers.proposal_type import EagleSpecProposal + + +def test_eagle3_reuses_attention_flow_and_maps_proposal_tokens(): + target_input_ids = torch.tensor([10, 20], dtype=torch.int64) + target_hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + + def forward(model_input): + assert model_input.input_ids is target_input_ids + assert model_input.mtp_draft_input_hiddens is target_hidden + return ModelOutput( + logits=torch.tensor([[1.0], [2.0]]), + mtp_collector=ModelMtpOutputCollector(spec_hidden=target_hidden + 10), + ) + + draft_model = SimpleNamespace( + forward=forward, + map_draft_vocab_to_main_vocab=lambda token_ids: token_ids + 100, + ) + proposer = Eagle3Proposer( + backend=SimpleNamespace( + draft_models=[draft_model], + _gen_argmax_token_ids=lambda output: output.logits[:, 0].long(), + ), + enable_dynmaic_mtp=False, + ) + target_input = SimpleNamespace( + is_prefill=False, + batch_size=2, + input_ids=None, + b_position_delta=torch.zeros(2, dtype=torch.int32), + mtp_draft_input_hiddens=None, + ) + + proposal = proposer.propose_next( + target_model_input=target_input, + target_model_output=SimpleNamespace(mtp_collector=SimpleNamespace(spec_hidden=target_hidden)), + target_next_token_ids=target_input_ids, + b_req_mtp_start_loc=torch.tensor([0, 1], dtype=torch.int32), + draft_step=1, + accept_len=torch.ones(2, dtype=torch.int32), + ) + + assert isinstance(proposal, EagleSpecProposal) + torch.testing.assert_close(proposal.token_ids, torch.tensor([[101], [102]])) + assert proposal.schedule_scores is None + assert target_input.input_ids is None + assert target_input.mtp_draft_input_hiddens is None + + +def test_eagle_with_att_rejects_zero_draft_steps(): + proposer = EagleWithAttProposer( + backend=SimpleNamespace(draft_models=[]), + enable_dynmaic_mtp=False, + ) + + with pytest.raises(AssertionError, match="requires draft_step to be greater than 0"): + proposer.propose_next( + target_model_input=None, + target_model_output=None, + target_next_token_ids=torch.tensor([10]), + b_req_mtp_start_loc=torch.tensor([0], dtype=torch.int32), + draft_step=0, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_eagle_with_att_prefill_builds_draft_kv_without_mutating_target(): + device = "cuda" + original_input_ids = torch.tensor([10, 11, 12, 20, 21, 22], dtype=torch.int64, device=device) + target_hidden = torch.arange(12, dtype=torch.float32, device=device).reshape(6, 2) + target_input = SimpleNamespace( + is_prefill=True, + batch_size=2, + input_ids=original_input_ids, + b_req_idx=torch.tensor([7, 9], dtype=torch.int32, device=device), + b_seq_len=torch.tensor([3, 3], dtype=torch.int32, device=device), + b_ready_cache_len=torch.zeros(2, dtype=torch.int32, device=device), + b_position_delta=None, + b_is_decode_req=None, + mtp_draft_input_hiddens=None, + ) + forwarded = [] + + def forward(model_input): + forwarded.append( + ( + model_input, + model_input.input_ids.clone(), + model_input.mtp_draft_input_hiddens, + model_input.b_is_decode_req, + ) + ) + return ModelOutput(logits=torch.empty((0,), device=device)) + + proposer = EagleWithAttProposer( + backend=SimpleNamespace(draft_models=[SimpleNamespace(forward=forward)]), + enable_dynmaic_mtp=False, + ) + proposer.fill_draft_model_kv_state( + target_model_input=target_input, + target_model_output=SimpleNamespace(mtp_collector=SimpleNamespace(spec_hidden=target_hidden)), + target_next_token_ids=torch.tensor([13, 23], dtype=torch.int64, device=device), + ) + + assert len(forwarded) == 1 + assert forwarded[0][0] is not target_input + torch.testing.assert_close( + forwarded[0][1], + torch.tensor([11, 12, 13, 21, 22, 23], dtype=torch.int64, device=device), + ) + assert forwarded[0][2] is target_hidden + assert not forwarded[0][3].any() + assert target_input.input_ids is original_input_ids + assert target_input.b_is_decode_req is None + assert target_input.mtp_draft_input_hiddens is None + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_eagle_with_att_commits_verify_kv_then_recurrently_decodes(monkeypatch): + device = "cuda" + original_input_ids = torch.tensor([10, 11, 12, 20, 21, 22], dtype=torch.int64, device=device) + target_hidden = torch.arange(12, dtype=torch.float32, device=device).reshape(6, 2) + extend_hidden = target_hidden + 100 + decode_hidden = torch.tensor([[201.0, 202.0], [203.0, 204.0]], device=device) + draft_calls = [] + + def forward(model_input): + draft_calls.append( + { + "model_input": model_input, + "batch_size": model_input.batch_size, + "input_ids": model_input.input_ids.clone(), + "draft_hidden": model_input.mtp_draft_input_hiddens.clone(), + "b_req_idx": model_input.b_req_idx.clone(), + "b_mtp_index": model_input.b_mtp_index.clone(), + "b_seq_len": model_input.b_seq_len.clone(), + "max_kv_seq_len": model_input.max_kv_seq_len, + "total_token_num": model_input.total_token_num, + } + ) + if len(draft_calls) == 1: + return ModelOutput( + logits=torch.arange(30, 36, dtype=torch.float32, device=device).unsqueeze(1), + mtp_collector=ModelMtpOutputCollector(spec_hidden=extend_hidden), + ) + return ModelOutput( + logits=torch.tensor([[40.0], [41.0]], device=device), + mtp_collector=ModelMtpOutputCollector(spec_hidden=decode_hidden), + ) + + backend = SimpleNamespace( + draft_models=[SimpleNamespace(forward=forward)], + _gen_argmax_token_ids=lambda output: output.logits[:, 0].long(), + _gen_argmax_token_ids_and_prob=lambda output: ( + output.logits[:, 0].long(), + output.logits[:, 0] / 100, + ), + ) + proposer = EagleWithAttProposer(backend=backend, enable_dynmaic_mtp=True) + target_input = SimpleNamespace( + is_prefill=False, + batch_size=6, + total_token_num=66, + max_kv_seq_len=22, + input_ids=original_input_ids, + b_req_idx=torch.tensor([7, 7, 7, 9, 9, 9], dtype=torch.int32, device=device), + b_mtp_index=torch.tensor([0, 1, 2, 0, 1, 2], dtype=torch.int32, device=device), + b_seq_len=torch.tensor([10, 11, 12, 20, 21, 22], dtype=torch.int32, device=device), + b_position_delta=torch.tensor([0, 1, 2, 3, 4, 5], dtype=torch.int32, device=device), + b_shared_seq_len=torch.tensor([8, 8, 8, 6, 6, 6], dtype=torch.int32, device=device), + b_shared_radix_node_id=torch.tensor([70, 70, 70, 90, 90, 90], dtype=torch.int64, device=device), + multimodal_params=[{"images": [], "audios": []} for _ in range(6)], + mtp_draft_input_hiddens=None, + ) + proposal = proposer.propose_next( + target_model_input=target_input, + target_model_output=SimpleNamespace(mtp_collector=SimpleNamespace(spec_hidden=target_hidden)), + target_next_token_ids=original_input_ids, + b_req_mtp_start_loc=torch.tensor([0, 3], dtype=torch.int32, device=device), + draft_step=2, + accept_len=torch.tensor([3, 2], dtype=torch.int32, device=device), + ) + + assert isinstance(proposal, EagleSpecProposal) + torch.testing.assert_close(proposal.token_ids, torch.tensor([[32, 40], [34, 41]], device=device)) + torch.testing.assert_close(proposal.schedule_scores, torch.tensor([[0.32, 0.40], [0.34, 0.41]], device=device)) + assert len(draft_calls) == 2 + assert draft_calls[0]["model_input"] is not target_input + assert draft_calls[0]["batch_size"] == 6 + torch.testing.assert_close(draft_calls[0]["input_ids"], original_input_ids) + torch.testing.assert_close(draft_calls[0]["draft_hidden"], target_hidden) + assert draft_calls[1]["batch_size"] == 2 + torch.testing.assert_close(draft_calls[1]["input_ids"], torch.tensor([32, 34], device=device)) + torch.testing.assert_close( + draft_calls[1]["draft_hidden"], + extend_hidden.index_select(0, torch.tensor([2, 4], device=device)), + ) + torch.testing.assert_close(draft_calls[1]["b_req_idx"], torch.tensor([7, 9], dtype=torch.int32, device=device)) + torch.testing.assert_close(draft_calls[1]["b_mtp_index"], torch.zeros(2, dtype=torch.int32, device=device)) + torch.testing.assert_close(draft_calls[1]["b_seq_len"], torch.tensor([13, 22], dtype=torch.int32, device=device)) + assert draft_calls[1]["max_kv_seq_len"] == 23 + assert draft_calls[1]["total_token_num"] == 46 + assert target_input.input_ids is original_input_ids + assert target_input.mtp_draft_input_hiddens is None diff --git a/unit_tests/server/router/model_infer/mtp_speculative/test_planner.py b/unit_tests/server/router/model_infer/mtp_speculative/test_planner.py new file mode 100644 index 0000000000..2e488ebb30 --- /dev/null +++ b/unit_tests/server/router/model_infer/mtp_speculative/test_planner.py @@ -0,0 +1,964 @@ +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +from lightllm.server.router.model_infer.mtp_speculative.engine import SpecEngine +from lightllm.server.router.model_infer.mtp_speculative.dp_overlap_proposers import ( + build_dp_overlap_spec_proposer, +) +from lightllm.server.router.model_infer.mtp_speculative.dp_overlap_proposers.base import ( + BaseDpOverlapProposer, +) +from lightllm.server.router.model_infer.mtp_speculative.dp_overlap_proposers.eagle3 import ( + DpOverlapEagle3Proposer, +) +from lightllm.server.router.model_infer.mtp_speculative.dp_overlap_proposers.eagle_no_att import ( + DpOverlapEagleNoAttProposer, +) +from lightllm.server.router.model_infer.mtp_speculative.dp_overlap_proposers.eagle_with_att import ( + DpOverlapEagleWithAttProposer, +) +from lightllm.server.router.model_infer.mtp_speculative.dp_overlap_proposers.vanilla_no_att import ( + DpOverlapVanillaNoAttProposer, +) +from lightllm.server.router.model_infer.mtp_speculative.dp_overlap_proposers.vanilla_with_att import ( + DpOverlapVanillaWithAttProposer, +) +from lightllm.server.router.model_infer.mtp_speculative.planner import ( + BaseMtpPlanner, + DSparkPlanner, + FixedSpecPlanner, + LightSpecPlanner, + SpecDecodePlan, +) +from lightllm.server.router.model_infer.mtp_speculative.planner.base import ( + _InferCostMsTable, +) +from lightllm.server.router.model_infer.mtp_speculative.proposers import ( + build_spec_proposer, +) +from lightllm.server.router.model_infer.mtp_speculative.proposers.base import ( + BaseSpecProposer, + SpecProposal, +) +from lightllm.server.router.model_infer.mtp_speculative.proposers.dflash import ( + DFlashProposer, +) +from lightllm.server.router.model_infer.mtp_speculative.proposers.dflash2 import ( + DFlash2Proposer, +) +from lightllm.server.router.model_infer.mtp_speculative.proposers.dspark import ( + DSparkProposer, +) +from lightllm.server.router.model_infer.mtp_speculative.proposers.eagle3 import ( + Eagle3Proposer, +) +from lightllm.server.router.model_infer.mtp_speculative.proposers.eagle_no_att import ( + EagleNoAttProposer, +) +from lightllm.server.router.model_infer.mtp_speculative.proposers.eagle_with_att import ( + EagleWithAttProposer, +) +from lightllm.server.router.model_infer.mtp_speculative.proposers.proposal_type import ( + DFlashSpecProposal, + DSparkSpecProposal, + EagleSpecProposal, + VanillaSpecProposal, +) +from lightllm.server.router.model_infer.mtp_speculative.proposers.vanilla_no_att import ( + VanillaNoAttProposer, +) +from lightllm.server.router.model_infer.mtp_speculative.proposers.vanilla_with_att import ( + VanillaWithAttProposer, +) +from lightllm.server.router.model_infer.mtp_speculative import utils as mtp_utils + + +def build_lightspec_planner( + max_draft_step: int = 3, + spec_mode: str = "vanilla_with_att", + block_size: int = 3, + enable_decode_microbatch_overlap: bool = False, +): + backend = SimpleNamespace( + args=SimpleNamespace(dp=1), + enable_decode_microbatch_overlap=enable_decode_microbatch_overlap, + max_draft_step=max_draft_step, + model=SimpleNamespace(graph=None), + draft_models=[SimpleNamespace(block_size=block_size, graph=None)], + ) + return LightSpecPlanner( + spec_mode=spec_mode, + backend=backend, + ) + + +def build_dspark_planner(max_draft_step: int = 3, block_size: int = 3): + backend = SimpleNamespace( + max_draft_step=max_draft_step, + model=SimpleNamespace(graph=None), + draft_models=[SimpleNamespace(block_size=block_size, graph=None)], + ) + return DSparkPlanner(backend=backend) + + +def build_decode_reqs(req_num: int, req_num_with_proposals: int | None = None): + if req_num_with_proposals is None: + req_num_with_proposals = req_num + return [SimpleNamespace(cur_output_len=2)] * req_num_with_proposals + [SimpleNamespace(cur_output_len=1)] * ( + req_num - req_num_with_proposals + ) + + +def build_planner(spec_mode: str, enable_dynmaic_mtp: bool = True): + engine = SpecEngine.__new__(SpecEngine) + engine.backend = SimpleNamespace( + args=SimpleNamespace(dp=1), + enable_decode_microbatch_overlap=False, + max_draft_step=3, + model=SimpleNamespace(graph=None), + draft_models=[SimpleNamespace(block_size=3, graph=None)], + ) + return engine._build_mtp_planner( + spec_mode=spec_mode, + enable_dynmaic_mtp=enable_dynmaic_mtp, + ) + + +def test_fixed_planner_returns_static_plan(): + planner = FixedSpecPlanner(max_draft_step=3) + plan = planner.plan(decode_reqs=build_decode_reqs(4), origin_batch_size=16) + + assert isinstance(planner, BaseMtpPlanner) + assert plan.origin_batch_size == 16 + assert plan.dynamic_batch_size == plan.origin_batch_size + assert plan.draft_step == plan.pre_draft_step == 3 + assert not plan.skip_verify_sync + + +def test_common_engine_accepts_empty_fixed_decode_batch(): + engine = SpecEngine.__new__(SpecEngine) + engine.planner = FixedSpecPlanner(max_draft_step=3) + + plan = engine.plan_decode(model_input=SimpleNamespace(batch_size=0), decode_reqs=[]) + + assert plan == SpecDecodePlan( + origin_batch_size=0, + dynamic_batch_size=0, + draft_step=3, + pre_draft_step=3, + ) + + +def test_common_engine_delegates_empty_dp_batch_to_lightspec_planner(): + engine = SpecEngine.__new__(SpecEngine) + engine.backend = SimpleNamespace(max_draft_step=3, dp_size=2) + engine.planner = build_lightspec_planner(spec_mode="eagle3") + engine.planner.backend.args.dp = 2 + engine.planner.pre_draft_step = 2 + + plan = engine.plan_decode(model_input=SimpleNamespace(batch_size=0), decode_reqs=[]) + + assert plan == SpecDecodePlan( + origin_batch_size=0, + dynamic_batch_size=0, + draft_step=1, + pre_draft_step=2, + ) + + +def test_spec_engine_only_exposes_planning_and_proposal_interfaces(): + public_methods = { + name for name, value in SpecEngine.__dict__.items() if callable(value) and not name.startswith("_") + } + + assert public_methods == { + "fill_draft_model_kv_state", + "plan_decode", + "prepare_decode_model_input", + "propose_next", + "update_planner_statics", + } + + +def test_mode_proposals_own_their_schedule_metadata(): + assert "schedule_scores" not in SpecProposal.__dataclass_fields__ + assert "schedule_scores_cpu" not in SpecProposal.__dataclass_fields__ + + for proposal_type in ( + VanillaSpecProposal, + EagleSpecProposal, + DFlashSpecProposal, + DSparkSpecProposal, + ): + assert "schedule_scores" in proposal_type.__dataclass_fields__ + assert "schedule_scores_cpu" not in VanillaSpecProposal.__dataclass_fields__ + assert "schedule_scores_cpu" not in EagleSpecProposal.__dataclass_fields__ + assert "schedule_scores_cpu" not in DFlashSpecProposal.__dataclass_fields__ + assert "schedule_scores_cpu" in DSparkSpecProposal.__dataclass_fields__ + + +def test_scatter_mtp_next_tokens_consumes_mode_proposal(monkeypatch): + scatter_args = {} + monkeypatch.setattr( + mtp_utils, + "mtp_scatter_next_token_ids", + lambda **kwargs: scatter_args.update(kwargs), + ) + req_to_next_token_ids = torch.empty((4, 2), dtype=torch.int64) + req_to_next_token_scores = torch.empty((4, 1), dtype=torch.float32) + backend = SimpleNamespace( + model=SimpleNamespace( + req_manager=SimpleNamespace( + req_sampling_params_manager=SimpleNamespace( + req_to_bitmask_enabled=None, + req_to_next_token_ids=req_to_next_token_ids, + req_to_next_token_scores=req_to_next_token_scores, + ) + ) + ) + ) + proposal = DFlashSpecProposal( + token_ids=torch.arange(2, dtype=torch.int64).view(2, 1), + schedule_scores=torch.arange(2, dtype=torch.float32).view(2, 1), + ) + next_token_ids = torch.tensor([10, 11], dtype=torch.int64) + + mtp_utils.scatter_mtp_next_tokens( + backend=backend, + proposal=proposal, + target_next_token_ids=next_token_ids, + b_req_mtp_start_loc=torch.tensor([0, 1], dtype=torch.int32), + b_req_idx=torch.tensor([0, 1], dtype=torch.int32), + mtp_accept_len=torch.ones(2, dtype=torch.int32), + ) + + assert scatter_args["req_to_next_token_ids"] is req_to_next_token_ids + assert scatter_args["req_to_next_token_scores"] is req_to_next_token_scores + assert torch.equal(scatter_args["target_next_token_ids"], next_token_ids) + assert torch.equal(scatter_args["draft_token_ids"], proposal.token_ids) + assert torch.equal(scatter_args["schedule_scores"], proposal.schedule_scores) + + +def test_scatter_mtp_next_tokens_ignores_empty_schedule_scores(monkeypatch): + scatter_args = {} + monkeypatch.setattr( + mtp_utils, + "mtp_scatter_next_token_ids", + lambda **kwargs: scatter_args.update(kwargs), + ) + backend = SimpleNamespace( + model=SimpleNamespace( + req_manager=SimpleNamespace( + req_sampling_params_manager=SimpleNamespace( + req_to_bitmask_enabled=None, + req_to_next_token_ids=torch.empty((4, 2), dtype=torch.int64), + req_to_next_token_scores=torch.empty((4, 2), dtype=torch.float32), + ) + ) + ) + ) + proposal = VanillaSpecProposal( + token_ids=torch.empty((2, 0), dtype=torch.int64), + schedule_scores=torch.empty((2, 0), dtype=torch.float32), + ) + + mtp_utils.scatter_mtp_next_tokens( + backend=backend, + proposal=proposal, + target_next_token_ids=torch.tensor([10, 11], dtype=torch.int64), + b_req_mtp_start_loc=torch.tensor([0, 1], dtype=torch.int32), + b_req_idx=torch.tensor([0, 1], dtype=torch.int32), + mtp_accept_len=torch.ones(2, dtype=torch.int32), + ) + + assert scatter_args["req_to_next_token_scores"] is None + assert scatter_args["schedule_scores"] is None + + +def test_infer_cost_candidates_include_feasible_boundaries(): + costs = _InferCostMsTable() + costs.update(batch_size=4, infer_cost_ms=1.0) + costs.update(batch_size=8, infer_cost_ms=2.0) + + assert costs.get_batch_size_keys_between(5, 10) == [5, 8, 10] + + +def test_engine_routes_only_dspark_to_the_confidence_planner(): + fixed_planner = build_planner("eagle3", enable_dynmaic_mtp=False) + dspark_planner = build_planner("dspark") + assert isinstance(fixed_planner, FixedSpecPlanner) + assert isinstance(dspark_planner, DSparkPlanner) + assert isinstance(fixed_planner, BaseMtpPlanner) + assert isinstance(dspark_planner, BaseMtpPlanner) + + vanilla_planner = build_planner("vanilla_with_att") + assert isinstance(vanilla_planner, LightSpecPlanner) + assert isinstance(vanilla_planner, BaseMtpPlanner) + assert vanilla_planner.draft_steps == (3,) + + vanilla_no_att_planner = build_planner("vanilla_no_att") + assert vanilla_no_att_planner.draft_steps == (0, 1, 2, 3) + + eagle_no_att_planner = build_planner("eagle_no_att") + assert eagle_no_att_planner.draft_steps == (0, 1, 2, 3) + + dflash_planner = build_planner("dflash") + assert isinstance(dflash_planner, LightSpecPlanner) + assert dflash_planner.draft_steps == (3,) + + dflash2_planner = build_planner("dflash2", enable_dynmaic_mtp=False) + assert isinstance(dflash2_planner, FixedSpecPlanner) + dflash2_dynamic_planner = build_planner("dflash2") + assert isinstance(dflash2_dynamic_planner, LightSpecPlanner) + assert dflash2_dynamic_planner.draft_steps == (3,) + + eagle_planner = build_planner("eagle3") + assert isinstance(eagle_planner, LightSpecPlanner) + assert eagle_planner.draft_steps == (1, 2, 3) + + eagle_with_att_planner = build_planner("eagle_with_att") + assert eagle_with_att_planner.draft_steps == (1, 2, 3) + + +def test_dynamic_planner_registers_cuda_graph_costs_from_backend(): + backend = SimpleNamespace( + args=SimpleNamespace(dp=1), + enable_decode_microbatch_overlap=False, + max_draft_step=3, + model=SimpleNamespace(graph=SimpleNamespace(infer_cost_ms_by_batch_size={4: 1.2})), + draft_models=[ + SimpleNamespace( + block_size=3, + graph=SimpleNamespace(infer_cost_ms_by_batch_size={4: 0.3}), + ) + ], + ) + + planner = LightSpecPlanner(spec_mode="vanilla_with_att", backend=backend) + + assert planner.target_infer_costs.estimate(4) == 1.2 + assert planner.draft_infer_costs.estimate(4) == 0.3 + + +def test_overlap_lightspec_consumes_combined_cuda_graph_batch_size(): + backend = SimpleNamespace( + args=SimpleNamespace(dp=1), + enable_decode_microbatch_overlap=True, + max_draft_step=3, + model=SimpleNamespace(graph=SimpleNamespace(infer_cost_ms_by_batch_size={8: 1.2})), + draft_models=[ + SimpleNamespace( + block_size=3, + graph=SimpleNamespace(infer_cost_ms_by_batch_size={8: 0.3}), + ) + ], + ) + + planner = LightSpecPlanner(spec_mode="vanilla_with_att", backend=backend) + + assert planner.target_infer_costs.estimate(7) == 1.2 + assert planner.target_infer_costs.estimate(8) == 1.2 + assert planner.draft_infer_costs.estimate(7) == 0.3 + assert planner.draft_infer_costs.estimate(8) == 0.3 + + +def test_each_mode_proposer_inherits_its_expected_implementation_base(): + proposer_types = ( + VanillaWithAttProposer, + VanillaNoAttProposer, + EagleWithAttProposer, + EagleNoAttProposer, + DFlashProposer, + DSparkProposer, + ) + dp_overlap_proposer_types = ( + DpOverlapVanillaWithAttProposer, + DpOverlapVanillaNoAttProposer, + DpOverlapEagleWithAttProposer, + DpOverlapEagleNoAttProposer, + ) + + for proposer_type in proposer_types: + assert proposer_type.__bases__ == (BaseSpecProposer,) + assert Eagle3Proposer.__bases__ == (EagleWithAttProposer,) + assert DFlash2Proposer.__bases__ == (DFlashProposer,) + for proposer_type in dp_overlap_proposer_types: + assert proposer_type.__bases__ == (BaseDpOverlapProposer,) + assert DpOverlapEagle3Proposer.__bases__ == (DpOverlapEagleWithAttProposer,) + + +def test_each_mtp_mode_builds_its_own_proposer(): + backend = SimpleNamespace() + proposer_types = { + "vanilla_with_att": VanillaWithAttProposer, + "vanilla_no_att": VanillaNoAttProposer, + "eagle_with_att": EagleWithAttProposer, + "eagle_no_att": EagleNoAttProposer, + "eagle3": Eagle3Proposer, + "dflash": DFlashProposer, + "dflash2": DFlash2Proposer, + "dspark": DSparkProposer, + } + + for spec_mode, proposer_type in proposer_types.items(): + proposer = build_spec_proposer(spec_mode=spec_mode, backend=backend, enable_dynmaic_mtp=False) + assert type(proposer) is proposer_type + + +def test_each_supported_dp_overlap_mtp_mode_builds_its_own_proposer(): + backend = SimpleNamespace() + proposer_types = { + "vanilla_with_att": DpOverlapVanillaWithAttProposer, + "vanilla_no_att": DpOverlapVanillaNoAttProposer, + "eagle_with_att": DpOverlapEagleWithAttProposer, + "eagle_no_att": DpOverlapEagleNoAttProposer, + "eagle3": DpOverlapEagle3Proposer, + } + + for spec_mode, proposer_type in proposer_types.items(): + proposer = build_dp_overlap_spec_proposer(spec_mode=spec_mode, backend=backend, enable_dynmaic_mtp=False) + assert type(proposer) is proposer_type + assert isinstance(proposer, BaseDpOverlapProposer) + + +def test_eagle_proposer_skips_draft_forward_for_zero_steps(): + proposer = EagleNoAttProposer( + backend=SimpleNamespace(draft_models=[]), + enable_dynmaic_mtp=True, + ) + next_token_ids = torch.tensor([10, 11], dtype=torch.int64) + + proposal = proposer.propose_next( + target_model_input=None, + target_model_output=None, + target_next_token_ids=next_token_ids, + b_req_mtp_start_loc=torch.tensor([0, 1], dtype=torch.int32), + draft_step=0, + ) + + assert isinstance(proposal, EagleSpecProposal) + assert proposal.token_ids.shape == (2, 0) + assert proposal.schedule_scores.shape == (2, 0) + + +def test_records_request_mtp_metrics_in_one_pass(): + class MetricReq: + def __init__(self, req_idx: int, mtp_step: int): + self.req_idx = req_idx + self.mtp_step = mtp_step + self.accepted = 0 + self.verified = 0 + self.verify_steps = 0 + + def update_mtp_accepted_token_num(self, accept_token_num: int): + self.accepted += accept_token_num + + def update_mtp_verify_token_num(self, verify_token_num: int): + self.verified += verify_token_num + + def update_mtp_verify_step_num(self, verify_step_num: int): + self.verify_steps += verify_step_num + + backend = SimpleNamespace(is_master_in_dp=True) + req0 = MetricReq(req_idx=10, mtp_step=3) + req1 = MetricReq(req_idx=11, mtp_step=3) + + mtp_utils.record_request_mtp_metrics( + backend=backend, + decode_reqs=[req0, req1], + accept_lengths_cpu=torch.tensor([2, 1], dtype=torch.int32), + verify_run_reqs=[req0, req0, req1], + ) + + assert (req0.accepted, req0.verified, req0.verify_steps) == (1, 2, 1) + assert (req1.accepted, req1.verified, req1.verify_steps) == (0, 1, 1) + + fixed_req = MetricReq(req_idx=12, mtp_step=3) + mtp_utils.record_request_mtp_metrics( + backend=backend, + decode_reqs=[fixed_req], + accept_lengths_cpu=torch.tensor([3], dtype=torch.int32), + verify_run_reqs=[fixed_req] * 4, + ) + assert (fixed_req.accepted, fixed_req.verified, fixed_req.verify_steps) == (2, 4, 1) + + +def test_lightspec_stays_full_width_until_costs_are_profiled(): + plan = build_lightspec_planner().plan( + decode_reqs=build_decode_reqs(2), + origin_batch_size=8, + ) + + assert plan.origin_batch_size == 8 + assert plan.dynamic_batch_size == 8 + assert plan.draft_step == plan.pre_draft_step == 3 + + +def test_lightspec_collects_full_width_progress_before_adapting(): + planner = build_lightspec_planner(spec_mode="eagle_with_att") + for batch_size in (2, 4, 8): + planner.target_infer_costs.update(batch_size=batch_size, infer_cost_ms=float(batch_size)) + planner.draft_infer_costs.update(batch_size=batch_size, infer_cost_ms=float(batch_size)) + + plan = planner.plan(decode_reqs=build_decode_reqs(2), origin_batch_size=8) + + assert plan.dynamic_batch_size == 8 + assert plan.draft_step == plan.pre_draft_step == 3 + + +def test_lightspec_selects_eagle_draft_depth_and_verify_capacity(): + planner = build_lightspec_planner(spec_mode="eagle_with_att") + for batch_size, target_cost in ((2, 1.0), (4, 1.1), (8, 10.0)): + planner.target_infer_costs.update(batch_size=batch_size, infer_cost_ms=target_cost) + for batch_size, draft_cost in ((2, 0.1), (4, 0.1), (8, 0.2)): + planner.draft_infer_costs.update(batch_size=batch_size, infer_cost_ms=draft_cost) + for _ in range(8): + planner._update_verified_batch( + accept_lengths=[4, 4], + req_num=2, + dynamic_batch_size=8, + verified_draft_step=3, + ) + + plan = planner.plan(decode_reqs=build_decode_reqs(2), origin_batch_size=8) + + assert plan.dynamic_batch_size == 4 + assert plan.pre_draft_step == 3 + assert plan.draft_step == 1 + + +def test_lightspec_overlap_prices_combined_batch_size(): + planner = build_lightspec_planner( + spec_mode="eagle_with_att", + enable_decode_microbatch_overlap=True, + ) + planner.target_infer_costs.update(batch_size=4, infer_cost_ms=1.0) + planner.target_infer_costs.update(batch_size=8, infer_cost_ms=3.0) + planner.draft_infer_costs.update(batch_size=4, infer_cost_ms=0.25) + planner.draft_infer_costs.update(batch_size=8, infer_cost_ms=1.0) + + assert planner.target_infer_costs.estimate(7) == 3.0 + assert planner.target_infer_costs.get_batch_size_keys_between(4, 8) == [4, 8] + assert planner._get_cost_ms(req_num=4, dynamic_batch_size=8, draft_step=1) == 0.5 + + +def test_lightspec_overlap_selects_dynamic_draft_step(): + planner = build_lightspec_planner( + spec_mode="eagle_with_att", + enable_decode_microbatch_overlap=True, + ) + for batch_size, target_cost in ((2, 1.0), (4, 1.1), (8, 10.0)): + planner.target_infer_costs.update(batch_size=batch_size, infer_cost_ms=target_cost) + for batch_size, draft_cost in ((2, 0.1), (4, 0.1), (8, 0.2)): + planner.draft_infer_costs.update(batch_size=batch_size, infer_cost_ms=draft_cost) + planner._update_verified_batch( + accept_lengths=[4, 4], + req_num=2, + dynamic_batch_size=8, + verified_draft_step=3, + ) + + plan = planner.plan(decode_reqs=build_decode_reqs(2), origin_batch_size=8) + + assert plan.dynamic_batch_size == 4 + assert plan.pre_draft_step == 3 + assert plan.draft_step == 1 + + +def test_lightspec_keeps_vanilla_attention_chained_depth_fixed(): + planner = build_lightspec_planner(spec_mode="vanilla_with_att") + for batch_size, target_cost in ((2, 1.0), (4, 1.1), (8, 10.0)): + planner.target_infer_costs.update(batch_size=batch_size, infer_cost_ms=target_cost) + for batch_size, draft_cost in ((2, 0.1), (4, 0.1), (8, 0.2)): + planner.draft_infer_costs.update(batch_size=batch_size, infer_cost_ms=draft_cost) + planner._update_verified_batch( + accept_lengths=[4, 4], + req_num=2, + dynamic_batch_size=8, + verified_draft_step=3, + ) + + plan = planner.plan(decode_reqs=build_decode_reqs(2), origin_batch_size=8) + + assert planner.draft_steps == (3,) + assert plan.draft_step == plan.pre_draft_step == 3 + + +@pytest.mark.parametrize("spec_mode,block_size", [("dflash", 7), ("dflash2", 8)]) +def test_lightspec_compacts_block_verify_without_changing_draft_shape(spec_mode, block_size): + planner = build_lightspec_planner( + max_draft_step=7, + spec_mode=spec_mode, + block_size=block_size, + ) + for batch_size, target_cost in ((2, 1.0), (4, 1.1), (8, 3.0), (16, 8.0)): + planner.target_infer_costs.update(batch_size=batch_size, infer_cost_ms=target_cost) + planner.draft_infer_costs.update(batch_size=14, infer_cost_ms=0.5) + for _ in range(8): + planner._update_verified_batch( + accept_lengths=[3, 3], + req_num=2, + dynamic_batch_size=16, + verified_draft_step=7, + ) + + plan = planner.plan(decode_reqs=build_decode_reqs(2), origin_batch_size=16) + + assert plan.dynamic_batch_size < 16 + assert plan.draft_step == plan.pre_draft_step == 7 + + +def test_lightspec_bounds_verify_to_existing_proposals(): + planner = build_lightspec_planner() + + cold_start = planner.plan(decode_reqs=build_decode_reqs(2, 0), origin_batch_size=8) + mixed_batch = planner.plan(decode_reqs=build_decode_reqs(2, 1), origin_batch_size=8) + ready_batch = planner.plan(decode_reqs=build_decode_reqs(2, 2), origin_batch_size=8) + + assert cold_start.dynamic_batch_size == 2 + assert mixed_batch.dynamic_batch_size == 5 + assert ready_batch.dynamic_batch_size == 8 + assert not cold_start.all_reqs_have_proposals + assert not mixed_batch.all_reqs_have_proposals + assert ready_batch.all_reqs_have_proposals + + +def test_engine_lets_planner_count_requests_with_a_previous_proposal(): + engine = SpecEngine.__new__(SpecEngine) + engine.planner = build_lightspec_planner() + model_input = SimpleNamespace(batch_size=8) + + mixed_plan = engine.plan_decode( + model_input=model_input, + decode_reqs=[ + SimpleNamespace(cur_output_len=1), + SimpleNamespace(cur_output_len=2), + ], + ) + ready_plan = engine.plan_decode( + model_input=model_input, + decode_reqs=[ + SimpleNamespace(cur_output_len=2), + SimpleNamespace(cur_output_len=2), + ], + ) + + assert mixed_plan.dynamic_batch_size == 5 + assert not mixed_plan.all_reqs_have_proposals + assert ready_plan.dynamic_batch_size == 8 + assert ready_plan.all_reqs_have_proposals + + +def test_dynamic_prepare_delegates_row_compaction(monkeypatch): + from lightllm.common.basemodel.triton_kernel import dynamic_mtp_utils + from lightllm.server.router.model_infer.mtp_speculative import ( + engine as engine_module, + ) + + selected_row_mask = torch.tensor([1, 0, 1, 0], dtype=torch.int32) + monkeypatch.setattr( + dynamic_mtp_utils, + "prepare_dynamic_mtp_model_input", + lambda model_input, **kwargs: (model_input, selected_row_mask), + ) + async_mask = object() + monkeypatch.setattr( + engine_module.g_pin_mem_manager, + "async_copy_from_gpu_tensor_with_event", + lambda **kwargs: async_mask, + ) + + engine = SpecEngine.__new__(SpecEngine) + engine.backend = SimpleNamespace( + model=SimpleNamespace( + req_manager=SimpleNamespace( + req_sampling_params_manager=SimpleNamespace(req_to_next_token_scores=torch.empty(0)) + ) + ) + ) + model_input = SimpleNamespace( + batch_size=4, + ) + plan = SpecDecodePlan(origin_batch_size=4, dynamic_batch_size=2, draft_step=1, pre_draft_step=1) + + compacted_input, selected_mask_cpu = engine.prepare_decode_model_input( + model_input=model_input, + req_num=1, + plan=plan, + ) + + assert compacted_input is model_input + assert selected_mask_cpu is async_mask + + +def test_lightspec_eagle_draft_always_keeps_the_extend_candidate(): + planner = build_lightspec_planner(spec_mode="eagle_with_att") + for batch_size in (2, 4, 8): + planner.target_infer_costs.update(batch_size=batch_size, infer_cost_ms=float(batch_size)) + planner.draft_infer_costs.update(batch_size=batch_size, infer_cost_ms=float(batch_size)) + planner._update_verified_batch( + accept_lengths=[2, 2], + req_num=2, + dynamic_batch_size=8, + verified_draft_step=3, + ) + + plan = planner.plan(decode_reqs=build_decode_reqs(2), origin_batch_size=8) + + assert planner.draft_steps == (1, 2, 3) + assert plan.draft_step >= 1 + + +def test_vanilla_with_attention_planner_prices_extend_then_normal_batches(): + planner = build_lightspec_planner(spec_mode="vanilla_with_att") + planner.draft_infer_costs.update(batch_size=2, infer_cost_ms=0.1) + + draft_cost_ms = planner._get_draft_cost_ms( + req_num=4, + verify_batch_size=8, + draft_step=3, + ) + + assert np.isclose(draft_cost_ms, 0.8) + with pytest.raises(AssertionError, match="requires draft_step to be greater than 0"): + planner._get_draft_cost_ms(req_num=4, verify_batch_size=8, draft_step=0) + + +def test_no_attention_planner_prices_only_normal_request_batches(): + for spec_mode in ("vanilla_no_att", "eagle_no_att"): + planner = build_lightspec_planner(spec_mode=spec_mode) + planner.draft_infer_costs.update(batch_size=2, infer_cost_ms=0.1) + + draft_cost_ms = planner._get_draft_cost_ms( + req_num=4, + verify_batch_size=8, + draft_step=3, + ) + + assert np.isclose(draft_cost_ms, 0.6) + + +def test_eagle_planner_prices_extend_and_decode_separately(): + planner = build_lightspec_planner(spec_mode="eagle_with_att") + planner.draft_infer_costs.update(batch_size=2, infer_cost_ms=0.25) + + draft_cost_ms = planner._get_draft_cost_ms( + req_num=2, + verify_batch_size=8, + draft_step=3, + ) + + assert draft_cost_ms == 1.5 + + +def test_autoregressive_eagle_planner_prices_extend_and_decode_rows(): + for spec_mode in ("eagle_with_att", "eagle3"): + planner = build_lightspec_planner( + max_draft_step=7, + spec_mode=spec_mode, + ) + for batch_size in (2, 4, 8, 16): + planner.draft_infer_costs.update(batch_size=batch_size, infer_cost_ms=float(batch_size)) + + draft_cost_ms = planner._get_draft_cost_ms( + req_num=8, + verify_batch_size=16, + draft_step=7, + ) + + assert draft_cost_ms == 64.0 + + +@pytest.mark.parametrize("spec_mode,block_size,expected_cost", [("dflash", 7, 1.5), ("dflash2", 8, 1.6)]) +def test_block_planner_prices_commit_and_complete_block(spec_mode, block_size, expected_cost): + planner = build_lightspec_planner( + max_draft_step=7, + spec_mode=spec_mode, + block_size=block_size, + ) + planner.draft_infer_costs.update(batch_size=8, infer_cost_ms=0.4) + + draft_cost_ms = planner._get_draft_cost_ms( + req_num=2, + verify_batch_size=16, + draft_step=7, + ) + + assert np.isclose(draft_cost_ms, expected_cost) + + +def test_dspark_planner_prices_commit_and_complete_block(): + planner = build_dspark_planner(max_draft_step=3, block_size=3) + planner.draft_infer_costs.update(batch_size=8, infer_cost_ms=0.8) + + draft_cost_ms = planner._get_draft_cost_ms( + req_num=4, + verify_batch_size=8, + draft_step=3, + ) + + assert np.isclose(draft_cost_ms, 2.0) + + +def test_lightspec_records_one_batch_observation_per_configuration(): + planner = build_lightspec_planner() + + planner._update_verified_batch( + accept_lengths=[2, 2], + req_num=2, + dynamic_batch_size=4, + verified_draft_step=1, + ) + planner._update_verified_batch( + accept_lengths=[1, 1], + req_num=2, + dynamic_batch_size=4, + verified_draft_step=3, + ) + + assert planner.progress_ema_by_config[(2, 4, 1)].get() == 1.0 + assert planner.progress_ema_by_config[(2, 4, 3)].get() == 0.5 + assert planner.progress_ema_by_config[(2, 4, 1)].get_count() == 1 + assert planner.progress_ema_by_config[(2, 4, 3)].get_count() == 1 + + +def test_lightspec_high_concurrency_does_not_multiply_ema_updates(): + planner = build_lightspec_planner() + + planner._update_verified_batch( + accept_lengths=[1] * 128, + req_num=128, + dynamic_batch_size=128, + verified_draft_step=1, + ) + + assert planner.progress_ema_by_config[(128, 128, 1)].get_count() == 1 + + +def test_lightspec_estimates_unseen_shapes_from_prefix_survival(): + planner = build_lightspec_planner() + planner._update_verified_batch( + accept_lengths=[2, 1], + req_num=2, + dynamic_batch_size=8, + verified_draft_step=3, + ) + + assert planner._estimate_progress(req_num=2, dynamic_batch_size=4, draft_step=3) == 0.75 + assert planner._estimate_progress(req_num=2, dynamic_batch_size=2, draft_step=3) == 1.0 + assert planner._estimate_progress(req_num=2, dynamic_batch_size=4, draft_step=2) == 0.75 + + +def test_lightspec_does_not_transfer_deep_progress_to_short_drafts(): + planner = build_lightspec_planner() + planner._update_verified_batch( + accept_lengths=[4, 2], + req_num=2, + dynamic_batch_size=8, + verified_draft_step=3, + ) + + assert np.isclose(planner._estimate_progress(req_num=2, dynamic_batch_size=6, draft_step=2), 5 / 6) + + +def test_lightspec_short_current_proposal_can_recover_to_a_deeper_draft(): + planner = build_lightspec_planner(spec_mode="eagle_with_att") + for batch_size, target_cost in ((2, 1.0), (4, 1.1), (6, 1.2), (8, 1.3)): + planner.target_infer_costs.update(batch_size=batch_size, infer_cost_ms=target_cost) + planner.draft_infer_costs.update(batch_size=batch_size, infer_cost_ms=0.01) + planner._update_verified_batch( + accept_lengths=[4, 4], + req_num=2, + dynamic_batch_size=8, + verified_draft_step=3, + ) + planner.pre_draft_step = 1 + + plan = planner.plan(decode_reqs=build_decode_reqs(2), origin_batch_size=8) + + assert plan.dynamic_batch_size <= 4 + assert plan.pre_draft_step == 1 + assert plan.draft_step == 3 + + +def test_engine_skips_feedback_for_a_mixed_proposal_batch(): + engine = SpecEngine.__new__(SpecEngine) + engine.planner = build_lightspec_planner() + plan = SpecDecodePlan( + origin_batch_size=8, + dynamic_batch_size=5, + draft_step=3, + pre_draft_step=3, + all_reqs_have_proposals=False, + ) + + engine.update_planner_statics( + plan=plan, + proposal=SpecProposal( + token_ids=torch.empty((0,), dtype=torch.int64), + ), + req_num=2, + accept_lengths_cpu=torch.tensor([1, 2], dtype=torch.int32), + ) + + assert not engine.planner.progress_ema_by_config + + +def test_dspark_applies_confidence_capacity_after_two_step_delay(): + planner = build_dspark_planner() + for batch_size, target_cost in ((2, 1.0), (4, 1.1), (8, 10.0)): + planner.target_infer_costs.update(batch_size=batch_size, infer_cost_ms=target_cost) + planner.draft_infer_costs.update(batch_size=6, infer_cost_ms=0.5) + confidence_probs = np.asarray([[0.9, 0.9, 0.9]] * 2, dtype=np.float64) + plan = SpecDecodePlan(origin_batch_size=8, dynamic_batch_size=8, draft_step=3, pre_draft_step=3) + proposal = DSparkSpecProposal( + token_ids=torch.empty((0,), dtype=torch.int64), + schedule_scores_cpu=torch.from_numpy(confidence_probs), + ) + engine = SpecEngine.__new__(SpecEngine) + engine.planner = planner + + engine.update_planner_statics( + plan=plan, + proposal=proposal, + req_num=2, + accept_lengths_cpu=torch.tensor([1, 1], dtype=torch.int32), + ) + first_plan = planner.plan(decode_reqs=build_decode_reqs(2), origin_batch_size=8) + engine.update_planner_statics( + plan=plan, + proposal=proposal, + req_num=2, + accept_lengths_cpu=torch.tensor([1, 1], dtype=torch.int32), + ) + second_plan = planner.plan(decode_reqs=build_decode_reqs(2), origin_batch_size=8) + + assert first_plan.dynamic_batch_size == 8 + assert second_plan.dynamic_batch_size == 4 + assert second_plan.draft_step == second_plan.pre_draft_step == 3 + + +def test_dspark_uses_confidence_scores_without_acceptance_ema(): + planner = build_dspark_planner() + for batch_size, target_cost in ((2, 1.0), (4, 1.2), (8, 10.0)): + planner.target_infer_costs.update(batch_size=batch_size, infer_cost_ms=target_cost) + planner.draft_infer_costs.update(batch_size=6, infer_cost_ms=0.5) + + low_confidence = np.cumprod(np.full((2, 3), 0.1), axis=1) + high_confidence = np.cumprod(np.full((2, 3), 0.9), axis=1) + + assert planner._select_dynamic_batch_size_from_survival_scores(2, low_confidence) == 2 + assert planner._select_dynamic_batch_size_from_survival_scores(2, high_confidence) == 4 + + +def test_topk_prefix_sums_only_computes_requested_counts(): + result = DSparkPlanner._topk_prefix_sums( + values=np.asarray([0.1, 0.9, 0.4, 0.7]), + counts=[0, 2, 4], + ) + + assert set(result) == {0, 2, 4} + np.testing.assert_allclose([result[0], result[2], result[4]], [0.0, 1.6, 2.1]) diff --git a/unit_tests/server/router/model_infer/mtp_speculative/test_vanilla_no_att.py b/unit_tests/server/router/model_infer/mtp_speculative/test_vanilla_no_att.py new file mode 100644 index 0000000000..bedf17a5b0 --- /dev/null +++ b/unit_tests/server/router/model_infer/mtp_speculative/test_vanilla_no_att.py @@ -0,0 +1,211 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.basemodel.triton_kernel.select_mtp_rows import select_accepted_tail_rows +from lightllm.server.router.model_infer.mtp_speculative.dp_overlap_proposers.vanilla_no_att import ( + DpOverlapVanillaNoAttProposer, +) +from lightllm.server.router.model_infer.mtp_speculative.proposers.vanilla_no_att import VanillaNoAttProposer + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_select_accepted_tail_rows_triton_matches_index_select(): + device = "cuda" + b_req_mtp_start_loc = torch.tensor([0, 3, 6], dtype=torch.int32, device=device) + accept_len = torch.tensor([2, 3, 2], dtype=torch.int32, device=device) + expected_rows = torch.tensor([1, 5, 7], dtype=torch.int64, device=device) + + # Exercise non-contiguous source strides as well as hidden widths spanning + # multiple Triton column blocks. + input_ids = torch.arange(16, dtype=torch.int64, device=device)[::2] + hidden = torch.arange(8 * 514, dtype=torch.float32, device=device).reshape(8, 514)[:, ::2] + b_req_idx = (torch.arange(16, dtype=torch.int32, device=device) + 10)[::2] + b_mtp_index = torch.arange(8, dtype=torch.int32, device=device) + b_seq_len = (torch.arange(16, dtype=torch.int32, device=device) + 20)[::2] + b_shared_seq_len = (torch.arange(16, dtype=torch.int32, device=device) + 30)[::2] + b_shared_radix_node_id = (torch.arange(16, dtype=torch.int64, device=device) + 1000)[::2] + b_position_delta = (torch.arange(16, dtype=torch.int32, device=device) - 8)[::2] + + selected = select_accepted_tail_rows( + b_req_mtp_start_loc=b_req_mtp_start_loc, + accept_len=accept_len, + input_ids=input_ids, + hidden=hidden, + b_req_idx=b_req_idx, + b_mtp_index=b_mtp_index, + b_seq_len=b_seq_len, + b_shared_seq_len=b_shared_seq_len, + b_shared_radix_node_id=b_shared_radix_node_id, + b_position_delta=b_position_delta, + ) + + torch.testing.assert_close(selected.input_ids, input_ids.index_select(0, expected_rows)) + torch.testing.assert_close(selected.hidden, hidden.index_select(0, expected_rows)) + torch.testing.assert_close(selected.b_req_idx, b_req_idx.index_select(0, expected_rows)) + torch.testing.assert_close(selected.b_mtp_index, b_mtp_index.index_select(0, expected_rows)) + torch.testing.assert_close(selected.b_seq_len, b_seq_len.index_select(0, expected_rows)) + torch.testing.assert_close(selected.b_shared_seq_len, b_shared_seq_len.index_select(0, expected_rows)) + torch.testing.assert_close( + selected.b_shared_radix_node_id, + b_shared_radix_node_id.index_select(0, expected_rows), + ) + torch.testing.assert_close(selected.b_position_delta, b_position_delta.index_select(0, expected_rows)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_vanilla_no_att_runs_all_draft_steps_for_empty_dp_batch(): + device = "cuda" + draft_batch_sizes = [] + + def draft_forward(model_input): + draft_batch_sizes.append(model_input.batch_size) + return SimpleNamespace( + token_ids=torch.empty((0,), dtype=torch.int64, device=device), + mtp_collector=SimpleNamespace( + spec_hidden=torch.empty((0, 2), dtype=torch.float32, device=device), + ), + ) + + backend = SimpleNamespace( + draft_models=[SimpleNamespace(forward=draft_forward) for _ in range(2)], + _gen_argmax_token_ids=lambda output: output.token_ids, + ) + proposer = VanillaNoAttProposer(backend=backend, enable_dynmaic_mtp=False) + empty_i32 = torch.empty((0,), dtype=torch.int32, device=device) + target_model_input = SimpleNamespace( + batch_size=0, + b_req_idx=empty_i32, + b_mtp_index=empty_i32, + b_seq_len=empty_i32, + b_position_delta=empty_i32, + b_shared_seq_len=empty_i32, + b_shared_radix_node_id=torch.empty((0,), dtype=torch.int64, device=device), + multimodal_params=[], + ) + target_model_output = SimpleNamespace( + mtp_collector=SimpleNamespace( + spec_hidden=torch.empty((0, 2), dtype=torch.float32, device=device), + ) + ) + + proposal = proposer.propose_next( + target_model_input=target_model_input, + target_model_output=target_model_output, + target_next_token_ids=torch.empty((0,), dtype=torch.int64, device=device), + b_req_mtp_start_loc=empty_i32, + draft_step=2, + accept_len=empty_i32, + ) + + assert draft_batch_sizes == [0, 0] + assert proposal.token_ids.shape == (0, 2) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_vanilla_no_att_proposes_from_one_accepted_tail_per_request(): + device = "cuda" + draft_calls = [] + draft_outputs = [ + SimpleNamespace( + token_ids=torch.tensor([21, 24]), + token_probs=torch.tensor([0.8, 0.7]), + mtp_collector=SimpleNamespace(spec_hidden=torch.tensor([[102.0, 103.0], [108.0, 109.0]])), + ), + SimpleNamespace( + token_ids=torch.tensor([31, 34]), + token_probs=torch.tensor([0.6, 0.5]), + mtp_collector=SimpleNamespace(spec_hidden=torch.tensor([[202.0, 203.0], [208.0, 209.0]])), + ), + ] + + def build_draft_model(step): + def forward(model_input): + draft_calls.append( + { + "batch_size": model_input.batch_size, + "input_ids": model_input.input_ids.cpu(), + "draft_hidden": model_input.mtp_draft_input_hiddens.cpu(), + "b_req_idx": model_input.b_req_idx.cpu(), + "b_mtp_index": model_input.b_mtp_index.cpu(), + "b_seq_len": model_input.b_seq_len.cpu(), + } + ) + return draft_outputs[step] + + return SimpleNamespace(forward=forward) + + backend = SimpleNamespace( + draft_models=[build_draft_model(0), build_draft_model(1)], + _gen_argmax_token_ids=lambda output: output.token_ids, + _gen_argmax_token_ids_and_prob=lambda output: (output.token_ids, output.token_probs), + ) + proposer = VanillaNoAttProposer(backend=backend, enable_dynmaic_mtp=True) + target_model_input = SimpleNamespace( + batch_size=5, + input_ids=torch.tensor([10, 11, 12, 13, 14], device=device), + b_req_idx=torch.tensor([7, 7, 7, 9, 9], dtype=torch.int32, device=device), + b_mtp_index=torch.tensor([0, 1, 2, 0, 1], dtype=torch.int32, device=device), + b_seq_len=torch.tensor([10, 11, 12, 20, 21], dtype=torch.int32, device=device), + b_position_delta=torch.tensor([0, 1, 2, 3, 4], dtype=torch.int32, device=device), + b_shared_seq_len=torch.tensor([8, 8, 8, 6, 6], dtype=torch.int32, device=device), + b_shared_radix_node_id=torch.tensor([70, 70, 70, 90, 90], dtype=torch.int64, device=device), + multimodal_params=[{"images": [], "audios": []} for _ in range(5)], + ) + target_model_output = SimpleNamespace( + mtp_collector=SimpleNamespace(spec_hidden=torch.arange(10, dtype=torch.float32, device=device).reshape(5, 2)) + ) + + proposal = proposer.propose_next( + target_model_input=target_model_input, + target_model_output=target_model_output, + target_next_token_ids=target_model_input.input_ids, + b_req_mtp_start_loc=torch.tensor([0, 3], dtype=torch.int32, device=device), + draft_step=2, + accept_len=torch.tensor([2, 2], dtype=torch.int32, device=device), + ) + + torch.testing.assert_close(proposal.token_ids, torch.tensor([[21, 31], [24, 34]])) + torch.testing.assert_close(proposal.schedule_scores, torch.tensor([[0.8, 0.6], [0.7, 0.5]])) + assert len(draft_calls) == 2 + assert draft_calls[0]["batch_size"] == 2 + torch.testing.assert_close(draft_calls[0]["input_ids"], torch.tensor([11, 14])) + torch.testing.assert_close(draft_calls[0]["draft_hidden"], torch.tensor([[2.0, 3.0], [8.0, 9.0]])) + torch.testing.assert_close(draft_calls[0]["b_req_idx"], torch.tensor([7, 9], dtype=torch.int32)) + torch.testing.assert_close(draft_calls[0]["b_mtp_index"], torch.tensor([1, 1], dtype=torch.int32)) + torch.testing.assert_close(draft_calls[0]["b_seq_len"], torch.tensor([11, 21], dtype=torch.int32)) + torch.testing.assert_close(draft_calls[1]["input_ids"], torch.tensor([21, 24])) + torch.testing.assert_close( + draft_calls[1]["draft_hidden"], + torch.tensor([[102.0, 103.0], [108.0, 109.0]]), + ) + assert target_model_input.batch_size == 5 + torch.testing.assert_close(target_model_input.input_ids.cpu(), torch.tensor([10, 11, 12, 13, 14])) + + +def test_vanilla_no_att_skips_draft_forward_for_zero_steps(): + proposer = VanillaNoAttProposer( + backend=SimpleNamespace(draft_models=[]), + enable_dynmaic_mtp=True, + ) + + proposal = proposer.propose_next( + target_model_input=None, + target_model_output=None, + target_next_token_ids=torch.tensor([10, 11]), + b_req_mtp_start_loc=torch.tensor([0, 1], dtype=torch.int32), + draft_step=0, + ) + + assert proposal.token_ids.shape == (2, 0) + assert proposal.schedule_scores.shape == (2, 0) + + +def test_vanilla_no_att_fill_hooks_are_noops(): + backend = SimpleNamespace(draft_models=[]) + proposer = VanillaNoAttProposer(backend=backend, enable_dynmaic_mtp=False) + overlap_proposer = DpOverlapVanillaNoAttProposer(backend=backend, enable_dynmaic_mtp=False) + + proposer.fill_draft_model_kv_state(None, None, None) + overlap_proposer.fill_draft_model_kv_state_overlap(None, None, None, None, None, None) diff --git a/unit_tests/server/router/model_infer/mtp_speculative/test_vanilla_overlap.py b/unit_tests/server/router/model_infer/mtp_speculative/test_vanilla_overlap.py new file mode 100644 index 0000000000..23a6672fd2 --- /dev/null +++ b/unit_tests/server/router/model_infer/mtp_speculative/test_vanilla_overlap.py @@ -0,0 +1,147 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.basemodel.batch_objs import ModelMtpOutputCollector, ModelOutput +from lightllm.server.router.model_infer.mtp_speculative.dp_overlap_proposers.vanilla_no_att import ( + DpOverlapVanillaNoAttProposer, +) +from lightllm.server.router.model_infer.mtp_speculative.dp_overlap_proposers.vanilla_with_att import ( + DpOverlapVanillaWithAttProposer, +) + + +class _DraftModel: + def __init__(self): + self.decode_batch_sizes = [] + + def _microbatch_overlap_decode_cuda(self, input0, input1): + self.decode_batch_sizes.append((input0.batch_size, input1.batch_size)) + return tuple( + ModelOutput( + logits=torch.arange( + model_input.batch_size, + dtype=torch.float32, + device=model_input.input_ids.device, + ).view(-1, 1), + mtp_collector=ModelMtpOutputCollector( + spec_hidden=torch.ones( + (model_input.batch_size, 2), + device=model_input.input_ids.device, + ) + ), + ) + for model_input in (input0, input1) + ) + + +def test_dp_vanilla_no_att_supports_zero_dynamic_draft_step(): + proposer = DpOverlapVanillaNoAttProposer( + backend=SimpleNamespace(draft_models=[]), + enable_dynmaic_mtp=True, + ) + target_next_token_ids0 = torch.tensor([10], dtype=torch.int64) + target_next_token_ids1 = torch.tensor([11], dtype=torch.int64) + + proposal = proposer.propose_next_overlap( + target_model_input0=None, + target_model_output0=None, + target_next_token_ids0=target_next_token_ids0, + accept_len0=torch.ones((1,), dtype=torch.int32), + target_model_input1=None, + target_model_output1=None, + target_next_token_ids1=target_next_token_ids1, + accept_len1=torch.ones((1,), dtype=torch.int32), + draft_step=0, + ) + + assert proposal.token_ids.shape == (2, 0) + assert proposal.schedule_scores.shape == (2, 0) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_dp_vanilla_proposer_owns_overlap_decode(): + device = "cuda" + draft_models = [_DraftModel(), _DraftModel()] + backend = SimpleNamespace( + max_draft_step=2, + draft_models=draft_models, + _gen_argmax_token_ids=lambda output: output.logits[:, 0].to(torch.int64), + ) + proposer = DpOverlapVanillaWithAttProposer(backend=backend, enable_dynmaic_mtp=False) + b_mtp_index = torch.arange(3, dtype=torch.int32, device=device) + model_input0 = SimpleNamespace( + batch_size=3, + b_req_idx=torch.arange(3, dtype=torch.int32, device=device), + b_mtp_index=b_mtp_index, + ) + model_input1 = SimpleNamespace( + batch_size=3, + b_req_idx=torch.arange(3, dtype=torch.int32, device=device), + b_mtp_index=b_mtp_index, + ) + + proposal = proposer.propose_next_overlap( + target_model_input0=model_input0, + target_model_output0=ModelOutput( + logits=torch.empty((6, 1), device=device), + mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.ones((6, 2), device=device)), + ), + target_next_token_ids0=torch.tensor([10, 11, 0], dtype=torch.int64, device=device), + accept_len0=torch.tensor([2], dtype=torch.int32, device=device), + target_model_input1=model_input1, + target_model_output1=ModelOutput( + logits=torch.empty((6, 1), device=device), + mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.ones((6, 2), device=device)), + ), + target_next_token_ids1=torch.tensor([20, 21, 22], dtype=torch.int64, device=device), + accept_len1=torch.tensor([1], dtype=torch.int32, device=device), + draft_step=2, + ) + + assert proposal.token_ids.tolist() == [ + [1, 1], + [0, 0], + ] + assert draft_models[0].decode_batch_sizes == [(3, 3)] + assert draft_models[1].decode_batch_sizes == [(3, 3)] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_dp_vanilla_proposer_builds_padded_inputs_for_empty_verify_rows(): + device = "cuda" + draft_models = [_DraftModel(), _DraftModel()] + backend = SimpleNamespace( + max_draft_step=2, + draft_models=draft_models, + _gen_argmax_token_ids=lambda output: output.logits[:, 0].to(torch.int64), + ) + proposer = DpOverlapVanillaWithAttProposer(backend=backend, enable_dynmaic_mtp=False) + empty_i32 = torch.empty((0,), dtype=torch.int32, device=device) + model_input0 = SimpleNamespace(batch_size=0, b_req_idx=empty_i32, b_mtp_index=empty_i32) + model_input1 = SimpleNamespace(batch_size=0, b_req_idx=empty_i32, b_mtp_index=empty_i32) + model_output0 = ModelOutput( + logits=torch.empty((0, 1), device=device), + mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.ones((0, 2), device=device)), + ) + model_output1 = ModelOutput( + logits=torch.empty((0, 1), device=device), + mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.ones((0, 2), device=device)), + ) + + proposal = proposer.propose_next_overlap( + target_model_input0=model_input0, + target_model_output0=model_output0, + target_next_token_ids0=torch.empty((0,), dtype=torch.int64, device=device), + accept_len0=torch.empty((0,), dtype=torch.int32, device=device), + target_model_input1=model_input1, + target_model_output1=model_output1, + target_next_token_ids1=torch.empty((0,), dtype=torch.int64, device=device), + accept_len1=torch.empty((0,), dtype=torch.int32, device=device), + draft_step=2, + ) + + assert proposal.token_ids.shape == (0, 2) + assert draft_models[0].decode_batch_sizes == [(0, 0)] + assert draft_models[1].decode_batch_sizes == [(0, 0)] diff --git a/unit_tests/server/router/model_infer/mtp_speculative/test_vanilla_prefill.py b/unit_tests/server/router/model_infer/mtp_speculative/test_vanilla_prefill.py new file mode 100644 index 0000000000..d22db894f7 --- /dev/null +++ b/unit_tests/server/router/model_infer/mtp_speculative/test_vanilla_prefill.py @@ -0,0 +1,237 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.server.router.model_infer.mtp_speculative.dp_overlap_proposers.vanilla_with_att import ( + DpOverlapVanillaWithAttProposer, +) +from lightllm.server.router.model_infer.mtp_speculative.proposers.vanilla_with_att import VanillaWithAttProposer + + +def _prefill_input(input_ids, batch_size): + return SimpleNamespace( + is_prefill=True, + b_position_delta=None, + batch_size=batch_size, + input_ids=input_ids, + b_req_idx=torch.arange(batch_size, dtype=torch.int32, device=input_ids.device), + b_seq_len=torch.full( + (batch_size,), input_ids.shape[0] // batch_size, dtype=torch.int32, device=input_ids.device + ), + b_ready_cache_len=torch.zeros(batch_size, dtype=torch.int32, device=input_ids.device), + mtp_draft_input_hiddens=None, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_chained_prefill_advances_local_input_without_mutating_target(): + device = "cuda" + original_input_ids = torch.tensor([10, 11, 12, 20, 21, 22], dtype=torch.int64, device=device) + target_input = _prefill_input(original_input_ids, batch_size=2) + target_hidden = torch.arange(12, dtype=torch.float32, device=device).reshape(6, 2) + forwarded = [] + + def draft_model(output_tokens, output_hidden): + def forward(model_input): + forwarded.append( + ( + model_input, + model_input.input_ids.clone(), + model_input.mtp_draft_input_hiddens, + model_input.b_is_decode_req, + ) + ) + return SimpleNamespace( + token_ids=output_tokens, + mtp_collector=SimpleNamespace(spec_hidden=output_hidden), + ) + + return SimpleNamespace(forward=forward) + + stage0_hidden = target_hidden + 100 + stage1_hidden = target_hidden + 200 + backend = SimpleNamespace( + draft_models=[ + draft_model(torch.tensor([30, 40], dtype=torch.int64, device=device), stage0_hidden), + draft_model(torch.tensor([31, 41], dtype=torch.int64, device=device), stage1_hidden), + ], + _gen_argmax_token_ids=lambda output: output.token_ids, + ) + proposer = VanillaWithAttProposer(backend=backend, enable_dynmaic_mtp=False) + + proposer.fill_draft_model_kv_state( + target_model_input=target_input, + target_model_output=SimpleNamespace(mtp_collector=SimpleNamespace(spec_hidden=target_hidden)), + target_next_token_ids=torch.tensor([13, 23], dtype=torch.int64, device=device), + ) + + assert forwarded[0][0] is forwarded[1][0] + assert forwarded[0][0] is not target_input + torch.testing.assert_close( + forwarded[0][1], + torch.tensor([11, 12, 13, 21, 22, 23], dtype=torch.int64, device=device), + ) + torch.testing.assert_close( + forwarded[1][1], + torch.tensor([12, 13, 30, 22, 23, 40], dtype=torch.int64, device=device), + ) + assert forwarded[0][2] is target_hidden + assert forwarded[1][2] is stage0_hidden + assert not forwarded[0][3].any() + assert forwarded[0][3].data_ptr() == forwarded[1][3].data_ptr() + assert target_input.input_ids is original_input_ids + assert target_input.mtp_draft_input_hiddens is None + + +def test_overlap_chained_prefill_uses_local_microbatch_inputs(monkeypatch): + def prepare(model_input, b_next_token_ids, mtp_draft_input_hiddens): + model_input.input_ids = model_input.input_ids + b_next_token_ids + model_input.mtp_draft_input_hiddens = mtp_draft_input_hiddens + return model_input + + target_input0 = _prefill_input(torch.tensor([1, 2], dtype=torch.int64), batch_size=2) + target_input1 = _prefill_input(torch.tensor([3, 4], dtype=torch.int64), batch_size=2) + target_hidden0 = torch.tensor([[1.0], [2.0]]) + target_hidden1 = torch.tensor([[3.0], [4.0]]) + forwarded = [] + + class DraftModel: + def __init__(self, token_offset): + self.token_offset = token_offset + + def _microbatch_overlap_prefill_cuda(self, input0, input1): + forwarded.append((input0, input1, input0.input_ids.clone(), input1.input_ids.clone())) + return tuple( + SimpleNamespace( + token_ids=torch.full((2,), self.token_offset + index, dtype=torch.int64), + mtp_collector=SimpleNamespace(spec_hidden=hidden + self.token_offset), + ) + for index, hidden in enumerate((target_hidden0, target_hidden1)) + ) + + backend = SimpleNamespace( + draft_models=[DraftModel(10), DraftModel(20)], + _gen_argmax_token_ids=lambda output: output.token_ids, + ) + proposer = DpOverlapVanillaWithAttProposer(backend=backend, enable_dynmaic_mtp=False) + monkeypatch.setattr(proposer, "_prepare_mtp_prefill_inputs", prepare) + + proposer.fill_draft_model_kv_state_overlap( + target_model_input0=target_input0, + target_model_output0=SimpleNamespace(mtp_collector=SimpleNamespace(spec_hidden=target_hidden0)), + target_next_token_ids0=torch.tensor([5, 6], dtype=torch.int64), + target_model_input1=target_input1, + target_model_output1=SimpleNamespace(mtp_collector=SimpleNamespace(spec_hidden=target_hidden1)), + target_next_token_ids1=torch.tensor([7, 8], dtype=torch.int64), + ) + + assert forwarded[0][0] is forwarded[1][0] + assert forwarded[0][1] is forwarded[1][1] + assert forwarded[0][0] is not target_input0 + assert forwarded[0][1] is not target_input1 + torch.testing.assert_close(forwarded[0][2], torch.tensor([6, 8], dtype=torch.int64)) + torch.testing.assert_close(forwarded[0][3], torch.tensor([10, 12], dtype=torch.int64)) + torch.testing.assert_close(forwarded[1][2], torch.tensor([16, 18], dtype=torch.int64)) + torch.testing.assert_close(forwarded[1][3], torch.tensor([21, 23], dtype=torch.int64)) + torch.testing.assert_close(target_input0.input_ids, torch.tensor([1, 2], dtype=torch.int64)) + torch.testing.assert_close(target_input1.input_ids, torch.tensor([3, 4], dtype=torch.int64)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_chained_decode_overlays_verified_tokens_for_the_next_level(): + device = "cuda" + original_input_ids = torch.tensor([10, 11, 12, 20, 21, 22], dtype=torch.int64, device=device) + target_input = SimpleNamespace( + is_prefill=False, + batch_size=6, + input_ids=original_input_ids, + mtp_draft_input_hiddens=None, + ) + target_hidden = torch.arange(12, dtype=torch.float32, device=device).reshape(6, 2) + forwarded = [] + + def draft_model(output_tokens, output_probs, output_hidden): + def forward(model_input): + forwarded.append( + ( + model_input, + model_input.input_ids.clone(), + model_input.mtp_draft_input_hiddens, + ) + ) + return SimpleNamespace( + token_ids=output_tokens, + token_probs=output_probs, + mtp_collector=SimpleNamespace(spec_hidden=output_hidden), + ) + + return SimpleNamespace(forward=forward) + + stage0_tokens = torch.tensor([30, 31, 32, 33, 34, 35], dtype=torch.int64, device=device) + stage1_tokens = torch.tensor([40, 41, 42, 43, 44, 45], dtype=torch.int64, device=device) + stage2_tokens = torch.tensor([50, 51, 52, 53, 54, 55], dtype=torch.int64, device=device) + stage0_probs = torch.tensor([0.30, 0.31, 0.32, 0.33, 0.34, 0.35], device=device) + stage1_probs = torch.tensor([0.40, 0.41, 0.42, 0.43, 0.44, 0.45], device=device) + stage2_probs = torch.tensor([0.50, 0.51, 0.52, 0.53, 0.54, 0.55], device=device) + stage0_hidden = target_hidden + 100 + stage1_hidden = target_hidden + 200 + stage2_hidden = target_hidden + 300 + backend = SimpleNamespace( + max_draft_step=3, + draft_models=[ + draft_model(stage0_tokens, stage0_probs, stage0_hidden), + draft_model(stage1_tokens, stage1_probs, stage1_hidden), + draft_model(stage2_tokens, stage2_probs, stage2_hidden), + ], + _gen_argmax_token_ids_and_prob=lambda output: (output.token_ids, output.token_probs), + ) + proposer = VanillaWithAttProposer(backend=backend, enable_dynmaic_mtp=True) + + proposal = proposer.propose_next( + target_model_input=target_input, + target_model_output=SimpleNamespace(mtp_collector=SimpleNamespace(spec_hidden=target_hidden)), + target_next_token_ids=original_input_ids, + b_req_mtp_start_loc=torch.tensor([0, 3], dtype=torch.int32, device=device), + draft_step=3, + accept_len=torch.tensor([3, 2], dtype=torch.int32, device=device), + ) + + assert forwarded[0][0] is forwarded[1][0] is forwarded[2][0] + assert forwarded[0][0] is not target_input + torch.testing.assert_close(forwarded[0][1], original_input_ids) + # 第一次覆盖：[10, 11, 12] -> [11, 12, 32]， + # [20, 21] -> [21, 34]。 + torch.testing.assert_close( + forwarded[1][1], + torch.tensor([11, 12, 32, 21, 34, 35], dtype=torch.int64, device=device), + ) + # 第二次覆盖：[11, 12, 32] -> [12, 32, 42]， + # [21, 34] -> [34, 44]。 + torch.testing.assert_close( + forwarded[2][1], + torch.tensor([12, 32, 42, 34, 44, 45], dtype=torch.int64, device=device), + ) + assert forwarded[0][2] is target_hidden + assert forwarded[1][2] is stage0_hidden + assert forwarded[2][2] is stage1_hidden + torch.testing.assert_close( + proposal.token_ids, + torch.tensor([[32, 42, 52], [34, 44, 54]], dtype=torch.int64, device=device), + ) + torch.testing.assert_close( + proposal.schedule_scores, + torch.tensor([[0.32, 0.42, 0.52], [0.34, 0.44, 0.54]], device=device), + ) + assert target_input.input_ids is original_input_ids + assert target_input.mtp_draft_input_hiddens is None + + with pytest.raises(AssertionError, match="requires the full chained draft depth"): + proposer.propose_next( + target_model_input=target_input, + target_model_output=SimpleNamespace(mtp_collector=SimpleNamespace(spec_hidden=target_hidden)), + target_next_token_ids=original_input_ids, + b_req_mtp_start_loc=torch.tensor([0, 3], dtype=torch.int32, device=device), + draft_step=2, + accept_len=torch.tensor([3, 2], dtype=torch.int32, device=device), + ) diff --git a/unit_tests/server/router/model_infer/structured_output/test_error_output.py b/unit_tests/server/router/model_infer/structured_output/test_error_output.py new file mode 100644 index 0000000000..e8285f555a --- /dev/null +++ b/unit_tests/server/router/model_infer/structured_output/test_error_output.py @@ -0,0 +1,66 @@ +from types import MethodType, SimpleNamespace +from unittest.mock import MagicMock + +import numpy as np +import pytest +import torch +import xgrammar as xgr + +from unit_tests.server.grammar_helpers import init_request, make_req +from lightllm.server.router.model_infer.infer_batch import InferReq, InferReqUpdatePack + + +@pytest.mark.parametrize("failure", ["mask", "commit", "rollback"]) +@pytest.mark.parametrize("exception_type", [RuntimeError, BaseException]) +def test_grammar_failure_finishes_after_publishing_inflight_token(compiler, monkeypatch, failure, exception_type): + req = make_req(regular_constraint="abc") + init_request(compiler, req) + shm_req = req.shm_req + shm_req.shm_prompt_ids.arr = np.array([ord("P")] + [0] * 8, dtype=np.int64) + shm_req.shm_logprobs = SimpleNamespace(arr=np.zeros(9, dtype=[("logprob", np.float32), ("rank", np.int32)])) + req.set_next_gen_token_id = MethodType(InferReq.set_next_gen_token_id, req) + req.update_finish_status = lambda **kwargs: None + InferReqUpdatePack(req, 1).handle(ord("a"), -0.1, -1, [256], is_master_in_dp=True) + + if failure == "mask": + monkeypatch.setattr( + xgr.GrammarMatcher, + "fill_next_token_bitmask", + MagicMock(side_effect=exception_type("injected mask failure")), + ) + assert not req.output_constraint.fill_masks(torch.empty((1, 9), dtype=torch.int32), []) + elif failure == "commit": + monkeypatch.setattr( + xgr.GrammarMatcher, "accept_token", MagicMock(side_effect=exception_type("injected commit failure")) + ) + else: + monkeypatch.setattr( + xgr.GrammarMatcher, "rollback", MagicMock(side_effect=exception_type("injected rollback failure")) + ) + assert req.output_constraint.fill_masks(torch.empty((2, 9), dtype=torch.int32), [ord("b")]) + + next_token = ord("b") if failure == "rollback" else ord("!") + InferReqUpdatePack(req, 2).handle(next_token, -0.2, -1, [256], is_master_in_dp=True) + InferReqUpdatePack(req, 3).handle(ord("z"), -0.3, -1, [256], is_master_in_dp=True) + + # In-flight tokens can update KV history, but output publication stops at the error. + assert req.finish_status.is_finished_error() and shm_req.finish_status.is_finished_error() + assert req.output_constraint.is_terminated() + assert shm_req.shm_prompt_ids.arr[1:4].tolist() == [ord("a"), next_token, ord("z")] + assert req.output_constraint.error == f"injected {failure} failure" + assert shm_req.candetoken_out_len == shm_req.shm_cur_output_len == 2 + assert shm_req.finish_token_index == 2 + + +def test_draft_length_failure_disables_mask(compiler): + req = make_req(regular_constraint="ab") + init_request(compiler, req) + state = req.output_constraint + bitmask = torch.empty((1, 9), dtype=torch.int32) + draft_tokens = MagicMock() + draft_tokens.__len__.side_effect = BaseException("injected draft_length failure") + + assert not state.fill_masks(bitmask, draft_tokens) + assert state.error == "injected draft_length failure" + assert state.is_terminated() + assert state.matcher.accept_token(ord("a")) diff --git a/unit_tests/server/router/model_infer/structured_output/test_mtp.py b/unit_tests/server/router/model_infer/structured_output/test_mtp.py new file mode 100644 index 0000000000..823f3d6049 --- /dev/null +++ b/unit_tests/server/router/model_infer/structured_output/test_mtp.py @@ -0,0 +1,444 @@ +from contextlib import nullcontext +from types import SimpleNamespace + +import pytest +import torch + +from unit_tests.server.grammar_helpers import ( + allowed, + commit_token, + build_mask, + make_mask_buffers, + apply_masks, + make_req, + init_request, + prepare_sampling_tensors, +) + + +def verify_mask(compiler, reqs, inputs): + mask = build_mask(reqs, draft_input_ids=inputs) + logits = torch.zeros(len(reqs), compiler.vocab_size) + apply_masks(reqs, logits, mask) + return torch.isfinite(logits) + + +def test_verify_masks_follow_each_draft_prefix_and_roll_back(compiler): + req = make_req(regular_constraint="abc") + init_request(compiler, req) + mask = verify_mask(compiler, [req] * 4, [0, ord("a"), ord("b"), ord("c")]) + assert [row.nonzero().flatten().tolist() for row in mask] == [[ord("a")], [ord("b")], [ord("c")], [256]] + assert allowed([req])[0].nonzero().flatten().tolist() == [ord("a")] + + # Only tokens accepted by target verification become permanent history. + commit_token(req, ord("a")) + commit_token(req, ord("b")) + mask = verify_mask(compiler, [req] * 2, [ord("b"), ord("c")]) + assert mask[0, ord("c")] and mask[1, 256] + assert allowed([req])[0, ord("c")] + + +@pytest.mark.parametrize("draft", [b"!bc", b"a!c", b"ab!"]) +def test_invalid_draft_is_rejected_without_failing_the_request(compiler, draft): + req = make_req(regular_constraint="abc") + init_request(compiler, req) + mask = verify_mask(compiler, [req] * 4, [0, *draft]) + mismatch = draft.index(ord("!")) + assert mask[mismatch, b"abc"[mismatch]] + assert not mask[mismatch, ord("!")] + assert mask[mismatch + 1 :].all() # Unreachable suffix rows remain harmless. + assert req.output_constraint.error is None + assert not req.output_constraint.is_terminated() + assert allowed([req])[0, ord("a")] + + +def test_mixed_compacted_requests_have_independent_verify_states(compiler): + first, second, plain = make_req(regular_constraint="ab"), make_req(regular_constraint="cd"), make_req() + for req in (first, second, plain): + init_request(compiler, req) + reqs = [plain, plain, second, second, second, first, first] + mask = verify_mask(compiler, reqs, [0, 0, 0, ord("c"), ord("d"), 0, ord("a")]) + assert mask[:2].all() + assert [row.nonzero().flatten().tolist() for row in mask[2:]] == [ + [ord("c")], + [ord("d")], + [256], + [ord("a")], + [ord("b")], + ] + assert allowed([first])[0, ord("a")] + assert allowed([second])[0, ord("c")] + + +@pytest.mark.parametrize("reasoning_end", [b"", b"]"]) +def test_stop_token_in_draft_does_not_terminate_the_real_matcher(compiler, reasoning_end): + req = make_req(regular_constraint="a", guided_reasoning_end=tuple(reasoning_end)) + init_request(compiler, req) + draft = [*reasoning_end, ord("a"), 256, ord("!")] + mask = verify_mask(compiler, [req] * (len(draft) + 1), [0, *draft]) + assert mask[: len(reasoning_end)].all() + mask = mask[len(reasoning_end) :] + assert mask[0, ord("a")] and mask[1, 256] and mask[2:].all() + assert not req.output_constraint.matcher.is_terminated() + assert req.output_constraint.error is None + for token in reasoning_end: + assert not commit_token(req, token) + assert allowed([req])[0, ord("a")] + + +def test_completed_grammar_does_not_enable_mask(compiler): + req = make_req(regular_constraint="a") + init_request(compiler, req) + state = req.output_constraint + state.commit(ord("a")) + state.commit(256) + assert state.is_terminated() + + bitmask = torch.zeros((1, 9), dtype=torch.int32) + assert not state.fill_masks(bitmask, []) + assert (bitmask == -1).all() + assert state.is_terminated() and state.error is None + + +def test_long_reasoning_tail_restored_after_repeated_speculation(compiler): + req = make_req(regular_constraint="ab", guided_reasoning_end=tuple(b"]>")) + init_request(compiler, req) + state = req.output_constraint + # 长思考内容只保留结束标记长度的后缀，标记仍能跨已提交和 draft token 匹配。 + assert verify_mask(compiler, [req] * 4, [0, *b"xxx"]).all() + for token in b"x" * 100 + b"]": + commit_token(req, token) + assert state.reasoning_tail == tuple(b"x]") + saved_tail = state.reasoning_tail + + for _ in range(3): + mask = verify_mask(compiler, [req] * 4, [0, *b">ab"]) + assert mask[0].all() + assert [row.nonzero().flatten().tolist() for row in mask[1:]] == [[ord("a")], [ord("b")], [256]] + assert state.in_reasoning and state.reasoning_tail[-1] == ord("]") + # 快照恢复被临时推进替换的后缀，多轮验证始终观察同一已提交状态。 + assert state.reasoning_tail == saved_tail + assert state.error is None + + for token in b">ab": + commit_token(req, token) + # 思考结束标记保留在 tail 中，grammar token 的提交进度由 matcher 持有。 + assert state.reasoning_tail == tuple(b"]>") + assert not state.in_reasoning + assert allowed([req])[0].nonzero().flatten().tolist() == [256] + + +def test_reasoning_boundary_inside_verify_is_temporary_until_commit(compiler): + req = make_req(regular_constraint="ab", guided_reasoning_end=tuple(b"]>")) + init_request(compiler, req) + commit_token(req, ord("]")) + mask = verify_mask(compiler, [req] * 4, [ord("]"), ord(">"), ord("a"), ord("b")]) + assert mask[0].all() + assert [row.nonzero().flatten().tolist() for row in mask[1:]] == [[ord("a")], [ord("b")], [256]] + state = req.output_constraint + assert state.in_reasoning and state.reasoning_tail == (ord("]"),) + matcher = state.matcher + assert matcher is not None + # Repeating an uncommitted draft reuses the matcher at the same prefix. + repeated_mask = verify_mask(compiler, [req] * 4, [ord("]"), ord(">"), ord("a"), ord("b")]) + assert torch.equal(repeated_mask, mask) + assert state.matcher is matcher + + # A single MTP post_handle can commit the delimiter and the whole answer. + for token in b">ab": + assert not commit_token(req, token) + assert not state.in_reasoning + assert allowed([req])[0, 256] + assert state.matcher is matcher + assert commit_token(req, 256) + + +@pytest.mark.parametrize( + "marker,committed,draft,expected", + [ + (b"]>", b"", b"]>ab", [None, None, ord("a"), ord("b"), 256]), + (b"]>", b"]", b"x]>a", [None, None, None, ord("a"), ord("b")]), + (b"]>", b"]", b">!b", [None, ord("a"), None, None]), + (b"]>", b"]", b"x]x", [None, None, None, None]), + (b"aba", b"ab", b"baba", [None, None, None, None, ord("a")]), + (b"aba", b"ab", b"aab", [None, ord("a"), ord("b"), 256]), + ], +) +def test_verify_locates_answer_without_mutating_reasoning_progress(compiler, marker, committed, draft, expected): + req = make_req(regular_constraint="ab", guided_reasoning_end=tuple(marker)) + init_request(compiler, req) + for token in committed: + commit_token(req, token) + state = req.output_constraint + original_tail = state.reasoning_tail + + mask = verify_mask(compiler, [req] * (len(draft) + 1), [0, *draft]) + + for prediction, token in zip(mask, expected): + if token is None: + assert prediction.all() + else: + assert prediction.nonzero().flatten().tolist() == [token] + assert state.reasoning_end_token_ids == tuple(marker) + assert state.reasoning_tail == original_tail + assert state.in_reasoning and state.error is None + # The target may choose to continue thinking instead of accepting the draft. + assert not commit_token(req, ord("?")) + assert allowed([req])[0].all() + # A later real transition must start from the unmodified grammar root. + for token in marker: + assert not commit_token(req, token) + assert allowed([req])[0].nonzero().flatten().tolist() == [ord("a")] + + +@pytest.mark.parametrize("failure", ["fill_next_token_bitmask", "rollback"]) +@pytest.mark.parametrize("exception_type", [RuntimeError, BaseException]) +def test_failure_after_speculative_reasoning_end_records_terminal_error(compiler, monkeypatch, failure, exception_type): + import xgrammar as xgr + + req = make_req(regular_constraint="ab", guided_reasoning_end=(ord("]"),)) + init_request(compiler, req) + original = getattr(xgr.GrammarMatcher, failure) + + def fail_speculative_mask(matcher, *args): + if failure == "rollback" or args[1] == 2: + raise exception_type("injected speculative mask failure") + return original(matcher, *args) + + with monkeypatch.context() as patcher: + patcher.setattr(xgr.GrammarMatcher, failure, fail_speculative_mask) + masks = build_mask([req] * 5, draft_input_ids=[0, *b"]ab", 256]) + if failure == "rollback": + # 回滚失败不影响此前已完整生成的掩码，本轮仍可使用。 + assert masks is not None + else: + assert masks is None + + state = req.output_constraint + assert state.in_reasoning and state.reasoning_tail == () + assert state.error == "injected speculative mask failure" + assert state.is_terminated() + # 后续批次不再使用错误请求的 matcher。 + assert build_mask([req] * 4, draft_input_ids=[0, *b"]ab"]) is None + + +@pytest.mark.parametrize("exception_type", [RuntimeError, BaseException]) +def test_temporary_state_marks_rollback_failure_as_terminal(compiler, monkeypatch, exception_type): + import xgrammar as xgr + + req = make_req(regular_constraint="ab", guided_reasoning_end=tuple(b"]>")) + init_request(compiler, req) + commit_token(req, ord("]")) + state = req.output_constraint + + def fail_rollback(matcher, token_count): + raise exception_type("injected rollback failure") + + monkeypatch.setattr(xgr.GrammarMatcher, "rollback", fail_rollback) + # 直接使用临时状态接口，验证错误标记不依赖 fill_masks 的外层捕获。 + with state._temporary_state(): + assert state._accept_token(ord(">")) + assert state._accept_token(ord("a")) + + assert state.error == "injected rollback failure" + assert state.is_terminated() + assert state.in_reasoning and state.reasoning_tail == (ord("]"),) + + +def test_thinking_verify_does_not_enable_mask_or_advance_matcher(compiler): + req = make_req(regular_constraint="ab", guided_reasoning_end=(ord("]"),)) + init_request(compiler, req) + assert build_mask([req] * 3, draft_input_ids=[0, ord("x"), ord("y")]) is None + assert req.output_constraint.matcher.accept_token(ord("a")) + + +@pytest.mark.parametrize("failure", ["fill_next_token_bitmask", "rollback"]) +def test_speculative_mask_failure_isolated_and_draft_tokens_rolled_back(compiler, monkeypatch, failure): + import xgrammar as xgr + + bad = make_req(regular_constraint="ab") + good = make_req(regular_constraint="cd") + for req in [bad, good]: + init_request(compiler, req) + allowed([bad, good]) + bad_matcher = bad.output_constraint.matcher + original = getattr(xgr.GrammarMatcher, failure) + + def fail_for_bad_request(matcher, *args): + if matcher is bad_matcher and (failure == "rollback" or args[1] == 1): + raise RuntimeError("injected grammar failure") + return original(matcher, *args) + + with monkeypatch.context() as patcher: + patcher.setattr(xgr.GrammarMatcher, failure, fail_for_bad_request) + mask = verify_mask(compiler, [bad, bad, good], [0, ord("a"), 0]) + + if failure == "rollback": + assert [row.nonzero().flatten().tolist() for row in mask[:2]] == [[ord("a")], [ord("b")]] + else: + assert mask[:2].all() + assert mask[2].nonzero().flatten().tolist() == [ord("c")] + assert bad.output_constraint.error == "injected grammar failure" + assert good.output_constraint.error is None + if failure == "fill_next_token_bitmask": + # The failure happened after accepting draft 'a'; cleanup must still + # leave the matcher at its original prefix, which accepts 'a' again. + assert bad_matcher.accept_token(ord("a")) + + +@pytest.mark.parametrize("dp,microbatch_side", [(False, None), (True, None), (True, 0), (True, 1)]) +@pytest.mark.parametrize("dynamic", [False, True]) +def test_mtp_decode_handoff_compaction_and_accepted_commits(compiler, monkeypatch, dp, microbatch_side, dynamic): + from lightllm.server.router.model_infer.mode_backend.chunked_prefill import impl as chunked_impl + from lightllm.server.router.model_infer.mode_backend.dp_backend import impl as dp_impl + + impl = dp_impl if dp else chunked_impl + backend_class = dp_impl.DPChunkedPrefillBackend if dp else chunked_impl.ChunkedPrefillBackend + backend = backend_class.__new__(backend_class) + backend.eos_id = [256] + backend.is_master_in_dp = False + backend.extra_post_req_handle_func = None + req = make_req(regular_constraint="abc") + req.req_idx = 5 + req.out_token_id_count = {} + init_request(compiler, req) + order = [] + + def make_input(tokens): + return SimpleNamespace( + batch_size=len(tokens), + input_ids=None, + b_req_idx=torch.full((len(tokens),), req.req_idx, dtype=torch.int32), + b_mtp_index=torch.arange(len(tokens), dtype=torch.int32), + ) + + model_input = make_input(list(b"abcd") if dynamic else list(b"abc")) + empty = make_input([]) + original_reqs = [req] * model_input.batch_size + + def compact(**kwargs): + value = kwargs["model_input"] + selection = None + if dynamic and value.batch_size: + value.batch_size = 3 + value.b_req_idx = value.b_req_idx[:3] + value.b_mtp_index = value.b_mtp_index[:3] + selection = SimpleNamespace(wait=lambda: None, tensor=torch.tensor([True, True, True, False])) + return value, selection + + def compact_pair(**kwargs): + first, selection0 = compact(model_input=kwargs["model_input0"]) + second, selection1 = compact(model_input=kwargs["model_input1"]) + return first, selection0, second, selection1 + + def forward(value): + assert value.input_ids is None, "Constraint preparation must not materialize model inputs" + order.append("forward") + logits = torch.zeros(value.batch_size, 257) + logits[:, ord("!")] = 10 # Every constrained row must override this. + return SimpleNamespace(logits=logits) + + def forward_pair(first, second): + value = model_input + output = forward(value) + padding = SimpleNamespace(logits=torch.zeros(0, 257)) + return (output, padding) if microbatch_side == 0 else (padding, output) + + def propose(**kwargs): + order.append("propose") + return SimpleNamespace(token_ids=torch.empty(1, 2, dtype=torch.int64)) + + engine = SimpleNamespace( + plan_decode=lambda **kwargs: SimpleNamespace(draft_step=2, skip_verify_sync=False), + prepare_decode_model_input=compact, + prepare_decode_model_inputs=compact_pair, + propose_next=propose, + propose_next_overlap=propose, + update_planner_statics=lambda **kwargs: None, + ) + backend.spec_engine = backend.decode_draft_engine = engine + backend.model = SimpleNamespace(forward=forward, microbatch_overlap_decode=forward_pair) + monkeypatch.setattr(impl, "prepare_decode_inputs", lambda req_objs: (model_input, original_reqs)) + if microbatch_side is not None: + layout = ( + (model_input, original_reqs, [req], empty, [], []) + if microbatch_side == 0 + else (empty, [], [], model_input, original_reqs, [req]) + ) + monkeypatch.setattr(impl, "overlap_prepare_decode_inputs", lambda req_objs: layout) + monkeypatch.setattr(torch.cuda, "stream", lambda stream: nullcontext()) + monkeypatch.setattr(torch.cuda, "Event", lambda: SimpleNamespace(record=lambda: None, synchronize=lambda: None)) + monkeypatch.setattr(impl.g_infer_context, "get_overlap_stream", lambda: None) + monkeypatch.setattr( + impl.g_infer_context, + "req_sampling_manager", + SimpleNamespace(update_reqs_out_token_counter_gpu=lambda **kwargs: None), + raising=False, + ) + monkeypatch.setattr( + impl.g_pin_mem_manager, "async_copy_from_gpu_tensor", lambda key, gpu_tensor: gpu_tensor.clone() + ) + monkeypatch.setattr(impl, "gen_b_req_mtp_start_loc", lambda *args, **kwargs: torch.tensor([0])) + + token_table = torch.full((8, 4), ord("!"), dtype=torch.int64) + + class ReadyTokenTable: + def __getitem__(self, index): + assert "handoff" in order, "Token table read before the existing post_handle handoff" + assert "propose" not in order, "Token table read after the next proposal could overwrite it" + return token_table[index] + + backend._get_next_token_ranks = lambda logits, token_ids: torch.zeros_like(token_ids) + backend._async_copy_next_token_infos_to_pin_mem = lambda next_token_ids, next_token_logprobs, next_token_ranks: ( + next_token_ids, + next_token_logprobs, + next_token_ranks, + ) + + def sample(logits, reqs, eos_ids, b_mtp_index): + assert len(reqs) == 3 + assert req.out_token_id_count == {ord("a"): 1}, "MTP read stale CPU counters" + manager = make_mask_buffers(reqs, 4) + manager.req_to_next_token_ids_cpu = ReadyTokenTable() + assert prepare_sampling_tensors(reqs, manager)[-1] + assert b_mtp_index.tolist() == [0, 1, 2] + apply_masks(reqs, logits, manager, b_mtp_index) + order.append("sample") + result = logits.argmax(-1) + assert result.tolist() == [ord("b"), ord("c"), 256] + return result, torch.zeros(3) + + monkeypatch.setattr(impl, "sample", sample) + + def verify(**kwargs): + order.append("verify") + assert allowed([req])[0, ord("b")], "Draft advanced permanent matcher state" + return torch.tensor([2], dtype=torch.int32), torch.tensor([1, 1, 0], dtype=torch.int32) + + monkeypatch.setattr(impl.mtp_utils, "verify_mtp_tokens", verify) + monkeypatch.setattr(impl.mtp_utils, "scatter_mtp_next_tokens", lambda **kwargs: None) + backend._pre_post_handle = lambda *args, **kwargs: [] + + def post(**kwargs): + assert kwargs["next_token_ids"].tolist() == list(b"bc") + for token in kwargs["next_token_ids"].tolist(): + commit_token(req, token) + order.append("commit") + + backend._post_handle = post + + def handoff(): + order.append("handoff") + token_table[req.req_idx] = torch.tensor(list(b"abcd")) + commit_token(req, ord("a")) + req.out_token_id_count[ord("a")] = 1 + + events = SimpleNamespace( + notify_post_handle_and_wait_pre_post_handle=handoff, + notify_forward_and_wait_post_handle=lambda: None, + notify_pre_post_handle=lambda: None, + ) + method = backend.decode_overlap_mtp if microbatch_side is not None else backend.decode_mtp + method(events, [req]) + assert order == ["forward", "handoff", "sample", "verify", "propose", "commit"] + assert allowed([req])[0, 256] diff --git a/unit_tests/server/router/model_infer/structured_output/test_overlap.py b/unit_tests/server/router/model_infer/structured_output/test_overlap.py new file mode 100644 index 0000000000..930cc5c1e8 --- /dev/null +++ b/unit_tests/server/router/model_infer/structured_output/test_overlap.py @@ -0,0 +1,152 @@ +from contextlib import nullcontext +from threading import Event, Thread +from types import SimpleNamespace + +import pytest +import torch + +from unit_tests.server.grammar_helpers import ( + init_request, + make_req, + make_mask_buffers, + apply_masks, + prepare_sampling_tensors, +) +from lightllm.server.router.model_infer.mode_backend.chunked_prefill import impl as chunked_impl +from lightllm.server.router.model_infer.mode_backend.dp_backend import impl as dp_impl +from lightllm.server.router.model_infer.mode_backend.overlap_events import OverlapEventManager + + +@pytest.mark.parametrize("prefill_first", [False, True]) +@pytest.mark.parametrize("constrained", [False, True, "reasoning"]) +@pytest.mark.parametrize( + "dp,microbatch,empty_second", + [(False, False, False), (True, False, False), (True, False, True), (True, True, False), (True, True, True)], +) +def test_overlap_handoff_completes_postprocessing_before_sampling( + compiler, monkeypatch, prefill_first, constrained, dp, microbatch, empty_second +): + """Exercise the real two-thread event protocol and chunked-prefill methods. + + Sampling must see the previous matcher and CPU token counters. + All paths still submit the next forward before releasing CPU postprocessing. + """ + impl = dp_impl if dp else chunked_impl + backend_class = dp_impl.DPChunkedPrefillBackend if dp else chunked_impl.ChunkedPrefillBackend + backend = object.__new__(backend_class) + backend.args = SimpleNamespace(penalty_counter_mode="cpu_counter") + backend.disable_chunked_prefill = False + backend.extra_post_req_handle_func = backend.pd_prefill_chunked_handle_func = None + backend._capture_prompt_logprobs_if_needed = lambda *args: None + req = make_req( + regular_constraint="b" if constrained == "reasoning" else "ab" if constrained else None, + guided_reasoning_end=(ord("a"),) if constrained == "reasoning" else (), + ) + req.out_token_id_count = {} + init_request(compiler, req) + mask_buffers = make_mask_buffers([req]) + model_input = SimpleNamespace( + b_req_idx=torch.tensor([0]), b_mtp_index=torch.tensor([0]), b_prefill_has_output_cpu=[True] + ) + empty_input = SimpleNamespace(b_req_idx=torch.empty(0), b_mtp_index=torch.empty(0), b_prefill_has_output_cpu=[]) + monkeypatch.setattr(impl, "prepare_decode_inputs", lambda req_objs: (model_input, req_objs)) + monkeypatch.setattr(impl, "prepare_prefill_inputs", lambda req_objs, **kwargs: (model_input, req_objs)) + if microbatch: + monkeypatch.setattr( + impl, + "overlap_prepare_prefill_inputs", + lambda reqs: (model_input if reqs else empty_input, reqs, empty_input, []), + ) + monkeypatch.setattr( + impl, + "overlap_prepare_decode_inputs", + lambda req_objs: (model_input if req_objs else empty_input, req_objs, req_objs, empty_input, [], []), + ) + monkeypatch.setattr(impl.torch.cuda, "stream", lambda stream: nullcontext()) + monkeypatch.setattr( + impl.torch.cuda, "Event", lambda: SimpleNamespace(record=lambda: None, synchronize=lambda: None) + ) + monkeypatch.setattr(impl.g_infer_context, "get_overlap_stream", lambda: None) + monkeypatch.setattr(impl.g_infer_context, "save_hybrid_state_to_cache", lambda **kwargs: None) + monkeypatch.setattr(impl.g_infer_context, "is_hybrid_att_model", False) + second_forward = Event() + order, errors = [], [] + + def forward(model_input): + step = sum(event.startswith("forward") for event in order) + order.append(f"forward{step}") + if step == 1: + second_forward.set() + logits = torch.zeros(model_input.b_req_idx.numel(), 257) + logits[:, ord("a") + step] = 1 + return SimpleNamespace(logits=logits, prompt_logics=None) + + def forward_microbatch(first, second): + return forward(first), SimpleNamespace(logits=torch.zeros(0, 257), prompt_logics=None) + + backend.model = SimpleNamespace( + forward=forward, microbatch_overlap_prefill=forward_microbatch, microbatch_overlap_decode=forward_microbatch + ) + + def sample(**kwargs): + if second_forward.is_set(): + assert req.out_token_id_count.get(ord("a"), 0) == 1, "Sampling read stale CPU token counters" + logits = kwargs["logits"] + reqs = kwargs["run_reqs"] + if prepare_sampling_tensors(reqs, mask_buffers)[-1]: + apply_masks(reqs, logits, mask_buffers) + token = logits.argmax(-1) + order.append(f"sample{chr(token.item())}") + return token, token, torch.zeros(1), torch.zeros(1, dtype=torch.int64) + + backend._sample_and_scatter_token = sample + + def pre_post(*args, **kwargs): + req.cur_output_len += 1 + return [] + + backend._pre_post_handle = pre_post + + def post(**kwargs): + token = kwargs["next_token_ids"].item() + if token == ord("a"): + assert second_forward.wait(3), "CPU postprocessing must overlap the next forward" + if req.output_constraint is not None: + req.output_constraint.commit(token) + req.out_token_id_count[token] = req.out_token_id_count.get(token, 0) + 1 + order.append(f"post{chr(token)}") + + backend._post_handle = post + events = OverlapEventManager() + packs = [events.get_overlap_event_pack() for _ in range(3)] + + def run(step): + try: + packs[step].wait_to_forward() + run_reqs = [] if step == 1 and empty_second else [req] + if prefill_first and (step == 0 or empty_second): + method = backend.prefill_overlap if microbatch else backend.prefill_normal + else: + method = backend.decode_overlap if microbatch else backend.decode_normal + method(packs[step], run_reqs) + except BaseException as exc: + errors.append(exc) + + def drain(): + packs[2].wait_to_forward() + packs[2].notify_post_handle_and_wait_pre_post_handle() + + threads = [Thread(target=run, args=(step,), daemon=True) for step in range(2)] + threads.append(Thread(target=drain, daemon=True)) + for thread in threads: + thread.start() + for thread in threads: + thread.join(5) + assert not errors + assert not any(thread.is_alive() for thread in threads), f"Overlap deadlocked: {order}" + expected = ["forward0", "samplea", "forward1"] + if empty_second: + expected += ["posta"] + else: + expected += ["posta", "sampleb", "postb"] + assert order == expected diff --git a/unit_tests/server/router/model_infer/structured_output/test_reasoning.py b/unit_tests/server/router/model_infer/structured_output/test_reasoning.py new file mode 100644 index 0000000000..8c265847db --- /dev/null +++ b/unit_tests/server/router/model_infer/structured_output/test_reasoning.py @@ -0,0 +1,194 @@ +import asyncio +import ctypes +import json +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest +import torch + +from lightllm.server import api_http, api_openai, build_prompt +from lightllm.server.api_models import ChatCompletionRequest +from lightllm.server.core.objs import FinishStatus +from lightllm.server.core.objs.sampling_params import SamplingParams +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.server.router.model_infer.infer_batch import InferSamplingParams +from lightllm.server.router.model_infer.structured_output.state import ConstraintState + + +@pytest.mark.parametrize( + "parser,template_kwargs,effort,marker", + [ + ("qwen3", None, None, ""), + ("qwen3", {"enable_thinking": True}, None, ""), + ("qwen3", {"enable_thinking": False}, None, None), + ("qwen3", {"thinking": False}, None, None), + ("qwen3", None, "none", None), + ("deepseek-v3", None, None, None), + ("deepseek-v3", {"thinking": True}, None, ""), + ("qwen3-thinking", None, None, ""), + ("kimi", None, None, "◁/think▷"), + ("gpt-oss", None, None, "<|channel|>final<|message|>"), + ("gemma4", None, None, "answer"), + (None, None, None, None), + ], +) +def test_chat_grammar_activation_matches_template_thinking(monkeypatch, parser, template_kwargs, effort, marker): + monkeypatch.setattr(api_openai, "get_env_start_args", lambda: StartArgs(reasoning_parser=parser)) + monkeypatch.setattr(build_prompt, "tokenizer_supports_force_thinking", lambda: True) + encoded = [] + + def encode(text, add_special_tokens): + assert not add_special_tokens + encoded.append(text) + return [7, 8] + + request = ChatCompletionRequest( + model="test", + messages=[], + chat_template_kwargs=template_kwargs, + reasoning_effort=effort, + ) + result = api_openai._get_guided_reasoning_end(request, SimpleNamespace(encode=encode)) + assert encoded == ([marker] if marker else []) + assert result == ([7, 8] if marker else []) + + +@pytest.mark.parametrize("response_type", ["json_schema", "json_object", "text"]) +@pytest.mark.parametrize("thinking", [False, True]) +def test_chat_passes_reasoning_marker_through_shared_sampling_params(monkeypatch, compiler, response_type, thinking): + monkeypatch.setattr(api_openai, "get_env_start_args", lambda: StartArgs(reasoning_parser="qwen3")) + monkeypatch.setattr(build_prompt, "tokenizer_supports_force_thinking", lambda: True) + + async def prompt(*args): + return "prompt" + + monkeypatch.setattr(api_openai, "build_prompt", prompt) + captured = [] + + class RequestCaptured(Exception): + pass + + async def generate(prompt, params, *args, **kwargs): + # Exercise the real API-to-shared-memory and grammar-state boundary. + copy = SamplingParams.from_buffer_copy(bytes(params)) + payload = await copy.get_compiled_grammar(compiler) + shared_req = SimpleNamespace( + sample_params=copy, + get_compiled_grammar=lambda: payload, + input_len=1, + shm_prompt_ids=SimpleNamespace(arr=[ord("P")]), + ) + infer_params = InferSamplingParams(shared_req, vocab_size=257) + captured.append(compiler.output_constraint_state_factory.create_state(shared_req, infer_params)) + raise RequestCaptured + yield + + tokenizer = SimpleNamespace(encode=lambda *args, **kwargs: [7, 8]) + manager = SimpleNamespace(tokenizer=tokenizer, generate=generate, output_grammar_compiler=compiler) + monkeypatch.setattr(api_http, "g_objs", SimpleNamespace(httpserver_manager=manager)) + response_format = {"type": response_type} + if response_type == "json_schema": + response_format["json_schema"] = {"name": "result", "schema": {"type": "object"}} + request = ChatCompletionRequest( + model="test", + messages=[], + response_format=response_format, + chat_template_kwargs={"enable_thinking": thinking}, + ) + with pytest.raises(RequestCaptured): + asyncio.run(api_openai.chat_completions_impl(request, None)) + expected = (7, 8) if thinking and response_type != "text" else () + state = captured[0] + if response_type == "text": + assert state is None + else: + assert state.error is None + assert state.reasoning_end_token_ids == expected + assert state.in_reasoning == bool(expected) + + +def test_reasoning_marker_survives_serialization_and_parameter_reuse(): + params = SamplingParams() + params.init(None, guided_reasoning_end=[11, 12]) + copy = SamplingParams.from_buffer_copy(ctypes.string_at(ctypes.addressof(params), ctypes.sizeof(params))) + round_trip = SamplingParams() + round_trip.init(None, **copy.to_origin_dict()) + assert round_trip.constraint_out_params.guided_reasoning_end.to_list() == [11, 12] + round_trip.init(None) + assert round_trip.constraint_out_params.guided_reasoning_end.to_list() == [] + + +@pytest.mark.parametrize( + "parser,marker", + [ + ("deepseek-r1", ""), + ("qwen3-thinking", ""), + ("gpt-oss", "<|channel|>final<|message|>"), + ("minimax", ""), + ("qwen3", None), + ], +) +def test_fixed_reasoning_does_not_require_a_template_thinking_switch(monkeypatch, parser, marker): + monkeypatch.setattr(api_openai, "get_env_start_args", lambda: StartArgs(reasoning_parser=parser)) + monkeypatch.setattr(build_prompt, "tokenizer_supports_force_thinking", lambda: False) + tokenizer = SimpleNamespace(encode=lambda text, **kwargs: list(text.encode())) + request = ChatCompletionRequest(model="test", messages=[]) + assert api_openai._is_force_thinking_mode(request) == (marker is not None) + assert api_openai._get_guided_reasoning_end(request, tokenizer) == (list(marker.encode()) if marker else []) + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("response_type", ["json_object", "json_schema"]) +def test_deepseek_r1_json_answer_stays_in_content(monkeypatch, compiler, stream, response_type): + monkeypatch.setattr(api_openai, "get_env_start_args", lambda: StartArgs(reasoning_parser="deepseek-r1")) + monkeypatch.setattr(build_prompt, "tokenizer_supports_force_thinking", lambda: False) + monkeypatch.setattr(api_openai, "build_prompt", AsyncMock(return_value="")) + monkeypatch.setattr("lightllm.server.reasoning_parser.get_token_id", lambda token: ord(">")) + tokenizer = SimpleNamespace(encode=lambda text, **kwargs: list(text.encode())) + reasoning = "Let me think." + prefix = reasoning + "" + answer = '{"ok":true}' + + async def generate(prompt, params, *args, **kwargs): + payload = await params.get_compiled_grammar(compiler) + state = ConstraintState.from_grammar( + compiler.output_constraint_state_factory.get_grammar(payload), params.constraint_out_params + ) + bitmask = torch.empty((1, 9), dtype=torch.int32) + for index, token_id in enumerate([*tokenizer.encode(prefix + answer), 256]): + # The real matcher must allow thinking and its delimiter before constraining JSON. + masked = state.fill_masks(bitmask, []) + assert masked == (index >= len(prefix)) + if masked: + assert (int(bitmask[0, token_id // 32]) >> (token_id % 32)) & 1 + state.commit(token_id) + assert state.error is None + finish = FinishStatus(FinishStatus.FINISHED_STOP) if token_id == 256 else FinishStatus() + yield 0, "" if token_id == 256 else chr(token_id), {"id": token_id, "prompt_tokens": 1}, finish + + manager = SimpleNamespace(tokenizer=tokenizer, generate=generate, output_grammar_compiler=compiler) + monkeypatch.setattr(api_http, "g_objs", SimpleNamespace(httpserver_manager=manager)) + response_format = {"type": response_type} + if response_type == "json_schema": + response_format["json_schema"] = {"name": "result", "schema": {"const": {"ok": True}}} + request = ChatCompletionRequest(model="test", messages=[], response_format=response_format, stream=stream) + + async def run(): + response = await api_openai.chat_completions_impl(request, None) + if not stream: + message = response.choices[0].message + return message.reasoning, message.content + deltas = [] + async for chunk in response.body_iterator: + chunk = chunk.decode() if isinstance(chunk, bytes) else chunk + if chunk.strip() != "data: [DONE]": + for choice in json.loads(chunk.removeprefix("data: "))["choices"]: + deltas.append(choice["delta"]) + return "".join(delta.get("reasoning", "") for delta in deltas), "".join( + delta.get("content", "") for delta in deltas + ) + + actual_reasoning, content = asyncio.run(run()) + assert actual_reasoning == reasoning + assert content == answer diff --git a/unit_tests/server/router/model_infer/structured_output/test_sampling.py b/unit_tests/server/router/model_infer/structured_output/test_sampling.py new file mode 100644 index 0000000000..8043b70ec9 --- /dev/null +++ b/unit_tests/server/router/model_infer/structured_output/test_sampling.py @@ -0,0 +1,320 @@ +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +import torch + +from unit_tests.server.grammar_helpers import ( + allowed, + commit_token, + make_req, + init_request, +) +from lightllm.common.req_manager import req_sampling_params +from lightllm.common.basemodel.triton_kernel.mtp_utils import mtp_verify +from lightllm.server.router.model_infer.mode_backend import generic_post_process as sampling +from lightllm.server.router.model_infer.mtp_speculative.utils import scatter_mtp_next_tokens + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="Sampling integration requires CUDA") + + +@pytest.fixture +def sampling_manager(request, monkeypatch): + args = SimpleNamespace( + penalty_counter_mode=getattr(request, "param", "cpu_counter"), + model_dir="unused", + mtp_step=3, + mtp_dynamic_verify=True, + output_constraint_mode="xgrammar", + ) + monkeypatch.setattr(req_sampling_params, "get_env_start_args", lambda: args) + monkeypatch.setattr(req_sampling_params, "get_vocab_size", lambda model_dir: 257) + result = req_sampling_params.ReqSamplingParamsManager(8) + assert result.req_to_next_token_ids.is_cuda + assert result.req_to_next_token_ids_cpu.device.type == "cpu" and result.req_to_next_token_ids_cpu.is_pinned() + assert result.req_to_next_token_scores.is_cuda + monkeypatch.setattr(sampling.g_infer_context, "req_manager", SimpleNamespace(req_sampling_params_manager=result)) + yield result + torch.cuda.synchronize() + + +def sampling_req(compiler, sampling_manager, req_idx=5, **constraints): + req = make_req(**({"regular_constraint": "abc"} | constraints)) + req.req_idx = req_idx + req.sampling_param.shm_param = SimpleNamespace( + constraint_out_params=req.sampling_param.shm_param.constraint_out_params, + presence_penalty=0.0, + frequency_penalty=0.0, + repetition_penalty=1.0, + exponential_decay_length_penalty=SimpleNamespace(to_tuple=lambda: (0, 1.1)), + input_penalty=False, + min_new_tokens=0, + temperature=1.0, + top_p=1.0, + top_k=1, + ) + req.get_last_gen_token = lambda: req.shm_req.shm_prompt_ids.arr[-1] + init_request(compiler, req) + sampling_manager.init_req_sampling_params(req) + return req + + +def test_sample_disables_stale_constraint_mask_for_ordinary_requests(compiler, sampling_manager, monkeypatch): + req = sampling_req(compiler, sampling_manager, regular_constraint=None) + sampling_manager.req_to_bitmask_enabled[req.req_idx] = True + sampling_manager.req_to_bitmask[req.req_idx].zero_() + + def unexpected_mask(*args, **kwargs): + pytest.fail("Ordinary requests must not launch the constraint mask kernel") + + monkeypatch.setattr(sampling, "apply_constraint_mask", unexpected_mask) + logits = torch.zeros(1, 257, device="cuda") + logits[:, ord("!")] = 10 + next_tokens, logprobs = sampling.sample(logits, [req], [256]) + + assert next_tokens.tolist() == [ord("!")] + assert torch.isfinite(logprobs).all() + assert not sampling_manager.req_to_bitmask_enabled[req.req_idx] + assert not compiler._cache + + +@pytest.mark.parametrize("sampling_manager", ["cpu_counter", "pin_mem_counter", "gpu_counter"], indirect=True) +def test_sample_reads_scattered_prefix_after_existing_handoff(compiler, sampling_manager, monkeypatch): + req = sampling_req(compiler, sampling_manager) + table = sampling_manager.req_to_next_token_ids + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + # This is the existing previous-step completion event, before post_handle. + previous_done = torch.cuda.Event() + with torch.cuda.stream(stream): + scatter_mtp_next_tokens( + backend=SimpleNamespace(model=SimpleNamespace(req_manager=sampling.g_infer_context.req_manager)), + proposal=SimpleNamespace(token_ids=torch.tensor([[ord("b"), ord("!")]], device="cuda")), + b_req_mtp_start_loc=torch.tensor([0], dtype=torch.int32, device="cuda"), + target_next_token_ids=torch.tensor([ord("a")], device="cuda"), + b_req_idx=torch.tensor([req.req_idx], dtype=torch.int32, device="cuda"), + mtp_accept_len=torch.tensor([1], dtype=torch.int32, device="cuda"), + ) + previous_done.record() + previous_done.synchronize() + commit_token(req, ord("a")) + + def unexpected_copy(*args, **kwargs): + pytest.fail("Sampling must read the ready CPU snapshot without starting another D2H") + + monkeypatch.setattr(sampling.g_pin_mem_manager, "async_copy_from_gpu_tensor", unexpected_copy) + monkeypatch.setattr(sampling.g_pin_mem_manager, "async_copy_from_gpu_tensor_with_event", unexpected_copy) + tensor_tolist = torch.Tensor.tolist + + def cpu_tolist(tensor): + assert not tensor.is_cuda, "Sampling synchronized a GPU tensor to read the prefix" + return tensor_tolist(tensor) + + logits = torch.zeros(3, 257, device="cuda") + logits[:, ord("!")] = 10 + # Three selected verify rows use a prefix of the four-column request table. + with monkeypatch.context() as patcher: + patcher.setattr(torch.Tensor, "tolist", cpu_tolist) + next_tokens, logprobs = sampling.sample( + logits, [req] * 3, [256], b_mtp_index=torch.arange(3, dtype=torch.int32, device="cuda") + ) + assert next_tokens.tolist() == [ord("b"), ord("c"), ord("!")] + assert torch.isfinite(logprobs).all() + assert allowed([req])[0].nonzero().flatten().tolist() == [ord("b")] + + accept_len, accepted = mtp_verify( + table, + torch.tensor([0], dtype=torch.int32, device="cuda"), + next_tokens, + torch.full((3,), req.req_idx, dtype=torch.int32, device="cuda"), + ) + assert accept_len.is_cuda and accepted.is_cuda + assert accept_len.tolist() == [2] and accepted.tolist() == [1, 1, 0] + assert table[req.req_idx].tolist() == [ord("a"), ord("b"), ord("!"), 1] + + +def test_sample_does_not_activate_constraints_for_partial_prefill(compiler, sampling_manager): + req = sampling_req(compiler, sampling_manager) + req.output_constraint.fill_masks = MagicMock(side_effect=AssertionError("Partial prefill must not fill masks")) + logits = torch.zeros(1, 257, device="cuda") + logits[:, ord("!")] = 10 + + next_tokens, logprobs = sampling.sample(logits, [req], [256], has_output=[False]) + + assert next_tokens.tolist() == [ord("!")] + assert torch.isfinite(logprobs).all() + req.output_constraint.fill_masks.assert_not_called() + assert req.output_constraint.error is None + + +@pytest.mark.parametrize("failure", ["fill_next_token_bitmask", "rollback"]) +def test_sample_preserves_valid_masks_and_isolates_failed_requests(compiler, sampling_manager, monkeypatch, failure): + import xgrammar as xgr + + bad = sampling_req(compiler, sampling_manager, req_idx=0) + good = sampling_req(compiler, sampling_manager, req_idx=1, regular_constraint="cd") + bad_matcher = bad.output_constraint.matcher + original = getattr(xgr.GrammarMatcher, failure) + + def fail_for_bad_request(matcher, *args): + if matcher is bad_matcher and (failure == "rollback" or args[1] == 1): + if failure == "fill_next_token_bitmask": + args[0].zero_() + raise RuntimeError("injected grammar failure") + return original(matcher, *args) + + monkeypatch.setattr(xgr.GrammarMatcher, failure, fail_for_bad_request) + + sampling_manager.req_to_next_token_ids[bad.req_idx] = torch.tensor([0, *b"ab!"]) + copy_snapshot_and_wait(sampling_manager) + sampling_manager.req_to_bitmask_enabled[bad.req_idx] = True + sampling_manager.req_to_bitmask[bad.req_idx].zero_() + logits = torch.zeros(3, 257, device="cuda") + logits[:, ord("!")] = 10 + tokens, logprobs = sampling.sample( + logits, + [bad, bad, good], + [256], + b_mtp_index=torch.tensor([0, 1, 0], dtype=torch.int32, device="cuda"), + ) + + expected_tokens = b"abc" if failure == "rollback" else b"!!c" + assert tokens.tolist() == list(expected_tokens) + assert torch.isfinite(logprobs).all() + assert bool(sampling_manager.req_to_bitmask_enabled[bad.req_idx]) == (failure == "rollback") + assert sampling_manager.req_to_bitmask_enabled[good.req_idx] + assert bad.output_constraint.error == "injected grammar failure" + assert good.output_constraint.error is None + + +def test_sample_groups_mixed_verify_rows_and_preserves_speculative_reasoning(compiler, sampling_manager): + active = sampling_req(compiler, sampling_manager, req_idx=0) + thinking = sampling_req(compiler, sampling_manager, req_idx=1, guided_reasoning_end=tuple(b"]>")) + plain = sampling_req(compiler, sampling_manager, req_idx=2, regular_constraint=None) + partial = sampling_req(compiler, sampling_manager, req_idx=3) + commit_token(active, ord("a")) + commit_token(thinking, ord("]")) + table = sampling_manager.req_to_next_token_ids + table[active.req_idx] = torch.tensor(list(b"ab!!")) + table[thinking.req_idx] = torch.tensor(list(b"]>a!")) + copy_snapshot_and_wait(sampling_manager) + + partial.output_constraint.fill_masks = MagicMock(side_effect=AssertionError("Partial prefill must not fill masks")) + # Dynamic verification can keep a different prefix length per request. + reqs = [plain, active, active, thinking, thinking, thinking, plain, partial] + logits = torch.zeros(len(reqs), 257, device="cuda") + logits[:, ord("!")] = 10 + next_tokens, logprobs = sampling.sample( + logits, + reqs, + [256], + has_output=[True] * 7 + [False], + b_mtp_index=torch.tensor([0, 0, 1, 0, 1, 2, 0, 0], dtype=torch.int32, device="cuda"), + ) + + assert next_tokens.tolist() == list(b"!bc!ab!!") + assert torch.isfinite(logprobs).all() + assert allowed([active])[0].nonzero().flatten().tolist() == [ord("b")] + assert thinking.output_constraint.in_reasoning + assert thinking.output_constraint.reasoning_tail == (ord("]"),) + partial.output_constraint.fill_masks.assert_not_called() + assert partial.output_constraint.error is None + assert not commit_token(thinking, ord(">")) + assert allowed([thinking])[0].nonzero().flatten().tolist() == [ord("a")] + + +def test_reused_mask_slots_clear_an_unreachable_draft_suffix(compiler, sampling_manager): + req = sampling_req(compiler, sampling_manager) + positions = torch.arange(4, dtype=torch.int32, device="cuda") + sampling_manager.req_to_next_token_ids[req.req_idx] = torch.tensor([0, *b"abc"]) + copy_snapshot_and_wait(sampling_manager) + logits = torch.zeros(4, 257, device="cuda") + logits[:, ord("!")] = 10 + tokens, _ = sampling.sample(logits, [req] * 4, [256], b_mtp_index=positions) + assert tokens.tolist() == [*b"abc", 256] + completed = torch.cuda.Event() + completed.record() + completed.synchronize() + + # The second prefix rejects its draft at '!', so masks from the previous + # iteration must not constrain that unreachable suffix. + sampling_manager.req_to_next_token_ids[req.req_idx] = torch.tensor([0, *b"a!c"]) + copy_snapshot_and_wait(sampling_manager) + logits.zero_() + logits[:, ord("!")] = 10 + tokens, _ = sampling.sample(logits, [req] * 4, [256], b_mtp_index=positions) + assert tokens.tolist() == list(b"ab!!") + assert allowed([req])[0].nonzero().flatten().tolist() == [ord("a")] + + +@pytest.mark.parametrize("replacement", ["plain", "partial", "thinking", "finished"]) +def test_reused_request_slot_disables_stale_mask(compiler, sampling_manager, replacement): + from lightllm.server.core.objs import FinishStatus + + old = sampling_req(compiler, sampling_manager, req_idx=0) + tokens, _ = sampling.sample(torch.zeros(1, 257, device="cuda"), [old], [256]) + assert tokens.tolist() == [ord("a")] + completed = torch.cuda.Event() + completed.record() + completed.synchronize() + assert sampling_manager.req_to_bitmask_enabled[0] + + constraints = {"regular_constraint": None} if replacement == "plain" else {} + if replacement == "thinking": + constraints["guided_reasoning_end"] = (ord("]"),) + new = sampling_req(compiler, sampling_manager, req_idx=0, **constraints) + if replacement == "finished": + new.finish_status.set_status(FinishStatus.FINISHED_STOP) + good = sampling_req(compiler, sampling_manager, req_idx=1, regular_constraint="xy") + logits = torch.zeros(2, 257, device="cuda") + logits[:, ord("!")] = 10 + tokens, _ = sampling.sample(logits, [new, good], [256], has_output=[replacement != "partial", True]) + + assert tokens.tolist() == list(b"!x") + assert not sampling_manager.req_to_bitmask_enabled[0] + assert sampling_manager.req_to_bitmask_enabled[1] + + +@pytest.mark.parametrize("constraints_enabled,draft_step", [(True, 2), (False, 2), (True, 0)]) +def test_mtp_scatter_reuses_snapshot_and_only_copies_when_needed(sampling_manager, constraints_enabled, draft_step): + table = sampling_manager.req_to_next_token_ids + snapshot = sampling_manager.req_to_next_token_ids_cpu + snapshot.fill_(-42) + snapshot_ptr = snapshot.data_ptr() + if not constraints_enabled: + sampling_manager.req_to_bitmask_enabled = None + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + backend = SimpleNamespace(model=SimpleNamespace(req_manager=sampling.g_infer_context.req_manager)) + for first_draft in [10, 20]: + with torch.cuda.stream(stream): + drafts = torch.arange(first_draft, first_draft + draft_step, device="cuda").reshape(1, draft_step) + scatter_mtp_next_tokens( + backend=backend, + proposal=SimpleNamespace(token_ids=drafts), + target_next_token_ids=torch.tensor([7], device="cuda"), + b_req_mtp_start_loc=torch.tensor([0], dtype=torch.int32, device="cuda"), + b_req_idx=torch.tensor([0], dtype=torch.int32, device="cuda"), + mtp_accept_len=torch.tensor([1], dtype=torch.int32, device="cuda"), + ) + completed = torch.cuda.Event() + completed.record() + completed.synchronize() + assert snapshot.data_ptr() == snapshot_ptr + if constraints_enabled: + expected = [7, first_draft, first_draft + 1, 1] if draft_step else [7, 1, 1, 1] + assert snapshot[0].tolist() == expected + # CPU snapshot is independent of subsequent writes to the GPU table. + table[0].fill_(99) + torch.cuda.current_stream().synchronize() + assert snapshot[0].tolist() == expected + else: + assert snapshot.eq(-42).all() + + +def copy_snapshot_and_wait(manager): + manager.req_to_next_token_ids_cpu.copy_(manager.req_to_next_token_ids, non_blocking=True) + completed = torch.cuda.Event() + completed.record() + completed.synchronize() diff --git a/unit_tests/server/router/model_infer/structured_output/test_state.py b/unit_tests/server/router/model_infer/structured_output/test_state.py new file mode 100644 index 0000000000..aa09c7df2e --- /dev/null +++ b/unit_tests/server/router/model_infer/structured_output/test_state.py @@ -0,0 +1,407 @@ +import asyncio +from itertools import count +from types import MethodType, SimpleNamespace +from weakref import ref +from unittest.mock import patch + +import pytest +import torch +import xgrammar as xgr + +from lightllm.server.core.objs import FinishStatus, SamplingParams +from lightllm.server.httpserver.grammar import OutputGrammarCompiler +from unit_tests.server.grammar_helpers import ( + make_req, + init_request, + make_mask_buffers, + prepare_sampling_tensors, + build_mask, + apply_masks, + allowed, + commit_token, +) +from lightllm.server.router.model_infer import infer_batch + + +def test_sampling_preparation_without_constraint_buffers(): + req = make_req() + manager = SimpleNamespace(req_to_bitmask=None, req_to_bitmask_enabled=None) + tensors = prepare_sampling_tensors([req, req], manager) + assert tensors[0].tolist() == [req.req_idx, req.req_idx] + assert tensors[1].tolist() == [1.0, 1.0] + assert not tensors[-1] + + +@pytest.mark.parametrize("mode", ["none", "xgrammar"]) +def test_context_initializes_state_factory_with_resolved_eos_at_startup(monkeypatch, mode): + args = SimpleNamespace( + output_constraint_mode=mode, model_dir="unused", tokenizer_mode="auto", trust_remote_code=False, eos_id=[255] + ) + monkeypatch.setattr(infer_batch, "get_env_start_args", lambda: args) + tokenizer = SimpleNamespace(eos_token_id=256) + tokenizer_info = xgr.TokenizerInfo([bytes([i]) for i in range(256)] + [b""], stop_token_ids=[255]) + payload = xgr.GrammarCompiler(tokenizer_info).compile_regex("a").serialize_json().encode() + context = infer_batch.InferenceContext() + with ( + patch("lightllm.server.tokenizer.get_tokenizer", return_value=tokenizer) as load_tokenizer, + patch.object(xgr.TokenizerInfo, "from_huggingface", return_value=tokenizer_info) as build_tokenizer_info, + patch.object(xgr, "GrammarCompiler", side_effect=AssertionError("Inference must not compile")), + ): + context.register( + backend=SimpleNamespace( + dp_rank_in_node=0, model=SimpleNamespace(att_state_manager=None), is_hybrid_att_model=False + ), + req_manager=SimpleNamespace(req_sampling_params_manager=None), + radix_cache=None, + shm_req_manager=None, + vocab_size=257, + ) + if mode == "none": + assert context.output_constraint_state_factory is None + load_tokenizer.assert_not_called() + build_tokenizer_info.assert_not_called() + return + load_tokenizer.assert_called_once_with("unused", "auto", trust_remote_code=False) + build_tokenizer_info.assert_called_once_with(tokenizer, vocab_size=257, stop_token_ids=[255]) + req = make_req(regular_constraint="a") + req.shm_req.get_compiled_grammar = lambda: payload + req.output_constraint = context.output_constraint_state_factory.create_state(req.shm_req, req.sampling_param) + commit_token(req, ord("a")) + assert req.output_constraint.matcher.accept_token(255) + assert req.output_constraint.is_terminated() + + +@pytest.mark.parametrize("previous_output_len", [-1, 2]) +def test_pd_history_outside_prompt_returns_constraint_error(compiler, previous_output_len): + req = make_req(regular_constraint="ab") + req.sampling_param.shm_param.constraint_out_params.pd_previous_output_len = previous_output_len + init_request(compiler, req) + assert "pd_previous_output_len" in req.output_constraint.error + assert req.output_constraint.matcher is None + assert req.output_constraint.is_terminated() + + +def test_thinking_does_not_advance_matcher_or_enable_mask(compiler): + req = make_req(regular_constraint="ab", guided_reasoning_end=(ord("]"),)) + init_request(compiler, req) + # The state owns its marker snapshot even if the shared parameters change. + req.sampling_param.shm_param.constraint_out_params.guided_reasoning_end.initialize([ord("?")]) + manager = make_mask_buffers([req]) + manager.req_to_bitmask.fill_(42) + assert not prepare_sampling_tensors([req], manager)[-1] + assert (manager.req_to_bitmask[req.req_idx] == -1).all() + assert not manager.req_to_bitmask_enabled[req.req_idx] + assert not commit_token(req, ord("?")) + assert req.output_constraint.error is None + assert not commit_token(req, ord("]")) + mask = allowed([req]) + assert mask[0, ord("a")] and not mask[0, ord("?")] + + +@pytest.mark.parametrize("marker", [b"]", b"aba"]) +def test_reasoning_boundary_mixed_batch_and_shared_schema(compiler, marker): + thinking = make_req(guided_json='{"const":{"ok":true}}', guided_reasoning_end=tuple(marker)) + direct = make_req(guided_json='{"const":{"ok":true}}') + for req in (thinking, direct): + init_request(compiler, req) + assert compiler.output_constraint_state_factory.get_grammar.cache_info().currsize == 1 + mask = allowed([thinking, direct]) + assert mask[0].all() and not mask[1, ord("?")] + assert thinking.output_constraint.matcher is not direct.output_constraint.matcher + + # A failed partial marker must not activate constraints. The prompt is ignored. + for token in b"thought aab?" + marker[:-1]: + assert not commit_token(thinking, token) + assert allowed([thinking, direct])[0].all() + # Sampling preparation observes the delimiter committed by post_handle. + assert not commit_token(thinking, marker[-1]) + mask = build_mask([thinking, direct]) + logits = torch.zeros(2, compiler.vocab_size) + apply_masks([thinking, direct], logits, mask) + assert torch.isfinite(logits[:, ord("{")]).all() + assert not torch.isfinite(logits[:, ord("?")]).any() + + for token in b'{"ok":true}': + assert allowed([thinking])[0, token] + assert not commit_token(thinking, token) + assert allowed([thinking])[0, 256] + assert commit_token(thinking, 256) + assert allowed([direct])[0, ord("{")] + + +def test_unconstrained_batch_does_not_compile_or_allocate_mask(compiler): + req = make_req() + init_request(compiler, req) + assert allowed([req]).all() + assert req.output_constraint is None + assert not compiler._cache + + +@pytest.mark.parametrize("current_dp_rank", [0, 1]) +@pytest.mark.parametrize("constraints_enabled", [False, True]) +def test_request_admission_loads_shared_artifacts_without_compilation( + compiler, monkeypatch, current_dp_rank, constraints_enabled +): + payload = asyncio.run(compiler.compile("regex", "ab")) + shared_reqs = {} + requests = [] + for index, req_id in enumerate((1, 2, 3, 4)): + owner_dp_rank = 1 - current_dp_rank if req_id == 4 else current_dp_rank + params = SamplingParams() + params.init(None, regular_constraint="" if req_id == 3 else "ab") + + def read_artifact(req_id=req_id, owner_dp_rank=owner_dp_rank): + # Disabled constraints and KV-only donor requests must not load grammar artifacts. + assert constraints_enabled and owner_dp_rank == current_dp_rank + return None if req_id == 3 else payload + + shared_reqs[index] = SimpleNamespace( + sample_params=params, + input_len=1, + shm_prompt_ids=SimpleNamespace(arr=[ord("P")]), + link_prompt_ids_shm_array=lambda: None, + link_logprobs_shm_array=lambda: None, + get_compiled_grammar=read_artifact, + ) + requests.append((req_id, index, SimpleNamespace(to_dict=lambda: {}), owner_dp_rank)) + monkeypatch.setattr(compiler, "_compile", lambda *args: pytest.fail("Inference must not compile")) + monkeypatch.setattr( + infer_batch, + "get_env_start_args", + lambda: SimpleNamespace( + diverse_mode=False, mtp_step=0, disable_chunked_mm_input=False, req_token_index_mode="pin_mem_cache" + ), + ) + context = infer_batch.InferenceContext( + req_manager=SimpleNamespace( + alloc=count().__next__, + req_to_token_indexs=torch.zeros((len(requests), 1), dtype=torch.int32), + req_sampling_params_manager=SimpleNamespace(init_req_sampling_params=lambda req: None), + ), + shm_req_manager=SimpleNamespace(get_req_obj_by_index=lambda index: shared_reqs[index]), + requests_mapping={}, + output_constraint_state_factory=compiler.output_constraint_state_factory if constraints_enabled else None, + ) + context.backend = SimpleNamespace(dp_rank_in_node=current_dp_rank) + context.infer_req_ids = [] + context.vocab_size = compiler.vocab_size + monkeypatch.setattr(infer_batch, "g_infer_context", context) + # Exercise the real InferReq constructor, including its internal grammar initialization. + context.add_reqs(requests) + reqs = context.requests_mapping + assert [reqs[req_id].dp_rank_in_node for req_id in (1, 2, 3, 4)] == [ + current_dp_rank, + current_dp_rank, + current_dp_rank, + 1 - current_dp_rank, + ] + if not constraints_enabled: + assert all(req.output_constraint is None for req in reqs.values()) + return + assert reqs[3].output_constraint is None and reqs[4].output_constraint is None + assert compiler.output_constraint_state_factory.get_grammar.cache_info().misses == 1 + assert reqs[1].output_constraint.matcher is not reqs[2].output_constraint.matcher + mask = allowed([reqs[1], reqs[2], reqs[3]]) + assert mask[:2, ord("a")].all() and mask[2].all() + commit_token(reqs[1], ord("a")) + mask = allowed([reqs[2], reqs[1]]) + assert mask[0, ord("a")] and not mask[0, ord("b")] + assert mask[1, ord("b")] and not mask[1, ord("a")] + + +@pytest.mark.parametrize("failure", ["read", "deserialize", "empty", "matcher"]) +def test_artifact_load_failure_finishes_only_its_request(compiler, monkeypatch, failure): + bad = make_req(regular_constraint="ab") + good = make_req(regular_constraint="cd") + init_request(compiler, good) + payload = asyncio.run(compiler.compile("regex", "ab")) + + def read_artifact(): + if failure == "read": + raise FileNotFoundError("missing compiled artifact") + if failure == "empty": + return b"" + return b"invalid artifact" if failure == "deserialize" else payload + + def fail_to_create_matcher(*args): + raise RuntimeError("matcher failed") + + bad.shm_req.get_compiled_grammar = read_artifact + with monkeypatch.context() as patcher: + if failure == "matcher": + patcher.setattr(xgr, "GrammarMatcher", fail_to_create_matcher) + bad.output_constraint = compiler.output_constraint_state_factory.create_state(bad.shm_req, bad.sampling_param) + assert bad.output_constraint.error is not None + assert build_mask([bad]) is None + bad.set_next_gen_token_id = lambda *args, **kwargs: None + bad.update_finish_status = lambda **kwargs: None + infer_batch.InferReqUpdatePack(bad, 1).handle(ord("!"), 0.0, -1, [256], is_master_in_dp=False) + assert bad.finish_status.is_finished_error() + assert allowed([good])[0].nonzero().flatten().tolist() == [ord("c")] + + +def test_constraint_state_lifetime_follows_request_cleanup(compiler): + req = make_req(regular_constraint="ab") + req.req_idx = 7 + init_request(compiler, req) + allowed([req]) + state_ref = ref(req.output_constraint) + payload = req.shm_req.get_compiled_grammar() + grammar = compiler.output_constraint_state_factory.get_grammar(payload) + context = infer_batch.InferenceContext( + req_manager=SimpleNamespace(free=lambda *args: None), + requests_mapping={1: req}, + shm_req_manager=SimpleNamespace(put_back_req_obj=lambda req: None), + ) + context.infer_req_ids = [1] + context.args = SimpleNamespace(diverse_mode=False) + context.backend = SimpleNamespace(is_master_in_dp=False) + context.free_a_req_mem = lambda *args: None + + context._filter([1], modify_shm_finish_state=False) + assert context.requests_mapping == {} and context.infer_req_ids == [] + del req + assert state_ref() is None + assert compiler.output_constraint_state_factory.get_grammar(payload) is grammar + + +def test_mixed_batch_partial_prefill_and_finished_requests(compiler): + partial = make_req(regular_constraint="ab") + active = make_req(regular_constraint="cd") + plain = make_req() + for req in [partial, active, plain]: + init_request(compiler, req) + mask = allowed([partial, active, plain], [False, True, True]) + assert mask[0].all() and mask[2].all() + assert mask[1, ord("c")] and not mask[1, ord("a")] + assert allowed([partial])[0].nonzero().flatten().tolist() == [ord("a")] + + active.finish_status.set_status(FinishStatus.FINISHED_STOP) + assert build_mask([active]) is None + + +@pytest.mark.parametrize("is_master_in_dp", [False, True]) +@pytest.mark.parametrize( + "reasoning_end,history,first_token,expected_next", + [ + (b"", b"", ord("a"), ord("b")), + (b"]", b"", ord("]"), ord("a")), + (b"]>", b"", ord("]"), None), + (b"", b"", ord("!"), None), + (b"", b"a", ord("b"), 256), + (b"]>", b"thought]", ord(">"), ord("a")), + ], +) +def test_pd_first_token_updates_constraint_before_sampling( + compiler, monkeypatch, is_master_in_dp, reasoning_end, history, first_token, expected_next +): + from lightllm.server.router.model_infer.mode_backend.base_backend import ModeBackend + + req = make_req(regular_constraint="ab", guided_reasoning_end=tuple(reasoning_end)) + req.shm_req.shm_prompt_ids.arr.extend(history) + req.shm_req.input_len += len(history) + req.sampling_param.shm_param.constraint_out_params.pd_previous_output_len = len(history) + good = make_req(regular_constraint="cd") + for request in (req, good): + init_request(compiler, request) + req.pd_task_success_num = 0 + req.shm_req.shm_prompt_ids.arr.append(0) + req.shm_req.shm_logprobs = SimpleNamespace(arr=[None] * (req.shm_req.input_len + 1)) + req.set_next_gen_token_id = MethodType(infer_batch.InferReq.set_next_gen_token_id, req) + req.update_finish_status = lambda **kwargs: None + sampling_params = make_mask_buffers([req, good]) + transfer = SimpleNamespace( + request_id=req.req_idx, + has_error=False, + first_gen_token_id=first_token, + first_gen_token_logprob=0.0, + ) + backend = SimpleNamespace( + shm_pd_trans_io_buffer=SimpleNamespace(read_obj=lambda: [transfer], sub_state=lambda: None), + is_pd_decode_mode=True, + model=SimpleNamespace(req_manager=SimpleNamespace(req_sampling_params_manager=sampling_params)), + eos_id=[256], + is_master_in_dp=is_master_in_dp, + ) + monkeypatch.setattr(infer_batch.g_infer_context, "requests_mapping", {req.req_idx: req}) + monkeypatch.setattr(torch.cuda, "current_stream", lambda: SimpleNamespace(synchronize=lambda: None)) + + # Exercise the real PD receive handler and InferReqUpdatePack.handle() + # before any sampling preparation can update the constraint. + ModeBackend._read_pd_trans_io_buffer_and_update_req_status(backend) + + assert req.cur_output_len == 1 and req.pd_task_success_num == 1 + assert req.shm_req.shm_prompt_ids.arr == [ord("P")] + list(history) + [first_token] + assert sampling_params.req_to_next_token_ids[req.req_idx, 0].item() == first_token + if is_master_in_dp: + assert req.shm_req.shm_cur_output_len == 1 + state = req.output_constraint + if first_token == ord("!"): + assert req.finish_status.status == FinishStatus.FINISHED_ERROR + assert "Grammar rejected sampled token" in state.error + assert allowed([good])[0].nonzero().flatten().tolist() == [ord("c")] + assert good.output_constraint.error is None + return + + assert state.error is None and not req.finish_status.is_finished() + if expected_next is None: + assert state.in_reasoning and state.reasoning_tail == (first_token,) + assert allowed([req])[0].all() + assert not commit_token(req, ord(">")) + expected_next = ord("a") + + # Repeated mask preparation must observe the same committed prefix. + for _ in range(2): + mask = allowed([req, good]) + assert [row.nonzero().flatten().tolist() for row in mask] == [[expected_next], [ord("c")]] + + +@pytest.mark.parametrize( + "constraints,text", + [ + ({"guided_grammar": 'root ::= "hello"'}, "hello"), + ( + { + "guided_json": '{"type":"object","properties":{"x":{"const":1}},' + '"required":["x"],"additionalProperties":false}' + }, + '{"x":1}', + ), + ({"guided_grammar": "json"}, '{"x":1}'), + ], +) +def test_grammar_and_schema_enforce_each_token(compiler, constraints, text): + req = make_req(**constraints) + init_request(compiler, req) + for token in text.encode(): + mask = allowed([req]) + assert mask[0, token] + req.output_constraint.commit(token) + assert not req.output_constraint.is_terminated() + assert allowed([req])[0, 256] + req.output_constraint.commit(256) + assert req.output_constraint.is_terminated() + + +def test_json_object_does_not_accept_array_root(compiler): + req = make_req(guided_grammar="json") + init_request(compiler, req) + mask = allowed([req]) + assert mask[0, ord("{")] + assert not mask[0, ord("[")] + + +def test_wrapped_tokenizer_uses_underlying_hf_tokenizer(monkeypatch): + tokenizer = object() + wrapper = SimpleNamespace(tokenizer=SimpleNamespace(tokenizer=tokenizer)) + received = [] + + def from_huggingface(value, **kwargs): + received.append(value) + return xgr.TokenizerInfo([bytes([i]) for i in range(256)] + [b""], stop_token_ids=[256]) + + monkeypatch.setattr(xgr.TokenizerInfo, "from_huggingface", from_huggingface) + compiler = OutputGrammarCompiler(wrapper, 257, [256]) + asyncio.run(compiler.compile("regex", "ab")) + compiler.shutdown() + assert received == [tokenizer] 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..21bbb06339 --- /dev/null +++ b/unit_tests/server/router/model_infer/test_cpu_resource_layout.py @@ -0,0 +1,513 @@ +"""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.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 只有一个自有层，不能沿用 main 的多层类型列表。 + layer_types=["deepseek_sparse_attention"] if kind == "glm" else ["full_attention"], + ) + 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, + 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 + prepare_model_states(draft, logical_layer_ids=(config["num_hidden_layers"] + step,)) + drafts.append(draft) + 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_state_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=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") + assert main is not draft + 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.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_cache_operator().get_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=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: + 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_cache_operator().get_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_runtime_to_pages( + req_indices=[0], + state_token_lens=[8], + page_kind="big", + page_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.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_cache_operator().get_page_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.get_cpu_cache_operator().read_req( + cpu_page_buffers=buffer, + cpu_page_indices=[2], + token_indices=torch.arange( + start=8, + end=16, + dtype=torch.long, + device="cuda", + ), + state_token_len=8, + layouts=resources.get_cpu_cache_operator().get_page_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() + + +@pytest.mark.parametrize("hybrid", [False, True]) +def test_production_offload_preserves_ready_pages_without_host_sync( + 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=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) + 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.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.get_cpu_cache_operator(), "write_req", 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_offload_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 + token_indexes = torch.arange(16, device="cuda") + req.token_index = SimpleNamespace(get_gpu_token_indexes=lambda start, end: token_indexes[start:end]) + module.backend = SimpleNamespace( + is_master_in_dp=True, + is_hybrid_att_model=hybrid, + dp_world_size=1, + model=SimpleNamespace( + att_state_manager=resources, + req_manager=SimpleNamespace(req_to_token_indexs=token_indexes.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) + module.get_cpu_offload_stream = lambda: copy_stream + task = module.write_req_to_cpu(req) + assert waited_streams == [] + task.sync_event.synchronize() + assert write_batches == [(pages, [True, False])] + assert req.cpu_cache_offload_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.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_cache_operator().get_page_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( + 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=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.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_cache_operator().get_page_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.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_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 new file mode 100644 index 0000000000..514aa19bf0 --- /dev/null +++ b/unit_tests/server/router/model_infer/test_hybrid_cpu_state_cache.py @@ -0,0 +1,362 @@ +from unit_tests.common.att_resources.model_helpers import build_hybrid_resources +from unit_tests.common.req_token_index_test_utils import attach_req_token_index +import dataclasses +import json +from collections import deque +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.attention_state.allocators import FullAttKVAllocator + +from lightllm.common.attention_state.full_att import ( + GQAFullAttentionState, + GQAAttentionConfig, +) +from lightllm.common.req_manager import ReqManager +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.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"), + pytest.mark.usefixtures("kv_allocator_environment"), +] + + +@pytest.fixture(params=[False, True], ids=["default_chunks", "whole_mm_blocks"]) +def linear_cache(monkeypatch, tmp_path, request): + args = StartArgs( + tp=1, + page_size=4, + max_req_total_len=512, + linear_att_hash_page_size=64, + linear_att_page_block_num=2, + cpu_cache_token_page_size=128, + disable_chunked_mm_input=request.param, + ) + monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(dataclasses.asdict(args))) + monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_NODE", "0") + monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") + monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_DP", "0") + monkeypatch.setenv("LIGHTLLM_DP_WORLD_SIZE", "1") + monkeypatch.setenv("LIGHTLLM_UNIQUE_SERVICE_NAME_ID", tmp_path.name) + 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 = Qwen3NextResourceConfig( + tp_world_size=1, + full_att_all_num_kv_heads=1, + full_att_head_dim=8, + global_linear_k_heads=1, + global_linear_v_heads=1, + head_linear_k_dim=8, + head_linear_v_dim=8, + 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( + GQAAttentionConfig( + logical_layer_groups=tuple((layer,) for layer in config.get_full_att_layer_indices()), + global_head_num=1, + head_dim=8, + dtype=torch.bfloat16, + ), + page_size=args.page_size, + ) + full.allocate_storage(384) + requests = ReqManager(2, 512) + 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, + ) + 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( + g_infer_context, + "radix_cache", + SimpleNamespace( + get_big_page_ids_by_node=lambda node: [], + free_one_small_page_buffer=lambda: None, + free_radix_cache_to_get_enough_token=lambda need: None, + ), + ) + yield full, resources, requests, config + get_env_start_args.cache_clear() + get_unique_server_name.cache_clear() + + +@pytest.mark.parametrize("token_num", [320, 384]) +@pytest.mark.parametrize("token_index_mode", ["pin_mem_cache", "pin_mem_page_cache", "pure_gpu"]) +def test_cpu_roundtrip_requires_all_linear_pages_and_restores_final_state( + monkeypatch, + linear_cache, + cpu_cache_factory, + token_num, + token_index_mode, +): + 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 + + 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_conv_bf16_ssm_fp32_state + big = linear.big_pages_pool + small = linear.small_pages_pool + big_states = {} + for end in range(128, token_num + 1, 128): + slot = resources.bounded_att_big_page_allocator.alloc(1)[0] + big_states[end] = slot + for tensor in (big.conv_state[slot], big.ssm_state[slot]): + tensor.normal_() + tail_slot = None + if token_num % 128: + 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 = (small.conv_state[tail_slot], small.ssm_state[tail_slot]) + else: + 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.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, resources.get_cpu_cache_operator().get_page_bytes()), + ) + indexes = torch.arange(token_num, dtype=torch.int32, device="cuda") + req = SimpleNamespace( + req_idx=0, + cache_tiers=(CacheTier.CPU,), + cpu_cache_offload_task_status=InferReq._CpuCacheTaskStatus.NOT_STARTED, + cur_kv_len=token_num, + hold_kv_len=token_num, + hybrid_cache_len=token_num, + hybrid_len_to_big_page_id=big_states, + tail_small_page_buffer_id=tail_slot, + shared_kv_node=None, + mm_input_spans=[(96, 160)] if disable_chunked_mm_input else [], + sampling_param=SimpleNamespace(shm_param=SimpleNamespace(prompt_logprobs=-1)), + shm_req=SimpleNamespace( + 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]), + ), + ) + + # 通过接收 InferReq 的生产入口完成页来源查询、CPU 分配与快照写出。 + attach_req_token_index(req, requests, indexes.cpu(), mode=token_index_mode) + module = MultiLevelKvCacheModule.__new__(MultiLevelKvCacheModule) + module.load_tasks = deque() + 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_state_manager=resources), + ) + monkeypatch.setattr(common.dist, "broadcast_object_list", lambda *args, **kwargs: None) + module.get_cpu_offload_stream().wait_stream(torch.cuda.current_stream()) + task = module.write_req_to_cpu(req) + assert task is not None + task.sync_event.synchronize() + pages = task.page_indexes.tolist() + client.update_pages_status_to_ready(pages, deref=False) + resources.bounded_att_big_page_allocator.free(list(big_states.values())) + if tail_slot is not None: + 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 + req.token_index.reset() + full.k_storage.zero_() + full.v_storage.zero_() + big.conv_state.zero_() + big.ssm_state.zero_() + + module.args = get_env_start_args() + module.backend = SimpleNamespace( + is_master_in_dp=True, + dp_world_size=1, + model=SimpleNamespace(req_manager=requests, att_state_manager=resources), + is_hybrid_att_model=True, + ) + module.init_sync_group = None + module.need_sync_compute_stream = lambda: False + monkeypatch.setattr(common.g_infer_context, "get_can_alloc_token_num", lambda: 384) + monkeypatch.setattr(common.dist, "barrier", lambda **kwargs: None) + module.load_cpu_cache_to_reqs([req]) + torch.cuda.synchronize() + + assert req.cur_kv_len == req.hold_kv_len == token_num + 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)) + # 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.conv_state[slot], big.ssm_state[slot]), expected_big_states[end]): + torch.testing.assert_close(actual, expected) + torch.testing.assert_close(linear.conv_state[:, 0].cpu(), expected_states[0]) + torch.testing.assert_close(linear.ssm_state[:, 0].cpu(), expected_states[1]) + + # Model data validity and multimodal resume boundaries are separate checks. + tree = HybridAttPagedRadixCache( + total_token_num=384, + rank_in_node=103, + hash_page_size=64, + big_page_num=2, + att_state_manager=resources, + ) + cache_len = max(req.hybrid_len_to_big_page_id) + tree.insert( + 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), + len_to_big_page_id=req.hybrid_len_to_big_page_id, + ) + for length, expected in [(128, 128 if first_page_has_state else 0), (256, 256)]: + node, matched, _ = tree.match_prefix( + key=torch.arange(length), + block_hashs=list(range(1, length // 64 + 1)), + update_refs=True, + mm_input_spans=req.mm_input_spans, + ) + assert matched == expected + tree.dec_node_ref_counter(node) + req.shared_kv_node, _, _ = tree.match_prefix( + key=torch.arange(cache_len), + block_hashs=list(range(1, cache_len // 64 + 1)), + update_refs=True, + 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, + hold_kv_len=256, + cache_tiers=(CacheTier.CPU,), + shm_req=SimpleNamespace( + input_len=257, + token_hash_list=SimpleNamespace(get_all=lambda: [11, 22]), + token_hash_page_len_list=SimpleNamespace(get_all=lambda: [128, 256]), + ), + ) + attach_req_token_index(req, requests, torch.arange(256, dtype=torch.int32)) + 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_state_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.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) + if allocated_pages == 0: + assert task is None + assert req.cpu_cache_offload_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) + 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 new file mode 100644 index 0000000000..eb64c6ed1d --- /dev/null +++ b/unit_tests/server/router/model_infer/test_multimodal_chunking.py @@ -0,0 +1,404 @@ +from types import SimpleNamespace + +import pytest +import torch + + +@pytest.mark.parametrize("hybrid", [False, True]) +@pytest.mark.parametrize("spans", [[], [(450, 600)], [(100, 800)], [(0, 384), (384, 768)]]) +def test_chunk_ids_and_lengths_use_the_same_complete_boundaries(hybrid, spans): + from lightllm.server.router.model_infer.infer_batch import InferReq + + req = InferReq.__new__(InferReq) + req.args = SimpleNamespace( + chunked_prefill_size=256, + batch_max_tokens=1024, + linear_att_hash_page_size=64, + linear_att_page_block_num=8, + ) + req.cur_kv_len = req.cur_output_len = 0 + req.mm_input_spans = spans + req.hybrid_cache_len = 960 + req.shm_req = SimpleNamespace(input_len=1000, shm_prompt_ids=SimpleNamespace(arr=list(range(1000)))) + if hybrid: + req.get_chuncked_input_token_len = req.get_chuncked_input_token_len_for_hybrid_att + req.get_chuncked_input_token_ids = req.get_chuncked_input_token_ids_for_hybrid_att + + while req.cur_kv_len < req.get_cur_total_len(): + end = req.get_chuncked_input_token_len() + assert req.get_chuncked_input_token_ids() == list(range(end)) + assert 0 < end - req.cur_kv_len <= req.args.batch_max_tokens + assert not any(start < end < stop for start, stop in spans) + req.cur_kv_len = end + + +@pytest.mark.usefixtures("kv_allocator_environment") +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for bounded transfers") +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 + + monkeypatch.setattr( + "lightllm.common.attention_state.bounded_att.utils.get_env_start_args", + lambda: SimpleNamespace(enable_dp_prompt_cache_fetch=False), + ) + resources = composition(device="cuda") + slots = init_source(resources) + resources.reset_bounded_att_state([1]) + resources.restore_bounded_att_runtime_from_pages( + req_indices=[1], + state_token_lens=[8], + page_kind="big", + 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_state_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_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 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", + [ + (2049, 8, [], [512, 1024, 1536, 2048], []), + (2049, 8, [(450, 600)], [1024, 1536, 2048], []), + (2049, 8, [(1100, 1700)], [512, 1024, 2048], []), + (2049, 8, [(0, 700)], [1024, 1536, 2048], []), + (2049, 8, [(450, 1700)], [2048], []), + (2049, 8, [(1900, 2049)], [512, 1024, 1536], []), + (2010, 8, [], [512, 1024, 1536], [1984]), + (2010, 8, [(1900, 2000)], [512, 1024, 1536], []), + (2010, 8, [(1900, 1984)], [512, 1024, 1536], [1984]), + (2010, 10000000, [], [], [1984]), + (2010, 10000000, [(1900, 2000)], [], []), + ], +) +def test_hybrid_chunks_save_state_only_at_original_boundaries( + input_len, + big_page_num, + spans, + expected_big, + expected_small, +): + from lightllm.server.router.model_infer.infer_batch import InferReq, InferenceContext + + req = InferReq.__new__(InferReq) + req.args = SimpleNamespace( + chunked_prefill_size=256, + batch_max_tokens=2048, + linear_att_hash_page_size=64, + linear_att_page_block_num=big_page_num, + disable_chunked_prefill=False, + ) + req.cur_kv_len = req.cur_output_len = req.req_idx = 0 + req.mm_input_spans = spans + req.hybrid_cache_len = (input_len - 1) // 64 * 64 + req.hybrid_len_to_big_page_id = {} + req.tail_small_page_buffer_id = None + req.shm_req = SimpleNamespace(input_len=input_len) + req.get_chuncked_input_token_len = req.get_chuncked_input_token_len_for_hybrid_att + big_snapshots, small_snapshots = [], [] + page_ends = list(range(64 * big_page_num, req.hybrid_cache_len + 1, 64 * big_page_num)) + slots = iter(range(len(page_ends))) + context = InferenceContext() + context.args = req.args + context.is_hybrid_att_model = True + context.radix_cache = SimpleNamespace(free_one_small_page_buffer=lambda: None) + 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_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), + ) + + 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() + + assert big_snapshots == expected_big + assert small_snapshots == expected_small + assert sorted(req.hybrid_len_to_big_page_id) == page_ends + assert len(set(req.hybrid_len_to_big_page_id.values())) == len(page_ends) + assert big_snapshots == [end for end in page_ends if not any(start < end < stop for start, stop in spans)] + + +@pytest.mark.parametrize( + "page_size,same_image,gpu_end,cpu_end", + [(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, +): + from lightllm.server.multi_level_kv_cache.manager import MultiLevelKVCacheManager + from lightllm.server.router.dynamic_prompt.radix_cache import RadixCache + from lightllm.server.router.model_infer.infer_batch import InferReq, g_infer_context + from lightllm.common.req_manager import ReqManager + from lightllm.utils import shm_utils + from lightllm.utils.kv_cache_utils import CpuKVCacheMeta, compute_token_list_hash + + # Neo layout: text + + independent virtual image IDs + + text. + prefix = [7] * 699 + [10] + source = prefix + list(range(1000000, 1000500)) + [11] + [20] * 400 + image_start_id = 1000000 if same_image else 2000000 + query = prefix + list(range(image_start_id, image_start_id + 500)) + [11] + [21] * 400 + monkeypatch.setattr(shm_utils, "get_unique_server_name", lambda: tmp_path.name) + radix_cache = RadixCache(4096, 0, page_size=page_size) + radix_cache.insert(torch.tensor(source)) + + monkeypatch.setattr(g_infer_context, "is_hybrid_att_model", False) + monkeypatch.setattr(g_infer_context, "radix_cache", radix_cache) + from lightllm.server.router.model_infer.req_token_index import ( + PinMemPageCacheReqTokenIndex, + create_req_token_index, + ) + + if ( + PinMemPageCacheReqTokenIndex._page_offsets is not None + and PinMemPageCacheReqTokenIndex._page_offsets.numel() != page_size + ): + PinMemPageCacheReqTokenIndex._page_offsets = None + + req_manager = ReqManager.__new__(ReqManager) + req_manager.page_size = page_size + req_manager.req_to_token_indexs = torch.zeros((1, len(query)), dtype=torch.int32) + monkeypatch.setattr(g_infer_context, "req_manager", req_manager) + req = InferReq.__new__(InferReq) + req.req_manager = req_manager + req.args = SimpleNamespace(page_size=page_size, req_token_index_mode="pin_mem_page_cache") + req.sampling_param = SimpleNamespace(disable_prompt_cache=False) + req.cur_kv_len = req.cur_output_len = req.hold_kv_len = req.req_idx = 0 + req.shared_kv_node = None + req.mm_input_spans = [(700, 1200)] + req.shm_req = SimpleNamespace(input_len=len(query), shm_prompt_ids=SimpleNamespace(arr=query)) + req.token_index = create_req_token_index(req, torch.cuda.current_stream()) + req._match_radix_cache() + assert req.cur_kv_len == req.hold_kv_len == gpu_end + if gpu_end: + got = req.token_index.get_cpu_token_indexes(0, gpu_end) + assert got[::page_size].tolist() == query[:gpu_end:page_size] + assert not any(start < gpu_end < end for start, end in req.mm_input_spans) + assert radix_cache.get_refed_tokens_num() == gpu_end + + hashes = compute_token_list_hash(source, 256) + 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) + manager.cpu_cache_client = client + matched = manager._cpu_cache_match(compute_token_list_hash(query, 256)) + assert len(matched) * 256 == cpu_end + # The page containing the image tail also contains the differing text. + assert any(start < cpu_end < end for start, end in req.mm_input_spans) == same_image + client.deref_pages(matched) + radix_cache.dec_node_ref_counter(req.shared_kv_node) + assert radix_cache.get_refed_tokens_num() == 0 + assert radix_cache.root_node.ref_counter == 1 + + +@pytest.mark.parametrize("disable_chunked_mm_input", [False, True]) +def test_mm_chunking_keeps_the_original_state_cache_position(monkeypatch, disable_chunked_mm_input): + from lightllm.server.router.model_infer import infer_batch + + shm_req = SimpleNamespace( + link_prompt_ids_shm_array=lambda: None, + link_logprobs_shm_array=lambda: None, + hybrid_token_hash_list=SimpleNamespace(size=3), + input_len=1000, + shm_prompt_ids=SimpleNamespace(arr=list(range(1000))), + ) + monkeypatch.setattr(infer_batch.g_infer_context, "is_hybrid_att_model", True) + monkeypatch.setattr( + infer_batch.g_infer_context, "shm_req_manager", SimpleNamespace(get_req_obj_by_index=lambda index: shm_req) + ) + monkeypatch.setattr( + infer_batch.g_infer_context, + "req_manager", + SimpleNamespace(req_sampling_params_manager=SimpleNamespace(init_req_sampling_params=lambda req: None)), + ) + monkeypatch.setattr( + infer_batch, + "InferSamplingParams", + lambda *args: SimpleNamespace( + pd_decode_node=None, + shm_param=SimpleNamespace(stop_sequences=SimpleNamespace(to_token_ids_list=lambda: [])), + ), + ) + monkeypatch.setattr(infer_batch, "PromptSelectedLogprobsExt", lambda req: None) + monkeypatch.setattr(infer_batch, "FinalTokenMetadataExt", lambda req: None) + req = infer_batch.InferReq.__new__(infer_batch.InferReq) + req.args = SimpleNamespace( + linear_att_hash_page_size=256, + linear_att_page_block_num=8, + disable_chunked_mm_input=disable_chunked_mm_input, + batch_max_tokens=1024, + ) + req.shm_index, req.vocab_size = 0, 100 + 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}, + ], + "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]) +@pytest.mark.parametrize( + "loaded_start,load_end,spans,disk_tokens,expected,cpu_tokens,remaining_disk_tokens", + [ + (0, 0, [], 0, 0, 0, 0), + (0, 4096, [], 2048, 4096, 2048, 2048), + (0, 2048, [(2048, 2500)], 0, 2048, 2048, 0), + (0, 4096, [(3500, 4096)], 2048, 4096, 2048, 2048), + (0, 4096, [(3500, 4500)], 0, 2048, 2048, 0), + (0, 4096, [(3500, 4500)], 2048, 2048, 2048, 0), + (0, 4096, [(3500, 4500)], 4096, 2048, 0, 2048), + (0, 4096, [(1000, 2500), (3500, 4500)], 2048, 0, 0, 0), + (2304, 4096, [(3000, 4500)], 2048, 2304, 0, 0), + (0, 6144, [(1000, 2500), (3500, 4500)], 2048, 6144, 4096, 2048), + (2304, 6144, [(3000, 4500)], 2048, 6144, 1792, 2048), + ], +) +def test_cpu_cache_match_and_load_use_a_complete_restore_point( + monkeypatch, + cpu_cache_factory, + hybrid, + loaded_start, + load_end, + spans, + disk_tokens, + expected, + cpu_tokens, + remaining_disk_tokens, +): + from lightllm.server.router.model_infer.mode_backend import multi_level_kv_cache as common + 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)), 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 = [] + req = SimpleNamespace( + req_idx=0, + cur_kv_len=loaded_start, + hold_kv_len=loaded_start, + mm_input_spans=spans, + sampling_param=SimpleNamespace(shm_param=SimpleNamespace(prompt_logprobs=-1)), + shm_req=SimpleNamespace( + input_len=8192, + disk_prompt_cache_len=disk_tokens, + cpu_cache_match_page_indexes=SimpleNamespace(get_all=lambda: page_ids), + token_hash_page_len_list=SimpleNamespace(get_all=lambda: ends), + ), + ) + table = torch.arange(8192, dtype=torch.int32).view(1, -1) + + def allocate(alloc_token_num): + if alloc_token_num == 0: + return None + start = req.cur_kv_len + req.hold_kv_len += alloc_token_num + return torch.arange(start, start + alloc_token_num, dtype=torch.int32) + + req.args = SimpleNamespace(page_size=256) + req.token_index = SimpleNamespace(get_gpu_token_indexes=lambda start, end: table[0, start:end]) + req.kv_cache_alloc_need = lambda target: max(0, target - req.hold_kv_len) + req.alloc_kv_mem = allocate + + from collections import deque + + module = common.MultiLevelKvCacheModule.__new__(common.MultiLevelKvCacheModule) + module.load_tasks = deque() + module.args = SimpleNamespace(page_size=256, cpu_cache_token_page_size=256) + module.backend = SimpleNamespace( + is_master_in_dp=True, + dp_world_size=1, + is_hybrid_att_model=hybrid, + model=SimpleNamespace( + req_manager=SimpleNamespace( + req_to_token_indexs=table, alloc_token_indexes=allocate, wait_token_indexes=lambda: None + ), + att_state_manager=SimpleNamespace( + 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))), + full_att_allocator=None, + ), + ), + ) + module.cpu_cache_client = client + module.init_sync_group = None + module.need_sync_compute_stream = lambda: False + monkeypatch.setattr( + common.g_infer_context, + "req_manager", + SimpleNamespace(alloc_token_indexes=allocate), + ) + monkeypatch.setattr(common.g_infer_context, "get_can_alloc_token_num", lambda: 8192) + monkeypatch.setattr(common.g_infer_context, "is_hybrid_att_model", hybrid) + monkeypatch.setattr(common.dist, "barrier", lambda **kwargs: None) + monkeypatch.setattr(torch.Tensor, "cuda", lambda self, **kwargs: self) + req.hybrid_len_to_big_page_id = {} + module.load_cpu_cache_to_reqs([req]) + for task in getattr(module, "load_tasks", ()): + task.event.synchronize() + module.update_cpu_cache_load_task_states() + assert req.cur_kv_len == req.hold_kv_len == expected + assert loaded == ([expected] if expected > loaded_start else []) + assert req.shm_req.cpu_prompt_cache_len == cpu_tokens + assert req.shm_req.disk_prompt_cache_len == remaining_disk_tokens + assert all(client.page_items.linked_items[page].ref_count == 0 for page in page_ids) + + +@pytest.mark.parametrize("state_slot,expected", [(None, 1), (3, 2)]) +def test_cpu_offload_requires_tail_state(monkeypatch, state_slot, expected): + from lightllm.server.router.model_infer.mode_backend import multi_level_kv_cache as common + + monkeypatch.setattr(common.g_infer_context, "is_hybrid_att_model", True) + module = common.MultiLevelKvCacheModule.__new__(common.MultiLevelKvCacheModule) + module.args = SimpleNamespace(cpu_cache_token_page_size=512, disable_linear_att_small_page_cpu_cache=False) + req = SimpleNamespace(tail_small_page_buffer_id=state_slot) + assert module._handle_hybrid_att_last_page(req, 2, [512, 960]) == expected diff --git a/unit_tests/server/router/req_queue/test_pd_queue_selection.py b/unit_tests/server/router/req_queue/test_pd_queue_selection.py new file mode 100644 index 0000000000..0aa1a45c9e --- /dev/null +++ b/unit_tests/server/router/req_queue/test_pd_queue_selection.py @@ -0,0 +1,108 @@ +from types import SimpleNamespace +from unittest.mock import MagicMock + +from lightllm.server.router.req_queue import _get_req_queue_class +from lightllm.server.router.batch import Batch +from lightllm.server.router.req_queue.chunked_prefill.impl_for_pd_decode import PDDecodeQueue +from lightllm.server.router.req_queue.chunked_prefill.impl_for_pd_prefill import PDPrefillQueue + + +def _make_args(run_mode: str): + return SimpleNamespace( + run_mode=run_mode, + diverse_mode=False, + output_constraint_mode="none", + disable_chunked_prefill=False, + ) + + +def test_pd_prefill_uses_prefill_queue(): + args = _make_args("prefill") + + assert _get_req_queue_class(args, router=None, dp_size_in_node=1) is PDPrefillQueue + + +def test_pd_decode_uses_decode_queue(): + args = _make_args("decode") + + assert _get_req_queue_class(args, router=None, dp_size_in_node=1) is PDDecodeQueue + + +def test_pd_prefill_peak_tokens_sum_request_tuples(): + queue = PDPrefillQueue.__new__(PDPrefillQueue) + queue.dp_index = 0 + queue.args = SimpleNamespace(page_size=16) + queue.is_busy = lambda: False + queue.router = SimpleNamespace(router_statics=SimpleNamespace(ema_req_out_len=10)) + reqs = [ + SimpleNamespace( + request_id="req-0", + input_len=10, + sample_params=SimpleNamespace(max_new_tokens=1, suggested_dp_index=0), + get_tuple_tokens=MagicMock(return_value=(11, 34)), + ), + SimpleNamespace( + request_id="req-1", + input_len=20, + sample_params=SimpleNamespace(max_new_tokens=1, suggested_dp_index=0), + get_tuple_tokens=MagicMock(return_value=(21, 34)), + ), + SimpleNamespace( + request_id="req-2", + input_len=100, + sample_params=SimpleNamespace(max_new_tokens=1, suggested_dp_index=1), + get_tuple_tokens=MagicMock(return_value=(101, 34)), + ), + ] + batch = Batch(batch_id=1, reqs=reqs, dp_size_in_node=2) + + assert queue._caclu_batch_estimated_peak_token_num(batch) == 100 + reqs[0].get_tuple_tokens.assert_called_once_with(False, 10) + reqs[1].get_tuple_tokens.assert_called_once_with(False, 10) + reqs[2].get_tuple_tokens.assert_not_called() + + +def test_pd_decode_sums_request_tuples_for_non_decode_requests(): + queue = PDDecodeQueue.__new__(PDDecodeQueue) + queue.args = SimpleNamespace(page_size=16) + queue.dp_index = 0 + queue.is_busy = lambda: False + queue.router = SimpleNamespace(router_statics=SimpleNamespace(ema_req_out_len=10)) + reqs = [ + SimpleNamespace( + request_id="req-0", + input_len=20, + sample_params=SimpleNamespace(max_new_tokens=1, suggested_dp_index=0), + is_infer_decode=lambda: False, + get_tuple_tokens=MagicMock(return_value=(21, 34)), + ), + SimpleNamespace( + request_id="req-1", + input_len=100, + sample_params=SimpleNamespace(max_new_tokens=1, suggested_dp_index=1), + is_infer_decode=lambda: False, + get_tuple_tokens=MagicMock(return_value=(101, 34)), + ), + ] + batch = Batch(batch_id=1, reqs=reqs, dp_size_in_node=2) + + assert queue._caclu_batch_estimated_peak_token_num(batch) == 55 + reqs[0].get_tuple_tokens.assert_called_once_with(False, 10) + reqs[1].get_tuple_tokens.assert_not_called() + + +def test_pd_decode_uses_tuple_estimation(): + queue = PDDecodeQueue.__new__(PDDecodeQueue) + queue.args = SimpleNamespace(page_size=16) + queue.dp_index = 0 + queue.is_busy = lambda: False + queue.router = SimpleNamespace(router_statics=SimpleNamespace(ema_req_out_len=10)) + req = SimpleNamespace( + request_id="req-0", + sample_params=SimpleNamespace(suggested_dp_index=0), + is_infer_decode=lambda: True, + get_tuple_tokens=lambda is_busy, ema_req_out_len: (16, 32), + ) + batch = Batch(batch_id=1, reqs=[req], dp_size_in_node=1) + + assert queue._caclu_batch_estimated_peak_token_num(batch) == 48 diff --git a/unit_tests/server/test_api_responses.py b/unit_tests/server/test_api_responses.py new file mode 100644 index 0000000000..498904b613 --- /dev/null +++ b/unit_tests/server/test_api_responses.py @@ -0,0 +1,313 @@ +import asyncio + +import pytest +import ujson as json + +from lightllm.server.api_responses import ( + _chat_response_to_responses, + _openai_sse_to_responses_events, + _responses_to_chat_request, +) +from lightllm.server.api_models import ChatCompletionRequest + + +async def _chunks(*payloads): + for payload in payloads: + yield f"data: {json.dumps(payload)}\n\n" + + +def _collect_events(*payloads): + async def collect(): + return [event async for event in _openai_sse_to_responses_events(_chunks(*payloads), {"input": "hi"})] + + raw_events = asyncio.run(collect()) + return [json.loads(event.decode().split("data: ", 1)[1]) for event in raw_events] + + +def test_function_call_arguments_done_includes_name(): + events = _collect_events( + { + "choices": [ + { + "delta": { + "tool_calls": [ + { + "id": "call_1", + "function": {"name": "get_weather", "arguments": '{"city":"Paris"}'}, + } + ] + } + } + ] + } + ) + + done = next(event for event in events if event["type"] == "response.function_call_arguments.done") + assert done["name"] == "get_weather" + + +def test_content_array_function_output_is_preserved(): + request = _responses_to_chat_request( + { + "input": [ + { + "type": "function_call_output", + "call_id": "call_1", + "output": [ + {"type": "input_text", "text": "sunny"}, + {"type": "input_image", "image_url": "https://example.com/weather.png"}, + ], + } + ] + } + ) + + assert request["messages"] == [ + { + "role": "tool", + "tool_call_id": "call_1", + "content": [ + {"type": "text", "text": "sunny"}, + {"type": "image_url", "image_url": {"url": "https://example.com/weather.png"}}, + ], + } + ] + + +def test_replayed_function_call_preserves_content_after_validation(): + request = _responses_to_chat_request( + { + "input": [ + { + "type": "function_call", + "call_id": "call_1", + "name": "get_weather", + "arguments": '{"city":"Paris"}', + } + ] + } + ) + + chat_request = ChatCompletionRequest(**request) + message = chat_request.messages[0].model_dump(exclude_none=True) + assert message["content"] == "" + + +def test_parallel_function_calls_share_one_assistant_message(): + request = _responses_to_chat_request( + { + "input": [ + { + "type": "function_call", + "call_id": "call_paris", + "name": "get_weather", + "arguments": '{"city":"Paris"}', + }, + { + "type": "function_call", + "call_id": "call_tokyo", + "name": "get_weather", + "arguments": '{"city":"Tokyo"}', + }, + {"type": "function_call_output", "call_id": "call_paris", "output": "sunny"}, + {"type": "function_call_output", "call_id": "call_tokyo", "output": "rainy"}, + ] + } + ) + + assert request["messages"] == [ + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_paris", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city":"Paris"}'}, + }, + { + "id": "call_tokyo", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city":"Tokyo"}'}, + }, + ], + }, + {"role": "tool", "tool_call_id": "call_paris", "content": "sunny"}, + {"role": "tool", "tool_call_id": "call_tokyo", "content": "rainy"}, + ] + + +@pytest.mark.parametrize( + ("field", "part_type"), + [("content", "reasoning_text"), ("summary", "summary_text")], +) +def test_replayed_reasoning_is_attached_to_assistant_tool_call(field, part_type): + request = _responses_to_chat_request( + { + "input": [ + { + "type": "reasoning", + field: [{"type": part_type, "text": "I should check the weather."}], + }, + { + "type": "function_call", + "call_id": "call_1", + "name": "get_weather", + "arguments": '{"city":"Paris"}', + }, + {"type": "function_call_output", "call_id": "call_1", "output": "sunny"}, + ] + } + ) + + assert request["messages"][0]["reasoning_content"] == "I should check the weather." + assert request["messages"][0]["tool_calls"][0]["id"] == "call_1" + assert len(request["messages"]) == 2 + + +def test_replayed_reasoning_is_attached_to_assistant_message(): + request = _responses_to_chat_request( + { + "input": [ + { + "type": "reasoning", + "content": [{"type": "reasoning_text", "text": "I found the answer."}], + }, + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "It is sunny."}], + }, + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "Thanks"}], + }, + ] + } + ) + + assert request["messages"] == [ + { + "role": "assistant", + "content": [{"type": "text", "text": "It is sunny."}], + "reasoning_content": "I found the answer.", + }, + {"role": "user", "content": [{"type": "text", "text": "Thanks"}]}, + ] + + +def test_non_streamed_truncated_function_call_is_incomplete(): + response = _chat_response_to_responses( + { + "choices": [ + { + "finish_reason": "length", + "message": { + "tool_calls": [ + { + "id": "call_1", + "function": {"name": "get_weather", "arguments": '{"city":'}, + } + ] + }, + } + ] + }, + {"input": "hi"}, + ) + + assert response["status"] == "incomplete" + assert response["output"][0]["status"] == "incomplete" + + +def test_streamed_truncated_function_call_is_incomplete(): + events = _collect_events( + { + "choices": [ + { + "delta": { + "tool_calls": [ + { + "id": "call_1", + "function": {"name": "get_weather", "arguments": '{"city":'}, + } + ] + } + } + ] + }, + {"choices": [{"delta": {}, "finish_reason": "length"}]}, + ) + + item_done = next(event for event in events if event["type"] == "response.output_item.done") + response_done = events[-1] + assert item_done["item"]["status"] == "incomplete" + assert response_done["type"] == "response.incomplete" + assert response_done["response"]["output"][0]["status"] == "incomplete" + + +def test_streamed_message_adds_text_content_once(): + events = _collect_events({"choices": [{"delta": {"content": "hello"}, "finish_reason": "stop"}]}) + + item_added = next(event for event in events if event["type"] == "response.output_item.added") + part_added = next(event for event in events if event["type"] == "response.content_part.added") + assert item_added["item"]["content"] == [] + assert part_added["part"] == {"type": "output_text", "text": "", "annotations": []} + + +def test_route_maps_downstream_value_error_to_bad_request(monkeypatch): + from types import SimpleNamespace + + from lightllm.server import api_http, api_responses + + async def raise_value_error(raw_request): + raise ValueError("Unrecognized image input.") + + monkeypatch.setattr(api_http, "get_env_start_args", lambda: SimpleNamespace(run_mode="normal")) + monkeypatch.setattr(api_http.g_objs, "metric_client", SimpleNamespace(counter_inc=lambda *args: None)) + monkeypatch.setattr(api_responses, "responses_impl", raise_value_error) + + response = asyncio.run(api_http.openai_responses(None)) + assert response.status_code == 400 + + +@pytest.mark.parametrize( + "partial_payload", + [ + {"choices": [{"delta": {"content": "partial"}}]}, + { + "choices": [ + { + "delta": { + "tool_calls": [ + { + "id": "call_1", + "function": {"name": "get_weather", "arguments": '{"city":'}, + } + ] + } + } + ] + }, + ], +) +def test_stream_failure_does_not_complete_partial_item(partial_payload): + events = _collect_events(partial_payload, {"error": {"message": "generation failed"}}) + event_types = [event["type"] for event in events] + + assert "response.output_text.done" not in event_types + assert "response.function_call_arguments.done" not in event_types + assert "response.output_item.done" not in event_types + assert event_types[-1] == "response.failed" + + +@pytest.mark.parametrize("effort", ["low", "medium", "high", "none", "minimal", "xhigh"]) +def test_reasoning_effort_is_rejected(effort): + with pytest.raises(ValueError, match="reasoning.effort is not supported"): + _responses_to_chat_request({"input": "hi", "reasoning": {"effort": effort}}) + + +def test_automatic_truncation_is_rejected(): + with pytest.raises(ValueError, match="truncation"): + _responses_to_chat_request({"input": "hi", "truncation": "auto"}) diff --git a/unit_tests/server/test_application_stop.py b/unit_tests/server/test_application_stop.py new file mode 100644 index 0000000000..0603c29c67 --- /dev/null +++ b/unit_tests/server/test_application_stop.py @@ -0,0 +1,200 @@ +import asyncio +import json +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from lightllm.server.api_lightllm import lightllm_generate, lightllm_generate_stream +from lightllm.server.api_tgi import tgi_generate_impl, tgi_generate_stream_impl +from lightllm.server.core.objs import FinishStatus, SamplingParams, StopStrMatchMode +from lightllm.server.httpserver.manager import HttpServerManager +from lightllm.server.pd_io_struct import NodeRole + + +def _manager(tokenizer, args, raw_results): + manager = HttpServerManager.__new__(HttpServerManager) + manager.output_grammar_compiler = None + manager.tokenizer = tokenizer + manager.args = args + manager.pd_mode = NodeRole.NORMAL + manager.is_multinode_tp_slave = False + manager._generate = raw_results + manager.stop_request = AsyncMock() + return manager + + +@pytest.mark.parametrize("stream", [False, True]) +def test_native_api_discards_output_after_application_stop(stream): + async def run(): + closed = [] + exhausted = [] + + async def raw_results(*args, **kwargs): + metadata = {"prompt_tokens": 2, "prompt_cache_len": 0, "logprobs": {}, "logprob": -0.1} + try: + yield 80, "helloE", {**metadata, "id": 10}, FinishStatus() + yield 80, "NDextra", {**metadata, "id": 11}, FinishStatus() + yield 80, "late", {**metadata, "id": 12}, FinishStatus(FinishStatus.FINISHED_LENGTH) + exhausted.append(True) + finally: + closed.append(True) + + manager = _manager( + SimpleNamespace(encode=lambda text, **kwargs: [ord(c) for c in text]), + SimpleNamespace(reasoning_parser=None, enable_rl=False, enable_return_routed_experts=False), + raw_results, + ) + request = SimpleNamespace( + json=AsyncMock( + return_value={ + "inputs": "prompt", + "parameters": {"stop_sequences": ["END"], "max_new_tokens": 10}, + } + ) + ) + response = await (lightllm_generate_stream if stream else lightllm_generate)(request, manager) + if stream: + events = [json.loads(chunk.decode().split(":", 1)[1]) async for chunk in response.body_iterator] + assert "".join(event["token"]["text"] for event in events) == "hello" + assert events[-1]["finished"] + assert events[-1]["finish_reason"] == "stop" + else: + data = json.loads(response.body) + assert data["generated_text"] == ["hello"] + assert data["count_output_tokens"] == 2 + assert data["finish_reason"] == "stop" + manager.stop_request.assert_awaited_once_with(80) + assert closed == [True] + assert exhausted == [True] + + asyncio.run(run()) + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("include", [False, True]) +def test_tgi_api_uses_shared_stop_wrapper(stream, include): + async def run(): + async def raw_results(*args, **kwargs): + metadata = {"prompt_tokens": 2, "logprobs": {}, "logprob": -0.1, "cumlogprob": -0.1} + yield 80, "helloE", {**metadata, "id": 10}, FinishStatus() + yield 80, "NDextra", {**metadata, "id": 11}, FinishStatus() + yield 80, "late", {**metadata, "id": 12}, FinishStatus(FinishStatus.FINISHED_LENGTH) + + manager = _manager( + SimpleNamespace(encode=lambda text, **kwargs: [ord(c) for c in text]), + SimpleNamespace(reasoning_parser=None, enable_rl=False, enable_return_routed_experts=False), + raw_results, + ) + request = SimpleNamespace( + json=AsyncMock( + return_value={ + "inputs": "prompt", + "parameters": {"stop": ["END"], "details": True, "include_stop_str_in_output": include}, + } + ) + ) + response = await (tgi_generate_stream_impl if stream else tgi_generate_impl)(request, manager) + expected = "helloEND" if include else "hello" + if stream: + events = [] + async for chunk in response.body_iterator: + chunk = chunk.decode() if isinstance(chunk, bytes) else chunk + events.append(json.loads(chunk.split(":", 1)[1])) + assert "".join(event["token"]["text"] for event in events) == expected + assert events[-1]["finish_reason"] == "stop" + else: + data = json.loads(response.body) + data = data[0] if isinstance(data, list) else data + assert data["generated_text"] == expected + assert data["details"]["finish_reason"] == "stop" + manager.stop_request.assert_awaited_once_with(80) + + asyncio.run(run()) + + +@pytest.mark.parametrize("mode", [NodeRole.NORMAL, NodeRole.P, NodeRole.D]) +@pytest.mark.parametrize("slave", [False, True]) +def test_http_generate_only_wraps_client_output(mode, slave): + async def run(): + tokenizer = SimpleNamespace(encode=lambda text, **kwargs: [ord(c) for c in text]) + params = SamplingParams() + params.init(tokenizer, stop_sequences="END", max_new_tokens=10) + websocket, pd_event = object(), object() + + async def raw_results(prompt, sampling_params, multimodal, request, upload, event): + assert upload is websocket + assert event is pd_event + yield 80, "helloENDextra", {"id": 10}, FinishStatus(FinishStatus.FINISHED_STOP) + + manager = _manager( + tokenizer, + SimpleNamespace(reasoning_parser=None, enable_rl=False, enable_return_routed_experts=False), + raw_results, + ) + manager.pd_mode = mode + manager.is_multinode_tp_slave = slave + results = [result async for result in manager.generate("prompt", params, None, None, websocket, pd_event)] + expected = "hello" if mode == NodeRole.NORMAL and not slave else "helloENDextra" + assert "".join(result[1] for result in results) == expected + manager.stop_request.assert_not_awaited() + + asyncio.run(run()) + + +@pytest.mark.parametrize("mode", [StopStrMatchMode.ALL_TEXT, StopStrMatchMode.CONTENT_ONLY]) +@pytest.mark.parametrize("feature", ["enable_rl", "enable_return_routed_experts", "prompt_logprobs"]) +@pytest.mark.parametrize("stop", [None, [], "END", [[10, 11]], ["END", [10, 11]]]) +@pytest.mark.parametrize("include", [False, True]) +@pytest.mark.parametrize("split", [False, True]) +def test_final_metadata_stop_forwards_late_tokens_and_backend_finish(mode, feature, stop, include, split): + async def run(): + tokenizer = SimpleNamespace(encode=lambda text, **kwargs: [ord(c) for c in text]) + params = SamplingParams() + params.init(tokenizer, stop_sequences=stop, stop_str_match_mode=mode, include_stop_str_in_output=include) + if feature == "prompt_logprobs": + params.prompt_logprobs = 0 + metadata = {"id": 13, "prompt_token_ids": [1, 2], "prompt_logprobs": [None, {}], "routed_experts": [[2]]} + + async def raw_results(*args, **kwargs): + yield 80, "helloE" if split else "hello", {"id": 10}, FinishStatus() + yield 80, "NDextra" if split else "ENDextra", {"id": 11}, FinishStatus() + yield 80, "lateEND", {"id": 12}, FinishStatus() + yield 80, "tail", metadata, FinishStatus(FinishStatus.FINISHED_ABORTED) + yield 80, "after finish", {"id": 14}, FinishStatus() + + manager = _manager( + tokenizer, + SimpleNamespace( + reasoning_parser=None, + enable_rl=feature == "enable_rl", + enable_return_routed_experts=feature == "enable_return_routed_experts", + ), + raw_results, + ) + results = [result async for result in manager.generate("prompt", params, None, None)] + first_text = "helloE" if split else "hello" + hit_text = "NDextra" if split else "ENDextra" + if stop: + if params.stop_sequences.has_stop_strs(): + first_text = first_text if include else "hello" + hit_text = ("ND" if split else "END") if include else "" + assert [result[1] for result in results] == [ + first_text, + hit_text, + "lateEND", + "tail", + ] + assert [result[2]["id"] for result in results] == [10, 11, 12, 13] + expected_stops = ( + 2 if params.stop_sequences.has_stop_strs() and params.stop_sequences.to_token_ids_list() else 1 + ) + assert [call.args for call in manager.stop_request.await_args_list] == [(80,)] * expected_stops + else: + assert [result[1] for result in results] == [first_text, hit_text, "lateEND", "tail", "after finish"] + manager.stop_request.assert_not_awaited() + assert all(not result[3].is_finished() for result in results[:3]) + assert results[3][2] is metadata + assert results[3][3].status == FinishStatus.FINISHED_ABORTED + + asyncio.run(run()) diff --git a/unit_tests/server/test_detokenization.py b/unit_tests/server/test_detokenization.py new file mode 100644 index 0000000000..4d12363290 --- /dev/null +++ b/unit_tests/server/test_detokenization.py @@ -0,0 +1,114 @@ +"""Incremental detokenization must ignore virtual image/audio cache IDs.""" + +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import numpy as np +from tokenizers import Tokenizer, decoders, models, pre_tokenizers +from transformers import PreTrainedTokenizerFast + +from lightllm.server.detokenization import decode_req as decode_req_module +from lightllm.server.detokenization.decode import decode_token +from lightllm.server.detokenization.decode_req import DecodeReq +from lightllm.server.embed_cache.utils import MULTIMODAL_TOKEN_ID_START + + +class TestDetokenization(unittest.TestCase): + def setUp(self): + vocab = {char: i for i, char in enumerate(sorted(pre_tokenizers.ByteLevel.alphabet()))} + backend = Tokenizer(models.BPE(vocab=vocab, merges=[])) + backend.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False) + backend.decoder = decoders.ByteLevel() + self.tokenizer = PreTrainedTokenizerFast(tokenizer_object=backend) + self.tokenizer.add_special_tokens({"additional_special_tokens": [""]}) + self.tokenizer.add_tokens(["added_text"]) + prefix_length = patch.object(decode_req_module, "LIGHTLLM_DECODE_PREFIX_LENGTH", 5) + prefix_length.start() + self.addCleanup(prefix_length.stop) + + def make_decode_req(self, prompt_ids, output_ids, **sampling_params): + params = dict(print_eos_token=False, skip_special_tokens=True, add_spaces_between_special_tokens=False) + params.update(sampling_params) + req = SimpleNamespace( + request_id=1, + group_req_id=1, + input_len=len(prompt_ids), + shm_prompt_ids=SimpleNamespace(arr=np.array(prompt_ids + output_ids, dtype=np.int64)), + sample_params=SimpleNamespace(**params), + ) + return DecodeReq(req) + + def decode_stream(self, decode_req, output_ids): + pieces = [] + for token_id in output_ids: + decode_req.output_ids.append(token_id) + pieces.append(decode_token(self.tokenizer, decode_req, token_id, [])) + return pieces + + def test_virtual_prompt_ids_at_boundary_and_large_values(self): + output_ids = self.tokenizer.encode("hello world", add_special_tokens=False) + for virtual_id in (MULTIMODAL_TOKEN_ID_START, 2 ** 31, 2 ** 32, 2 ** 63 - 2): + with self.subTest(virtual_id=virtual_id): + prompt_ids = self.tokenizer.encode("prompt", add_special_tokens=False) + [virtual_id] + decode_req = self.make_decode_req(prompt_ids, output_ids) + self.assertEqual("".join(self.decode_stream(decode_req, output_ids)), "hello world") + + def test_five_token_prompt_tail_and_partial_utf8_keep_original_positions(self): + prompt_ids = self.tokenizer.encode("prompt", add_special_tokens=False) + [2 ** 40 + i for i in range(5)] + output_ids = self.tokenizer.encode("你 world", add_special_tokens=False) + decode_req = self.make_decode_req(prompt_ids, output_ids) + original_tokens = decode_req.req.shm_prompt_ids.arr.copy() + initial_offsets = (len(prompt_ids) - 5, len(prompt_ids)) + self.assertEqual((decode_req.prefix_offset, decode_req.read_offset), initial_offsets) + + pieces = [] + for index, token_id in enumerate(output_ids): + shared_token_id, shared_index = decode_req.get_next_token_id_and_index() + self.assertEqual(shared_token_id, token_id) + self.assertEqual(shared_index, len(prompt_ids) + index) + decode_req.output_ids.append(shared_token_id) + pieces.append(decode_token(self.tokenizer, decode_req, token_id, [])) + if index < 2: + self.assertEqual((decode_req.prefix_offset, decode_req.read_offset), initial_offsets) + elif index == 2: + self.assertEqual( + (decode_req.prefix_offset, decode_req.read_offset), (len(prompt_ids), len(prompt_ids) + 3) + ) + + self.assertEqual(pieces[:2], ["", ""]) + self.assertEqual("".join(pieces), "你 world") + self.assertEqual(decode_req.prompt_ids, prompt_ids) + self.assertEqual(decode_req.output_ids, output_ids) + np.testing.assert_array_equal(decode_req.req.shm_prompt_ids.arr, original_tokens) + self.assertEqual(decode_req.read_offset, len(original_tokens)) + self.assertEqual(decode_req.prefix_offset, len(original_tokens) - 1) + + def test_added_and_special_tokens_keep_tokenizer_options(self): + output_ids = self.tokenizer.encode("hello", add_special_tokens=False) + output_ids += self.tokenizer.convert_tokens_to_ids(["", "added_text"]) + for skip_special_tokens in (False, True): + for add_spaces in (False, True): + with self.subTest(skip_special_tokens=skip_special_tokens, add_spaces=add_spaces): + decode_req = self.make_decode_req( + [2 ** 40], + output_ids, + skip_special_tokens=skip_special_tokens, + add_spaces_between_special_tokens=add_spaces, + ) + expected = self.tokenizer.decode( + output_ids, + skip_special_tokens=skip_special_tokens, + spaces_between_special_tokens=add_spaces, + ) + self.assertEqual("".join(self.decode_stream(decode_req, output_ids)), expected) + + def test_text_only_prompt(self): + prompt_ids = self.tokenizer.encode("prompt", add_special_tokens=False) + output_ids = self.tokenizer.encode("你 world", add_special_tokens=False) + decode_req = self.make_decode_req(prompt_ids, output_ids) + self.assertEqual("".join(self.decode_stream(decode_req, output_ids)), "你 world") + + +if __name__ == "__main__": + unittest.main() diff --git a/unit_tests/server/test_dp_request_capacity.py b/unit_tests/server/test_dp_request_capacity.py new file mode 100644 index 0000000000..800f0cdda3 --- /dev/null +++ b/unit_tests/server/test_dp_request_capacity.py @@ -0,0 +1,109 @@ +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +from lightllm.server.core.objs.shm_req_manager import ShmReqManager +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.server.router.batch import Batch +from lightllm.server.router.req_queue.chunked_prefill.impl import ChunkedPrefillQueue +from lightllm.server.router.req_queue.dp_base_queue import DpQueue +from lightllm.utils.config_utils import get_running_max_req_size_per_dp + + +@pytest.mark.parametrize( + "global_capacity,dp,nnodes,expected", + [(32, 8, 1, 4), (33, 8, 1, 5), (4, 8, 1, 1), (32, 8, 2, 8), (32, 1, 2, 32), (32, 1, 1, 32)], +) +def test_hybrid_capacity_uses_local_dp_count(monkeypatch, global_capacity, dp, nnodes, expected): + monkeypatch.setattr("lightllm.utils.config_utils.is_hybrid_att_model", lambda _: True) + args = StartArgs(running_max_req_size=global_capacity, dp=dp, nnodes=nnodes) + assert get_running_max_req_size_per_dp(args) == expected + # HTTP/SHM indexes remain global even when GPU request indexes are local. + monkeypatch.setattr("lightllm.server.core.objs.shm_req_manager.get_env_start_args", lambda: args) + assert ShmReqManager.get_max_req_num(None) == global_capacity + + +@pytest.mark.parametrize("mode", ["full_attention", "diverse", "cache_fetch"]) +def test_modes_requiring_existing_capacity_keep_global_default(monkeypatch, mode): + monkeypatch.setattr("lightllm.utils.config_utils.is_hybrid_att_model", lambda _: mode != "full_attention") + args = StartArgs( + running_max_req_size=32, + dp=8, + diverse_mode=mode == "diverse", + enable_dp_prompt_cache_fetch=mode == "cache_fetch", + ) + assert get_running_max_req_size_per_dp(args) == 32 + + +@pytest.mark.parametrize("capacity", [0, -1]) +def test_reject_invalid_global_capacity(capacity): + args = StartArgs(running_max_req_size=capacity, dp=8) + with pytest.raises(ValueError, match="running_max_req_size"): + get_running_max_req_size_per_dp(args) + + +def _make_dp_queue(monkeypatch, balancer): + monkeypatch.setattr("lightllm.utils.config_utils.is_hybrid_att_model", lambda _: True) + args = StartArgs( + running_max_req_size=32, + dp=8, + max_total_token_num=1048576, + batch_max_tokens=4096, + router_token_ratio=0.85, + dp_balancer=balancer, + ) + router = SimpleNamespace( + router_statics=SimpleNamespace(ema_req_out_len=16), + get_used_tokens=lambda dp: 0, + shared_token_load=MagicMock(), + ) + return DpQueue(args, router, ChunkedPrefillQueue, dp_size_in_node=8) + + +def _req(request_id, dp=-1): + return SimpleNamespace( + request_id=request_id, + sample_params=SimpleNamespace(suggested_dp_index=dp, high_priority_request=False), + is_aborted=False, + get_tuple_tokens=lambda busy, ema: (64, 16), + get_first_router_need_tokens=lambda: 64, + get_decode_need_tokens=lambda: 3, + ) + + +@pytest.mark.parametrize("balancer", ["bs_balancer", "round_robin"]) +def test_global_32_requests_run_as_four_per_rank(monkeypatch, balancer): + queue = _make_dp_queue(monkeypatch, balancer) + for i in range(32): + queue.extend([_req(i)]) + batch = queue.generate_new_batch(None) + assert batch.get_all_dp_req_num() == [4] * 8 + assert queue.get_wait_req_num() == 0 + + +def test_pinned_overflow_waits_until_a_local_slot_is_released(monkeypatch): + queue = _make_dp_queue(monkeypatch, "bs_balancer") + # A skewed client can use any global SHM slot, but only four may enter + # inference on this rank. The fifth must not exhaust its GPU request pool. + pinned = [_req(i, dp=2) for i in range(6)] + for req in pinned: + queue.extend([req]) + queue.extend([_req(6, dp=3)]) + batch = queue.generate_new_batch(None) + assert batch.get_all_dp_req_num() == [0, 0, 4, 1, 0, 0, 0, 0] + assert queue.get_wait_req_num() == 2 + assert queue.generate_new_batch(batch) is None + # Releasing a request on a different rank must not open a slot on rank 2. + batch.pop_req(6) + assert queue.generate_new_batch(batch) is None + batch.pop_req(0) + resumed = queue.generate_new_batch(batch) + assert [req.request_id for req in resumed.reqs] == [4] + batch.merge(resumed) + assert batch.get_all_dp_req_num()[2] == 4 + assert queue.get_wait_req_num() == 1 + assert queue.generate_new_batch(batch) is None + rest = queue.generate_new_batch(Batch(-1, [], dp_size_in_node=8)) + assert [req.request_id for req in rest.reqs] == [5] + assert queue.get_wait_req_num() == 0 diff --git a/unit_tests/server/test_generation_error_response.py b/unit_tests/server/test_generation_error_response.py new file mode 100644 index 0000000000..87670d5bf0 --- /dev/null +++ b/unit_tests/server/test_generation_error_response.py @@ -0,0 +1,91 @@ +import json +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import pytest + +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from lightllm.server import api_openai, api_stream_obj +from lightllm.server.api_models import ChatCompletionRequest +from lightllm.server.core.objs import FinishStatus, SamplingParams, StartArgs +from lightllm.server.api_http import g_objs +from lightllm.server.api_http import generation_exception_handler +from lightllm.utils.error_utils import GenerationError + + +def test_generation_failure_before_first_chunk_returns_http_500(monkeypatch): + monkeypatch.setattr(g_objs, "metric_client", Mock()) + monkeypatch.setattr(api_stream_obj, "get_env_start_args", lambda: StartArgs()) + app = FastAPI() + app.add_exception_handler(GenerationError, generation_exception_handler) + + @app.get("/generate") + async def generate(): + async def failed_stream(): + raise GenerationError("Generation failed before producing output") + yield + + return api_stream_obj.CustomStreamingResponse(failed_stream(), media_type="text/event-stream") + + response = TestClient(app).get("/generate") + assert response.status_code == 500 + assert "Generation failed" in response.json()["error"]["message"] + + +@pytest.mark.parametrize( + "early_headers, finish_status, expected_status", + [ + (False, FinishStatus.FINISHED_ERROR, 500), + (True, FinishStatus.FINISHED_ERROR, 200), + (False, FinishStatus.FINISHED_STOP, 200), + (True, FinishStatus.FINISHED_STOP, 200), + ], +) +def test_chat_first_result_sets_http_status(monkeypatch, early_headers, finish_status, expected_status): + args = StartArgs(disable_delay_response_start=early_headers) + + async def generate(*_args, **_kwargs): + yield 0, "answer", {"prompt_tokens": 2, "prompt_cache_len": 0}, FinishStatus(finish_status) + + monkeypatch.setattr(g_objs, "metric_client", Mock()) + monkeypatch.setattr(g_objs, "args", args) + monkeypatch.setattr( + g_objs, + "httpserver_manager", + SimpleNamespace(tokenizer=Mock(), output_grammar_compiler=None, generate=generate), + ) + monkeypatch.setattr(api_openai, "get_env_start_args", lambda: args) + monkeypatch.setattr(api_stream_obj, "get_env_start_args", lambda: args) + monkeypatch.setattr(api_openai, "build_prompt", AsyncMock(return_value="prompt")) + monkeypatch.setattr(SamplingParams, "init", lambda *_args, **_kwargs: None) + app = FastAPI() + app.add_exception_handler(GenerationError, generation_exception_handler) + + @app.post("/v1/chat/completions") + async def chat(request: ChatCompletionRequest): + return await api_openai.chat_completions_impl(request, Mock()) + + response = TestClient(app).post( + "/v1/chat/completions", + json={"model": "test", "messages": [{"role": "user", "content": "hi"}], "stream": True}, + ) + assert response.status_code == expected_status + if expected_status == 500: + assert "Generation failed" in response.json()["error"]["message"] + elif finish_status == FinishStatus.FINISHED_ERROR: + assert response.headers["content-type"].startswith("text/event-stream") + events = response.text.strip().split("\n\n") + assert len(events) == 2 + assert json.loads(events[0].removeprefix("data: ")) == { + "error": { + "message": "Generation failed before producing output", + "type": "server_error", + "code": "stream_error", + } + } + assert events[1] == "data: [DONE]" + else: + assert "answer" in response.text + assert "data: [DONE]" in response.text diff --git a/unit_tests/server/test_hybrid_start_args.py b/unit_tests/server/test_hybrid_start_args.py new file mode 100644 index 0000000000..1c391c09cd --- /dev/null +++ b/unit_tests/server/test_hybrid_start_args.py @@ -0,0 +1,138 @@ +import json + +import pytest + +from lightllm.server.api_start import _launch_subprocesses +from lightllm.server.core.objs.start_args_type import StartArgs + + +@pytest.fixture +def launch_args(monkeypatch, tmp_path): + for name in ( + "_set_envs_and_config", + "auto_set_max_req_total_len", + "auto_set_fused_shared_experts", + "set_unique_server_name", + ): + monkeypatch.setattr(f"lightllm.server.api_start.{name}", lambda args: None) + + def validation_finished(args): + raise RuntimeError("startup validation passed") + + monkeypatch.setattr("lightllm.server.api_start.auto_set_response_parsers", validation_finished) + + def make_args(model_type, **kwargs): + (tmp_path / "config.json").write_text(json.dumps({"model_type": model_type})) + return StartArgs( + model_dir=str(tmp_path), + max_req_total_len=2048, + eos_id=[2], + router_token_ratio=0.0, + disable_vision=True, + disable_audio=True, + disable_shm_warning=True, + **kwargs, + ) + + return make_args + + +@pytest.mark.parametrize("run_mode", ["normal", "prefill", "decode"]) +@pytest.mark.parametrize("output_constraint_mode", ["none", "xgrammar"]) +@pytest.mark.parametrize( + "options, message", + [ + ({"disable_dynamic_prompt_cache": True}, "hybrid attention models require dynamic prompt cache"), + ({"disable_chunked_prefill": True}, "hybrid attention models require chunked prefill"), + ({"diverse_mode": True}, "--diverse_mode does not support hybrid attention models"), + ], +) +def test_hybrid_rejects_unsupported_modes(launch_args, run_mode, output_constraint_mode, options, message): + args = launch_args("qwen3_5_text", run_mode=run_mode, output_constraint_mode=output_constraint_mode, **options) + with pytest.raises(AssertionError, match=message): + _launch_subprocesses(args) + + +@pytest.mark.parametrize("output_constraint_mode", ["none", "xgrammar"]) +@pytest.mark.parametrize( + "model_type, options", + [ + ("qwen3_5_text", {}), + ("qwen3", {"disable_dynamic_prompt_cache": True}), + ("qwen3", {"disable_chunked_prefill": True}), + ("qwen3", {"diverse_mode": True}), + ], +) +def test_supported_models_pass_startup_validation(launch_args, output_constraint_mode, model_type, options): + args = launch_args(model_type, output_constraint_mode=output_constraint_mode, **options) + with pytest.raises(RuntimeError, match="startup validation passed"): + _launch_subprocesses(args) + + +@pytest.mark.parametrize( + "model_type,disable_mm,disable_vision,batch_tokens,passes", + [ + ("qwen3", True, False, 512, False), + ("qwen3", True, False, 1024, True), + ("qwen3", True, False, 2048, True), + ("qwen3", True, False, None, True), + ("qwen3", False, False, 512, True), + ("qwen3", True, True, 512, True), + ("neo_chat", False, False, 512, False), + ("deepseek_v4", False, False, 512, False), + ("gemma4", False, False, 512, False), + ("neo_chat", False, False, 1024, True), + ("qwen3_vl", False, False, 512, True), + ("qwen3_5", False, False, 512, True), + ], +) +def test_whole_image_blocks_fit_batch_at_startup( + launch_args, model_type, disable_mm, disable_vision, batch_tokens, passes +): + args = launch_args( + model_type, + disable_chunked_mm_input=disable_mm, + batch_max_tokens=batch_tokens, + max_image_token_count=1024, + ) + args.disable_vision = disable_vision + if passes: + with pytest.raises(RuntimeError, match="startup validation passed"): + _launch_subprocesses(args) + else: + with pytest.raises(AssertionError, match="batch_max_tokens >= max_image_token_count"): + _launch_subprocesses(args) + + +@pytest.mark.parametrize("mode,step", [("native", 0), ("replay", 0), ("replay", 4)]) +def test_s3_configuration_is_prepared_before_ssm_validation(launch_args, monkeypatch, tmp_path, mode, step): + from pathlib import Path + from lightllm.utils.petrel_helper import PetrelHelper + + monkeypatch.chdir(tmp_path) + model_dir = "s3://test-bucket/qwen" + config = {"model_type": "qwen3_5_text"} + downloads = [] + + def list_files(path, extension): + assert path == model_dir + return ["config.json"] if extension == "json" else [] + + def download_json(remote, local): + downloads.append(remote) + Path(local).write_text(json.dumps(config)) + + monkeypatch.setattr(PetrelHelper, "list", list_files) + monkeypatch.setattr(PetrelHelper, "download_json", download_json) + args = launch_args("qwen3_5_text", ssm_state_mode=mode, mtp_step=step, mtp_mode="eagle_with_att" if step else None) + args.model_dir = model_dir + assert not Path(model_dir, "config.json").exists() + with pytest.raises(RuntimeError, match="startup validation passed"): + _launch_subprocesses(args) + assert downloads == [model_dir + "/config.json"] + if mode == "replay": + expected_compute_mode = "precompute" if step > 0 else "inline" + assert args.ssm_replay_config == { + "cache_len": 16, + "compute_mode": expected_compute_mode, + } diff --git a/unit_tests/server/test_mtp_start_args.py b/unit_tests/server/test_mtp_start_args.py new file mode 100644 index 0000000000..3aa269bb0e --- /dev/null +++ b/unit_tests/server/test_mtp_start_args.py @@ -0,0 +1,55 @@ +import pytest + +from lightllm.server.api_start import _launch_subprocesses +from lightllm.server.core.objs.start_args_type import StartArgs + + +class _StartupStopped(Exception): + pass + + +def test_mtp_requires_cuda_graph(monkeypatch): + monkeypatch.setattr("lightllm.server.api_start._set_envs_and_config", lambda args: None) + args = StartArgs(mtp_mode="vanilla_no_att", disable_cudagraph=True, output_constraint_mode="none") + + with pytest.raises(AssertionError, match="only supported on Prefill nodes"): + _launch_subprocesses(args) + + +def test_mtp_prefill_still_requires_positive_step(monkeypatch): + monkeypatch.setattr("lightllm.server.api_start._set_envs_and_config", lambda args: None) + monkeypatch.setattr("lightllm.server.api_start.auto_set_max_req_total_len", lambda args: None) + monkeypatch.setattr("lightllm.server.api_start.auto_set_fused_shared_experts", lambda args: None) + monkeypatch.setattr("lightllm.server.api_start.set_unique_server_name", lambda args: None) + args = StartArgs( + run_mode="prefill", + mtp_mode="dspark", + output_constraint_mode="none", + mtp_draft_model_dir=["draft"], + mtp_step=0, + disable_cudagraph=True, + disable_vision=True, + disable_audio=True, + disable_shm_warning=True, + ) + + with pytest.raises(AssertionError): + _launch_subprocesses(args) + + +def test_dsv4_cpu_cache_uses_common_adaptive_placement(monkeypatch): + monkeypatch.setattr("lightllm.server.api_start._set_envs_and_config", lambda args: None) + monkeypatch.setattr("lightllm.server.api_start.auto_set_max_req_total_len", lambda args: None) + monkeypatch.setattr("lightllm.server.api_start.auto_set_fused_shared_experts", lambda args: None) + monkeypatch.setattr("lightllm.server.api_start.set_unique_server_name", lambda args: None) + monkeypatch.setattr("lightllm.utils.config_utils.get_model_type", lambda model_dir: "deepseek_v4") + monkeypatch.setattr( + "lightllm.server.api_start.has_vision_module", + lambda model_dir: (_ for _ in ()).throw(_StartupStopped), + ) + args = StartArgs(enable_cpu_cache=True, disable_audio=True) + + with pytest.raises(_StartupStopped): + _launch_subprocesses(args) + + assert args.cache_placement_strategy == "adaptive" diff --git a/unit_tests/server/test_pd_cache_aware.py b/unit_tests/server/test_pd_cache_aware.py new file mode 100644 index 0000000000..f71105c0d7 --- /dev/null +++ b/unit_tests/server/test_pd_cache_aware.py @@ -0,0 +1,82 @@ +import pytest + +from lightllm.server.httpserver_for_pd_master.pd_selector.prompt_cache_tree import PromptCacheTree + + +@pytest.mark.parametrize("operation", ["insert", "prefix_match"]) +def test_prompt_cache_tree_rejects_mismatched_sample_stride(operation): + tree = PromptCacheTree(sample_stride=4) + + with pytest.raises(AssertionError, match="sample_stride mismatch"): + if operation == "insert": + tree.insert_sampled("po", "prefill", 2) + else: + tree.prefix_match_sampled("po", len("prompt"), 2) + + +def test_prompt_cache_tree_handles_key_deeper_than_python_recursion_limit(): + tree = PromptCacheTree(sample_stride=1) + key = "x" * 5000 + + tree.insert_sampled(key, "prefill", sample_stride=1) + result = tree.prefix_match_sampled(key, input_char_count=len(key), sample_stride=1) + + assert result.prefill_node == "prefill" + assert result.matched_char_count == len(key) + + +def test_prompt_cache_tree_empty_key_does_not_create_nodes(): + tree = PromptCacheTree(sample_stride=1) + + tree.insert_sampled("", "prefill", sample_stride=1) + result = tree.prefix_match_sampled("", input_char_count=0, sample_stride=1) + + assert tree._node_count == 0 + assert not tree.root.children + assert result.prefill_node is None + assert result.matched_char_count == 0 + assert result.input_char_count == 0 + assert result.last_insert_time is None + + +def test_prompt_cache_tree_zero_capacity_evicts_the_complete_path(): + tree = PromptCacheTree(sample_stride=1, max_node_count=0, evict_node_batch=1) + + tree.insert_sampled("你好🙂", "prefill", sample_stride=1) + result = tree.prefix_match_sampled("你好🙂", input_char_count=3, sample_stride=1) + + assert tree._node_count == 0 + assert not tree.root.children + assert result.prefill_node is None + assert result.matched_char_count == 0 + assert result.input_char_count == 3 + assert result.last_insert_time is None + + +def test_prompt_cache_tree_matches_unicode_partial_prefix(): + tree = PromptCacheTree(sample_stride=1) + tree.insert_sampled("你好🙂", "prefill", sample_stride=1) + + result = tree.prefix_match_sampled("你好🙃", input_char_count=3, sample_stride=1) + + assert result.prefill_node == "prefill" + assert result.matched_char_count == 2 + assert result.input_char_count == 3 + assert result.last_insert_time is not None + + +def test_prompt_cache_tree_iterative_insert_keeps_lru_eviction_order(): + tree = PromptCacheTree(sample_stride=1, max_node_count=3, evict_node_batch=1) + tree.insert_sampled("ab", "prefill-ab", sample_stride=1) + tree.insert_sampled("ac", "prefill-ac", sample_stride=1) + + # 再次访问 ab，使 ac 成为最早未使用的叶节点。 + tree.insert_sampled("ab", "prefill-ab", sample_stride=1) + # 插入 d 后超出容量，依次驱逐 ac 和 ab；共同父节点 a 随后成为叶节点。 + tree.insert_sampled("d", "prefill-d", sample_stride=1) + + assert tree._node_count == 2 + assert set(tree.root.children) == {"a", "d"} + result = tree.prefix_match_sampled("ab", input_char_count=2, sample_stride=1) + assert result.prefill_node == "prefill-ab" + assert result.matched_char_count == 1 diff --git a/unit_tests/server/test_pd_local_selection.py b/unit_tests/server/test_pd_local_selection.py new file mode 100644 index 0000000000..669be609b6 --- /dev/null +++ b/unit_tests/server/test_pd_local_selection.py @@ -0,0 +1,919 @@ +"""Local PD policies use shared cache/load snapshots without remote selection.""" + +import asyncio +import random +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock, call + +import pytest + +from lightllm.server.core.objs import FinishStatus, SamplingParams +from lightllm.server.httpserver_for_pd_master.manager import ( + HttpServerManagerForPDMaster, + PDManager, +) +from lightllm.server.httpserver_for_pd_master.pd_selector import pd_node_stats_client +from lightllm.server.httpserver_for_pd_master.pd_selector import pd_selector +from lightllm.server.httpserver_for_pd_master.pd_selector import prompt_cache_tree_client +from lightllm.server.httpserver_for_pd_master.pd_selector.prompt_cache_tree_client import ( + LocalPromptCacheTreeClient, + RemotePromptCacheTreeClient, +) +from lightllm.server.httpserver_for_pd_master.pd_selector.pd_selector import ( + AdaptiveLoadSelector, + LoadBalancedCacheAwareSelector, + PDSelectionExtraInfo, + RandomSelector, + RoundRobinSelector, +) +from lightllm.server.httpserver_for_pd_master.pd_selector.prompt_cache_tree import ( + PromptCacheMatchResult, +) +from lightllm.server.httpserver_for_pd_master.pd_selector.pd_node_stats_server import ( + PDNodeStatsState, +) +from lightllm.server.httpserver_for_pd_master.pd_selector.prompt_cache_tree_server import ( + PromptCacheTreeState, +) +from lightllm.server.pd_io_struct import PD_Client_Obj +from lightllm.utils.error_utils import ServerBusyError + + +PROMPT = "x" * 2048 +COUNTER_METHODS = ("add_dispatched_info", "sub_dispatched_info") + + +@pytest.fixture(autouse=True) +def local_pd_master_mode(monkeypatch): + start_args = SimpleNamespace(httpserver_workers=1) + monkeypatch.setattr(pd_node_stats_client, "get_env_start_args", lambda: start_args) + monkeypatch.setattr(prompt_cache_tree_client, "get_env_start_args", lambda: start_args) + + +def _node( + node_id, + address, + *, + mode="prefill", + chars=0, + requests=0, + load=0.0, + connection_id=None, +): + node = PD_Client_Obj( + node_id=node_id, + client_ip_port=address, + mode=mode, + start_args={}, + connection_id=connection_id or f"connection-{node_id}", + dispatched_prompt_chars=chars, + dispatched_req_num=requests, + ) + node.run_status.total_token_usage_rate = load + return node + + +def _match(node=None, *, matched=1500): + return PromptCacheMatchResult( + prefill_node=None if node is None else node.connection_key, + matched_char_count=0 if node is None else matched, + input_char_count=len(PROMPT), + last_insert_time=None if node is None else 123.0, + ) + + +def _assert_backend_client(selector): + assert isinstance( + selector.policy.prompt_cache_tree, + (LocalPromptCacheTreeClient, RemotePromptCacheTreeClient), + ) + + +def _set_nodes(manager, prefill_nodes, decode_nodes): + manager.pd_manager.prefill_nodes = prefill_nodes + manager.pd_manager.decode_nodes = decode_nodes + manager.pd_manager.url_to_pd_nodes = {node.connection_key: node for node in [*prefill_nodes, *decode_nodes]} + manager.pd_manager.selector.update_nodes(prefill_nodes, decode_nodes) + + +def _manager(strategy="cache_aware", *, prefill_nodes=None, decode_nodes=None, mock_prefix=True): + manager = HttpServerManagerForPDMaster.__new__(HttpServerManagerForPDMaster) + manager.args = SimpleNamespace(select_p_d_node_strategy=strategy, pd_master_mode="elastic") + manager.pd_manager = PDManager(manager.args) + _set_nodes( + manager, + prefill_nodes if prefill_nodes is not None else [_node(100, "p1:8000"), _node(101, "p2:8000")], + decode_nodes if decode_nodes is not None else [_node(200, "d:8000", mode="decode")], + ) + if mock_prefix: + selector = manager.pd_manager.selector + selector.pd_node_stats.update_dispatched_info = AsyncMock(wraps=selector.pd_node_stats.update_dispatched_info) + if strategy == "cache_aware": + selector.policy.prompt_cache_tree.prefix_match_sampled = AsyncMock(return_value=_match()) + manager.pd_node_stats = manager.pd_manager.selector.pd_node_stats + return manager + + +@pytest.mark.parametrize("busy", [False, True]) +def test_shared_tree_is_used_for_selection_and_cache_metadata(monkeypatch, busy): + selector = LoadBalancedCacheAwareSelector(None) + cached = _node(100, "p2:8000", chars=110, requests=int(busy)) + other = _node(101, "p1:8000", chars=100, requests=int(busy)) + decode = _node(200, "d:8000", mode="decode") + selector.update_nodes([other, cached], [decode]) + _assert_backend_client(selector) + selector.policy.prompt_cache_tree.prefix_match_sampled = AsyncMock(return_value=_match(cached)) + + p_node, d_node, extra = asyncio.run(selector.select_p_d_node(PROMPT, None, None)) + + assert p_node is cached + assert d_node is decode + assert extra == PDSelectionExtraInfo(1500 / len(PROMPT), 123.0) + + +def test_cache_metadata_comes_from_the_selection_query(monkeypatch): + selector = LoadBalancedCacheAwareSelector(None) + cached, other = _node(100, "p2:8000"), _node(101, "p1:8000") + selector.update_nodes([other, cached], [_node(200, "d:8000", mode="decode")]) + _assert_backend_client(selector) + selector.policy.prompt_cache_tree.prefix_match_sampled = AsyncMock(return_value=_match(cached)) + + selected, _, extra = asyncio.run(selector.select_p_d_node(PROMPT, None, None)) + + assert selected is cached + assert extra == PDSelectionExtraInfo(1500 / len(PROMPT), 123.0) + selector.policy.prompt_cache_tree.prefix_match_sampled.assert_awaited_once_with( + PROMPT[:: selector.policy.config.sample_stride], + len(PROMPT), + selector.policy.config.sample_stride, + ) + + +@pytest.mark.parametrize("restarted", [False, True]) +def test_cache_hint_survives_endpoint_reconnect_or_process_restart(monkeypatch, restarted): + original = _node(100, "p2:8000", connection_id="old-socket") + current = _node(300 if restarted else 100, "p2:8000", connection_id="new-socket") + other = _node(101, "p1:8000") + selector = LoadBalancedCacheAwareSelector(None) + selector.update_nodes([other, current], [_node(200, "d:8000", mode="decode")]) + _assert_backend_client(selector) + selector.policy.prompt_cache_tree.prefix_match_sampled = AsyncMock(return_value=_match(original)) + + selected, _, extra = asyncio.run(selector.select_p_d_node(PROMPT, None, None)) + + assert original.connection_key == current.connection_key + assert selected is current + assert extra == PDSelectionExtraInfo(1500 / len(PROMPT), 123.0) + + +@pytest.mark.parametrize("match_kind", ["miss", "remote_worker"]) +def test_shared_match_never_falls_back_to_local_tree_or_remote_candidates(monkeypatch, match_kind): + busy = _node(100, "p1:8000", chars=2000, requests=1) + available = _node(101, "p2:8000", chars=100, requests=1) + selector = LoadBalancedCacheAwareSelector(None) + selector.update_nodes([busy, available], [_node(200, "d:8000", mode="decode")]) + asyncio.run( + selector.policy.prompt_cache_tree.insert_sampled( + PROMPT[:: selector.policy.config.sample_stride], + busy.connection_key, + selector.policy.config.sample_stride, + ) + ) + _assert_backend_client(selector) + shared_match = _match() if match_kind == "miss" else _match(_node(999, "remote:8000")) + selector.policy.prompt_cache_tree.prefix_match_sampled = AsyncMock(return_value=shared_match) + + selected, _, extra = asyncio.run(selector.select_p_d_node(PROMPT, None, None)) + + assert selected is available + assert extra == PDSelectionExtraInfo() + + +def test_default_cache_selection_uses_connection_keys(): + selector = LoadBalancedCacheAwareSelector(None) + cached, other = _node(100, "p2:8000"), _node(101, "p1:8000") + selector.update_nodes([other, cached], [_node(200, "d:8000", mode="decode")]) + asyncio.run( + selector.policy.prompt_cache_tree.insert_sampled( + PROMPT[:: selector.policy.config.sample_stride], + cached.connection_key, + selector.policy.config.sample_stride, + ) + ) + + selected, _, extra = asyncio.run(selector.select_p_d_node(PROMPT, None, None)) + + assert selected is cached + assert extra.estimated_cache_hit_rate > 0.5 + assert ( + asyncio.run( + selector.policy.prompt_cache_tree.prefix_match_sampled( + PROMPT[:: selector.policy.config.sample_stride], + len(PROMPT), + selector.policy.config.sample_stride, + ) + ).prefill_node + == cached.connection_key + ) + + +def test_round_robin_progress_is_local_to_each_worker(): + nodes = [_node(100, "p1:8000"), _node(101, "p2:8000")] + decode = [ + _node(200, "d1:8000", mode="decode"), + _node(201, "d2:8000", mode="decode"), + ] + first, second = RoundRobinSelector(None), RoundRobinSelector(None) + for selector in (first, second): + selector.update_nodes(nodes, decode) + + async def run(): + assert (await first.select_p_d_node(PROMPT, None, None))[:2] == ( + nodes[0], + decode[0], + ) + assert (await first.select_p_d_node(PROMPT, None, None))[:2] == ( + nodes[1], + decode[1], + ) + assert (await second.select_p_d_node(PROMPT, None, None))[:2] == ( + nodes[0], + decode[0], + ) + + asyncio.run(run()) + + +@pytest.mark.parametrize("selector_type", [RandomSelector, AdaptiveLoadSelector]) +def test_random_and_adaptive_keep_their_local_load_preferences(monkeypatch, selector_type): + monkeypatch.setattr(pd_selector, "random", random.Random(23)) + busy, available = _node(100, "p1:8000", load=1.0), _node(101, "p2:8000", load=0.0) + selector = selector_type(None) + selector.update_nodes([busy, available], [_node(200, "d:8000", mode="decode")]) + + async def select_many(): + return [(await selector.select_p_d_node(PROMPT, None, None))[0] for _ in range(100)] + + selections = asyncio.run(select_many()) + + assert all(node is busy or node is available for node in selections) + available_count = sum(node is available for node in selections) + if selector_type is AdaptiveLoadSelector: + assert available_count > 90 + else: + assert 20 < available_count < 80 + + +@pytest.mark.parametrize("strategy", ["random", "round_robin", "adaptive_load", "cache_aware"]) +def test_only_cache_aware_policy_queries_prefix_data_and_all_policies_choose_local_nodes(monkeypatch, strategy): + manager = _manager(strategy) + local_nodes = manager.pd_manager.prefill_nodes + if strategy == "cache_aware": + _assert_backend_client(manager.pd_manager.selector) + manager.pd_manager.selector.policy.prompt_cache_tree.prefix_match_sampled.return_value = _match(local_nodes[-1]) + + selected, decode, _ = asyncio.run(manager.select_p_d_node(PROMPT, None, None)) + + assert any(selected is node for node in local_nodes) + assert decode is manager.pd_manager.decode_nodes[0] + if strategy == "cache_aware": + assert manager.pd_manager.selector.policy.prompt_cache_tree.prefix_match_sampled.await_count == 1 + manager.pd_manager.selector.policy.prompt_cache_tree.prefix_match_sampled.assert_awaited_with( + PROMPT[:: manager.pd_manager.selector.policy.config.sample_stride], + len(PROMPT), + manager.pd_manager.selector.policy.config.sample_stride, + ) + manager.pd_manager.selector.pd_node_stats.update_dispatched_info.assert_awaited_once_with(p_nodes=local_nodes) + else: + manager.pd_manager.selector.pd_node_stats.update_dispatched_info.assert_not_awaited() + + +def test_manager_uses_shared_node_load_and_local_balancing_threshold_for_selection( + monkeypatch, +): + cached, available = _node(100, "p1:8000"), _node(101, "p2:8000") + manager = _manager(prefill_nodes=[cached, available]) + _assert_backend_client(manager.pd_manager.selector) + + async def update_dispatched_info(p_nodes): + assert p_nodes == [cached, available] + cached.dispatched_prompt_chars, cached.dispatched_req_num = 500, 1 + available.dispatched_prompt_chars, available.dispatched_req_num = 100, 1 + + manager.pd_manager.selector.pd_node_stats.update_dispatched_info.side_effect = update_dispatched_info + manager.pd_manager.selector.policy.prompt_cache_tree.prefix_match_sampled.return_value = _match(cached) + + manager.pd_manager.selector.policy.config.balance_rel_threshold = 1.1 + first = asyncio.run(manager.select_p_d_node(PROMPT, None, None)) + manager.pd_manager.selector.policy.config.balance_rel_threshold = 2.0 + second = asyncio.run(manager.select_p_d_node(PROMPT, None, None)) + + assert first[0] is available + assert first[2] == PDSelectionExtraInfo() + assert second[0] is cached + assert second[2].estimated_cache_hit_rate == 1500 / len(PROMPT) + + +@pytest.mark.parametrize("replacement_load", [(0, 0), (999, 9)]) +def test_query_completion_uses_current_connections_without_overwriting_their_load(monkeypatch, replacement_load): + async def run(): + old_p, old_d = _node(100, "p:8000"), _node(200, "d:8000", mode="decode") + new_p = _node(300, "p:8000", chars=replacement_load[0], requests=replacement_load[1]) + new_d = _node(400, "d:8000", mode="decode") + manager = _manager(prefill_nodes=[old_p], decode_nodes=[old_d]) + _assert_backend_client(manager.pd_manager.selector) + started, complete = asyncio.Event(), asyncio.Event() + + async def update_dispatched_info(p_nodes): + assert len(p_nodes) == 1 + if not started.is_set(): + assert p_nodes[0] is old_p + started.set() + await complete.wait() + old_p.dispatched_prompt_chars, old_p.dispatched_req_num = 800, 2 + + manager.pd_manager.selector.pd_node_stats.update_dispatched_info.side_effect = update_dispatched_info + manager.pd_manager.selector.policy.prompt_cache_tree.prefix_match_sampled.return_value = _match(old_p) + task = asyncio.create_task(manager.select_p_d_node(PROMPT, None, None)) + await asyncio.wait_for(started.wait(), timeout=2) + _set_nodes(manager, [new_p], [new_d]) + complete.set() + selected, decode, extra = await task + + assert selected is new_p and decode is new_d + assert ( + new_p.dispatched_prompt_chars, + new_p.dispatched_req_num, + ) == replacement_load + assert extra == PDSelectionExtraInfo(1500 / len(PROMPT), 123.0) + assert manager.pd_manager.selector.policy.prompt_cache_tree.prefix_match_sampled.await_count == 1 + manager.pd_manager.selector.policy.prompt_cache_tree.prefix_match_sampled.assert_awaited_with( + PROMPT[:: manager.pd_manager.selector.policy.config.sample_stride], + len(PROMPT), + manager.pd_manager.selector.policy.config.sample_stride, + ) + + asyncio.run(run()) + + +@pytest.mark.parametrize("replaced", ["prefill", "decode"]) +def test_selection_query_uses_current_decode_and_rejects_a_replaced_prefill(replaced): + async def run(): + manager = _manager() + selector = manager.pd_manager.selector + prefill, decode = ( + manager.pd_manager.prefill_nodes, + manager.pd_manager.decode_nodes, + ) + selected_prefill = prefill[0] + replacement = _node(300, "p1:8000") if replaced == "prefill" else _node(400, "d:8000", mode="decode") + started, complete = asyncio.Event(), asyncio.Event() + + async def query(*_args): + started.set() + await complete.wait() + return _match(selected_prefill) + + selector.policy.prompt_cache_tree.prefix_match_sampled.side_effect = query + task = asyncio.create_task(manager.select_p_d_node(PROMPT, None, None)) + await asyncio.wait_for(started.wait(), timeout=2) + _set_nodes( + manager, + [replacement, *prefill[1:]] if replaced == "prefill" else prefill, + [replacement] if replaced == "decode" else decode, + ) + complete.set() + if replaced == "prefill": + with pytest.raises(ServerBusyError, match="connection changed"): + await task + else: + selected, current_decode, _ = await task + assert selected is selected_prefill + assert current_decode is replacement + + asyncio.run(run()) + + +@pytest.fixture(params=["local", "rpyc"]) +def backend(request, monkeypatch): + start_args = SimpleNamespace(httpserver_workers=2 if request.param == "rpyc" else 1) + monkeypatch.setattr(pd_node_stats_client, "get_env_start_args", lambda: start_args) + monkeypatch.setattr(prompt_cache_tree_client, "get_env_start_args", lambda: start_args) + monkeypatch.setattr( + pd_node_stats_client, + "get_pd_master_unix_socket_address", + lambda service_name: f"\0lightllm_test-pd-info_{service_name}".encode(), + ) + monkeypatch.setattr( + prompt_cache_tree_client, + "get_pd_master_unix_socket_address", + lambda service_name: f"\0lightllm_test-pd-info_{service_name}".encode(), + ) + if request.param == "rpyc": + monkeypatch.setattr(pd_node_stats_client, "unix_connect", Mock()) + monkeypatch.setattr(prompt_cache_tree_client, "unix_connect", Mock()) + return request.param + + +def _generator_manager(backend, strategy="round_robin"): + manager = _manager( + strategy, + prefill_nodes=[_node(100, "p:8000")], + mock_prefix=False, + ) + states = SimpleNamespace(stats=None, tree=None) + if backend != "local": + states.stats = PDNodeStatsState() + stats = manager.pd_manager.selector.pd_node_stats + stats._get_dispatched_info = states.stats.get_dispatched_info + stats._add_dispatched_info = states.stats.add_dispatched_info + stats._sub_dispatched_info = states.stats.sub_dispatched_info + if strategy == "cache_aware": + states.tree = PromptCacheTreeState() + tree = manager.pd_manager.selector.policy.prompt_cache_tree + + def prefix_match_sampled(*args): + return PromptCacheMatchResult(*states.tree.prefix_match_sampled(*args)) + + tree._prefix_match_sampled = prefix_match_sampled + tree._insert_sampled = states.tree.insert_sampled + for name in COUNTER_METHODS: + stats = manager.pd_manager.selector.pd_node_stats + setattr(stats, name, AsyncMock(wraps=getattr(stats, name))) + manager.pd_manager.selector.record_prompt_cache_hit_rate = Mock( + wraps=manager.pd_manager.selector.record_prompt_cache_hit_rate + ) + if strategy == "cache_aware": + tree = manager.pd_manager.selector.policy.prompt_cache_tree + tree.insert_sampled = AsyncMock(wraps=tree.insert_sampled) + manager.active_req_ids = set() + manager.disable_pd_cache_high_priority = True + manager.enable_pd_node_self_request_limit = False + manager.remove_req = AsyncMock() + manager.abort = AsyncMock() + ids = iter(range(8, 80, 8)) + manager.id_gen = SimpleNamespace(generate_id=lambda: next(ids)) + return manager, states + + +def _pending(state, node): + if state.stats is None: + return node.dispatched_prompt_chars, node.dispatched_req_num + return state.stats.get_dispatched_info((node.connection_key,))[0] + + +def test_backend_tree_queries_preserve_cache_hint_for_restarted_endpoint(monkeypatch, backend): + async def run(): + manager, state = _generator_manager(backend, "cache_aware") + cached, other = _node(100, "p2:8000"), _node(101, "p1:8000") + decode = manager.pd_manager.decode_nodes + _set_nodes(manager, [other, cached], decode) + selector = manager.pd_manager.selector + await selector.policy.prompt_cache_tree.insert_sampled( + PROMPT[:: selector.policy.config.sample_stride], + cached.connection_key, + selector.policy.config.sample_stride, + ) + injected_tree = selector.policy.prompt_cache_tree + tree = ( + selector.policy.prompt_cache_tree.prompt_cache_tree if state.tree is None else state.tree.prompt_cache_tree + ) + prefix_match = Mock(wraps=tree.prefix_match_sampled) + monkeypatch.setattr(tree, "prefix_match_sampled", prefix_match) + assert isinstance( + injected_tree, + (LocalPromptCacheTreeClient, RemotePromptCacheTreeClient), + ) + + selected, _, extra = await manager.select_p_d_node(PROMPT, None, None) + + assert selected is cached + assert extra.estimated_cache_hit_rate > 0.5 + assert extra.cache_last_insert_time is not None + original_extra = extra + assert prefix_match.call_count == 1 + + restarted = _node(300, cached.client_ip_port) + _set_nodes(manager, [other, restarted], decode) + selected, _, extra = await manager.select_p_d_node(PROMPT, None, None) + assert selected is restarted + assert extra == original_extra + assert prefix_match.call_count == 2 + assert selector.policy.prompt_cache_tree is injected_tree + + asyncio.run(run()) + + +def test_cache_feedback_and_dynamic_thresholds_stay_private_to_each_worker(backend): + first, state = _generator_manager(backend, "cache_aware") + second, _ = _generator_manager(backend, "cache_aware") + stats_calls = tree_calls = None + if state.stats is not None: + stats_calls = Mock(wraps=state.stats.get_dispatched_info) + + def prefix_match_sampled(*args): + return PromptCacheMatchResult(*state.tree.prefix_match_sampled(*args)) + + tree_calls = Mock(wraps=prefix_match_sampled) + for manager in (first, second): + selector = manager.pd_manager.selector + selector.pd_node_stats._get_dispatched_info = stats_calls + selector.policy.prompt_cache_tree._prefix_match_sampled = tree_calls + + for manager, rates in ( + (first, (0.2, 0.2, 0.8, 0.8)), + (second, (0.8, 0.8, 0.2, 0.2)), + ): + selector = manager.pd_manager.selector + selector.policy.config.cache_hit_rate_window_size = 2 + for rate in rates: + assert selector.record_prompt_cache_hit_rate(rate) is None + + first_policy, second_policy = ( + first.pd_manager.selector.policy, + second.pd_manager.selector.policy, + ) + assert first_policy.config.balance_rel_threshold == pytest.approx(1.55) + assert second_policy.config.balance_rel_threshold == pytest.approx(1.45) + assert first_policy.balance_rel_threshold_controller is not second_policy.balance_rel_threshold_controller + if stats_calls is not None: + stats_calls.assert_not_called() + tree_calls.assert_not_called() + + async def query_both_workers(): + for manager in (first, second): + await manager.select_p_d_node(PROMPT, None, None) + + asyncio.run(query_both_workers()) + + assert first_policy.config.balance_rel_threshold == pytest.approx(1.55) + assert second_policy.config.balance_rel_threshold == pytest.approx(1.45) + if stats_calls is not None: + assert stats_calls.call_count == 2 + assert tree_calls.call_count == 2 + + +def _params(max_new_tokens=3): + params = SamplingParams() + params.init(None, max_new_tokens=max_new_tokens) + return params + + +async def _consume(manager, params): + return [item async for item in manager._generate_one_attempt("prompt", params, None, None, 0, 80, 2)] + + +def test_cache_aware_selection_and_initial_reservation_are_serialized_per_worker(backend): + async def run(): + manager, state = _generator_manager(backend, "cache_aware") + cached = _node(100, "p1:8000") + other = _node(101, "p2:8000") + _set_nodes(manager, [cached, other], manager.pd_manager.decode_nodes) + manager.pd_manager.selector.policy.prompt_cache_tree.prefix_match_sampled = AsyncMock( + return_value=_match(cached) + ) + + active_selections = 0 + max_active_selections = 0 + original_select = manager.select_p_d_node + + async def select_p_d_node(*args): + nonlocal active_selections, max_active_selections + active_selections += 1 + max_active_selections = max(max_active_selections, active_selections) + try: + # 主动让出事件循环，验证并发请求仍无法同时进入选点区间。 + await asyncio.sleep(0) + return await original_select(*args) + finally: + active_selections -= 1 + + manager.select_p_d_node = select_p_d_node + selected_nodes = [] + all_reserved = asyncio.Event() + release_requests = asyncio.Event() + request_count = 16 + request_ids = iter(range(8, 8 * (request_count + 1), 8)) + manager.id_gen = SimpleNamespace(generate_id=lambda: next(request_ids)) + + async def generate_segment(p_node, _d_node, _start_time, _prompt, params, _multimodal, _request): + selected_nodes.append(p_node) + if len(selected_nodes) == request_count: + all_reserved.set() + await release_requests.wait() + yield ( + params.group_request_id, + "x", + { + "prompt_tokens": len(PROMPT), + "prompt_cache_len": 0, + "needs_prefill_first_token": False, + }, + FinishStatus(FinishStatus.FINISHED_STOP), + ) + + manager._wait_to_token_package = generate_segment + + async def consume(origin_request_id): + return [ + item + async for item in manager._generate_one_attempt( + PROMPT, _params(max_new_tokens=1), None, None, 0, origin_request_id, 2 + ) + ] + + tasks = [asyncio.create_task(consume(1000 + index * 8)) for index in range(request_count)] + try: + await asyncio.wait_for(all_reserved.wait(), timeout=5) + assert max_active_selections == 1 + assert cached in selected_nodes + assert other in selected_nodes + finally: + release_requests.set() + await asyncio.gather(*tasks) + + assert _pending(state, cached) == (0, 0) + assert _pending(state, other) == (0, 0) + + asyncio.run(run()) + + +@pytest.mark.parametrize("strategy", ["random", "round_robin", "adaptive_load", "cache_aware"]) +@pytest.mark.parametrize("needs_prefill", [True, False]) +def test_segment_counters_clear_at_first_token_and_restart_for_continuation(backend, strategy, needs_prefill): + manager, state = _generator_manager(backend, strategy) + prefill = manager.pd_manager.prefill_nodes[0] + segments = [] + + async def generate_segment(p_node, d_node, start_time, prompt, params, multimodal, request): + segments.append(prompt) + assert _pending(state, p_node) == (len(prompt), 1) + yield ( + params.group_request_id, + "x", + { + "prompt_tokens": 100, + "prompt_cache_len": 20, + "needs_prefill_first_token": needs_prefill, + }, + FinishStatus(), + ) + assert _pending(state, p_node) == (0, 0) + finish = FinishStatus.FINISHED_PD_DECODE_CAPACITY if len(segments) == 1 else FinishStatus.FINISHED_STOP + yield params.group_request_id, "", {"prompt_tokens": 100}, FinishStatus(finish) + + manager._wait_to_token_package = generate_segment + results = asyncio.run(_consume(manager, _params())) + + assert segments == ["prompt", "promptx"] + assert len(results) == 3 + assert _pending(state, prefill) == (0, 0) + assert not manager.active_req_ids + stats = manager.pd_manager.selector.pd_node_stats + assert stats.add_dispatched_info.await_args_list == [ + call(p_node=prefill, chars=len("prompt"), req_num=1), + call(p_node=prefill, chars=len("promptx"), req_num=1), + ] + assert stats.sub_dispatched_info.await_args_list == [ + call(p_node=prefill, chars=len("prompt"), req_num=1), + call(p_node=prefill, chars=len("promptx"), req_num=1), + ] + selector = manager.pd_manager.selector + selector.record_prompt_cache_hit_rate.assert_called_once_with(0.2) + if strategy == "cache_aware": + assert selector.policy.balance_rel_threshold_controller._cache_hit_rates == [0.2] + if needs_prefill: + sample_stride = manager.pd_manager.selector.policy.config.sample_stride + selector.policy.prompt_cache_tree.insert_sampled.assert_awaited_once_with( + "prompt"[::sample_stride], prefill.connection_key, sample_stride + ) + else: + selector.policy.prompt_cache_tree.insert_sampled.assert_not_awaited() + + +def test_continuation_tokenization_failure_clears_reserved_counter(backend): + manager, state = _generator_manager(backend) + prefill = manager.pd_manager.prefill_nodes[0] + params = SamplingParams() + params.init(None, max_new_tokens=3, regular_constraint="ab") + + async def generate_first_segment(_p_node, _d_node, _start_time, _prompt, params, _multimodal, _request): + yield ( + params.group_request_id, + "x", + { + "prompt_tokens": 100, + "prompt_cache_len": 0, + "needs_prefill_first_token": False, + }, + FinishStatus(), + ) + yield ( + params.group_request_id, + "", + {"prompt_tokens": 100}, + FinishStatus(FinishStatus.FINISHED_PD_DECODE_CAPACITY), + ) + + manager._wait_to_token_package = generate_first_segment + manager.tokens = Mock(side_effect=RuntimeError("continuation tokenization failed")) + + with pytest.raises(RuntimeError, match="continuation tokenization failed"): + asyncio.run(_consume(manager, params)) + + assert manager.pd_node_stats.add_dispatched_info.await_args_list == [ + call(p_node=prefill, chars=len("prompt"), req_num=1), + call(p_node=prefill, chars=len("promptx"), req_num=1), + ] + assert manager.pd_node_stats.sub_dispatched_info.await_args_list == [ + call(p_node=prefill, chars=len("prompt"), req_num=1), + call(p_node=prefill, chars=len("promptx"), req_num=1), + ] + assert _pending(state, prefill) == (0, 0) + + +def test_error_output_records_hit_rate_without_inserting_cache(backend): + manager, state = _generator_manager(backend, "cache_aware") + prefill = manager.pd_manager.prefill_nodes[0] + + async def generate_segment(p_node, d_node, start_time, prompt, params, multimodal, request): + yield ( + params.group_request_id, + "", + { + "prompt_tokens": 100, + "prompt_cache_len": 20, + "needs_prefill_first_token": True, + }, + FinishStatus(FinishStatus.FINISHED_ERROR), + ) + + manager._wait_to_token_package = generate_segment + results = asyncio.run(_consume(manager, _params())) + + assert len(results) == 1 and results[0][3].is_error_finished() + assert _pending(state, prefill) == (0, 0) + selector = manager.pd_manager.selector + selector.record_prompt_cache_hit_rate.assert_called_once_with(0.2) + assert selector.policy.balance_rel_threshold_controller._cache_hit_rates == [0.2] + selector.policy.prompt_cache_tree.insert_sampled.assert_not_awaited() + + +def test_error_before_first_token_clears_segment_counters(backend): + manager, state = _generator_manager(backend) + prefill = manager.pd_manager.prefill_nodes[0] + + async def generate_segment(*_args): + assert _pending(state, prefill) == (len("prompt"), 1) + raise RuntimeError("before first token") + yield + + manager._wait_to_token_package = generate_segment + with pytest.raises(RuntimeError, match="before first token"): + asyncio.run(_consume(manager, _params())) + + assert _pending(state, prefill) == (0, 0) + assert not manager.active_req_ids + manager.abort.assert_awaited_once() + manager.remove_req.assert_awaited_once() + manager.pd_node_stats.sub_dispatched_info.assert_awaited_once_with(p_node=prefill, chars=len("prompt"), req_num=1) + + +def test_sampling_params_copy_failure_clears_initial_reservation(monkeypatch, backend): + manager, state = _generator_manager(backend) + prefill = manager.pd_manager.prefill_nodes[0] + manager._wait_to_token_package = Mock() + params = _params() + + def fail_copy(*_args, **_kwargs): + raise RuntimeError("segment initialization failed") + + monkeypatch.setattr(SamplingParams, "from_buffer_copy", fail_copy) + with pytest.raises(RuntimeError, match="segment initialization"): + asyncio.run(_consume(manager, params)) + manager.pd_node_stats.add_dispatched_info.assert_awaited_once_with(p_node=prefill, chars=len("prompt"), req_num=1) + manager.pd_node_stats.sub_dispatched_info.assert_awaited_once_with(p_node=prefill, chars=len("prompt"), req_num=1) + assert _pending(state, prefill) == (0, 0) + manager._wait_to_token_package.assert_not_called() + manager.remove_req.assert_awaited_once() + assert not manager.active_req_ids + + +@pytest.mark.parametrize("replaced", ["prefill", "decode"]) +def test_connection_closed_during_counter_update_fails_at_dispatch_and_clears_load(backend, replaced): + async def run(): + manager, state = _generator_manager(backend) + prefill, decode = ( + manager.pd_manager.prefill_nodes[0], + manager.pd_manager.decode_nodes[0], + ) + created, complete = asyncio.Event(), asyncio.Event() + original_add_info = manager.pd_node_stats.add_dispatched_info + + async def add_info(**kwargs): + await original_add_info(**kwargs) + created.set() + await complete.wait() + + async def generate_segment(p_node, d_node, *_args): + assert p_node is prefill and d_node is decode + raise ConnectionError("selected PD connection closed") + yield + + manager.pd_node_stats.add_dispatched_info = add_info + manager._wait_to_token_package = Mock(side_effect=generate_segment) + task = asyncio.create_task(_consume(manager, _params())) + await asyncio.wait_for(created.wait(), timeout=2) + previous = prefill if replaced == "prefill" else decode + replacement = _node( + previous.node_id, + previous.client_ip_port, + mode=previous.mode, + chars=23, + requests=2, + ) + _set_nodes( + manager, + [replacement] if replaced == "prefill" else [prefill], + [replacement] if replaced == "decode" else [decode], + ) + complete.set() + with pytest.raises(ConnectionError, match="selected PD connection closed"): + await task + + manager._wait_to_token_package.assert_called_once() + manager.abort.assert_awaited_once() + manager.remove_req.assert_awaited_once() + assert _pending(state, prefill) == (0, 0) + assert ( + replacement.dispatched_prompt_chars, + replacement.dispatched_req_num, + ) == (23, 2) + assert not manager.active_req_ids + + asyncio.run(run()) + + +def test_cancelled_request_clears_counters_before_removing_request(backend): + manager, state = _generator_manager(backend) + prefill = manager.pd_manager.prefill_nodes[0] + + async def run(): + started = asyncio.Event() + + async def remove_req(*_args): + assert _pending(state, prefill) == (0, 0) + + manager.remove_req.side_effect = remove_req + + async def generate_segment(*_args): + started.set() + await asyncio.Event().wait() + yield + + manager._wait_to_token_package = generate_segment + task = asyncio.create_task(_consume(manager, _params())) + try: + await asyncio.wait_for(started.wait(), timeout=2) + assert _pending(state, prefill) == (len("prompt"), 1) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + finally: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + asyncio.run(run()) + assert _pending(state, prefill) == (0, 0) + assert not manager.active_req_ids + manager.remove_req.assert_awaited_once() + + +def test_cancellation_after_counter_release_does_not_release_twice(backend): + manager, state = _generator_manager(backend) + prefill = manager.pd_manager.prefill_nodes[0] + + async def run(): + released = asyncio.Event() + original_sub_dispatched_info = manager.pd_node_stats.sub_dispatched_info + + async def sub_dispatched_info(**kwargs): + await original_sub_dispatched_info(**kwargs) + released.set() + await asyncio.Event().wait() + + manager.pd_node_stats.sub_dispatched_info = AsyncMock(side_effect=sub_dispatched_info) + + async def generate_segment(_p_node, _d_node, _start_time, _prompt, params, _multimodal, _request): + yield params.group_request_id, "x", {"prompt_tokens": 1}, FinishStatus() + + manager._wait_to_token_package = generate_segment + task = asyncio.create_task(_consume(manager, _params())) + await asyncio.wait_for(released.wait(), timeout=2) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + asyncio.run(run()) + assert _pending(state, prefill) == (0, 0) + manager.pd_node_stats.sub_dispatched_info.assert_awaited_once() diff --git a/unit_tests/server/test_pd_master_haproxy_proxy.py b/unit_tests/server/test_pd_master_haproxy_proxy.py new file mode 100644 index 0000000000..3642bc69e3 --- /dev/null +++ b/unit_tests/server/test_pd_master_haproxy_proxy.py @@ -0,0 +1,434 @@ +"""Exercise the real streaming proxy bundled by ray-haproxy.""" + +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +import json +import os +from pathlib import Path +import signal +import socket +import subprocess +import threading +import time +from urllib.error import HTTPError, URLError +from urllib.request import ProxyHandler, Request, build_opener + +import pytest +import psutil + +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.server.httpserver_for_pd_master.multiprocess import ( + _build_haproxy_config, + _build_initial_health_state, +) +from lightllm.utils import start_utils +from lightllm.utils.start_utils import kill_recursive +from unit_tests.server.pd_master_test_utils import get_test_haproxy_binary + + +def test_real_haproxy_starts_with_a_low_fd_hard_limit(tmp_path): + binary = str(Path(get_test_haproxy_binary()).resolve()) + with socket.socket() as listener: + listener.bind(("127.0.0.1", 0)) + public_port = listener.getsockname()[1] + + args = StartArgs( + run_mode="pd_master", + host="127.0.0.1", + port=public_port, + httpserver_workers=2, + ) + state_file = tmp_path / "haproxy-low-fd.state" + state_file.write_text(_build_initial_health_state([1, 2]), encoding="utf-8") + config = tmp_path / "haproxy-low-fd.cfg" + config.write_text(_build_haproxy_config(args, [1, 2], state_file), encoding="utf-8") + error_log_path = tmp_path / "haproxy-low-fd.stderr" + # 同时降低 soft/hard limit，验证 HAProxy 会按实际 FD 上限自动缩小 maxconn。 + command = [ + "/bin/bash", + "-c", + 'ulimit -n 1024 && exec "$@"', + "lightllm-haproxy-low-fd", + binary, + "-db", + "-f", + str(config), + ] + with error_log_path.open("wb") as error_log: + haproxy = subprocess.Popen(command, stdout=subprocess.DEVNULL, stderr=error_log) + try: + deadline = time.monotonic() + 5 + while True: + assert haproxy.poll() is None, error_log_path.read_text() + try: + with socket.create_connection(("127.0.0.1", public_port), timeout=0.2): + break + except OSError: + assert time.monotonic() < deadline, error_log_path.read_text() + time.sleep(0.02) + finally: + if haproxy.poll() is None: + haproxy.send_signal(signal.SIGUSR1) + haproxy.wait(timeout=5) + + assert haproxy.returncode == 0 + + +@pytest.fixture +def haproxy_proxy(tmp_path, request): + binary = get_test_haproxy_binary() + state = { + "requests": [], + "lock": threading.Lock(), + "release_stream": threading.Event(), + "release_initial_probes": threading.Event(), + "stream_disconnected": threading.Event(), + "runtime": tmp_path, + "probe_status": {worker: {path: 200 for path in ("/liveness", "/readiness", "/health")} for worker in range(2)}, + "probe_delays": {}, + "probe_observations": [], + } + if getattr(request, "param", None) != "block_initial_health_checks": + state["release_initial_probes"].set() + + class Handler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def log_message(self, *args): + pass + + def do_GET(self): + if self.path in ("/stream", "/stream_until_disconnect"): + self.send_response(200) + self.send_header("Content-Type", "text/event-stream") + self.send_header("Connection", "close") + self.end_headers() + self.close_connection = True + try: + self.wfile.write(b"data: first\n\n") + self.wfile.flush() + if self.path == "/stream_until_disconnect": + for _ in range(500): + self.wfile.write(b"data: " + b"x" * (64 * 1024) + b"\n\n") + self.wfile.flush() + time.sleep(0.01) + else: + state["release_stream"].wait(timeout=5) + self.wfile.write(b"data: last\n\n") + self.wfile.flush() + except (BrokenPipeError, ConnectionResetError): + state["stream_disconnected"].set() + return + with state["lock"]: + status = state["probe_status"][self.server.worker_id].get(self.path, 200) + delay = state["probe_delays"].get((self.server.worker_id, self.path), 0) + if self.path in ("/liveness", "/readiness", "/health"): + state["release_initial_probes"].wait(timeout=15) + if delay: + time.sleep(delay) + with state["lock"]: + if self.path in ("/liveness", "/readiness", "/health"): + state["probe_observations"].append((self.server.worker_id, self.path, status)) + body = json.dumps({"worker_id": self.server.worker_id}).encode() + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + try: + self.wfile.write(body) + except (BrokenPipeError, ConnectionResetError): + pass + + def do_POST(self): + body = self.rfile.read(int(self.headers["Content-Length"])) + with state["lock"]: + state["requests"].append((self.server.worker_id, body)) + # Simulate a worker dying after accepting a non-idempotent request. + # A proxy retry would submit this same body to the other worker. + self.close_connection = True + self.connection.shutdown(socket.SHUT_RDWR) + self.connection.close() + + servers = [] + threads = [] + haproxy = None + try: + for worker_id in range(2): + server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + server.daemon_threads = True + server.worker_id = worker_id + thread = threading.Thread(target=server.serve_forever, kwargs={"poll_interval": 0.05}, daemon=True) + thread.start() + servers.append(server) + threads.append(thread) + with socket.socket() as listener: + listener.bind(("127.0.0.1", 0)) + public_port = listener.getsockname()[1] + args = StartArgs( + run_mode="pd_master", + host="127.0.0.1", + port=public_port, + httpserver_workers=2, + ) + worker_ports = [server.server_port for server in servers] + state_file = tmp_path / "haproxy-health.state" + state_file.write_text(_build_initial_health_state(worker_ports), encoding="utf-8") + config = tmp_path / "haproxy.cfg" + config.write_text(_build_haproxy_config(args, worker_ports, state_file), encoding="utf-8") + executable = str(Path(binary).resolve()) + check = subprocess.run([executable, "-c", "-f", str(config)], capture_output=True, text=True, timeout=5) + assert check.returncode == 0, check.stderr + command = [executable, "-db", "-f", str(config)] + with (tmp_path / "haproxy.stderr").open("wb") as error_log: + haproxy = subprocess.Popen(command, stdout=subprocess.DEVNULL, stderr=error_log, start_new_session=True) + state["haproxy"] = haproxy + opener = build_opener(ProxyHandler({})) + base_url = f"http://127.0.0.1:{public_port}" + deadline = time.monotonic() + 5 + while True: + assert haproxy.poll() is None, (tmp_path / "haproxy.stderr").read_text() + try: + with opener.open(base_url + "/worker", timeout=0.5) as response: + assert response.status == 200 + break + except (OSError, URLError): + if time.monotonic() >= deadline: + pytest.fail("HAProxy did not start: " + (tmp_path / "haproxy.stderr").read_text()) + time.sleep(0.02) + yield base_url, opener, state + finally: + state["release_stream"].set() + state["release_initial_probes"].set() + if haproxy is not None: + if haproxy.poll() is None: + haproxy.send_signal(signal.SIGUSR1) + try: + haproxy.wait(timeout=5) + except subprocess.TimeoutExpired: + kill_recursive(haproxy) + haproxy.wait(timeout=5) + try: + os.killpg(haproxy.pid, signal.SIGKILL) + except ProcessLookupError: + pass + for server in servers: + server.shutdown() + server.server_close() + for thread in threads: + thread.join(timeout=1) + + +def test_real_haproxy_balances_workers_and_delivers_sse_before_stream_completion(haproxy_proxy): + base_url, opener, state = haproxy_proxy + proxy_process = psutil.Process(state["haproxy"].pid) + assert proxy_process.children(recursive=True) == [] + assert proxy_process.num_threads() >= 2 + workers = set() + for _ in range(8): + with opener.open(base_url + "/worker", timeout=3) as response: + workers.add(json.load(response)["worker_id"]) + assert workers == {0, 1} + + # The backend waits for this client to read the first event before it sends + # the last one. Buffering the complete stream would deadlock until timeout. + with opener.open(base_url + "/stream", timeout=3) as response: + assert response.headers["Content-Type"] == "text/event-stream" + assert response.readline() == b"data: first\n" + assert response.readline() == b"\n" + assert not state["release_stream"].is_set() + state["release_stream"].set() + assert response.read() == b"data: last\n\n" + + +def test_real_haproxy_never_replays_a_post_and_keeps_pd_websockets_direct(haproxy_proxy): + base_url, opener, state = haproxy_proxy + body = b'{"prompt":"submit once"}' + with pytest.raises(HTTPError) as response: + opener.open(Request(base_url + "/generate", data=body, method="POST"), timeout=3) + assert response.value.code == 502 + response.value.close() + assert [request_body for _worker_id, request_body in state["requests"]] == [body] + + websocket_headers = { + "Connection": "Upgrade", + "Upgrade": "websocket", + "Sec-WebSocket-Version": "13", + "Sec-WebSocket-Key": "dGhlIHNhbXBsZSBub25jZQ==", + } + for path in ("/pd_register", "/kv_move_status"): + for headers in ({}, websocket_headers): + with pytest.raises(HTTPError) as response: + opener.open(Request(base_url + path, headers=headers), timeout=3) + assert response.value.code == 426 + response.value.close() + + +def test_real_haproxy_propagates_client_disconnect_to_the_request_worker(haproxy_proxy): + base_url, opener, state = haproxy_proxy + response = opener.open(base_url + "/stream_until_disconnect", timeout=3) + assert response.readline() == b"data: first\n" + response.close() + assert state["stream_disconnected"].wait(timeout=3), "Backend continued producing after the client disconnected" + + +def test_launcher_sigterm_drains_an_active_sse_response(haproxy_proxy, monkeypatch): + base_url, opener, state = haproxy_proxy + registered_handlers = {} + shutdown_result = [] + process_manager = start_utils.SubmoduleManager() + monkeypatch.setattr( + start_utils.signal, + "signal", + lambda sig, handler: registered_handlers.__setitem__(sig, handler), + ) + monkeypatch.setattr(process_manager, "terminate_all_processes", lambda: None) + process_manager.setup_signal_handlers( + state["haproxy"], + graceful_shutdown_signal=signal.SIGUSR1, + ) + + def stop_launcher(): + try: + registered_handlers[signal.SIGTERM](signal.SIGTERM, None) + except SystemExit as error: + shutdown_result.append(error.code) + + with opener.open(base_url + "/stream", timeout=3) as response: + assert response.readline() == b"data: first\n" + assert response.readline() == b"\n" + + shutdown_thread = threading.Thread(target=stop_launcher) + shutdown_thread.start() + time.sleep(0.1) + assert state["haproxy"].poll() is None, "HAProxy exited before the active stream completed" + assert state["runtime"].is_dir() + + state["release_stream"].set() + assert response.read() == b"data: last\n\n" + shutdown_thread.join(timeout=5) + assert not shutdown_thread.is_alive() + assert shutdown_result == [0] + assert state["haproxy"].returncode == 0 + + +def _probe_response(opener, base_url, path, method="GET"): + request = Request(base_url + path, data=b"" if method == "POST" else None, method=method) + try: + response = opener.open(request, timeout=3) + except HTTPError as error: + response = error + with response: + return response.status, response.read(), response.headers.get_content_type() + + +def _wait_for_probe_response(opener, base_url, path, status, timeout=15): + deadline = time.monotonic() + timeout + while True: + response = _probe_response(opener, base_url, path) + if response[0] == status: + return response + assert time.monotonic() < deadline, response + time.sleep(0.05) + + +@pytest.mark.parametrize("haproxy_proxy", ["block_initial_health_checks"], indirect=True) +def test_real_haproxy_reports_unready_until_every_initial_probe_succeeds(haproxy_proxy): + base_url, opener, state = haproxy_proxy + + # 后端首次探测仍在处理中时，三类健康接口及 healthz 别名必须全部返回 503。 + for path in ("/liveness", "/readiness", "/health", "/healthz"): + assert _probe_response(opener, base_url, path)[0] == 503 + with opener.open(base_url + "/worker", timeout=3) as response: + assert response.status == 200 + + state["release_initial_probes"].set() + for path in ("/liveness", "/readiness", "/health", "/healthz"): + assert _wait_for_probe_response(opener, base_url, path, 200)[0] == 200 + + +@pytest.mark.parametrize( + "path,failure,success_body,error_body", + [ + ("/liveness", 503, {"status": "ok"}, {"status": "not alive"}), + ("/readiness", 302, {"status": "ok"}, {"status": "not ready"}), + ("/health", 500, {"message": "Ok"}, {"message": "Error"}), + ], +) +def test_real_haproxy_requires_every_worker_for_each_probe(haproxy_proxy, path, failure, success_body, error_body): + base_url, opener, state = haproxy_proxy + for initial_path in ("/liveness", "/readiness", "/health"): + _wait_for_probe_response(opener, base_url, initial_path, 200) + assert json.loads(_probe_response(opener, base_url, path)[1]) == success_body + + def change_worker_status(worker, status): + with state["lock"]: + before = len(state["probe_observations"]) + state["probe_status"][worker][path] = status + # Wait for a new probe of this worker, rather than treating HAProxy's + # initial UP state or an old response as evidence of recovery. + deadline = time.monotonic() + 15 + while True: + with state["lock"]: + observed = (worker, path, status) in state["probe_observations"][before:] + if observed: + return + assert time.monotonic() < deadline, (worker, path, status) + time.sleep(0.05) + + change_worker_status(1, failure) + status, body, content_type = _wait_for_probe_response(opener, base_url, path, 503) + assert content_type == "application/json" + assert json.loads(body) == error_body + for other in ("/liveness", "/readiness", "/health"): + if other != path: + assert _probe_response(opener, base_url, other)[0] == 200 + + # Probe failure must not filter ordinary traffic. + for ordinary_path in ("/worker", "/metrics"): + workers = set() + for _ in range(6): + with opener.open(base_url + ordinary_path, timeout=3) as response: + workers.add(json.load(response)["worker_id"]) + assert workers == {0, 1} + aliases = ("/health", "/healthz") if path == "/health" else (path,) + for alias in aliases: + assert _probe_response(opener, base_url, alias + "?detail=1")[0] == 503 + method = "HEAD" if path == "/health" else "POST" + method_status, method_body, _ = _probe_response(opener, base_url, path, method) + assert method_status == 503 + if method == "HEAD": + assert method_body == b"" + else: + assert json.loads(method_body) == error_body + # Unsupported methods still reach the original backend handler. + assert _probe_response(opener, base_url, path, "HEAD")[0] == 501 + + change_worker_status(0, failure) + assert _wait_for_probe_response(opener, base_url, path, 503)[0] == 503 + change_worker_status(1, 200) + assert _wait_for_probe_response(opener, base_url, path, 503)[0] == 503 + change_worker_status(0, 200) + assert json.loads(_wait_for_probe_response(opener, base_url, path, 200)[1]) == success_body + if method == "HEAD": + assert _probe_response(opener, base_url, path, method)[:2] == (200, b"") + else: + assert json.loads(_probe_response(opener, base_url, path, method)[1]) == success_body + + +def test_real_haproxy_allows_slow_health_checks_and_rejects_a_timed_out_worker(haproxy_proxy): + base_url, opener, state = haproxy_proxy + with state["lock"]: + state["probe_status"][1]["/health"] = 503 + _wait_for_probe_response(opener, base_url, "/health", 503) + + # Recover from a known DOWN state so the initial UP default cannot hide a + # check timeout that is too short for the fixed-topology eight-second probe. + with state["lock"]: + state["probe_status"][1]["/health"] = 200 + state["probe_delays"][(1, "/health")] = 8.5 + _wait_for_probe_response(opener, base_url, "/health", 200, timeout=25) + + with state["lock"]: + state["probe_delays"][(1, "/health")] = 13 + _wait_for_probe_response(opener, base_url, "/health", 503, timeout=30) + assert _probe_response(opener, base_url, "/liveness")[0] == 200 + assert _probe_response(opener, base_url, "/readiness")[0] == 200 diff --git a/unit_tests/server/test_pd_master_mode.py b/unit_tests/server/test_pd_master_mode.py new file mode 100644 index 0000000000..d9163366f0 --- /dev/null +++ b/unit_tests/server/test_pd_master_mode.py @@ -0,0 +1,612 @@ +import asyncio +import json +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +from easydict import EasyDict + +from lightllm.server.api_cli import make_argument_parser +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.server.core.objs import StopStrMatchMode, FinishStatus, SamplingParams +from lightllm.server.httpserver_for_pd_master.pd_selector import pd_node_stats_client +from lightllm.server.httpserver_for_pd_master.manager import HttpServerManagerForPDMaster, PDManager +from lightllm.server.httpserver_for_pd_master.pd_selector import prompt_cache_tree_client + + +@pytest.fixture(autouse=True) +def local_pd_master_clients(monkeypatch): + start_args = SimpleNamespace(httpserver_workers=1) + monkeypatch.setattr(pd_node_stats_client, "get_env_start_args", lambda: start_args) + monkeypatch.setattr( + prompt_cache_tree_client, + "get_env_start_args", + lambda: start_args, + ) + + +def test_pd_node_self_request_limit_cli_defaults_to_enabled_and_can_be_disabled(): + parser = make_argument_parser() + + assert parser.parse_args([]).disable_pd_node_self_request_limit is False + assert parser.parse_args(["--disable_pd_node_self_request_limit"]).disable_pd_node_self_request_limit is True + assert StartArgs().disable_pd_node_self_request_limit is False + + +def test_pd_cache_high_priority_cli_defaults_to_enabled_and_can_be_disabled(): + parser = make_argument_parser() + + assert parser.parse_args([]).disable_pd_cache_high_priority is False + assert parser.parse_args(["--disable_pd_cache_high_priority"]).disable_pd_cache_high_priority is True + assert StartArgs().disable_pd_cache_high_priority is False + + +def test_pd_master_request_slot_is_reserved_noop(): + manager = HttpServerManagerForPDMaster.__new__(HttpServerManagerForPDMaster) + + asyncio.run(manager._wait_for_pd_master_request_slot()) + + +def test_auto_set_response_parsers_from_qwen35_model_config(tmp_path): + from lightllm.utils.config_utils import auto_set_response_parsers + + (tmp_path / "config.json").write_text('{"model_type": "qwen3_5"}', encoding="utf-8") + args = StartArgs(model_dir=str(tmp_path)) + + auto_set_response_parsers(args) + + assert args.tool_call_parser == "qwen3_coder" + assert args.reasoning_parser == "qwen3" + + +def test_auto_set_response_parsers_preserves_explicit_values(tmp_path): + from lightllm.utils.config_utils import auto_set_response_parsers + + (tmp_path / "config.json").write_text('{"model_type": "qwen3_5"}', encoding="utf-8") + args = StartArgs( + model_dir=str(tmp_path), + tool_call_parser="llama3", + reasoning_parser="deepseek-r1", + ) + + auto_set_response_parsers(args) + + assert args.tool_call_parser == "llama3" + assert args.reasoning_parser == "deepseek-r1" + + +def test_get_server_info_serializes_runtime_easydict(monkeypatch): + from lightllm.server import api_http + + runtime_args = EasyDict( + model_name="qwen35_0.8b", + tool_call_parser="qwen3_coder", + reasoning_parser="qwen3", + ) + monkeypatch.setattr(api_http.g_objs, "args", runtime_args) + + assert api_http.get_server_info() == dict(runtime_args) + + +@pytest.mark.parametrize( + "runtime,worker_ports", + [(False, None), (True, None), (True, [8100, 8101])], +) +def test_pd_master_server_info_reports_worker_ports_without_changing_start_args(monkeypatch, runtime, worker_ports): + from lightllm.server import api_http + + worker_count = 1 if worker_ports is None else len(worker_ports) + start_args = StartArgs( + run_mode="pd_master", + httpserver_workers=worker_count, + pd_node_id=100, + ) + original = vars(start_args).copy() + args = EasyDict(original) if runtime else start_args + monkeypatch.setattr(api_http.g_objs, "args", args) + if worker_ports is not None: + monkeypatch.setattr( + api_http, + "get_shm_port_args", + lambda: SimpleNamespace(pd_master_multi_ports=worker_ports), + ) + + expected_ports = worker_ports or [start_args.port] + assert api_http.get_server_info() == { + **original, + "pd_master_multi_ports": expected_ports, + } + assert (dict(args) if runtime else vars(args)) == original + + +def test_get_server_info_serializes_start_args_dataclass(monkeypatch): + from lightllm.server import api_http + + start_args = StartArgs(model_name="qwen35_0.8b") + monkeypatch.setattr(api_http.g_objs, "args", start_args) + + server_info = api_http.get_server_info() + assert server_info["model_name"] == "qwen35_0.8b" + assert server_info["run_mode"] == "normal" + + +def test_pd_master_models_endpoint_has_created_timestamp(monkeypatch): + from lightllm.server import api_http + + args = StartArgs(run_mode="pd_master", model_dir="/tmp/test-model", model_name="test-model") + manager = SimpleNamespace(get_real_supported_max_req_total_len=lambda: 1024, metric_client=object()) + monkeypatch.setattr(api_http, "init_tokenizer", lambda *_args, **_kwargs: None) + monkeypatch.setattr(api_http.SamplingParams, "load_generation_cfg", lambda _model_dir: None) + monkeypatch.setattr(api_http.CompletionRequest, "load_generation_cfg", lambda _model_dir: None) + monkeypatch.setattr(api_http.ChatCompletionRequest, "load_generation_cfg", lambda _model_dir: None) + monkeypatch.setattr(api_http, "MetricClient", lambda _port: object()) + monkeypatch.setattr(api_http, "get_shm_port_args", lambda: SimpleNamespace(metric_port=1234)) + monkeypatch.setattr(api_http, "HttpServerManagerForPDMaster", lambda args: manager) + monkeypatch.setattr(api_http, "get_unique_server_name", lambda: "test-server") + monkeypatch.setattr(api_http.setproctitle, "setproctitle", lambda _title: None) + monkeypatch.setattr(api_http.time, "time", lambda: 1234.9) + + global_objs = api_http.G_Objs() + global_objs.set_args(args) + monkeypatch.setattr(api_http, "g_objs", global_objs) + + response = asyncio.run(api_http.get_models(None)) + assert response.data[0].created == 1234 + + +def test_elastic_pd_nodes_are_ready_with_at_least_one_node_of_each_role(): + manager = PDManager(StartArgs(pd_master_mode="elastic")) + assert manager.is_pd_nodes_ready() is False + + manager.prefill_nodes = [object()] + assert manager.is_pd_nodes_ready() is False + + manager.decode_nodes = [object()] + assert manager.is_pd_nodes_ready() is True + + manager.prefill_nodes.append(object()) + manager.decode_nodes.append(object()) + assert manager.is_pd_nodes_ready() is True + + +def test_fixed_pd_nodes_are_ready_only_with_exact_node_counts(): + manager = PDManager(StartArgs(pd_master_mode="2p1d")) + assert manager.is_pd_nodes_ready() is False + + manager.prefill_nodes = [SimpleNamespace(client_ip_port="p1:8000"), SimpleNamespace(client_ip_port="p2:8000")] + manager.decode_nodes = [object()] + assert manager.is_pd_nodes_ready() is True + + manager.decode_nodes.append(object()) + assert manager.is_pd_nodes_ready() is False + + +@pytest.mark.parametrize("mode", ["flex", "xpxd", "1p1"]) +def test_malformed_pd_master_mode_is_not_ready(mode): + manager = PDManager(StartArgs(pd_master_mode=mode)) + assert manager.is_pd_nodes_ready() is False + + +def test_pd_master_readiness_endpoint_uses_pd_node_readiness(monkeypatch): + from lightllm.server import api_http + + args = StartArgs(run_mode="pd_master", pd_master_mode="elastic") + pd_manager = PDManager(args) + monkeypatch.setattr(api_http.g_objs, "args", args) + monkeypatch.setattr(api_http.g_objs, "httpserver_manager", SimpleNamespace(pd_manager=pd_manager)) + + response = api_http.readiness() + assert response.status_code == 503 + assert json.loads(response.body) == {"status": "not ready"} + + pd_manager.prefill_nodes = [object()] + pd_manager.decode_nodes = [object()] + response = api_http.readiness() + assert response.status_code == 200 + assert json.loads(response.body) == {"status": "ok"} + + +def test_pd_manager_checks_all_connected_node_health(monkeypatch): + from lightllm.server.httpserver_for_pd_master import manager as manager_module + + requested_urls = [] + responses = { + "http://10.0.0.1:8000/health": SimpleNamespace(status_code=200), + "http://10.0.0.2:8000/health": SimpleNamespace(status_code=200), + } + + class FakeAsyncClient: + def __init__(self, **kwargs): + assert kwargs == {"timeout": 8, "trust_env": False} + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc_value, traceback): + return False + + async def get(self, url): + requested_urls.append(url) + result = responses[url] + if isinstance(result, BaseException): + raise result + return result + + monkeypatch.setattr(manager_module.httpx, "AsyncClient", FakeAsyncClient) + + manager = PDManager(StartArgs()) + manager.prefill_nodes = [SimpleNamespace(mode="prefill", client_ip_port="10.0.0.1:8000")] + manager.decode_nodes = [SimpleNamespace(mode="decode", client_ip_port="10.0.0.2:8000")] + + assert asyncio.run(manager.check_pd_nodes_health()) is True + assert set(requested_urls) == set(responses) + + responses["http://10.0.0.2:8000/health"] = SimpleNamespace(status_code=503) + assert asyncio.run(manager.check_pd_nodes_health()) is False + + responses["http://10.0.0.2:8000/health"] = RuntimeError("connection failed") + assert asyncio.run(manager.check_pd_nodes_health()) is False + + +def test_pd_manager_without_connected_nodes_is_healthy(): + manager = PDManager(StartArgs()) + assert asyncio.run(manager.check_pd_nodes_health()) is True + + +def test_prefill_registration_preserves_existing_inflight_prompt_chars(): + args = StartArgs() + manager = PDManager(args) + + def register_prefill(node_id, client_ip_port): + manager.register_pd( + { + "node_id": node_id, + "client_ip_port": client_ip_port, + "mode": "prefill", + "start_args": { + "max_req_total_len": args.max_req_total_len, + "max_image_pixels": args.max_image_pixels, + "disable_image_resize": args.disable_image_resize, + }, + }, + websocket=object(), + ) + + register_prefill(1, "10.0.0.1:8000") + manager.prefill_nodes[0].dispatched_prompt_chars = 1234 + manager.prefill_nodes[0].dispatched_req_num = 12 + + register_prefill(2, "10.0.0.2:8000") + + assert [node.dispatched_prompt_chars for node in manager.prefill_nodes] == [1234, 0] + assert [node.dispatched_req_num for node in manager.prefill_nodes] == [12, 0] + + +def test_prefill_reconnection_preserves_other_nodes_inflight_prompt_chars(): + args = StartArgs() + manager = PDManager(args) + + def pd_info(node_id, client_ip_port): + return { + "node_id": node_id, + "client_ip_port": client_ip_port, + "mode": "prefill", + "start_args": { + "max_req_total_len": args.max_req_total_len, + "max_image_pixels": args.max_image_pixels, + "disable_image_resize": args.disable_image_resize, + }, + } + + manager.register_pd(pd_info(1, "10.0.0.1:8000"), websocket=object()) + manager.register_pd(pd_info(2, "10.0.0.2:8000"), websocket=object()) + manager.prefill_nodes[0].dispatched_prompt_chars = 100 + manager.prefill_nodes[1].dispatched_prompt_chars = 200 + manager.prefill_nodes[0].dispatched_req_num = 1 + manager.prefill_nodes[1].dispatched_req_num = 2 + + manager.register_pd(pd_info(3, "10.0.0.2:8000"), websocket=object()) + + assert [node.client_ip_port for node in manager.prefill_nodes] == ["10.0.0.1:8000", "10.0.0.2:8000"] + assert [node.dispatched_prompt_chars for node in manager.prefill_nodes] == [100, 0] + assert [node.dispatched_req_num for node in manager.prefill_nodes] == [1, 0] + + +def test_pd_master_inference_health_matches_normal_node_semantics(monkeypatch): + monkeypatch.setattr("lightllm.server.httpserver_for_pd_master.manager.time.time", lambda: 1000) + manager = SimpleNamespace( + health_timeout=200, + latest_success_infer_time=900, + running_request_count=1, + req_id_to_out_inf={1: object()}, + ) + + assert HttpServerManagerForPDMaster.is_healthy(manager) is True + + manager.latest_success_infer_time = 700 + assert HttpServerManagerForPDMaster.is_healthy(manager) is False + + manager.running_request_count = 0 + manager.req_id_to_out_inf.clear() + assert HttpServerManagerForPDMaster.is_healthy(manager) is True + + +def test_pd_master_restores_request_count_when_preload_fails(): + class FailingMultimodalParams: + async def verify_and_preload(self, request): + raise RuntimeError("preload failed") + + manager = HttpServerManagerForPDMaster.__new__(HttpServerManagerForPDMaster) + manager.args = SimpleNamespace(reasoning_parser=None, enable_rl=False, enable_return_routed_experts=False) + manager.running_request_count = 0 + + params = SamplingParams() + params.init(None) + + async def consume_generate(): + async for _ in manager.generate("prompt", params, FailingMultimodalParams(), None): + pass + + with pytest.raises(RuntimeError, match="preload failed"): + asyncio.run(consume_generate()) + + assert manager.running_request_count == 0 + + +@pytest.mark.parametrize("parser", [None, "qwen3"]) +@pytest.mark.parametrize("enabled", [False, True]) +def test_pd_master_request_count_covers_async_generator_lifecycle(monkeypatch, parser, enabled): + monkeypatch.setenv("LIGHTLLM_STOP_IN_REASONING", "1" if enabled else "0") + manager = HttpServerManagerForPDMaster.__new__(HttpServerManagerForPDMaster) + manager.args = SimpleNamespace(reasoning_parser=None, enable_rl=False, enable_return_routed_experts=False) + manager.running_request_count = 0 + manager.args.reasoning_parser = parser + inner_generator_closed = False + + async def fake_generate(prompt, sampling_params, multimodal_params, request): + nonlocal inner_generator_closed + try: + assert manager.running_request_count == 1 + yield "result" + finally: + inner_generator_closed = True + + manager._generate_choices = fake_generate + + async def consume_one_result_and_close(): + params = SamplingParams() + params.init(None) + results_generator = manager.generate(None, params, None, None) + assert await results_generator.__anext__() == "result" + assert manager.running_request_count == 1 + await results_generator.aclose() + + asyncio.run(consume_one_result_and_close()) + assert manager.running_request_count == 0 + assert inner_generator_closed is True + + +@pytest.mark.parametrize("include", [False, True]) +def test_pd_master_forwards_filtered_choices(include): + manager = HttpServerManagerForPDMaster.__new__(HttpServerManagerForPDMaster) + manager.args = SimpleNamespace(reasoning_parser=None, enable_rl=False, enable_return_routed_experts=False) + manager.running_request_count = 0 + manager.tokenizer = SimpleNamespace(encode=lambda text, **kwargs: [ord(c) for c in text]) + params = SamplingParams() + params.init(manager.tokenizer, stop_sequences="END", n=2, best_of=2, include_stop_str_in_output=include) + + async def fake_generate(*args): + # Prefill's first token and subsequent decode tokens arrive in one stream. + yield 80, "first", {"id": 1}, FinishStatus() + yield 81, "second", {"id": 2}, FinishStatus() + yield 80, "END" if include else "", {"id": 3}, FinishStatus(FinishStatus.FINISHED_STOP) + yield 81, "EN", {"id": 4}, FinishStatus(FinishStatus.FINISHED_LENGTH) + + manager._generate_choices = fake_generate + + async def collect(): + return [result async for result in manager.generate("prompt", params, None, None)] + + results = asyncio.run(collect()) + assert "".join(result[1] for result in results if result[0] == 80) == ("firstEND" if include else "first") + assert "".join(result[1] for result in results if result[0] == 81) == "secondEN" + assert sorted(result[2]["id"] for result in results) == [1, 2, 3, 4] + assert manager.running_request_count == 0 + + +def test_fixed_pd_master_health_endpoint_combines_ready_and_health_status(monkeypatch): + from lightllm.server import api_http + + args = StartArgs(run_mode="pd_master", pd_master_mode="2p1d") + pd_manager = PDManager(args) + pd_node_health_results = iter([True, False]) + pd_node_health_check_count = 0 + + async def check_pd_nodes_health(): + nonlocal pd_node_health_check_count + pd_node_health_check_count += 1 + return next(pd_node_health_results) + + pd_manager.check_pd_nodes_health = check_pd_nodes_health + httpserver_manager = SimpleNamespace(pd_manager=pd_manager, is_healthy=lambda: True) + monkeypatch.setattr(api_http.g_objs, "args", args) + monkeypatch.setattr(api_http.g_objs, "httpserver_manager", httpserver_manager) + + response = asyncio.run(api_http.healthcheck(None)) + assert response.status_code == 503 + assert json.loads(response.body) == { + "message": "Error", + "pd_master_mode": "2p1d", + "inference_healthy": True, + "pd_nodes_ready": False, + "pd_nodes_healthy": False, + "registered_prefill_nodes": 0, + "registered_decode_nodes": 0, + } + assert pd_node_health_check_count == 0 + + pd_manager.prefill_nodes = [SimpleNamespace(client_ip_port="p1:8000"), SimpleNamespace(client_ip_port="p2:8000")] + pd_manager.decode_nodes = [object()] + response = asyncio.run(api_http.healthcheck(None)) + assert response.status_code == 200 + assert json.loads(response.body)["message"] == "Ok" + assert json.loads(response.body)["pd_nodes_ready"] is True + assert pd_node_health_check_count == 1 + + response = asyncio.run(api_http.healthcheck(None)) + assert response.status_code == 503 + assert json.loads(response.body)["pd_nodes_healthy"] is False + assert pd_node_health_check_count == 2 + + httpserver_manager.is_healthy = lambda: False + response = asyncio.run(api_http.healthcheck(None)) + assert response.status_code == 503 + assert json.loads(response.body)["inference_healthy"] is False + assert pd_node_health_check_count == 2 + + +def test_elastic_pd_master_health_does_not_check_individual_pd_nodes(monkeypatch): + from lightllm.server import api_http + + args = StartArgs(run_mode="pd_master", pd_master_mode="elastic") + pd_manager = PDManager(args) + pd_manager.prefill_nodes = [object()] + pd_manager.decode_nodes = [object()] + + async def unexpected_pd_node_health_check(): + raise AssertionError("elastic mode must not check individual PD node health") + + pd_manager.check_pd_nodes_health = unexpected_pd_node_health_check + httpserver_manager = SimpleNamespace(pd_manager=pd_manager, is_healthy=lambda: True) + monkeypatch.setattr(api_http.g_objs, "args", args) + monkeypatch.setattr(api_http.g_objs, "httpserver_manager", httpserver_manager) + + response = asyncio.run(api_http.healthcheck(None)) + response_body = json.loads(response.body) + assert response.status_code == 200 + assert response_body["message"] == "Ok" + assert "pd_nodes_healthy" not in response_body + + +@pytest.mark.parametrize( + "run_mode,start_name,workers", + [ + ("pd_master", "pd_master_start", 1), + ("pd_master", "pd_master_multi_worker_start", 2), + ("prefill", "_launch_subprocesses", 1), + ("decode", "_launch_subprocesses", 1), + ], +) +@pytest.mark.parametrize("feature", ["enable_rl", "enable_return_routed_experts", "enable_prompt_logprobs"]) +def test_pd_rejects_final_metadata_features_before_startup(monkeypatch, run_mode, start_name, workers, feature): + from lightllm.server import api_start + + def unexpected_startup(*args, **kwargs): + pytest.fail("Unsupported PD features must be rejected before process setup") + + monkeypatch.setattr(api_start.mp, "set_start_method", unexpected_startup) + start = getattr(api_start, start_name) + with pytest.raises(ValueError, match="PD mode does not support"): + start(StartArgs(run_mode=run_mode, httpserver_workers=workers, **{feature: True})) + + +@pytest.mark.parametrize("workers", [1, 2, 0]) +def test_pd_master_cli_dispatches_to_the_requested_startup_mode(monkeypatch, workers): + from lightllm.server import api_server, api_start + + single, multiple = Mock(), Mock() + monkeypatch.setattr(api_start, "pd_master_start", single) + monkeypatch.setattr(api_start, "pd_master_multi_worker_start", multiple) + monkeypatch.setattr(api_server, "set_log_node_role", Mock()) + monkeypatch.setattr(api_server.torch.multiprocessing, "set_start_method", Mock()) + args = StartArgs(run_mode="pd_master", httpserver_workers=workers) + + if workers == 0: + with pytest.raises(ValueError, match="httpserver_workers.*at least 1"): + api_server.launch_server(args) + single.assert_not_called() + multiple.assert_not_called() + else: + api_server.launch_server(args) + expected, unexpected = (single, multiple) if workers == 1 else (multiple, single) + expected.assert_called_once_with(args) + unexpected.assert_not_called() + + +def test_multiworker_preflight_failure_precedes_environment_and_process_initialization(monkeypatch): + from lightllm.server import api_start + from lightllm.server.httpserver_for_pd_master import multiprocess + + preflight = Mock(side_effect=ValueError("preflight rejected")) + configure = Mock(side_effect=AssertionError("Environment setup must follow preflight")) + launch = Mock(side_effect=AssertionError("Process startup must follow preflight")) + monkeypatch.setattr(multiprocess, "validate_multiprocess_pd_master_args", preflight) + monkeypatch.setattr(multiprocess, "start_multiprocess_pd_master", launch) + monkeypatch.setattr(api_start, "_set_envs_and_config", configure) + args = StartArgs(run_mode="pd_master", httpserver_workers=2) + + with pytest.raises(ValueError, match="preflight rejected"): + api_start.pd_master_multi_worker_start(args) + + preflight.assert_called_once_with(args) + configure.assert_not_called() + launch.assert_not_called() + + +def test_multiworker_starts_required_services_in_order(monkeypatch): + from lightllm.server import api_start + from lightllm.server.httpserver_for_pd_master import multiprocess + from lightllm.server.httpserver_for_pd_master.pd_selector import pd_node_stats_server + from lightllm.server.httpserver_for_pd_master.pd_selector import prompt_cache_tree_server + + args = StartArgs(run_mode="pd_master", port=8000, httpserver_workers=2) + ports = SimpleNamespace( + close=Mock(), + metric_port=8104, + pd_master_multi_ports=[8100, 8101], + ) + metric_server = Mock(name="metric_server") + request_id_server = Mock(name="request_id_server") + stats_server = Mock(name="stats_server") + tree_server = Mock(name="tree_server") + started_services = [] + manager = Mock() + manager.start_submodule_processes.side_effect = lambda *, start_funcs, start_args: started_services.append( + start_funcs[0] + ) + launch = Mock(return_value="completed") + publish = Mock() + read_start_args = Mock() + validate_ports = Mock() + + monkeypatch.setattr(api_start, "process_manager", manager) + monkeypatch.setattr(api_start, "start_metric_manager", metric_server) + monkeypatch.setattr(api_start, "start_req_id_generator_server", request_id_server) + monkeypatch.setattr(api_start, "_set_envs_and_config", Mock()) + monkeypatch.setattr(api_start, "set_unique_server_name", Mock()) + monkeypatch.setattr(api_start, "auto_set_max_req_total_len", Mock()) + monkeypatch.setattr(api_start, "auto_set_response_parsers", Mock()) + monkeypatch.setattr(api_start, "set_env_start_args", publish) + monkeypatch.setattr(api_start, "get_env_start_args", read_start_args) + monkeypatch.setattr(api_start, "get_shm_port_args", Mock(return_value=ports)) + monkeypatch.setattr(api_start, "validate_ports", validate_ports) + monkeypatch.setattr(multiprocess, "validate_multiprocess_pd_master_args", Mock(return_value="haproxy")) + monkeypatch.setattr(multiprocess, "start_multiprocess_pd_master", launch) + monkeypatch.setattr(pd_node_stats_server, "start_pd_node_stats_server", stats_server) + monkeypatch.setattr(prompt_cache_tree_server, "start_prompt_cache_tree_server", tree_server) + + assert api_start.pd_master_multi_worker_start(args) == "completed" + assert started_services == [metric_server, request_id_server, stats_server, tree_server] + manager.setup_exit_controller.assert_called_once_with() + manager.terminate_all_processes.assert_not_called() + ports.close.assert_not_called() + multiprocess.validate_multiprocess_pd_master_args.assert_called_once_with(args) + publish.assert_called_once_with(args) + read_start_args.cache_clear.assert_called_once_with() + api_start.get_shm_port_args.assert_called_once_with(create=True) + validate_ports.assert_called_once_with([8000, 8104, 8100, 8101]) + launch.assert_called_once_with( + args, + ports=ports, + haproxy_binary="haproxy", + ) diff --git a/unit_tests/server/test_pd_master_multiprocess_e2e.py b/unit_tests/server/test_pd_master_multiprocess_e2e.py new file mode 100644 index 0000000000..cc4037a90b --- /dev/null +++ b/unit_tests/server/test_pd_master_multiprocess_e2e.py @@ -0,0 +1,496 @@ +"""Real PD master service tests with CPU-only P/D protocol peers. + +The functional cases cover worker discovery, HTTP generation, SSE, busy replies, +cancellation, shared IDs and metrics, health aggregation, and private channel +isolation. The opt-in performance case compares the same end-to-end protocol path +with one worker and multiple workers without loading model weights. +""" + +import asyncio +from contextlib import AsyncExitStack +import json +import os +import pickle +import statistics +import time +from types import SimpleNamespace +from urllib.parse import urlsplit + +import httpx +import psutil +import pytest +import websockets +from prometheus_client.parser import text_string_to_metric_families + +from lightllm.server.core.objs import FinishStatus +from lightllm.server.httpserver import pd_loop +from lightllm.server.pd_io_struct import ObjType, PDDecodeNodeInfo, PDUpKVStatus +from unit_tests.server.pd_master_test_utils import running_pd_master + + +class FakePDPeers: + """Drive the actual registration, prompt-ID, KV status, and token channels.""" + + def __init__(self, workers, decode_delay=0.02): + self.workers = workers + self.decode_delay = decode_delay + self.stack = AsyncExitStack() + self.tasks = [] + self.requests = [] + self.prefill_done = {} + self.generations = {} + self.generation_errors = [] + self.aborted = set() + self.blocked_prefills = [] + self.hold_prompt_ids = False + self.hold_decode = False + self.busy_prefill_requests = 0 + self.node_addresses = {"prefill": "prefill.internal:18001", "decode": "decode.internal:18002"} + + async def __aenter__(self): + async with httpx.AsyncClient(trust_env=False, timeout=5) as client: + for endpoint in self.workers: + owner = endpoint["node_id"] + response = await client.get("http://" + endpoint["host_ip_port"] + "/get_server_info") + response.raise_for_status() + start_args = response.json() + kv = await self.stack.enter_async_context( + websockets.connect("ws://" + endpoint["host_ip_port"] + "/kv_move_status", proxy=None) + ) + for mode, node_id in (("prefill", 501), ("decode", 502)): + websocket = await self.stack.enter_async_context( + websockets.connect("ws://" + endpoint["host_ip_port"] + "/pd_register", proxy=None) + ) + await websocket.send( + json.dumps( + { + "node_id": node_id, + "client_ip_port": self.node_addresses[mode], + "mode": mode, + "start_args": dict(start_args, run_mode=mode), + "connection_id": f"{mode}-{owner}", + } + ) + ) + self.tasks.append(asyncio.create_task(self._run(mode, owner, websocket, kv))) + return self + + async def __aexit__(self, *exc): + tasks = [*self.tasks, *self.generations.values()] + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + await self.stack.aclose() + + def _event(self, request_id): + return self.prefill_done.setdefault(request_id, asyncio.Event()) + + async def _token(self, websocket, mode, owner, request_id, index, text, finish): + metadata = { + "id": 2 if index == 1 else 3, + "logprob": -0.1, + "logprobs": None, + "prompt_tokens": 3, + "count_output_tokens": index, + "node_mode": mode, + "prompt_cache_len": 0, + } + await websocket.send( + pickle.dumps( + ( + ObjType.TOKEN_PACKS, + [(request_id, text, metadata, FinishStatus(finish))], + { + "connection_key": self.node_addresses[mode] + "/dpNone", + "total_token_usage_rate": 0.1, + }, + ) + ) + ) + + async def _run(self, mode, owner, websocket, kv): + async for data in websocket: + message = pickle.loads(data) + if message[0] == ObjType.ABORT: + request_id = message[1] + self.aborted.add((mode, owner, request_id)) + task = self.generations.get((owner, request_id)) + if mode == "decode" and task is not None: + task.cancel() + continue + if mode == "prefill": + if message[0] == ObjType.REQ: + _, params, _ = message[1] + assert params.pd_master_node_id.get() == owner + if self.busy_prefill_requests: + self.busy_prefill_requests -= 1 + await websocket.send( + pickle.dumps( + (ObjType.PD_UPLOAD_SERVER_BUSY, params.group_request_id, "test capacity exhausted") + ) + ) + continue + if self.hold_prompt_ids: + self.blocked_prefills.append((owner, params.group_request_id)) + continue + await websocket.send( + pickle.dumps((ObjType.PD_UPLOAD_PREFILL_PROMPT_IDS, params.group_request_id, [2, 3, 2])) + ) + elif message[0] == ObjType.PD_REQ_DECODE_NODE_INFO: + request_id = message[1] + await self._token(websocket, mode, owner, request_id, 1, "hello", FinishStatus.FINISHED_LENGTH) + self._event(request_id).set() + elif message[0] == ObjType.REQ: + task = asyncio.create_task(self._decode(owner, websocket, kv, message[1])) + self.generations[(owner, message[1][1].group_request_id)] = task + task.add_done_callback(self._generation_finished) + + def _generation_finished(self, task): + if not task.cancelled() and task.exception() is not None: + self.generation_errors.append(task.exception()) + + async def _decode(self, owner, websocket, kv, request): + prompt_ids, params, _ = request + request_id = params.group_request_id + assert params.pd_master_node_id.get() == owner + self.requests.append((owner, request_id)) + info = PDDecodeNodeInfo( + decode_node_id=502, + pd_master_node_id=owner, + agent_name="test", + agent_metadata=b"", + num_pages=0, + page_reg_desc=b"", + request_id=request_id, + ready_kv_len=0, + ) + assert len(prompt_ids) == 3 + await kv.send(pickle.dumps(PDUpKVStatus(request_id, owner, pickle.dumps(info)))) + await asyncio.wait_for(self._event(request_id).wait(), timeout=5) + # Decode also reports the transferred first token. P/D WebSocket delivery + # order is independent; the master deduplicates their index-one outputs. + await self._token(websocket, "decode", owner, request_id, 1, "hello", FinishStatus.NO_FINISH) + if self.hold_decode: + await asyncio.Event().wait() + # Keep enough requests in flight to exercise load balancing. + if self.decode_delay: + await asyncio.sleep(self.decode_delay) + await self._token(websocket, "decode", owner, request_id, 2, " world", FinishStatus.FINISHED_STOP) + + +async def _wait_until(predicate, timeout=5): + async def poll(): + while not predicate(): + await asyncio.sleep(0.01) + + await asyncio.wait_for(poll(), timeout=timeout) + + +async def _wait_http_status(client, url, expected_status, timeout=20): + deadline = asyncio.get_running_loop().time() + timeout + while True: + response = await client.get(url) + if response.status_code == expected_status: + return response + assert asyncio.get_running_loop().time() < deadline, response.text + await asyncio.sleep(0.1) + + +@pytest.mark.parametrize("worker_count,strategy", [(1, "cache_aware"), (2, "round_robin"), (4, "cache_aware")]) +def test_real_pd_master_generation_sse_busy_and_cancellation(tmp_path, monkeypatch, worker_count, strategy): + monkeypatch.setenv("LIGHTLLM_PD_NODE_BUSY_RETRY_TIMEOUT_SECONDS", "0") + + async def run(master): + public = urlsplit(master.public_url) + monkeypatch.setattr( + pd_loop, + "get_shm_port_args", + lambda: SimpleNamespace(config_server_port=public.port, pd_master_port=public.port), + ) + discovery_args = SimpleNamespace( + config_server_host=None, + config_server_port=None, + pd_master_ip=public.hostname, + ) + discovered = await pd_loop._get_pd_master_objs(discovery_args) + assert discovered is not None + assert {(node.node_id, node.host_ip_port) for node in discovered.values()} == { + (worker["node_id"], worker["host_ip_port"]) for worker in master.workers + } + async with FakePDPeers(master.workers) as peers: + async with httpx.AsyncClient(trust_env=False, timeout=15) as client: + deadline = asyncio.get_running_loop().time() + 10 + while True: + response = await client.get(master.public_url + "/readiness") + if response.status_code == 200: + break + assert asyncio.get_running_loop().time() < deadline, response.text + await asyncio.sleep(0.1) + + async def generate(): + response = await client.post( + master.public_url + "/generate", + json={ + "inputs": "hello world", + "parameters": {"max_new_tokens": 2, "n": 2, "best_of": 2}, + }, + ) + assert response.status_code == 200, response.text + assert response.json()["generated_text"] == ["hello world", "hello world"] + + await asyncio.gather(*(generate() for _ in range(32))) + assert len(peers.requests) == 64 + assert {owner for owner, _ in peers.requests} == {worker["node_id"] for worker in master.workers} + ids = [request_id for _, request_id in peers.requests] + assert len(set(ids)) == len(ids) + assert all(request_id % 8 == 0 for request_id in ids) + + # Scraping through each worker shares its manager's metrics client + # and exposes the same service-wide request counter. + for worker in master.workers: + deadline = asyncio.get_running_loop().time() + 5 + while True: + response = await client.get("http://" + worker["host_ip_port"] + "/metrics") + assert response.status_code == 200 + counts = [ + sample.value + for family in text_string_to_metric_families(response.text) + for sample in family.samples + if sample.name == "lightllm_request_count_total" + ] + if counts == [32]: + break + assert asyncio.get_running_loop().time() < deadline, response.text + await asyncio.sleep(0.05) + + async with client.stream( + "POST", + master.public_url + "/generate_stream", + json={ + "inputs": "hello world", + "parameters": {"max_new_tokens": 2}, + }, + ) as response: + assert response.status_code == 200 + events = [json.loads(line[5:]) async for line in response.aiter_lines() if line.startswith("data:")] + assert [event["token"]["text"] for event in events] == ["hello", " world"] + assert [event["finished"] for event in events] == [False, True] + + # A rejected Prefill must release its reservation and preserve the + # public busy envelope, without killing a control connection. + peers.busy_prefill_requests = 1 + response = await client.post( + master.public_url + "/generate", + json={ + "inputs": "hello world", + "parameters": {"max_new_tokens": 2}, + }, + ) + assert response.status_code == 429, response.text + assert response.json()["error"]["type"] == "RateLimitError" + + # Closing SSE after the first token must reach both P/D actors. + peers.hold_decode = True + async with client.stream( + "POST", + master.public_url + "/generate_stream", + json={ + "inputs": "hello world", + "parameters": {"max_new_tokens": 2}, + }, + ) as response: + assert response.status_code == 200 + async for line in response.aiter_lines(): + if line.startswith("data:"): + assert json.loads(line[5:])["token"]["text"] == "hello" + break + owner, request_id = peers.requests[-1] + await _wait_until( + lambda: all((mode, owner, request_id) in peers.aborted for mode in ("prefill", "decode")) + ) + peers.hold_decode = False + + # Cancellation before prompt IDs must clean up a still-held + # Prefill reservation, rather than only the output iterator. + peers.hold_prompt_ids = True + pending = asyncio.create_task( + client.post( + master.public_url + "/generate", + json={ + "inputs": "hello world", + "parameters": {"max_new_tokens": 2}, + }, + ) + ) + await _wait_until(lambda: bool(peers.blocked_prefills)) + owner, request_id = peers.blocked_prefills[-1] + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await pending + await _wait_until( + lambda: all((mode, owner, request_id) in peers.aborted for mode in ("prefill", "decode")) + ) + peers.hold_prompt_ids = False + + await generate() + assert not peers.generation_errors + for task in peers.tasks: + assert not task.done(), task.exception() + + with running_pd_master( + tmp_path, extra_args=["--httpserver_workers", str(worker_count), "--select_p_d_node_strategy", strategy] + ) as master: + assert len(master.workers) == worker_count + try: + asyncio.run(run(master)) + except BaseException: + print(master.log_path.read_text()[-12000:]) + raise + + +def test_real_health_aggregation_and_private_channel_isolation(tmp_path): + """Public health succeeds only while every real worker has usable P/D peers.""" + + async def run(master): + async with httpx.AsyncClient(trust_env=False, timeout=15) as client: + assert (await client.get(master.public_url + "/liveness")).status_code == 200 + assert (await client.get(master.public_url + "/pd_register")).status_code == 426 + await _wait_http_status(client, master.public_url + "/readiness", 503) + await _wait_http_status(client, master.public_url + "/health", 503) + + async with FakePDPeers(master.workers): + readiness = await _wait_http_status(client, master.public_url + "/readiness", 200) + health = await _wait_http_status(client, master.public_url + "/health", 200) + assert readiness.json() == {"status": "ok"} + assert health.json() == {"message": "Ok"} + + readiness = await _wait_http_status(client, master.public_url + "/readiness", 503) + health = await _wait_http_status(client, master.public_url + "/health", 503) + assert readiness.json() == {"status": "not ready"} + assert health.json() == {"message": "Error"} + + with running_pd_master(tmp_path, extra_args=["--httpserver_workers", "2"]) as master: + try: + asyncio.run(run(master)) + except BaseException: + print(master.log_path.read_text()[-12000:]) + raise + + +def _service_cpu_seconds(master): + processes = [psutil.Process(master.process.pid), *master.descendants()] + cpu_seconds = 0.0 + for process in {process.pid: process for process in processes}.values(): + try: + cpu_time = process.cpu_times() + except (psutil.NoSuchProcess, psutil.AccessDenied): + continue + cpu_seconds += cpu_time.user + cpu_time.system + return cpu_seconds + + +async def _run_high_load(master, request_count, concurrency): + prompt = "hello world " * 48 + payload = { + "inputs": prompt, + "parameters": {"max_new_tokens": 2}, + } + limits = httpx.Limits(max_connections=concurrency, max_keepalive_connections=concurrency) + timeout = httpx.Timeout(60) + + async with FakePDPeers(master.workers, decode_delay=0) as peers, httpx.AsyncClient( + trust_env=False, + timeout=timeout, + limits=limits, + ) as client: + await _wait_http_status(client, master.public_url + "/readiness", 200) + semaphore = asyncio.Semaphore(concurrency) + + async def generate(latencies=None): + async with semaphore: + start = time.perf_counter() + response = await client.post(master.public_url + "/generate", json=payload) + if latencies is not None: + latencies.append(time.perf_counter() - start) + assert response.status_code == 200, response.text + assert response.json()["generated_text"] == ["hello world"] + + # 预热连接池、tokenizer 和请求 ID 客户端，避免一次性初始化污染结果。 + await asyncio.gather(*(generate() for _ in range(concurrency))) + peers.requests.clear() + cpu_start = _service_cpu_seconds(master) + latencies = [] + start = time.perf_counter() + await asyncio.gather(*(generate(latencies) for _ in range(request_count))) + elapsed = time.perf_counter() - start + cpu_seconds = _service_cpu_seconds(master) - cpu_start + + assert len(peers.requests) == request_count + request_ids = [request_id for _, request_id in peers.requests] + assert len(set(request_ids)) == request_count + assert {owner for owner, _ in peers.requests} == {worker["node_id"] for worker in master.workers} + + sorted_latencies = sorted(latencies) + return { + "requests": request_count, + "concurrency": concurrency, + "elapsed_seconds": elapsed, + "requests_per_second": request_count / elapsed, + "p50_ms": statistics.median(sorted_latencies) * 1000, + "p95_ms": sorted_latencies[int((len(sorted_latencies) - 1) * 0.95)] * 1000, + "service_cpu_cores": cpu_seconds / elapsed, + } + + +@pytest.mark.skipif( + os.getenv("LIGHTLLM_RUN_PD_MASTER_PERF") != "1", + reason="set LIGHTLLM_RUN_PD_MASTER_PERF=1 to run the real PD master load benchmark", +) +def test_multiprocess_pd_master_high_load_speedup(tmp_path): + """Compare the complete HTTP -> PD master -> fake P/D protocol path.""" + available_cpus = len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else (os.cpu_count() or 1) + if available_cpus < 2: + pytest.skip("the multiprocess benchmark requires at least two available CPUs") + + request_count = int(os.getenv("LIGHTLLM_PD_MASTER_PERF_REQUESTS", "2000")) + concurrency = int(os.getenv("LIGHTLLM_PD_MASTER_PERF_CONCURRENCY", "128")) + multi_worker_count = int(os.getenv("LIGHTLLM_PD_MASTER_PERF_WORKERS", str(min(4, available_cpus)))) + if not 2 <= multi_worker_count <= available_cpus: + raise ValueError("LIGHTLLM_PD_MASTER_PERF_WORKERS must be between 2 and the available CPU count") + results = {} + + for worker_count in (1, multi_worker_count): + runtime_dir = tmp_path / f"workers-{worker_count}" + with running_pd_master( + runtime_dir, + extra_args=[ + "--httpserver_workers", + str(worker_count), + "--select_p_d_node_strategy", + "round_robin", + ], + ) as master: + try: + results[worker_count] = asyncio.run(_run_high_load(master, request_count, concurrency)) + except BaseException: + print(master.log_path.read_text()[-12000:]) + raise + + single_qps = results[1]["requests_per_second"] + multi_qps = results[multi_worker_count]["requests_per_second"] + speedup = multi_qps / single_qps + print( + "PD master high-load benchmark: " + + json.dumps( + { + "single_worker": results[1], + f"{multi_worker_count}_workers": results[multi_worker_count], + "throughput_speedup": speedup, + }, + sort_keys=True, + ) + ) + + minimum_speedup = float(os.getenv("LIGHTLLM_PD_MASTER_MIN_SPEEDUP", "1.0")) + if minimum_speedup: + assert speedup >= minimum_speedup diff --git a/unit_tests/server/test_pd_master_multiprocess_launcher.py b/unit_tests/server/test_pd_master_multiprocess_launcher.py new file mode 100644 index 0000000000..18282463c7 --- /dev/null +++ b/unit_tests/server/test_pd_master_multiprocess_launcher.py @@ -0,0 +1,399 @@ +import argparse +import json +from pathlib import Path +import signal +import sys +import uuid +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.server.httpserver_for_pd_master import multiprocess +from lightllm.utils.envs_utils import get_unique_server_name +from lightllm.utils.net_utils import get_pd_master_unix_socket_address + + +def _args(**kwargs): + options = dict( + run_mode="pd_master", + host="10.0.0.1", + port=8000, + httpserver_workers=2, + ) + options.update(kwargs) + return StartArgs(**options) + + +def _worker_ports(): + return [8100, 8101] + + +def _ports(**kwargs): + values = dict( + port=8000, + metric_port=8104, + pd_master_multi_ports=[8100, 8101], + ) + values.update(kwargs) + return SimpleNamespace(**values) + + +def test_workers_use_the_shared_service_ipc_address(monkeypatch): + monkeypatch.setenv("LIGHTLLM_UNIQUE_SERVICE_NAME_ID", "shared-service") + get_unique_server_name.cache_clear() + launcher_address = get_pd_master_unix_socket_address("pd_node_stats") + + worker_address = get_pd_master_unix_socket_address("pd_node_stats") + + assert worker_address == launcher_address + get_unique_server_name.cache_clear() + + +@pytest.mark.parametrize("exited_worker", [None, 1]) +def test_backend_launch_tracks_processes_and_passes_worker_index(monkeypatch, exited_worker): + children = [Mock(pid=2101), Mock(pid=2102)] + managed = [object(), object()] + manager = SimpleNamespace(processes=[], process_names={}) + process_factory = Mock(side_effect=children) + process_context = Mock(Process=process_factory) + monkeypatch.setattr( + multiprocess.multiprocessing, + "get_context", + Mock(return_value=process_context), + ) + monkeypatch.setattr(multiprocess, "process_manager", manager) + + def lookup(pid): + index = [child.pid for child in children].index(pid) + if index == exited_worker: + raise multiprocess.psutil.NoSuchProcess(pid) + return managed[index] + + monkeypatch.setattr(multiprocess.psutil, "Process", lookup) + if exited_worker is None: + multiprocess._start_backends(2) + else: + with pytest.raises(multiprocess.psutil.NoSuchProcess): + multiprocess._start_backends(2) + + registered_count = len(children) if exited_worker is None else exited_worker + assert manager.processes == managed[:registered_count] + assert manager.process_names == { + process: f"pd_master_worker_{index}" for index, process in enumerate(managed[:registered_count]) + } + multiprocess.multiprocessing.get_context.assert_called_once_with("spawn") + for index, call in enumerate(process_factory.call_args_list): + assert call.kwargs == { + "target": multiprocess._run_pd_master_worker, + "args": (index,), + } + children[index].start.assert_called_once_with() + + +def test_worker_bootstrap_uses_its_assigned_port(monkeypatch): + import hypercorn.__main__ + + args = _args(port=8000, httpserver_workers=2) + ports = SimpleNamespace(pd_master_multi_ports=[8100, 8101]) + create_ports = Mock(return_value=ports) + monkeypatch.setenv("LIGHTLLM_PD_MASTER_WORKER_INDEX", "1") + monkeypatch.setattr(multiprocess, "get_env_start_args", lambda: args) + monkeypatch.setattr(multiprocess, "get_shm_port_args", create_ports) + + def start_hypercorn(argv): + assert args.port == 8000 + assert args.pd_node_id == 8101 + assert argv[argv.index("--workers") + 1] == "0" + assert argv[argv.index("--bind") + 1] == "10.0.0.1:8101" + assert argv[-1] == "lightllm.server.api_http:app" + + monkeypatch.setattr(hypercorn.__main__, "main", start_hypercorn) + multiprocess._start_pd_master_worker() + create_ports.assert_called_once_with() + + +def test_worker_port_table_preserves_dynamically_allocated_internal_ports(monkeypatch): + from lightllm.utils.envs_utils import get_env_start_args, get_unique_server_name + from lightllm.utils.shm_port_args import ShmPortArgs + + args = _args(port=8000, pd_master_port=8001, config_server_port=8103) + monkeypatch.setenv("LIGHTLLM_UNIQUE_SERVICE_NAME_ID", "test_pd_worker_ports_" + uuid.uuid4().hex) + monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(vars(args))) + get_unique_server_name.cache_clear() + get_env_start_args.cache_clear() + ports = ShmPortArgs(create=True) + linked = None + try: + allocator = Mock(side_effect=[8120, 8104]) + monkeypatch.setattr(ports, "_alloc_free_port", allocator) + assert ports.router_port == 8120 + assert ports.metric_port == 8104 + assert ports.port == 8000 + assert ports.config_server_port == 8103 + linked = ShmPortArgs() + assert linked.router_port == 8120 + assert linked.metric_port == 8104 + assert linked.port == 8000 + assert linked.config_server_port == 8103 + assert allocator.call_count == 2 + finally: + if linked is not None: + linked.close() + ports.shm.unlink() + ports.close() + get_unique_server_name.cache_clear() + get_env_start_args.cache_clear() + + +def test_haproxy_config_streams_without_request_replay(): + state_path = Path("/tmp/lightllm-test-health.state") + config = multiprocess._build_haproxy_config(_args(), _worker_ports(), state_path) + + assert "balance leastconn" in config + assert "nbthread 2" in config + assert "fd-hard-limit 131072" in config + assert f"server-state-file {state_path}" in config + assert config.count("load-server-state-from-file global") == 3 + assert "maxconn 65536" not in config + assert "server worker_0 10.0.0.1:8100" in config + assert "server worker_1 10.0.0.1:8101" in config + assert "bind 10.0.0.1:8000" in config + assert "no option http-buffer-request" in config + assert "option http-server-close" in config + assert "retries 0" in config + assert "retry-on none" in config + assert "timeout client 3600s" in config + assert "timeout server 3600s" in config + assert "acl private_pd_channel path -i /pd_register /kv_move_status" in config + assert "http-request return status 426 if private_pd_channel" in config + assert "option forwardfor" in config + assert "registered_objects" not in config + assert "allocate_global_unique_multimodal_id_range" not in config + assert "pd_master_config_server" not in config + + +@pytest.mark.parametrize( + ("host", "expected_address"), + [("0.0.0.0", "0.0.0.0:8100"), ("::", "[::]:8100")], +) +def test_haproxy_worker_addresses_follow_the_configured_listen_host(host, expected_address): + config = multiprocess._build_haproxy_config( + _args(host=host), _worker_ports(), Path("/tmp/lightllm-test-health.state") + ) + + assert f"server worker_0 {expected_address}" in config + + +def test_initial_health_state_only_contains_down_health_backends(): + state = multiprocess._build_initial_health_state(_worker_ports()) + state_lines = state.splitlines() + + assert state_lines[0] == "1" + assert "pd_master_workers" not in state + assert len(state_lines) == 2 + 3 * len(_worker_ports()) + for backend_name in ("pd_master_liveness", "pd_master_readiness", "pd_master_health"): + backend_lines = [line for line in state_lines if f" {backend_name} " in line] + assert len(backend_lines) == len(_worker_ports()) + for line in backend_lines: + fields = line.split() + assert fields[4] == "-" # srv_addr: keep the configured server address + assert fields[5] == "0" # srv_op_state: DOWN + assert fields[12] == "0" # srv_check_health: no successful check + assert fields[13] == "6" # srv_check_state: configured and enabled + + +@pytest.mark.parametrize("config", [{"config_server_host": "10.0.0.2"}, {"config_server_port": 9000}]) +def test_multiworker_rejects_config_server_before_haproxy_lookup(monkeypatch, config): + lookup = Mock() + monkeypatch.setitem( + sys.modules, + "ray_haproxy", + SimpleNamespace(get_haproxy_binary=lookup), + ) + + with pytest.raises(ValueError, match="does not yet support"): + multiprocess.validate_multiprocess_pd_master_args(_args(**config)) + lookup.assert_not_called() + + +@pytest.mark.parametrize("port", [None, 0, 65536]) +def test_multiworker_requires_a_valid_public_port_before_haproxy_lookup(monkeypatch, port): + lookup = Mock() + monkeypatch.setitem( + sys.modules, + "ray_haproxy", + SimpleNamespace(get_haproxy_binary=lookup), + ) + + with pytest.raises(ValueError, match="between 1 and 65535"): + multiprocess.validate_multiprocess_pd_master_args(_args(port=port)) + lookup.assert_not_called() + + +def test_cli_rejects_the_removed_worker_ports_option(capsys): + from lightllm.server.api_cli import add_cli_args + + parser = add_cli_args(argparse.ArgumentParser()) + with pytest.raises(SystemExit) as error: + parser.parse_args(["--pd_master_worker_ports", "8100", "8101"]) + assert error.value.code == 2 + assert "unrecognized arguments: --pd_master_worker_ports" in capsys.readouterr().err + + +def test_cli_rejects_the_removed_multiprocess_ports_option(capsys): + from lightllm.server.api_cli import add_cli_args + + parser = add_cli_args(argparse.ArgumentParser()) + with pytest.raises(SystemExit) as error: + parser.parse_args(["--pd_master_multi_ports", "8100", "8101"]) + assert error.value.code == 2 + assert "unrecognized arguments: --pd_master_multi_ports" in capsys.readouterr().err + + +def test_missing_haproxy_package_reports_the_install_command(monkeypatch): + monkeypatch.setitem(sys.modules, "ray_haproxy", None) + with pytest.raises(RuntimeError, match=r"pip install ray-haproxy==2\.8\.25"): + multiprocess.validate_multiprocess_pd_master_args(_args()) + + +def test_haproxy_uses_the_native_package(monkeypatch): + package = SimpleNamespace(get_haproxy_binary=Mock(return_value="/tmp/native-wheel/haproxy")) + monkeypatch.setitem(sys.modules, "ray_haproxy", package) + args = _args() + assert multiprocess.validate_multiprocess_pd_master_args(args) == "/tmp/native-wheel/haproxy" + package.get_haproxy_binary.assert_called_once_with() + + +def test_package_resolution_failure_remains_an_error(monkeypatch): + error = FileNotFoundError("bundled executable unavailable") + package = SimpleNamespace(get_haproxy_binary=Mock(side_effect=error)) + monkeypatch.setitem(sys.modules, "ray_haproxy", package) + with pytest.raises(RuntimeError, match="Cannot locate the packaged HAProxy executable"): + multiprocess.validate_multiprocess_pd_master_args(_args()) + + +def test_worker_ports_are_allocated_as_one_launcher_owned_list(monkeypatch): + from lightllm.utils.envs_utils import get_env_start_args, get_unique_server_name + from lightllm.utils.shm_port_args import ShmPortArgs + + args = _args() + monkeypatch.setenv( + "LIGHTLLM_UNIQUE_SERVICE_NAME_ID", + "test_pd_master_multi_ports_" + uuid.uuid4().hex, + ) + monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(vars(args))) + get_unique_server_name.cache_clear() + get_env_start_args.cache_clear() + ports = ShmPortArgs(create=True) + linked = None + try: + allocator = Mock(side_effect=[18100, 18101]) + monkeypatch.setattr(ports, "_alloc_free_port", allocator) + + assert ports.pd_master_multi_ports == [18100, 18101] + linked = ShmPortArgs() + assert linked.pd_master_multi_ports == [18100, 18101] + assert allocator.call_count == 2 + finally: + if linked is not None: + linked.close() + ports.shm.unlink() + ports.close() + get_unique_server_name.cache_clear() + get_env_start_args.cache_clear() + + +@pytest.mark.parametrize("run_mode", ["pd_master", "prefill", "config_server"]) +def test_config_server_port_is_never_allocated_dynamically(monkeypatch, run_mode): + from lightllm.utils.envs_utils import get_env_start_args, get_unique_server_name + from lightllm.utils.shm_port_args import ShmPortArgs + + args = _args(run_mode=run_mode, config_server_port=None) + monkeypatch.setenv("LIGHTLLM_UNIQUE_SERVICE_NAME_ID", "test_pd_config_port_" + uuid.uuid4().hex) + monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(vars(args))) + get_unique_server_name.cache_clear() + get_env_start_args.cache_clear() + ports = ShmPortArgs(create=True) + try: + with pytest.raises(RuntimeError, match="config_server_port.*None"): + _ = ports.config_server_port + finally: + ports.shm.unlink() + ports.close() + get_unique_server_name.cache_clear() + get_env_start_args.cache_clear() + + +@pytest.mark.parametrize("failure_stage", [None, "backend", "proxy"]) +def test_launcher_uses_process_manager_before_releasing_runtime_and_ports(monkeypatch, tmp_path, failure_stage): + args, ports = _args(), _ports(close=Mock()) + proxy = Mock(pid=3103) + managed_proxy = object() + runtime_paths, events = [], [] + create_runtime = multiprocess.tempfile.TemporaryDirectory + + def runtime_directory(**kwargs): + runtime = create_runtime(dir=tmp_path, **kwargs) + runtime_paths.append(multiprocess.Path(runtime.name)) + return runtime + + def start_backends(worker_count): + assert worker_count == 2 + if failure_stage == "backend": + raise RuntimeError("injected backend startup failure") + + def terminate(): + assert runtime_paths[0].is_dir() + ports.close.assert_not_called() + events.append("stop") + + manager = SimpleNamespace( + processes=[], + process_names={}, + terminate_all_processes=Mock(side_effect=terminate), + setup_signal_handlers=Mock(), + supervise_processes=Mock(), + ) + ports.close.side_effect = lambda: events.append("close ports") + monkeypatch.setattr(multiprocess, "process_manager", manager) + monkeypatch.setattr(multiprocess, "get_unique_server_name", lambda: "test-cleanup-order") + monkeypatch.setattr(multiprocess.tempfile, "TemporaryDirectory", runtime_directory) + monkeypatch.setattr(multiprocess, "_start_backends", start_backends) + config_check = Mock() + monkeypatch.setattr(multiprocess.subprocess, "run", config_check) + monkeypatch.setattr(multiprocess.psutil, "Process", Mock(return_value=managed_proxy)) + launch_proxy = Mock(return_value=proxy) + if failure_stage == "proxy": + launch_proxy.side_effect = OSError("injected proxy startup failure") + monkeypatch.setattr(multiprocess.subprocess, "Popen", launch_proxy) + + if failure_stage is None: + multiprocess.start_multiprocess_pd_master( + args, + ports=ports, + haproxy_binary="/bundled/haproxy", + ) + manager.supervise_processes.assert_called_once_with(proxy) + manager.setup_signal_handlers.assert_called_once_with( + proxy, + graceful_shutdown_signal=signal.SIGUSR1, + ) + assert manager.processes == [managed_proxy] + assert manager.process_names == {managed_proxy: "pd_master_haproxy"} + else: + with pytest.raises((RuntimeError, OSError), match=f"injected {failure_stage} startup failure"): + multiprocess.start_multiprocess_pd_master( + args, + ports=ports, + haproxy_binary="/bundled/haproxy", + ) + manager.supervise_processes.assert_not_called() + + manager.terminate_all_processes.assert_called_once_with() + assert events == ["stop", "close ports"] + assert not runtime_paths[0].exists() + ports.close.assert_called_once_with() + if failure_stage != "backend": + assert config_check.call_args.args[0][0:3] == ["/bundled/haproxy", "-c", "-f"] diff --git a/unit_tests/server/test_pd_master_req_ids.py b/unit_tests/server/test_pd_master_req_ids.py new file mode 100644 index 0000000000..d4f80acb03 --- /dev/null +++ b/unit_tests/server/test_pd_master_req_ids.py @@ -0,0 +1,149 @@ +"""Request IDs are owned by one RPyC service instead of shared memory.""" + +import socket +import threading +import uuid +from types import SimpleNamespace +from unittest.mock import Mock, patch + +from rpyc.utils.server import ThreadedServer + +from lightllm.server.httpserver import req_id_generator_client +from lightllm.server.httpserver.req_id_generator_server import ( + CLIENT_REQUEST_ID_RANGE_SIZE, + CONFIG_SERVER_REQUEST_TIMEOUT_SECONDS, + MAX_BEST_OF, + get_req_id_generator_socket_address, +) +from lightllm.utils.net_utils import get_abstract_unix_socket_address +from lightllm.server.httpserver.req_id_generator_client import ReqIDGenerator +from lightllm.server.httpserver.req_id_generator_server import ( + ConfigServerReqIDGeneratorState, + StandaloneReqIDGeneratorState, + _create_req_id_generator_state, +) + + +def _args(**overrides): + values = { + "run_mode": "normal", + "config_server_host": None, + "config_server_port": None, + } + values.update(overrides) + return SimpleNamespace(**values) + + +def _remote_client(ipc_name): + socket_address = get_abstract_unix_socket_address("req_id_generator", ipc_name) + with patch.object( + req_id_generator_client, + "get_req_id_generator_socket_address", + return_value=socket_address, + ): + return ReqIDGenerator() + + +def test_server_allocates_aligned_request_id_ranges(): + state = StandaloneReqIDGeneratorState() + + first = state.allocate_id_range(4) + second = state.allocate_id_range(2) + + assert first == (8, 40) + assert second == (40, 56) + + +def test_client_generates_locally_until_its_range_is_exhausted(monkeypatch): + connection = Mock() + connection.root.allocate_id_range.side_effect = [(8, 8200), (8200, 16392)] + monkeypatch.setattr(req_id_generator_client, "unix_connect", Mock(return_value=connection)) + client = _remote_client("request-id-client") + + first_range = [client.generate_id() for _ in range(CLIENT_REQUEST_ID_RANGE_SIZE)] + + assert first_range == list(range(8, 8200, MAX_BEST_OF)) + connection.root.allocate_id_range.assert_called_once_with(CLIENT_REQUEST_ID_RANGE_SIZE) + assert client.generate_id() == 8200 + assert connection.root.allocate_id_range.call_count == 2 + + +def test_config_server_range_is_used_and_renewed(monkeypatch): + responses = [ + Mock(status_code=200, json=Mock(return_value={"start_id": 0, "end_id": 800000})), + Mock( + status_code=200, + json=Mock(return_value={"start_id": 800000, "end_id": 1600000}), + ), + ] + get = Mock(side_effect=responses) + monkeypatch.setattr("lightllm.server.httpserver.req_id_generator_server.requests.get", get) + state = ConfigServerReqIDGeneratorState("10.0.0.1", 9000) + + assert state.allocate_id_range(1) == (8, 16) + state.next_id = state.range_end - MAX_BEST_OF + assert state.allocate_id_range(1) == (799992, 800000) + assert state.allocate_id_range(2) == (800000, 800016) + assert [call.args for call in get.call_args_list] == [ + ("http://10.0.0.1:9000/allocate_global_unique_id_range",), + ("http://10.0.0.1:9000/allocate_global_unique_id_range",), + ] + assert [call.kwargs for call in get.call_args_list] == [ + {"timeout": CONFIG_SERVER_REQUEST_TIMEOUT_SECONDS}, + {"timeout": CONFIG_SERVER_REQUEST_TIMEOUT_SECONDS}, + ] + + +def test_request_id_socket_uses_configured_ipc_name(): + with patch( + "lightllm.server.httpserver.req_id_generator_server.get_pd_master_unix_socket_address", + return_value=get_abstract_unix_socket_address("req_id_generator", "launcher-name"), + ): + assert get_req_id_generator_socket_address().endswith(b"launcher-name_req_id_generator") + + +def test_state_type_is_selected_once_during_server_startup(): + assert isinstance(_create_req_id_generator_state(_args()), StandaloneReqIDGeneratorState) + assert isinstance( + _create_req_id_generator_state(_args(config_server_host="127.0.0.1")), + StandaloneReqIDGeneratorState, + ) + assert isinstance( + _create_req_id_generator_state(_args(config_server_port=9000)), + StandaloneReqIDGeneratorState, + ) + assert isinstance( + _create_req_id_generator_state( + _args( + run_mode="normal", + config_server_host="127.0.0.1", + config_server_port=9000, + ) + ), + ConfigServerReqIDGeneratorState, + ) + + +def test_real_rpc_clients_share_one_sequence(): + ipc_name = uuid.uuid4().hex + state = StandaloneReqIDGeneratorState() + server = ThreadedServer( + state, + socket_path=get_abstract_unix_socket_address("req_id_generator", ipc_name), + protocol_config={"allow_pickle": True}, + ) + thread = threading.Thread(target=server.start, daemon=True) + thread.start() + clients = [_remote_client(ipc_name) for _ in range(4)] + try: + ids = [clients[index % len(clients)].generate_id() for index in range(200)] + assert len(set(ids)) == len(ids) + assert all(request_id % MAX_BEST_OF == 0 for request_id in ids) + assert state.next_id == 8 + 4 * CLIENT_REQUEST_ID_RANGE_SIZE * MAX_BEST_OF + assert clients[0].connection._channel.stream.sock.family == socket.AF_UNIX + finally: + for client in clients: + client.connection.close() + server.close() + thread.join(timeout=5) + assert not thread.is_alive() diff --git a/unit_tests/server/test_pd_master_topology_api.py b/unit_tests/server/test_pd_master_topology_api.py new file mode 100644 index 0000000000..de92995ddc --- /dev/null +++ b/unit_tests/server/test_pd_master_topology_api.py @@ -0,0 +1,96 @@ +import asyncio +import json +import pickle +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import pytest +from fastapi.testclient import TestClient + +from lightllm.server import api_http +from lightllm.server.core.objs import StartArgs +from lightllm.server.httpserver_for_pd_master.manager import HttpServerManagerForPDMaster +from lightllm.server.pd_io_struct import PDUpKVStatus + + +def _install(monkeypatch, *, ready=False, healthy=False): + # Exercise the real HTTP/WS handlers without loading a tokenizer/model. + monkeypatch.setattr(api_http.app.router, "on_startup", []) + monkeypatch.setattr(api_http.app.router, "on_shutdown", []) + args = StartArgs( + run_mode="pd_master", + host="10.0.0.1", + port=18090, + pd_node_id=100, + ) + manager = HttpServerManagerForPDMaster.__new__(HttpServerManagerForPDMaster) + manager.__dict__.update( + args=args, + pd_manager=SimpleNamespace( + prefill_nodes=[SimpleNamespace(node_id=199)], + decode_nodes=[SimpleNamespace(node_id=200)], + is_pd_nodes_ready=Mock(return_value=ready), + check_pd_nodes_health=AsyncMock(return_value=True), + ), + is_healthy=Mock(return_value=healthy), + update_req_status=AsyncMock(), + ) + monkeypatch.setattr(api_http, "g_objs", SimpleNamespace(args=args, httpserver_manager=manager)) + return args, manager + + +@pytest.mark.parametrize("ready,healthy", [(False, False), (True, False), (True, True)]) +def test_local_readiness_and_health_preserve_distinct_checks(monkeypatch, ready, healthy): + _, manager = _install(monkeypatch, ready=ready, healthy=healthy) + readiness = api_http.readiness() + assert readiness.status_code == (200 if ready else 503) + assert json.loads(readiness.body) == {"status": "ok" if ready else "not ready"} + response = asyncio.run(api_http.healthcheck(None)) + assert response.status_code == (200 if ready and healthy else 503) + assert json.loads(response.body) == { + "inference_healthy": healthy, + "pd_nodes_ready": ready, + "message": "Ok" if ready and healthy else "Error", + "pd_master_mode": "elastic", + "registered_prefill_nodes": 1, + "registered_decode_nodes": 1, + } + manager.pd_manager.check_pd_nodes_health.assert_not_awaited() + + +async def _wait_until(predicate): + async def poll(): + while not predicate(): + await asyncio.sleep(0) + + await asyncio.wait_for(poll(), timeout=1) + + +def test_local_disconnect_makes_readiness_and_health_unavailable(monkeypatch): + _, manager = _install(monkeypatch, ready=True, healthy=True) + manager.pd_manager.is_pd_nodes_ready.return_value = False + readiness = api_http.readiness() + health = asyncio.run(api_http.healthcheck(None)) + assert readiness.status_code == health.status_code == 503 + assert json.loads(readiness.body) == {"status": "not ready"} + assert json.loads(health.body)["pd_nodes_ready"] is False + assert json.loads(health.body)["message"] == "Error" + + +def test_fixed_topology_health_failure_has_consistent_payload(monkeypatch): + args, manager = _install(monkeypatch, ready=True, healthy=True) + args.pd_master_mode = "1p1d" + manager.pd_manager.check_pd_nodes_health.return_value = False + response = asyncio.run(api_http.healthcheck(None)) + assert response.status_code == 503 + assert json.loads(response.body)["pd_nodes_healthy"] is False + assert json.loads(response.body)["message"] == "Error" + manager.pd_manager.check_pd_nodes_health.assert_awaited_once() + + +def test_kv_status_keeps_existing_wire_protocol(monkeypatch): + _, manager = _install(monkeypatch) + with TestClient(api_http.app) as client: + with client.websocket_connect("/kv_move_status") as websocket: + websocket.send_bytes(pickle.dumps(PDUpKVStatus(8, 100, b"metadata"))) + client.portal.call(_wait_until, lambda: manager.update_req_status.await_count == 1) diff --git a/unit_tests/server/test_pd_node_stats.py b/unit_tests/server/test_pd_node_stats.py new file mode 100644 index 0000000000..d8a536ce4b --- /dev/null +++ b/unit_tests/server/test_pd_node_stats.py @@ -0,0 +1,321 @@ +"""PD node dispatch statistics stay independent from prompt-cache state.""" + +import asyncio +import errno +import os +import signal +import socket +import threading +import uuid +from unittest.mock import Mock, patch + +import pytest +from rpyc.utils.server import ThreadedServer + +from lightllm.server.httpserver_for_pd_master.pd_selector import pd_node_stats_client +from lightllm.server.httpserver_for_pd_master.pd_selector import pd_node_stats_server +from lightllm.server.httpserver_for_pd_master.pd_selector.pd_node_stats_client import ( + LocalPDNodeStats, + RemotePDNodeStats, + create_pd_node_stats_client, +) +from lightllm.server.httpserver_for_pd_master.pd_selector.pd_node_stats_server import PDNodeStatsState +from lightllm.server.pd_io_struct import PD_Client_Obj +from lightllm.utils.net_utils import get_abstract_unix_socket_address + + +def _node(address="p", node_id=1): + return PD_Client_Obj( + node_id=node_id, + client_ip_port=address, + mode="prefill", + start_args={}, + dp_index=0, + ) + + +def _counts(node): + return node.dispatched_prompt_chars, node.dispatched_req_num + + +def _query(state, *keys): + return dict(zip(keys, state.get_dispatched_info(tuple(keys)))) + + +def _remote_client(ipc_name="test-pd-node-stats"): + with patch.object( + pd_node_stats_client, + "get_pd_master_unix_socket_address", + return_value=get_abstract_unix_socket_address("pd_node_stats", ipc_name), + ): + return RemotePDNodeStats() + + +def test_ipc_address_uses_the_launcher_ipc_name(monkeypatch): + monkeypatch.setattr(pd_node_stats_client, "unix_connect", Mock()) + client = _remote_client("deployment-a") + assert client.socket_path == get_abstract_unix_socket_address("pd_node_stats", "deployment-a") + assert client.socket_path != get_abstract_unix_socket_address("pd_node_stats", "deployment-b") + + +def test_local_client_updates_node_objects_without_rpc(monkeypatch): + monkeypatch.setattr( + pd_node_stats_client, + "get_env_start_args", + lambda: Mock(httpserver_workers=1), + ) + client = create_pd_node_stats_client() + assert isinstance(client, LocalPDNodeStats) + node = _node() + + async def run(): + await client.add_dispatched_info(p_node=node, chars=10, req_num=2) + assert _counts(node) == (10, 2) + await client.sub_dispatched_info(p_node=node, chars=3, req_num=1) + assert _counts(node) == (7, 1) + await client.sub_dispatched_info(p_node=node, chars=100, req_num=100) + assert _counts(node) == (0, 0) + assert await client.update_dispatched_info([node]) is None + + asyncio.run(run()) + + +def test_remote_client_sends_only_connection_keys(monkeypatch): + monkeypatch.setattr(pd_node_stats_client, "unix_connect", Mock()) + client = _remote_client() + client._add_dispatched_info = Mock() + client._sub_dispatched_info = Mock() + client._get_dispatched_info = Mock(return_value=((17, 1),)) + node = _node() + + async def run(): + await client.add_dispatched_info(p_node=node, chars=17, req_num=1) + await client.sub_dispatched_info(p_node=node, chars=3, req_num=1) + await client.update_dispatched_info([node]) + + asyncio.run(run()) + client._add_dispatched_info.assert_called_once_with(node.connection_key, 17, 1) + client._sub_dispatched_info.assert_called_once_with(node.connection_key, 3, 1) + client._get_dispatched_info.assert_called_once_with((node.connection_key,)) + assert _counts(node) == (17, 1) + + +def test_state_creates_keys_lazily_and_clamps_subtraction(): + state = PDNodeStatsState() + assert _query(state, "a", "b") == {"a": (0, 0), "b": (0, 0)} + state.add_dispatched_info("a", 10, 2) + state.sub_dispatched_info("a", 100, 1) + assert _query(state, "a")["a"] == (0, 1) + state.sub_dispatched_info("a", 0, 100) + assert _query(state, "a")["a"] == (0, 0) + + +@pytest.mark.parametrize("key", ["", None, 1]) +def test_invalid_key_does_not_change_statistics(key): + state = PDNodeStatsState() + with pytest.raises(ValueError, match="nonempty string"): + state.add_dispatched_info(key, 1, 1) + assert not state.stats + + +def test_idle_and_capacity_eviction_are_lru(monkeypatch): + now = [1000.0] + monkeypatch.setattr(pd_node_stats_server.time, "monotonic", lambda: now[0]) + monkeypatch.setattr(pd_node_stats_server, "MAX_NODE_STATS", 2) + state = PDNodeStatsState() + _query(state, "a", "b") + _query(state, "a") + _query(state, "c") + assert list(state.stats) == ["a", "c"] + now[0] += pd_node_stats_server.NODE_STATS_IDLE_SECONDS + _query(state, "new") + assert list(state.stats) == ["new"] + + +def test_cancelled_update_finishes_before_following_update(monkeypatch): + monkeypatch.setattr(pd_node_stats_client, "unix_connect", Mock()) + state = PDNodeStatsState() + client = _remote_client() + node = _node() + started, finish = threading.Event(), threading.Event() + operations = [] + + def add_dispatched_info(*args): + started.set() + assert finish.wait(5) + state.add_dispatched_info(*args) + operations.append("add_dispatched_info") + + def sub_dispatched_info(*args): + state.sub_dispatched_info(*args) + operations.append("sub_dispatched_info") + + client._add_dispatched_info = add_dispatched_info + client._sub_dispatched_info = sub_dispatched_info + + async def run(): + first = asyncio.create_task(client.add_dispatched_info(p_node=node, chars=10, req_num=1)) + assert await asyncio.to_thread(started.wait, 5) + second = asyncio.create_task(client.sub_dispatched_info(p_node=node, chars=10, req_num=1)) + first.cancel() + await asyncio.sleep(0) + assert not first.done() and not second.done() + finish.set() + with pytest.raises(asyncio.CancelledError): + await first + await second + + asyncio.run(run()) + assert operations == ["add_dispatched_info", "sub_dispatched_info"] + assert _query(state, node.connection_key)[node.connection_key] == (0, 0) + + +def test_cancelled_update_waits_for_the_client_lock_and_still_executes(monkeypatch): + monkeypatch.setattr(pd_node_stats_client, "unix_connect", Mock()) + state = PDNodeStatsState() + client = _remote_client() + client._add_dispatched_info = state.add_dispatched_info + node = _node() + + async def run(): + await client.async_lock.acquire() + task = asyncio.create_task(client.add_dispatched_info(p_node=node, chars=10, req_num=1)) + await asyncio.sleep(0) + task.cancel() + await asyncio.sleep(0) + assert not task.done() + client.async_lock.release() + with pytest.raises(asyncio.CancelledError): + await task + + asyncio.run(run()) + assert _query(state, node.connection_key)[node.connection_key] == (10, 1) + + +def test_rpc_failure_terminates_the_worker(monkeypatch): + monkeypatch.setattr(pd_node_stats_client, "unix_connect", Mock()) + terminate = Mock() + monkeypatch.setattr(pd_node_stats_client.os, "kill", terminate) + client = _remote_client() + client._add_dispatched_info = Mock(side_effect=ConnectionError("connection lost")) + + with pytest.raises(ConnectionError, match="connection lost"): + asyncio.run(client.add_dispatched_info(p_node=_node(), chars=10, req_num=1)) + + terminate.assert_called_once_with(os.getpid(), signal.SIGTERM) + + +def test_rpc_failure_after_cancellation_still_terminates_the_worker(monkeypatch): + monkeypatch.setattr(pd_node_stats_client, "unix_connect", Mock()) + terminate = Mock() + monkeypatch.setattr(pd_node_stats_client.os, "kill", terminate) + client = _remote_client() + client._add_dispatched_info = Mock(side_effect=ConnectionError("connection lost")) + + async def run(): + await client.async_lock.acquire() + task = asyncio.create_task(client.add_dispatched_info(p_node=_node(), chars=10, req_num=1)) + await asyncio.sleep(0) + task.cancel() + await asyncio.sleep(0) + client.async_lock.release() + with pytest.raises(ConnectionError, match="connection lost"): + await task + + asyncio.run(run()) + terminate.assert_called_once_with(os.getpid(), signal.SIGTERM) + + +@pytest.fixture +def rpc_stats_cluster(): + state = PDNodeStatsState() + ipc_name = uuid.uuid4().hex + server = ThreadedServer( + state, + socket_path=get_abstract_unix_socket_address("pd_node_stats", ipc_name), + protocol_config={"allow_pickle": True}, + ) + thread = threading.Thread(target=server.start, daemon=True) + thread.start() + clients = [_remote_client(ipc_name) for _ in range(4)] + try: + yield state, clients + finally: + for client in clients: + connection = client.connection + if connection is not None: + connection.close() + server.close() + thread.join(timeout=5) + assert not thread.is_alive() + + +def test_real_rpc_clients_share_statistics(rpc_stats_cluster): + _, clients = rpc_stats_cluster + node = _node() + node.websocket = threading.Lock() + + async def run(): + await clients[0].add_dispatched_info(p_node=node, chars=17, req_num=1) + await clients[1].update_dispatched_info([node]) + assert _counts(node) == (17, 1) + transport = clients[0].connection._channel.stream.sock + assert transport.family == socket.AF_UNIX + await clients[1].sub_dispatched_info(p_node=node, chars=17, req_num=1) + await clients[0].update_dispatched_info([node]) + assert _counts(node) == (0, 0) + + asyncio.run(run()) + + +def test_abstract_socket_is_released_after_server_close(): + address = get_abstract_unix_socket_address("pd_node_stats", uuid.uuid4().hex) + server = ThreadedServer( + PDNodeStatsState(), + socket_path=address, + protocol_config={"allow_pickle": True}, + ) + try: + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as duplicate: + with pytest.raises(OSError) as conflict: + duplicate.bind(address) + assert conflict.value.errno == errno.EADDRINUSE + finally: + server.close() + ThreadedServer( + PDNodeStatsState(), + socket_path=address, + protocol_config={"allow_pickle": True}, + ).close() + + +def test_initial_connection_failure_is_reported_during_construction(monkeypatch): + attempts = [] + + def connect(*args, **kwargs): + attempts.append(1) + raise ConnectionRefusedError("not listening") + + monkeypatch.setattr(pd_node_stats_client, "unix_connect", connect) + + for _ in range(2): + with pytest.raises(ConnectionRefusedError): + _remote_client() + + assert len(attempts) == 2 + + +def test_real_rpc_combined_updates_are_atomic(rpc_stats_cluster): + state, clients = rpc_stats_cluster + node = _node() + + async def run(): + await asyncio.gather( + *( + clients[index % len(clients)].add_dispatched_info(p_node=node, chars=1, req_num=1) + for index in range(200) + ) + ) + + asyncio.run(run()) + assert _query(state, node.connection_key)[node.connection_key] == (200, 200) diff --git a/unit_tests/server/test_pd_start_args.py b/unit_tests/server/test_pd_start_args.py new file mode 100644 index 0000000000..3109fe11b9 --- /dev/null +++ b/unit_tests/server/test_pd_start_args.py @@ -0,0 +1,311 @@ +import json + +import pytest + +from lightllm.server.api_start import _launch_subprocesses +from lightllm.server.core.objs.start_args_type import StartArgs + + +@pytest.fixture +def startup_args(monkeypatch): + for name in ( + "_set_envs_and_config", + "auto_set_max_req_total_len", + "auto_set_fused_shared_experts", + "set_unique_server_name", + ): + monkeypatch.setattr("lightllm.server.api_start." + name, lambda args: None) + + def validation_finished(args): + raise RuntimeError("startup validation complete") + + monkeypatch.setattr("lightllm.server.api_start.auto_set_response_parsers", validation_finished) + return StartArgs( + model_dir="unused", + max_req_total_len=8192, + eos_id=[2], + disable_vision=True, + disable_audio=True, + disable_shm_warning=True, + ) + + +@pytest.mark.parametrize( + "model_type,block_num,cpu_page_size,expected", + [ + ("deepseek_v4", 8, None, 4096), + ("deepseek_v4", 8, 4096, 4096), + ("deepseek_v4", 10000000, None, 512 * 10000000), + ("deepseek_v4", 10000000, 4096, 512 * 10000000), + ("glm5_next", 8, None, 4096), + ("glm5_next", 8, 256, 4096), + ("llama", 8, None, 256), + ("llama", 8, 1024, 1024), + ], +) +def test_cpu_cache_page_defaults_preserve_model_contracts( + startup_args, monkeypatch, model_type, block_num, cpu_page_size, expected +): + monkeypatch.setattr("lightllm.utils.config_utils.get_model_type", lambda _: model_type) + monkeypatch.setattr( + "lightllm.server.api_start.is_hybrid_att_model", lambda _: model_type in ("deepseek_v4", "glm5_next") + ) + args = startup_args + args.enable_cpu_cache = True + args.linear_att_page_block_num = block_num + args.cpu_cache_token_page_size = cpu_page_size + with pytest.raises(RuntimeError, match="startup validation complete"): + _launch_subprocesses(args) + assert args.cpu_cache_token_page_size == expected + assert args.linear_att_page_block_num == block_num + + +@pytest.mark.parametrize("model_type", ["deepseek_v4", "llama"]) +@pytest.mark.parametrize("enable_cpu_cache", [False, True]) +@pytest.mark.parametrize("disable_chunked_prefill", [False, True]) +def test_model_kv_defaults_preserve_hash_pages_kv_selection_and_explicit_pd_tuning( + startup_args, monkeypatch, model_type, enable_cpu_cache, disable_chunked_prefill +): + monkeypatch.setattr("lightllm.utils.config_utils.get_model_type", lambda _: model_type) + monkeypatch.setattr("lightllm.server.api_start.is_hybrid_att_model", lambda _: model_type == "deepseek_v4") + args = startup_args + args.enable_cpu_cache = enable_cpu_cache + args.disable_chunked_prefill = disable_chunked_prefill + args.pd_kv_page_num, args.pd_kv_page_size = 8, 2048 + if model_type == "deepseek_v4" and disable_chunked_prefill: + with pytest.raises(AssertionError, match="hybrid attention models require chunked prefill"): + _launch_subprocesses(args) + return + with pytest.raises(RuntimeError, match="startup validation complete"): + _launch_subprocesses(args) + assert args.page_size == (256 if model_type == "deepseek_v4" else 1) + assert args.linear_att_hash_page_size == 512 + assert args.llm_kv_type == "None" + assert (args.pd_kv_page_num, args.pd_kv_page_size) == (8, 2048) + assert args.cache_placement_strategy == "adaptive" + + +def test_disabled_cpu_cache_keeps_page_size_unset(startup_args): + with pytest.raises(RuntimeError, match="startup validation complete"): + _launch_subprocesses(startup_args) + assert startup_args.cpu_cache_token_page_size is None + + +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", + ) + assert isinstance(state, GQAFullAttentionState) + + +def test_dsv4_cpu_cache_factory_resolves_packed_layout_without_overriding_kv_type(startup_args, monkeypatch, tmp_path): + 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 + + (tmp_path / "config.json").write_text( + json.dumps( + { + "model_type": "deepseek_v4", + "num_hidden_layers": 2, + "head_dim": 512, + "index_head_dim": 128, + "qk_rope_head_dim": 64, + "sliding_window": 128, + "compress_ratios": [4, 128], + } + ) + ) + args = startup_args + args.model_dir = str(tmp_path) + args.enable_cpu_cache = True + args.linear_att_page_block_num = 8 + args.data_type = "bf16" + with pytest.raises(RuntimeError, match="startup validation complete"): + _launch_subprocesses(args) + monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(vars(args))) + cached_functions = ( + get_env_start_args, + get_added_mtp_kv_layer_num, + get_llm_data_type, + get_llm_model_class, + calcu_cpu_cache_meta, + ) + for function in cached_functions: + function.cache_clear() + try: + 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 args.cpu_cache_token_page_size == args.linear_att_hash_page_size * args.linear_att_page_block_num + 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: + function.cache_clear() + + +@pytest.mark.parametrize( + "page_size,hash_page_size,cpu_page_size", + [ + (1, 256, None), + (256, 512, None), + (256, 256, 4096), + (256, 256, None), + (256, 512, 4096), + (512, 512, None), + (768, 768, None), + (1024, 1024, None), + ], +) +def test_dsv4_page_and_checkpoint_configuration(startup_args, monkeypatch, page_size, hash_page_size, cpu_page_size): + monkeypatch.setattr("lightllm.utils.config_utils.get_model_type", lambda model_dir: "deepseek_v4") + monkeypatch.setattr("lightllm.server.api_start.is_hybrid_att_model", lambda model_dir: True) + + def validation_finished(args): + assert args.page_size == (256 if page_size == 1 else page_size) + assert args.linear_att_hash_page_size == hash_page_size + assert args.cpu_cache_token_page_size == hash_page_size * 8 + assert args.llm_kv_type == "None" + assert (args.pd_kv_page_num, args.pd_kv_page_size) == (16, 1024) + raise RuntimeError("DSV4 page-size validation passed") + + monkeypatch.setattr("lightllm.server.api_start.auto_set_response_parsers", validation_finished) + args = startup_args + args.page_size = page_size + args.linear_att_hash_page_size = hash_page_size + args.linear_att_page_block_num = 8 + args.enable_cpu_cache = True + args.cpu_cache_token_page_size = cpu_page_size + with pytest.raises(RuntimeError, match="DSV4 page-size validation passed"): + _launch_subprocesses(args) + + +@pytest.mark.parametrize( + "page_size,hash_page_size,error", + [ + (256, 128, "--linear_att_hash_page_size must be divisible by --page_size"), + (256, 384, "--linear_att_hash_page_size must be divisible by --page_size"), + (512, 256, "--linear_att_hash_page_size must be divisible by --page_size"), + (0, 512, "--page_size must be >= 1"), + ], +) +def test_dsv4_rejects_unaligned_cache_pages(startup_args, monkeypatch, page_size, hash_page_size, error): + monkeypatch.setattr("lightllm.utils.config_utils.get_model_type", lambda _: "deepseek_v4") + monkeypatch.setattr("lightllm.server.api_start.is_hybrid_att_model", lambda _: True) + args = startup_args + args.page_size = page_size + args.linear_att_hash_page_size = hash_page_size + with pytest.raises(ValueError, match=error): + _launch_subprocesses(args) + + +def test_pd_cli_and_dataclass_keep_main_defaults(): + import argparse + from lightllm.server.api_cli import add_cli_args + + parsed = add_cli_args(argparse.ArgumentParser()).parse_args([]) + args = StartArgs() + assert (parsed.pd_kv_page_num, parsed.pd_kv_page_size) == (16, 1024) + assert (args.pd_kv_page_num, args.pd_kv_page_size) == (16, 1024) + + +def test_pd_master_does_not_apply_model_local_cpu_cache_restrictions(monkeypatch): + from lightllm.server.api_start import pd_master_start + + monkeypatch.setattr("lightllm.server.api_start._set_envs_and_config", lambda args: None) + monkeypatch.setattr("lightllm.server.api_start.set_unique_server_name", lambda args: None) + + def validation_finished(args): + raise RuntimeError("PD master uses no local CPU cache") + + monkeypatch.setattr("lightllm.server.api_start.auto_set_max_req_total_len", validation_finished) + pd_master_args = StartArgs(run_mode="pd_master", model_dir="unused", enable_cpu_cache=True) + with pytest.raises(RuntimeError, match="PD master uses no local CPU cache"): + pd_master_start(pd_master_args) + + +def test_pd_kv_page_size_must_be_divisible_by_model_page_size(monkeypatch): + monkeypatch.setattr("lightllm.server.api_start._set_envs_and_config", lambda args: None) + monkeypatch.setattr("lightllm.server.api_start.auto_set_max_req_total_len", lambda args: None) + monkeypatch.setattr("lightllm.server.api_start.auto_set_fused_shared_experts", lambda args: None) + monkeypatch.setattr("lightllm.server.api_start.set_unique_server_name", lambda args: None) + args = StartArgs( + run_mode="decode", + page_size=4, + pd_kv_page_size=6, + disable_vision=True, + disable_audio=True, + disable_shm_warning=True, + ) + + with pytest.raises(AssertionError, match="--pd_kv_page_size must be divisible by --page_size"): + _launch_subprocesses(args) + + +def test_normal_mode_does_not_validate_pd_kv_page_size(monkeypatch): + monkeypatch.setattr("lightllm.server.api_start._set_envs_and_config", lambda args: None) + monkeypatch.setattr("lightllm.server.api_start.auto_set_max_req_total_len", lambda args: None) + monkeypatch.setattr("lightllm.server.api_start.auto_set_fused_shared_experts", lambda args: None) + monkeypatch.setattr("lightllm.server.api_start.set_unique_server_name", lambda args: None) + monkeypatch.setattr("lightllm.server.api_start.is_hybrid_att_model", lambda model_dir: False) + + def validation_finished(args): + raise RuntimeError("PD page-size validation passed") + + monkeypatch.setattr("lightllm.server.api_start.auto_set_response_parsers", validation_finished) + args = StartArgs( + run_mode="normal", + model_dir="unused", + page_size=4, + pd_kv_page_size=6, + eos_id=[2], + disable_vision=True, + disable_audio=True, + disable_shm_warning=True, + ) + + with pytest.raises(RuntimeError, match="PD page-size validation passed"): + _launch_subprocesses(args) + + +@pytest.mark.parametrize("llm_kv_type", ["fp8kv_sph", "fp8kv_spt"]) +def test_fp8_kv_cache_allows_multi_token_model_pages(monkeypatch, llm_kv_type): + monkeypatch.setattr("lightllm.server.api_start._set_envs_and_config", lambda args: None) + monkeypatch.setattr("lightllm.server.api_start.auto_set_max_req_total_len", lambda args: None) + monkeypatch.setattr("lightllm.server.api_start.auto_set_fused_shared_experts", lambda args: None) + monkeypatch.setattr("lightllm.server.api_start.set_unique_server_name", lambda args: None) + monkeypatch.setattr("lightllm.server.api_start.is_hybrid_att_model", lambda model_dir: False) + + def validation_finished(args): + raise RuntimeError("FP8 page-size validation passed") + + monkeypatch.setattr("lightllm.server.api_start.auto_set_response_parsers", validation_finished) + args = StartArgs( + llm_kv_type=llm_kv_type, + kv_quant_calibration_config_path="unused", + model_dir="unused", + page_size=16, + eos_id=[2], + disable_vision=True, + disable_audio=True, + disable_shm_warning=True, + ) + + with pytest.raises(RuntimeError, match="FP8 page-size validation passed"): + _launch_subprocesses(args) diff --git a/unit_tests/server/test_reasoning_token_usage.py b/unit_tests/server/test_reasoning_token_usage.py new file mode 100644 index 0000000000..9be5611922 --- /dev/null +++ b/unit_tests/server/test_reasoning_token_usage.py @@ -0,0 +1,222 @@ +from unittest.mock import patch + +import pytest + +from lightllm.server.reasoning_parser import ReasoningParser +from lightllm.server.httpserver.stop_seq_match import ReasoningStateMachine + + +def _create_parser(model_type: str, force_reasoning: bool) -> ReasoningParser: + with patch("lightllm.server.reasoning_parser.get_token_id", return_value=99): + return ReasoningParser(model_type, force_reasoning=force_reasoning) + + +def _count_tokens(parser: ReasoningParser, token_ids: list[int]) -> int: + for token_id in token_ids: + parser.update_reasoning_token_count(token_id) + return parser.reasoning_tokens + + +def test_counts_tokens_until_single_token_closing_marker(): + parser = _create_parser("qwen3", force_reasoning=True) + + reasoning_tokens = _count_tokens(parser, [1, 2, 3, 99, 4]) + + assert reasoning_tokens == 3 + + +def test_counts_all_tokens_when_generation_is_truncated_before_closing_marker(): + parser = _create_parser("qwen3", force_reasoning=True) + + reasoning_tokens = _count_tokens(parser, [1, 2]) + + assert reasoning_tokens == 2 + + +def test_does_not_count_when_reasoning_is_disabled(): + parser = _create_parser("qwen3", force_reasoning=False) + + reasoning_tokens = _count_tokens(parser, [1, 2, 99]) + + assert reasoning_tokens == 0 + + +def test_counts_for_always_reasoning_detector(): + parser = _create_parser("deepseek-r1", force_reasoning=False) + + reasoning_tokens = _count_tokens(parser, [1, 2, 99]) + + assert reasoning_tokens == 2 + + +def test_minimax_append_think_output_is_not_counted_as_reasoning(): + parser = _create_parser("minimax-append-think", force_reasoning=True) + + reasoning_tokens = _count_tokens(parser, [1, 2, 99]) + + assert reasoning_tokens == 0 + + +class _StopTokenizer: + def __init__(self, end_ids=(92, 93), start_ids=(90, 91)): + self.start_ids = list(start_ids) + self.end_ids = list(end_ids) + + def encode(self, text, **kwargs): + return {"": self.start_ids, "": self.end_ids}[text] + + +def test_stop_state_tracks_multi_token_delimiters(): + state = ReasoningStateMachine("qwen3", _StopTokenizer(), 0, []) + tokens = [1, 90, 91, 2, 92, 93, 3] + assert [state.handle(token) for token in tokens] == [True, True, False, False, False, False, True] + + +@pytest.mark.parametrize( + "tail_prompt_ids,in_reasoning", + [([90, 91], True), ([1, 90, 91], True), ([90, 91, 1], False), ([90], False), ([], False)], +) +def test_prompt_complete_start_initializes_reasoning(tail_prompt_ids, in_reasoning): + state = ReasoningStateMachine("qwen3", _StopTokenizer(), 0, tail_prompt_ids) + assert state.in_reasoning is in_reasoning + assert state.handle(2) is (not in_reasoning) + + +@pytest.mark.parametrize("model,force", [("qwen3", 1), ("deepseek-r1", -1)]) +def test_prompt_end_marker_preserves_initial_reasoning_hint(model, force): + state = ReasoningStateMachine(model, _StopTokenizer(), force, [90, 91, 1, 92, 93]) + assert state.state is state.State.REASONING + assert state.handle(2) is False + + +@pytest.mark.parametrize( + "force,tail_prompt_ids,in_reasoning", + [(0, [90, 91, 1, 2], False), (1, [92, 93, 1, 2], True), (0, [90, 91, 92, 93, 1, 2], False)], +) +def test_prompt_marker_before_trailing_whitespace_does_not_pretransition(force, tail_prompt_ids, in_reasoning): + state = ReasoningStateMachine("qwen3", _StopTokenizer(), force, tail_prompt_ids) + assert state.in_reasoning is in_reasoning + assert state.handle(3) is (not in_reasoning) + + +@pytest.mark.parametrize("force,marker_ids", [(0, [90, 91]), (1, [92, 93])]) +def test_prompt_phase_hint_is_preserved_when_marker_is_outside_tail(force, marker_ids): + state = ReasoningStateMachine("qwen3", _StopTokenizer(), force, marker_ids + [1] * 20) + assert state.in_reasoning is bool(force) + assert state.handle(2) is (not bool(force)) + + +def test_short_prompt_retains_complete_tail_for_long_delimiters(): + tokenizer = _StopTokenizer(end_ids=[94, 95, 96, 97], start_ids=[90, 91, 92]) + state = ReasoningStateMachine("qwen3", tokenizer, 0, [90, 91]) + assert not state.in_reasoning + assert state.handle(92) is False + assert state.in_reasoning + for token in [94, 95, 96, 97]: + assert state.handle(token) is False + assert not state.in_reasoning + + +@pytest.mark.parametrize("model,force", [("qwen3", 1), ("deepseek-r1", -1)]) +def test_stop_state_starts_in_forced_reasoning(model, force): + state = ReasoningStateMachine(model, _StopTokenizer(), force, []) + assert state.state is state.State.REASONING + assert state.handle(1) is False + assert state.handle(92) is False + assert state.state is state.State.REASONING + assert state.in_reasoning + assert state.handle(93) is False + assert state.handle(2) is True + + +def test_stop_state_carries_a_partial_delimiter_from_pd_prompt(): + state = ReasoningStateMachine("qwen3", _StopTokenizer(), 1, [1, 92]) + assert state.state is state.State.REASONING + assert state.in_reasoning + assert state.handle(93) is False + assert state.in_reasoning is False + assert state.handle(2) is True + + +def test_stop_state_can_reenter_reasoning_after_answer_content(): + state = ReasoningStateMachine("qwen3", _StopTokenizer(), 0, []) + assert state.state is state.State.CONTENT + assert state.handle(90) is True + assert state.state is state.State.CONTENT + assert state.handle(1) is True # 尚未完整命中思维链 start 标记，仍处于正文。 + assert state.state is state.State.CONTENT + for _ in range(2): + assert state.handle(90) is True + assert state.handle(91) is False + assert state.in_reasoning is True + assert state.handle(1) is False + assert state.handle(92) is False + assert state.state is state.State.REASONING + assert state.in_reasoning + assert state.handle(93) is False + assert state.in_reasoning is False + assert state.handle(2) is True + + +@pytest.mark.parametrize( + "token,can_match,next_state", + [(91, False, "REASONING"), (90, True, "CONTENT"), (1, True, "CONTENT")], +) +def test_stop_state_requires_complete_start_from_prompt(token, can_match, next_state): + state = ReasoningStateMachine("qwen3", _StopTokenizer(), 0, [1, 90]) + assert state.state is state.State.CONTENT + assert not state.in_reasoning + assert state.handle(token) is can_match + assert state.state is state.State[next_state] + + +@pytest.mark.parametrize("end_ids,can_match,next_state", [([90, 92], True, "CONTENT"), ([90, 91], False, "REASONING")]) +def test_start_matching_only_recognizes_start_marker(end_ids, can_match, next_state): + state = ReasoningStateMachine("qwen3", _StopTokenizer(end_ids), 0, []) + assert state.handle(90) is True + assert state.state is state.State.CONTENT + assert state.handle(end_ids[-1]) is can_match + assert state.state is state.State[next_state] + + +def test_end_marker_outside_reasoning_is_content(): + state = ReasoningStateMachine("qwen3", _StopTokenizer(), 0, []) + assert state.handle(92) is True + assert state.handle(93) is True + assert state.state is state.State.CONTENT + + +@pytest.mark.parametrize("token,next_state", [(93, "CONTENT"), (92, "REASONING"), (1, "REASONING")]) +def test_stop_state_requires_complete_end_from_prompt(token, next_state): + state = ReasoningStateMachine("qwen3", _StopTokenizer(), 1, [1, 92]) + assert state.state is state.State.REASONING + assert state.handle(token) is False + assert state.state is state.State[next_state] + assert state.in_reasoning is (next_state != "CONTENT") + assert state.handle(2) is (next_state == "CONTENT") + + +def test_start_marker_inside_reasoning_is_reasoning_content(): + state = ReasoningStateMachine("qwen3", _StopTokenizer([90, 92]), 1, []) + assert state.handle(90) is False + assert state.state is state.State.REASONING + assert state.handle(91) is False + assert state.state is state.State.REASONING + assert state.in_reasoning + + +@pytest.mark.parametrize( + "end_ids,tokens,states", + [ + ([90, 91], [90, 91], ["REASONING", "CONTENT"]), + ([91], [90, 91], ["REASONING", "CONTENT"]), + ([91, 92], [90, 91, 92], ["REASONING", "REASONING", "CONTENT"]), + ], +) +def test_reasoning_matches_end_even_when_it_overlaps_start(end_ids, tokens, states): + state = ReasoningStateMachine("qwen3", _StopTokenizer(end_ids), 1, []) + for token, expected_state in zip(tokens, states): + assert state.handle(token) is False + assert state.state is state.State[expected_state] + assert not state.in_reasoning + assert state.handle(1) is True diff --git a/unit_tests/server/test_ssm_state_mode.py b/unit_tests/server/test_ssm_state_mode.py new file mode 100644 index 0000000000..e60e71028d --- /dev/null +++ b/unit_tests/server/test_ssm_state_mode.py @@ -0,0 +1,165 @@ +import json + +import pytest + +from lightllm.server.api_cli import make_argument_parser +from lightllm.server.api_start import _validate_and_normalize_ssm_config +from lightllm.server.core.objs.start_args_type import StartArgs + + +@pytest.mark.parametrize( + "mode,dtype,step,model,capacity,error", + [ + ("native", "bfloat16", 0, "llama", 16, None), + ("native", "bfloat16", 32, "llama", 16, None), + ("replay", "float32", 0, "sensenova_loop", 16, None), + ("replay", "bfloat16", 2, "sensenova_loop", 8, None), + ("replay", "float32", 0, "glm5_next", 8, None), + ("replay", "bfloat16", 2, "glm5_next_text", 8, None), + ("replay", "float32", 0, "qwen3_next", 16, None), + ("replay", "float32", 15, "qwen3_5", 16, None), + ("replay", "bfloat16", 2, "qwen3_5", 16, None), + ("replay", "bfloat16", 2, "qwen3_5", 4, None), + ("replay", "bfloat16", 0, "qwen3_5", 16, None), + ("replay", "float32", 4, "qwen3_5", 4, "capacity"), + ("replay", "float32", 2, "llama", 16, "require a GDN"), + ], +) +def test_ssm_mode_validation(monkeypatch, mode, dtype, step, model, capacity, error): + monkeypatch.setattr("lightllm.utils.config_utils.get_model_type", lambda _: model) + args = make_argument_parser().parse_args( + [ + "--ssm_state_mode", + mode, + "--linear_att_ssm_data_type", + dtype, + "--mtp_step", + str(step), + "--ssm_replay_config", + json.dumps({"cache_len": capacity} if mode == "replay" else {}), + ] + ) + if error: + with pytest.raises(AssertionError, match=error): + _validate_and_normalize_ssm_config(args) + else: + _validate_and_normalize_ssm_config(args) + + +@pytest.mark.parametrize("model", ["qwen3_5", "glm5_next"]) +@pytest.mark.parametrize("mode", ["native", "replay"]) +@pytest.mark.parametrize("step", [15, 16]) +def test_gdn_kda_mtp_width_limit(monkeypatch, model, mode, step): + monkeypatch.setattr("lightllm.utils.config_utils.get_model_type", lambda _: model) + args = StartArgs(ssm_state_mode=mode, mtp_step=step) + if step <= 15: + _validate_and_normalize_ssm_config(args) + else: + with pytest.raises(AssertionError, match="GDN/KDA requires mtp_step <= 15"): + _validate_and_normalize_ssm_config(args) + + +@pytest.mark.parametrize("model", ["qwen3_5", "glm5_next"]) +@pytest.mark.parametrize("compute_mode", ["inline", "precompute"]) +@pytest.mark.parametrize("cache_len", [32, 64]) +def test_replay_large_cache_does_not_bypass_mtp_width_limit(monkeypatch, model, compute_mode, cache_len): + monkeypatch.setattr("lightllm.utils.config_utils.get_model_type", lambda _: model) + args = StartArgs( + ssm_state_mode="replay", + ssm_replay_config={"cache_len": cache_len, "compute_mode": compute_mode}, + linear_att_ssm_data_type="bfloat16", + mtp_step=16, + ) + with pytest.raises(AssertionError, match="GDN/KDA requires mtp_step <= 15"): + _validate_and_normalize_ssm_config(args) + + +def test_ssm_mode_defaults_and_removed_flag(): + parser = make_argument_parser() + assert parser.parse_args([]).ssm_state_mode == StartArgs().ssm_state_mode == "native" + assert parser.parse_args([]).ssm_replay_config == StartArgs().ssm_replay_config == {} + with pytest.raises(SystemExit): + parser.parse_args(["--enable_replayssm"]) + with pytest.raises(SystemExit): + parser.parse_args(["--ssm_state_mode", "auto"]) + + +@pytest.mark.parametrize( + "overrides,error", + [ + ({"ssm_replay_config": {"cache_len": 2}}, "power of two and at least 4"), + ({"ssm_replay_config": {"cache_len": 6}}, "power of two and at least 4"), + ({"mtp_step": -1}, "mtp_step >= 0"), + ({"linear_att_ssm_data_type": "float16"}, "FP32 or BF16"), + ], +) +def test_programmatic_replay_config_rejected_before_model_loading(monkeypatch, overrides, error): + monkeypatch.setattr("lightllm.utils.config_utils.get_model_type", lambda _: "qwen3_5") + args = StartArgs(ssm_state_mode="replay") + for key, value in overrides.items(): + setattr(args, key, value) + with pytest.raises(AssertionError, match=error): + _validate_and_normalize_ssm_config(args) + + +def test_replay_compute_mode_requires_replay(monkeypatch): + monkeypatch.setattr("lightllm.utils.config_utils.get_model_type", lambda _: "qwen3_5") + parser = make_argument_parser() + args = parser.parse_args(["--ssm_state_mode", "replay", "--ssm_replay_config", '{"compute_mode":"precompute"}']) + _validate_and_normalize_ssm_config(args) + + +@pytest.mark.parametrize( + "config,error", + [ + ([], "JSON object"), + (None, "JSON object"), + ({"projection_mode": "inline"}, "Unknown"), + ({"cache_len": True}, "capacity"), + ({"cache_len": 16.0}, "capacity"), + ({"cache_len": "16"}, "capacity"), + ({"cache_len": 128}, "capacity"), + ({"compute_mode": "auto"}, "compute_mode"), + ({"fold_method": "raw"}, "Unknown"), + ], +) +def test_invalid_replay_config(config, error): + args = StartArgs(ssm_state_mode="replay", ssm_replay_config=config) + with pytest.raises(AssertionError, match=error): + _validate_and_normalize_ssm_config(args) + + +@pytest.mark.parametrize("mode", ["native"]) +def test_replay_config_rejects_other_modes(mode): + args = StartArgs(ssm_state_mode=mode, ssm_replay_config={"cache_len": 16}) + with pytest.raises(AssertionError, match="requires ssm_state_mode=replay"): + _validate_and_normalize_ssm_config(args) + + +def test_partial_replay_config_defaults_and_serialization(monkeypatch): + from dataclasses import asdict + + monkeypatch.setattr("lightllm.utils.config_utils.get_model_type", lambda _: "qwen3_5") + args = make_argument_parser().parse_args( + ["--ssm_state_mode", "replay", "--ssm_replay_config", '{"compute_mode":"precompute"}'] + ) + _validate_and_normalize_ssm_config(args) + assert args.ssm_replay_config == {"cache_len": 16, "compute_mode": "precompute"} + original = StartArgs(ssm_state_mode="replay", ssm_replay_config=args.ssm_replay_config) + restored = StartArgs(**json.loads(json.dumps(asdict(original)))) + _validate_and_normalize_ssm_config(restored) + assert restored.ssm_replay_config == args.ssm_replay_config + assert StartArgs().ssm_replay_config == {} + + +@pytest.mark.parametrize( + "option,value", + [ + ("--replayssm_cache_len", "16"), + ("--replayssm_projection_mode", "inline"), + ("--ssm_replay_config", "{invalid}"), + ], +) +def test_removed_flags_and_malformed_json(option, value): + with pytest.raises(SystemExit): + make_argument_parser().parse_args([option, value]) diff --git a/unit_tests/utils/test_config_utils.py b/unit_tests/utils/test_config_utils.py new file mode 100644 index 0000000000..ae3e093ca0 --- /dev/null +++ b/unit_tests/utils/test_config_utils.py @@ -0,0 +1,53 @@ +import pytest + +from lightllm.utils import config_utils + + +@pytest.fixture +def fake_config(monkeypatch): + def install(cfg_json): + monkeypatch.setattr(config_utils, "get_config_json", lambda model_path: cfg_json) + return "dummy/model/dir" + + return install + + +def test_get_keyvalue_reads_top_level_key(fake_config): + model_path = fake_config({"model_type": "llama", "hidden_size": 4096}) + + assert config_utils._get_config_llm_keyvalue(model_path, ["hidden_size"]) == 4096 + + +def test_get_keyvalue_does_not_crash_when_thinker_config_lacks_text_config(fake_config): + cfg = {"model_type": "qwen3_omni_moe", "hidden_size": 4096, "thinker_config": {"model_type": "text"}} + model_path = fake_config(cfg) + + assert config_utils._get_config_llm_keyvalue(model_path, ["hidden_size"]) == 4096 + + +def test_get_keyvalue_keeps_top_level_value_when_thinker_text_config_lacks_key(fake_config): + cfg = { + "model_type": "qwen3_omni_moe", + "hidden_size": 4096, + "thinker_config": {"text_config": {"model_type": "text"}}, + } + model_path = fake_config(cfg) + + assert config_utils._get_config_llm_keyvalue(model_path, ["hidden_size"]) == 4096 + + +def test_get_keyvalue_prefers_thinker_text_config_value_when_present(fake_config): + cfg = { + "model_type": "qwen3_omni_moe", + "hidden_size": 4096, + "thinker_config": {"text_config": {"hidden_size": 2048}}, + } + model_path = fake_config(cfg) + + assert config_utils._get_config_llm_keyvalue(model_path, ["hidden_size"]) == 2048 + + +def test_get_keyvalue_returns_none_when_key_missing_everywhere(fake_config): + model_path = fake_config({"model_type": "llama", "hidden_size": 4096}) + + assert config_utils._get_config_llm_keyvalue(model_path, ["head_dim"]) is None diff --git a/unit_tests/utils/test_envs_utils.py b/unit_tests/utils/test_envs_utils.py new file mode 100644 index 0000000000..18d0072a7a --- /dev/null +++ b/unit_tests/utils/test_envs_utils.py @@ -0,0 +1,68 @@ +from types import SimpleNamespace + +from lightllm.utils import envs_utils +from lightllm.utils.envs_utils import ( + get_deepep_num_max_dispatch_tokens_per_rank, + get_deepep_num_max_dispatch_tokens_per_rank_decode, + get_deepep_num_max_dispatch_tokens_per_rank_prefill, +) + + +def test_deepep_decode_buffer_capacity_covers_mtp_rows_and_safety_margin(monkeypatch): + monkeypatch.delenv("NUM_MAX_DISPATCH_TOKENS_PER_RANK_DECODE", raising=False) + monkeypatch.setattr( + envs_utils, + "get_env_start_args", + lambda: SimpleNamespace(running_max_req_size=256, mtp_step=2), + ) + get_deepep_num_max_dispatch_tokens_per_rank_decode.cache_clear() + + assert get_deepep_num_max_dispatch_tokens_per_rank_decode() == 784 + + get_deepep_num_max_dispatch_tokens_per_rank_decode.cache_clear() + + +def test_deepep_prefill_buffer_capacity_adds_safety_margin(monkeypatch): + monkeypatch.delenv("NUM_MAX_DISPATCH_TOKENS_PER_RANK_PREFILL", raising=False) + monkeypatch.setattr(envs_utils, "get_env_start_args", lambda: SimpleNamespace(batch_max_tokens=6000)) + get_deepep_num_max_dispatch_tokens_per_rank_prefill.cache_clear() + + assert get_deepep_num_max_dispatch_tokens_per_rank_prefill() == 6128 + + get_deepep_num_max_dispatch_tokens_per_rank_prefill.cache_clear() + + +def test_deepep_buffer_capacity_uses_larger_phase_capacity(monkeypatch): + cases = [ + (4096, 256, 4096), + (256, 512, 512), + ] + + for prefill_capacity, decode_capacity, expected_capacity in cases: + monkeypatch.setattr( + envs_utils, + "get_deepep_num_max_dispatch_tokens_per_rank_prefill", + lambda prefill_capacity=prefill_capacity: prefill_capacity, + ) + monkeypatch.setattr( + envs_utils, + "get_deepep_num_max_dispatch_tokens_per_rank_decode", + lambda decode_capacity=decode_capacity: decode_capacity, + ) + get_deepep_num_max_dispatch_tokens_per_rank.cache_clear() + assert get_deepep_num_max_dispatch_tokens_per_rank() == expected_capacity + + get_deepep_num_max_dispatch_tokens_per_rank.cache_clear() + + +def test_deepep_buffer_capacity_keeps_explicit_environment_override(monkeypatch): + monkeypatch.setenv("NUM_MAX_DISPATCH_TOKENS_PER_RANK_PREFILL", "1024") + monkeypatch.setenv("NUM_MAX_DISPATCH_TOKENS_PER_RANK_DECODE", "128") + get_deepep_num_max_dispatch_tokens_per_rank_prefill.cache_clear() + get_deepep_num_max_dispatch_tokens_per_rank_decode.cache_clear() + + assert get_deepep_num_max_dispatch_tokens_per_rank_prefill() == 1024 + assert get_deepep_num_max_dispatch_tokens_per_rank_decode() == 128 + + get_deepep_num_max_dispatch_tokens_per_rank_prefill.cache_clear() + get_deepep_num_max_dispatch_tokens_per_rank_decode.cache_clear() diff --git a/unit_tests/utils/test_service_shm_cleanup.py b/unit_tests/utils/test_service_shm_cleanup.py new file mode 100644 index 0000000000..d28434b995 --- /dev/null +++ b/unit_tests/utils/test_service_shm_cleanup.py @@ -0,0 +1,280 @@ +import ctypes +import json +import os +from pathlib import Path +import signal +import socket +import subprocess +import sys +import time +from types import SimpleNamespace +import uuid + +import psutil +import pytest + +from lightllm.utils import service_shm_cleanup + + +def test_cleanup_service_shm_only_removes_matching_service(monkeypatch, tmp_path): + shm_dir = tmp_path / "shm" + shm_dir.mkdir() + matching_names = ["service_0_req_pool", "service_0_token_load"] + for name in [*matching_names, "service_1_req_pool", "other_service_0_value"]: + (shm_dir / name).touch() + monkeypatch.setattr(service_shm_cleanup, "SHM_DIR", shm_dir) + + service_shm_cleanup.ServiceShmCleanup.cleanup_posix_shm("service_0") + + assert all(not (shm_dir / name).exists() for name in matching_names) + assert (shm_dir / "service_1_req_pool").exists() + assert (shm_dir / "other_service_0_value").exists() + + +def test_cleanup_ipc_files_only_removes_matching_service(monkeypatch, tmp_path): + service_socket = tmp_path / "_service_0_router" + other_socket = tmp_path / "_service_1_router" + sockets = [] + for path in [service_socket, other_socket]: + sock = socket.socket(socket.AF_UNIX) + sock.bind(str(path)) + sockets.append(sock) + service_lock = tmp_path / "service_0_shm_port_args.lock" + other_lock = tmp_path / "service_1_shm_port_args.lock" + service_lock.touch() + other_lock.touch() + monkeypatch.setattr(service_shm_cleanup, "IPC_DIR", tmp_path) + + try: + assert service_shm_cleanup.ServiceShmCleanup.cleanup_ipc_files("service_0") == 2 + assert not service_socket.exists() + assert not service_lock.exists() + assert other_socket.exists() + assert other_lock.exists() + finally: + for sock in sockets: + sock.close() + + +def test_system_v_shm_keys_follow_feature_switches(monkeypatch): + monkeypatch.setenv( + "LIGHTLLM_START_ARGS", + json.dumps( + { + "run_mode": "prefill", + "enable_cpu_cache": True, + "enable_multimodal": False, + "cpu_kv_cache_shm_id": 21, + "multi_modal_cache_shm_id": 22, + } + ), + ) + removed_system_v_keys = [] + monkeypatch.setattr( + service_shm_cleanup.ServiceShmCleanup, + "cleanup_posix_shm", + staticmethod(lambda service_name: 0), + ) + monkeypatch.setattr( + service_shm_cleanup.ServiceShmCleanup, + "cleanup_ipc_files", + staticmethod(lambda service_name: 0), + ) + monkeypatch.setattr( + service_shm_cleanup.ServiceShmCleanup, + "cleanup_system_v_shm", + staticmethod(lambda keys: removed_system_v_keys.extend(keys) or len(keys)), + ) + + service_shm_cleanup.ServiceShmCleanup("current_service_0").cleanup_service_resources() + + assert removed_system_v_keys == [21] + + +def test_non_inference_mode_skips_system_v_shm_cleanup(monkeypatch): + monkeypatch.setenv( + "LIGHTLLM_START_ARGS", + json.dumps( + { + "run_mode": "visual_only", + "enable_cpu_cache": True, + "enable_multimodal": True, + "cpu_kv_cache_shm_id": 21, + "multi_modal_cache_shm_id": 22, + } + ), + ) + system_v_cleanup_calls = [] + monkeypatch.setattr( + service_shm_cleanup.ServiceShmCleanup, + "cleanup_posix_shm", + staticmethod(lambda service_name: 0), + ) + monkeypatch.setattr( + service_shm_cleanup.ServiceShmCleanup, + "cleanup_ipc_files", + staticmethod(lambda service_name: 0), + ) + monkeypatch.setattr( + service_shm_cleanup.ServiceShmCleanup, + "cleanup_system_v_shm", + staticmethod(lambda keys: system_v_cleanup_calls.append(keys) or 0), + ) + + service_shm_cleanup.ServiceShmCleanup("current_service_0").cleanup_service_resources() + + assert system_v_cleanup_calls == [] + + +def test_active_process_rejects_missing_pid_and_zombie(monkeypatch): + monkeypatch.setattr( + service_shm_cleanup.psutil, + "Process", + lambda pid: SimpleNamespace(status=lambda: psutil.STATUS_RUNNING), + ) + assert service_shm_cleanup.is_process_active(123) + monkeypatch.setattr( + service_shm_cleanup.psutil, + "Process", + lambda pid: SimpleNamespace(status=lambda: psutil.STATUS_ZOMBIE), + ) + assert not service_shm_cleanup.is_process_active(123) + + +def test_cleanup_process_checks_parent_every_two_seconds(monkeypatch): + states = iter([object(), object(), None]) + sleeps = [] + cleanup_calls = [] + signal_calls = [] + monkeypatch.setattr(service_shm_cleanup, "is_process_active", lambda pid: next(states)) + monkeypatch.setattr(service_shm_cleanup.time, "sleep", sleeps.append) + monkeypatch.setattr(service_shm_cleanup.signal, "signal", lambda sig, handler: signal_calls.append((sig, handler))) + monkeypatch.setattr( + service_shm_cleanup.ServiceShmCleanup, + "cleanup_service_resources", + lambda self: cleanup_calls.append(self.service_name), + ) + + service_shm_cleanup.run_launcher_shm_cleanup_process("service_0", 123) + + assert sleeps == [2.0, 2.0] + assert cleanup_calls == ["service_0"] + assert signal_calls == [ + (signal.SIGINT, signal.SIG_IGN), + (signal.SIGTERM, signal.SIG_IGN), + (signal.SIGHUP, signal.SIG_IGN), + ] + + +def test_start_cleanup_process_is_independent_session(monkeypatch): + popen_calls = [] + monkeypatch.setattr(service_shm_cleanup.os, "getpid", lambda: 123) + monkeypatch.setattr( + service_shm_cleanup.subprocess, "Popen", lambda *args, **kwargs: popen_calls.append((args, kwargs)) + ) + + assert service_shm_cleanup.start_launcher_shm_cleanup_process("service_0") is None + + command = popen_calls[0][0][0] + assert command[-2:] == ["service_0", "123"] + assert popen_calls[0][1]["start_new_session"] is True + assert "env" not in popen_calls[0][1] + + +LAUNCHER_CODE = r""" +import json, os, psutil, sys, time +from pathlib import Path +from lightllm.utils.service_shm_cleanup import start_launcher_shm_cleanup_process +service, directory, exit_mode = sys.argv[1:] +directory = Path(directory) +start_launcher_shm_cleanup_process(service) +watcher = psutil.Process().children()[0] +(Path('/dev/shm')/(service+'_req_pool')).write_bytes(b'test') +lock = Path('/tmp')/(service+'_shm_port_args.lock') +lock.touch() +(directory/'launcher.json').write_text(json.dumps({'watcher':watcher.pid})) +while not (directory/'exit').exists(): time.sleep(.01) +if exit_mode == 'normal': sys.exit(0) +if exit_mode == 'exception': raise RuntimeError('test launcher failure') +if exit_mode == 'os_exit': os._exit(7) +while True: time.sleep(1) +""" + + +def _wait_until(predicate, timeout=15): + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if predicate(): + return + time.sleep(0.05) + assert predicate(), "Timed out waiting for resource cleanup" + + +def _process_stopped(pid): + try: + return psutil.Process(pid).status() == psutil.STATUS_ZOMBIE + except psutil.NoSuchProcess: + return True + + +def _cleanup_signals_are_ignored(pid): + status = Path(f"/proc/{pid}/status").read_text() + ignored = int(next(line.split()[1] for line in status.splitlines() if line.startswith("SigIgn:")), 16) + return all(ignored & (1 << (sig.value - 1)) for sig in [signal.SIGINT, signal.SIGTERM, signal.SIGHUP]) + + +@pytest.mark.parametrize("exit_mode", ["normal", "exception", "os_exit", "SIGKILL", "SIGINT", "SIGTERM", "SIGHUP"]) +def test_cleanup_process_after_real_launcher_exit(tmp_path, exit_mode): + service = "cleanup_test_" + uuid.uuid4().hex + libc = ctypes.CDLL(None) + while True: + key = uuid.uuid4().int % (2 ** 30) + 1 + shmid = libc.shmget(key, 4096, 0o1000 | 0o2000 | 0o600) + if shmid >= 0: + break + environment = os.environ.copy() + environment["LIGHTLLM_START_ARGS"] = json.dumps( + {"run_mode": "normal", "enable_cpu_cache": True, "cpu_kv_cache_shm_id": key} + ) + process = None + watcher_pid = None + try: + process = subprocess.Popen( + [sys.executable, "-c", LAUNCHER_CODE, service, str(tmp_path), exit_mode], + env=environment, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + start_new_session=True, + ) + _wait_until(lambda: (tmp_path / "launcher.json").exists()) + watcher_pid = json.loads((tmp_path / "launcher.json").read_text())["watcher"] + assert os.getsid(watcher_pid) == watcher_pid + assert os.getsid(watcher_pid) != os.getsid(process.pid) + _wait_until(lambda: _cleanup_signals_are_ignored(watcher_pid)) + assert (Path("/dev/shm") / (service + "_req_pool")).exists() + assert libc.shmget(key, 0, 0) == shmid + if exit_mode == "normal": + os.kill(watcher_pid, signal.SIGTERM) + time.sleep(0.1) + assert not _process_stopped(watcher_pid) + if exit_mode.startswith("SIG"): + if exit_mode == "SIGKILL": + process.kill() + else: + os.killpg(process.pid, getattr(signal, exit_mode)) + else: + (tmp_path / "exit").touch() + process.wait(timeout=10) + _wait_until(lambda: _process_stopped(watcher_pid)) + assert not (Path("/dev/shm") / (service + "_req_pool")).exists() + assert not (Path("/tmp") / (service + "_shm_port_args.lock")).exists() + assert libc.shmget(key, 0, 0) == -1 + finally: + if process is not None and process.poll() is None: + process.kill() + process.wait() + if watcher_pid and not _process_stopped(watcher_pid): + psutil.Process(watcher_pid).kill() + libc.shmctl(shmid, 0, None) + (Path("/dev/shm") / (service + "_req_pool")).unlink(missing_ok=True) + (Path("/tmp") / (service + "_shm_port_args.lock")).unlink(missing_ok=True) diff --git a/unit_tests/utils/test_sgl_utils.py b/unit_tests/utils/test_sgl_utils.py new file mode 100644 index 0000000000..85bb041e68 --- /dev/null +++ b/unit_tests/utils/test_sgl_utils.py @@ -0,0 +1,288 @@ +import weakref +import collections +import inspect +import math +from types import SimpleNamespace + +import pytest +import torch +from frozendict import frozendict + +from lightllm.common.basemodel.infer_struct import InferStateInfo +from lightllm.common.basemodel.attention.fa3 import fp as fa3_module +from lightllm.common.triton_utils import autotuner as autotuner_module +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, AutotuneLevel, Autotuner +from lightllm.utils import sgl_utils +from lightllm.utils.envs_utils import get_decode_attn_autotune_seq_len + + +@pytest.fixture(autouse=True) +def autotune_seq_len_environment(monkeypatch): + monkeypatch.delenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", raising=False) + monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", None) + get_decode_attn_autotune_seq_len.cache_clear() + yield + get_decode_attn_autotune_seq_len.cache_clear() + + +@pytest.mark.parametrize("seq_len, expected_len", [(None, 32768), ("8192", 8192), ("16384", 16384), ("8193", 8704)]) +@pytest.mark.parametrize("level", [AutotuneLevel.ADAPTIVE_AUTOTUNE, AutotuneLevel.FORCE_AUTOTUNE]) +def test_fa3_run_key_uses_configured_length_during_tuning(monkeypatch, seq_len, expected_len, level): + if seq_len is not None: + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", seq_len) + monkeypatch.setattr(sgl_utils, "get_triton_autotune_level", lambda: level) + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + for graph_capacity in [16, 32768]: + page_table = torch.empty(2, graph_capacity, dtype=torch.int32) + assert ( + sgl_utils._flash_attn_kvcache_run_key(page_table, 3, 2) + == 2 * 10_000_000_000_000 + 3 * 10_000_000 + expected_len + ) + + +@pytest.mark.parametrize( + "phase, level", + [ + (None, AutotuneLevel.USE_AUTOTUNE_HIS_CONFIG), + (None, AutotuneLevel.ADAPTIVE_AUTOTUNE), + (None, AutotuneLevel.FORCE_AUTOTUNE), + (AutotuneKernelType.GENERAL, AutotuneLevel.ADAPTIVE_AUTOTUNE), + (AutotuneKernelType.GENERAL, AutotuneLevel.FORCE_AUTOTUNE), + (AutotuneKernelType.DECODE_ATTENTION, AutotuneLevel.USE_AUTOTUNE_HIS_CONFIG), + ], +) +def test_fa3_run_key_uses_actual_length_without_tuning(monkeypatch, phase, level): + # 未调优时不应读取该环境变量，即使它无效也不影响正常配置查找。 + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", "invalid") + monkeypatch.setattr(sgl_utils, "get_triton_autotune_level", lambda: level) + monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", phase) + page_table = torch.empty(2, 32768, dtype=torch.int32) + for actual_len, expected_len in [(1, 512), (512, 512), (513, 1024), (8192, 8192), (16384, 16384)]: + assert ( + sgl_utils._flash_attn_kvcache_run_key(page_table, 3, actual_len) + == 2 * 10_000_000_000_000 + 3 * 10_000_000 + expected_len + ) + + +def test_fa3_runtime_matches_cached_config_by_actual_length(monkeypatch): + kernel = sgl_utils.flash_attn_with_kvcache_autotune + monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: AutotuneLevel.USE_AUTOTUNE_HIS_CONFIG) + monkeypatch.setattr(autotuner_module.dist, "is_initialized", lambda: False) + monkeypatch.setattr(kernel, "fn", lambda **kwargs: kwargs["run_config"]) + monkeypatch.setattr(kernel, "fast_match_configs", collections.defaultdict(dict)) + inputs = dict( + q=torch.empty(2, 4, 8), + k_cache=torch.empty(1, 1, 2, 8), + v_cache=torch.empty(1, 1, 2, 8), + page_table=torch.empty(2, 32768, dtype=torch.int32), + max_seqlen_q=1, + causal=True, + window_size=(-1, -1), + softcap=0.0, + sinks=None, + k_descale=None, + v_descale=None, + ) + static_key = frozendict(kernel._static_key(**inputs)) + base_key = 2 * 10_000_000_000_000 + 10_000_000 + monkeypatch.setattr( + kernel, + "cached_configs", + {static_key: {str(base_key + 8192): {"num_splits": 16}, str(base_key + 16384): {"num_splits": 32}}}, + ) + + assert kernel(**inputs, max_seqlen_k=8192) == {"num_splits": 16} + assert kernel(**inputs, max_seqlen_k=8191) == {"num_splits": 16} + assert kernel(**inputs, max_seqlen_k=15000) == {"num_splits": 32} + assert set(kernel.fast_match_configs[static_key]) == {str(base_key + 8192), str(base_key + 15360)} + + +@pytest.mark.parametrize("page_size, actual_kv_len", [(1, 8192), (16, 8191), (256, 8192)]) +def test_fa3_decode_preserves_actual_length_before_graph_capture(monkeypatch, page_size, actual_kv_len): + graph = SimpleNamespace(can_run=lambda **kwargs: True, graph_max_len_in_batch=32768) + model = SimpleNamespace( + graph=graph, + req_manager=SimpleNamespace(req_to_token_indexs=torch.zeros(2, 8192, dtype=torch.int32)), + is_mtp_draft_model=False, + mtp_manager=SimpleNamespace(get_decode_draft_step=lambda _: 0), + ) + infer_state = InferStateInfo() + infer_state.batch_size = 2 + infer_state.max_kv_seq_len = actual_kv_len + infer_state.microbatch_index = 0 + infer_state.b_req_idx = torch.tensor([0, 1], dtype=torch.int32) + infer_state.b_seq_len = torch.tensor([4096, actual_kv_len], dtype=torch.int32) + infer_state.b1_cu_q_seq_len = torch.tensor([0, 1, 2], dtype=torch.int32) + infer_state.b1_cu_kv_seq_len = torch.tensor([0, 4096, 4096 + actual_kv_len], dtype=torch.int32) + state = fa3_module.Fa3DecodeAttState( + backend=SimpleNamespace( + model=model, + infer_page_size=page_size, + uses_causal_attention=lambda: True, + uses_dynamic_spec_verify_layout=lambda: False, + get_page_table_view=lambda att_batch_size, max_kv_len, microbatch_index: torch.zeros( + att_batch_size, math.ceil(max_kv_len / page_size), dtype=torch.int32 + ), + ), + _infer_state_weak_ref=weakref.ref(infer_state), + ) + copy_calls = [] + monkeypatch.setattr(fa3_module, "page_table_copy", lambda **kwargs: copy_calls.append(kwargs)) + state.init_state() + # 模拟 CudaGraph._capture_decode 对 infer_state 的修改。 + state.infer_state.max_kv_seq_len = graph.graph_max_len_in_batch + calls = [] + + def attention(**kwargs): + calls.append(kwargs) + return kwargs["q"] + + monkeypatch.setattr(fa3_module, "flash_attn_with_kvcache_autotune", attention) + state.decode_att(torch.empty(2, 4, 8), torch.empty(page_size, 2, 8), torch.empty(page_size, 2, 8)) + assert calls[0]["max_seqlen_k"] == actual_kv_len + assert calls[0]["page_table"].shape == (2, 32768 // page_size) + + assert copy_calls[0]["page_table"].shape == (2, math.ceil(actual_kv_len / page_size)) + assert copy_calls[0]["page_size"] == page_size + assert calls[0]["k_cache"].shape == (1, page_size, 2, 8) + assert calls[0]["v_cache"].shape == (1, page_size, 2, 8) + + +@pytest.mark.parametrize("seq_len", ["0", "-1", "invalid"]) +def test_invalid_decode_autotune_seq_len_is_rejected(monkeypatch, seq_len): + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", seq_len) + with pytest.raises(ValueError): + get_decode_attn_autotune_seq_len() + + +@pytest.mark.parametrize("page_size, num_pages", [(1, 64), (1, 3), (256, 3)]) +@pytest.mark.parametrize("query_lengths", [[1, 1], [3, 3], [0, 3]]) +def test_fa3_rebuilds_valid_kv_metadata_without_changing_originals(monkeypatch, page_size, num_pages, query_lengths): + kv_len = 17 + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", str(kv_len)) + max_pages = (kv_len + page_size - 1) // page_size + q = torch.randn(sum(query_lengths), 4, 8) + k = torch.randn(num_pages, page_size, 2, 8) + v = torch.randn_like(k) + page_table = torch.full((2, 8), -1, dtype=torch.int32) + seq_lens = torch.full((2,), 2, dtype=torch.int32) + cu_q = torch.tensor([0, query_lengths[0], sum(query_lengths)], dtype=torch.int32) + original = (q, k, v, seq_lens, page_table, cu_q, max(query_lengths), 2) + snapshots = [tensor.clone() for tensor in original[:6]] + options = { + "cu_seqlens_k_new": None, + "causal": True, + "window_size": (16, 0), + "softmax_scale": 0.25, + "return_softmax_lse": False, + } + + args, kwargs = sgl_utils._flash_attn_kvcache_rebuild_inputs(*original, **options) + + assert all(args[i] is original[i] for i in [0, 1, 2, 5]) + assert args[6:] == (max(query_lengths), kv_len) + assert kwargs == options + assert args[4].shape == (2, max_pages) + assert args[4].min() >= 0 and args[4].max() < num_pages + if num_pages >= 2 * max_pages: + assert args[4].unique().numel() == 2 * max_pages + torch.testing.assert_close(args[3], torch.full_like(seq_lens, kv_len)) + for tensor, snapshot in zip(original[:6], snapshots): + torch.testing.assert_close(tensor, snapshot) + + # Batched Q does not require cumulative query or KV lengths. + batched_q = torch.randn(2, 1, 4, 8) + args, kwargs = sgl_utils._flash_attn_kvcache_rebuild_inputs( + q=batched_q, + k_cache=k, + v_cache=v, + page_table=page_table, + cache_seqlens=seq_lens, + cu_seqlens_q=None, + max_seqlen_q=1, + max_seqlen_k=2, + ) + assert args[0] is batched_q + assert args[5:] == (None, 1, kv_len) + + +@pytest.mark.parametrize("kv_len", [8192, 16384]) +@pytest.mark.parametrize("query_lengths", [[1, 1], [3, 3], [0, 3]]) +def test_fa3_autotunes_long_kv_then_captures_original_inputs(tmp_path, monkeypatch, kv_len, query_lengths): + if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 9: + pytest.skip("FA3 requires a Hopper GPU") + if sgl_utils.flash_attn_with_kvcache is None: + pytest.skip("sgl_kernel FA3 is unavailable") + + monkeypatch.setenv("LIGHTLLM_DECODE_ATTN_AUTOTUNE_SEQ_LEN", str(kv_len)) + kernel = sgl_utils.flash_attn_with_kvcache_autotune + monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", None) + monkeypatch.setattr(autotuner_module, "get_triton_autotune_level", lambda: AutotuneLevel.ADAPTIVE_AUTOTUNE) + monkeypatch.setattr(sgl_utils, "get_triton_autotune_level", lambda: AutotuneLevel.ADAPTIVE_AUTOTUNE) + monkeypatch.setattr(autotuner_module.dist, "is_initialized", lambda: False) + monkeypatch.setattr(kernel, "_cache_dir", str(tmp_path), raising=False) + monkeypatch.setattr(kernel, "cached_configs", {}) + monkeypatch.setattr(kernel, "fast_match_configs", collections.defaultdict(dict)) + monkeypatch.setattr(kernel, "warmuped_configs_set", set()) + + q = torch.randn(sum(query_lengths), 4, 64, device="cuda", dtype=torch.bfloat16) + k = torch.randn(2 * kv_len, 1, 2, 64, device="cuda", dtype=q.dtype) + v = torch.randn_like(k) + inputs = dict( + q=q, + k_cache=k, + v_cache=v, + page_table=torch.zeros(2, 32768, device="cuda", dtype=torch.int32), + cache_seqlens=torch.full((2,), 2, device="cuda", dtype=torch.int32), + cu_seqlens_q=torch.tensor([0, query_lengths[0], sum(query_lengths)], device="cuda", dtype=torch.int32), + cu_seqlens_k_new=None, + max_seqlen_q=max(query_lengths), + max_seqlen_k=2, + causal=True, + window_size=(-1, -1), + softcap=0.0, + sinks=None, + k_descale=None, + v_descale=None, + ) + snapshots = {name: value.clone() for name, value in inputs.items() if isinstance(value, torch.Tensor)} + reference = kernel.fn(**inputs) + benchmark = kernel._bench + timings = [] + + def checked_bench(*args, **kwargs): + bound = inspect.signature(kernel.fn).bind(*args, **kwargs).arguments + assert bound["cache_seqlens"].tolist() == [kv_len, kv_len] + assert bound["page_table"].shape == (2, kv_len) + assert bound["cu_seqlens_k_new"] is None + assert bound["max_seqlen_k"] == kv_len + assert bound["q"] is q + elapsed = benchmark(*args, **kwargs) + assert math.isfinite(elapsed), f"FA3 benchmark failed for {kwargs['run_config']}" + timings.append(elapsed) + return elapsed + + monkeypatch.setattr(kernel, "_bench", checked_bench) + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + output = kernel(**inputs) + assert len(timings) == 3 + expected_key = str(2 * 10_000_000_000_000 + max(query_lengths) * 10_000_000 + kv_len) + assert kernel._run_key(**inputs) == 2 * 10_000_000_000_000 + max(query_lengths) * 10_000_000 + 512 + assert all(list(configs) == [expected_key] for configs in kernel.cached_configs.values()) + torch.testing.assert_close(output, reference) + + def unexpected_rebuild(*args, **kwargs): + pytest.fail("Cached execution and CUDA Graph capture must use the original inputs") + + monkeypatch.setattr(kernel, "rebuild_input_func", unexpected_rebuild) + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + kernel(**inputs) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured_output = kernel(**inputs) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(captured_output, reference) + for name, snapshot in snapshots.items(): + torch.testing.assert_close(inputs[name], snapshot) diff --git a/unit_tests/utils/test_shm_utils.py b/unit_tests/utils/test_shm_utils.py new file mode 100644 index 0000000000..00203aba3d --- /dev/null +++ b/unit_tests/utils/test_shm_utils.py @@ -0,0 +1,88 @@ +from concurrent.futures import ThreadPoolExecutor +from multiprocessing import resource_tracker, shared_memory +from threading import Event + +import pytest + +from lightllm.utils import shm_utils + + +def test_get_service_shm_name_requires_service_name(monkeypatch): + monkeypatch.setattr(shm_utils, "get_unique_server_name", lambda: None) + + with pytest.raises(RuntimeError, match="LIGHTLLM_UNIQUE_SERVICE_NAME_ID is unset"): + shm_utils.get_service_shm_name("req_pool") + + +def test_get_service_shm_name_adds_prefix_once(monkeypatch): + monkeypatch.setattr(shm_utils, "get_unique_server_name", lambda: "service_uuid_0") + + assert shm_utils.get_service_shm_name("req_pool") == "service_uuid_0_req_pool" + assert shm_utils.get_service_shm_name("service_uuid_0_req_pool") == "service_uuid_0_req_pool" + + +def test_create_or_link_shm_passes_scoped_name_to_shared_memory_layer(monkeypatch): + created_names = [] + monkeypatch.setattr(shm_utils, "get_unique_server_name", lambda: "service_uuid_1") + monkeypatch.setattr( + shm_utils, + "_force_create_shm", + lambda name, expected_size: created_names.append((name, expected_size)) or object(), + ) + + shm_utils.create_or_link_shm("token_load", 128, force_mode="create") + + assert created_names == [("service_uuid_1_token_load", 128)] + + +@pytest.mark.parametrize("method, tracker_method", [("__init__", "register"), ("unlink", "unregister")]) +def test_service_shm_concurrent_calls_restore_tracker(monkeypatch, method, tracker_method): + original_tracker_method = getattr(resource_tracker, tracker_method) + # 即使回归导致 patch 恢复错误，也不污染后续测试。 + monkeypatch.setattr(resource_tracker, tracker_method, original_tracker_method) + first_entered = Event() + second_started = Event() + second_entered = Event() + release_first = Event() + first_finished = Event() + + def ordered_operation(self, *args, **kwargs): + if self._test_index == 0: + first_entered.set() + assert release_first.wait(5) + else: + second_entered.set() + # 无锁时强制先退出第一个 patch，再退出第二个，复现恢复顺序错误。 + assert first_finished.wait(5) + + monkeypatch.setattr(shared_memory.SharedMemory, method, ordered_operation) + + def worker(index): + # 只测试 patch 生命周期，不申请真实共享内存。 + shm = shm_utils.ServiceSharedMemory.__new__(shm_utils.ServiceSharedMemory) + shm._test_index = index + if index == 1: + second_started.set() + try: + if method == "__init__": + shm.__init__(name="test_tracker") + else: + shm.unlink() + finally: + if index == 0: + first_finished.set() + + with ThreadPoolExecutor(max_workers=2) as executor: + first = executor.submit(worker, 0) + try: + assert first_entered.wait(5) + second = executor.submit(worker, 1) + assert second_started.wait(5) + overlapped = second_entered.wait(0.2) + finally: + release_first.set() + first.result(timeout=5) + second.result(timeout=5) + + assert not overlapped, "Concurrent calls entered the patched operation together" + assert getattr(resource_tracker, tracker_method) is original_tracker_method diff --git a/unit_tests/utils/test_speculative_utils.py b/unit_tests/utils/test_speculative_utils.py new file mode 100644 index 0000000000..a58b35198d --- /dev/null +++ b/unit_tests/utils/test_speculative_utils.py @@ -0,0 +1,413 @@ +import weakref +import json +from importlib import import_module +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.basemodel.infer_struct import InferStateInfo +import lightllm.common.basemodel.attention.base_att as base_att_module +import lightllm.common.basemodel.hidden_collector as hidden_collector_module +from lightllm.common.basemodel.attention.base_att import BaseAttBackend +from lightllm.common.basemodel.attention.fa3.fp import Fa3DecodeAttState, Fa3PrefillAttState +from lightllm.common.basemodel.attention.fa3.mla import MlaFa3DecodeAttState, MlaFa3PrefillAttState +from lightllm.common.basemodel.attention.triton.fp import TritonDecodeAttState +from lightllm.models import get_draft_model_class +from lightllm.models.qwen3_eagle.layer_weights.transformer_layer_weight import Qwen3EagleTransformerLayerWeight +from lightllm.utils import envs_utils + + +@pytest.mark.parametrize( + "module_name, class_name", + [ + ("lightllm.models.deepseek_mtp.model", "Deepseek3MTPModel"), + ("lightllm.models.glm4_moe_lite_mtp.model", "Glm4MoeLiteMTPModel"), + ("lightllm.models.mistral_mtp.model", "MistralMTPModel"), + ("lightllm.models.qwen3_moe_mtp.model", "Qwen3MOEMTPModel"), + ("lightllm.models.qwen3_5_mtp.model", "Qwen3_5MTPModel"), + ("lightllm.models.qwen3_5_moe_mtp.model", "Qwen3_5MoeMTPModel"), + ("lightllm.models.qwen3_eagle.model", "Qwen3EagleModel"), + ("lightllm.models.qwen3_dflash.model", "Qwen3DFlashModel"), + ("lightllm.models.qwen3_dflash2.model", "Qwen3DFlash2Model"), + ("lightllm.models.qwen3_5_dflash.model", "Qwen3_5DFlashModel"), + ("lightllm.models.qwen3_dspark.model", "Qwen3DSparkModel"), + ("lightllm.models.qwen3_5_dspark.model", "Qwen3_5DSparkModel"), + ], +) +def test_spec_draft_model_class_is_marked(module_name, class_name): + model_class = getattr(import_module(module_name), class_name) + assert model_class.is_mtp_draft_model is True + + +def test_qwen3_eagle_uses_layers_checkpoint_prefix(): + layer_weight = Qwen3EagleTransformerLayerWeight.__new__(Qwen3EagleTransformerLayerWeight) + layer_weight.layer_num_ = 2 + + layer_weight._init_weight_names() + + assert layer_weight._q_weight_name == "layers.2.self_attn.q_proj.weight" + assert layer_weight._hidden_norm_weight_name == "layers.2.hidden_norm.weight" + + +@pytest.mark.parametrize( + "mtp_mode, is_draft_model, draft_step, dynamic_spec, expected", + [ + ("dspark", False, 7, True, True), + ("dspark", True, 7, True, False), + ("dflash", False, 7, True, True), + ("dflash", True, 7, True, False), + ("dflash2", False, 7, False, False), + ("dflash2", True, 7, False, False), + ("dflash2", False, 7, True, True), + ("dflash2", True, 7, True, False), + ("vanilla_with_att", True, 7, True, False), + ("vanilla_with_att", True, 0, True, False), + ("eagle3", True, 0, True, False), + ("eagle3", False, 7, True, True), + ("eagle3", False, 7, False, False), + ], +) +def test_attention_backend_selects_dynamic_spec_layout( + monkeypatch, + mtp_mode, + is_draft_model, + draft_step, + dynamic_spec, + expected, +): + monkeypatch.setattr( + base_att_module, + "get_env_start_args", + lambda: SimpleNamespace(mtp_mode=mtp_mode, mtp_dynamic_verify=dynamic_spec), + ) + backend = SimpleNamespace( + model=SimpleNamespace( + is_mtp_draft_model=is_draft_model, + mtp_manager=SimpleNamespace(get_decode_draft_step=lambda _: draft_step), + ) + ) + + assert BaseAttBackend.uses_dynamic_spec_verify_layout(backend) is expected + + +@pytest.mark.parametrize( + "mtp_mode, is_draft_model, expected", + [ + (None, False, True), + ("dflash", False, True), + ("dflash", True, False), + ("dflash2", True, False), + ("dspark", True, False), + ("eagle3", True, True), + ("vanilla_with_att", True, True), + ], +) +def test_attention_backend_selects_causality(monkeypatch, mtp_mode, is_draft_model, expected): + monkeypatch.setattr( + base_att_module, + "get_env_start_args", + lambda: SimpleNamespace(mtp_mode=mtp_mode), + ) + backend = SimpleNamespace(model=SimpleNamespace(is_mtp_draft_model=is_draft_model)) + + assert BaseAttBackend.uses_causal_attention(backend) is expected + + +@pytest.mark.parametrize("state_class", [Fa3PrefillAttState, MlaFa3PrefillAttState]) +def test_fa3_prefill_state_owns_causality(monkeypatch, state_class): + monkeypatch.setattr(import_module(Fa3PrefillAttState.__module__), "page_table_copy", lambda **kwargs: None) + infer_state = InferStateInfo() + infer_state.b1_cu_q_seq_len = torch.tensor([0, 1, 2], dtype=torch.int32) + infer_state.b1_cu_kv_seq_len = torch.tensor([0, 3, 7], dtype=torch.int32) + infer_state.b_req_idx = torch.tensor([0, 1], dtype=torch.int32) + infer_state.batch_size = 2 + infer_state.max_kv_seq_len = 4 + infer_state.input_ids = torch.empty(2, dtype=torch.int64) + infer_state.req_manager = SimpleNamespace(req_to_token_indexs=torch.arange(8, dtype=torch.int32).reshape(2, 4)) + state = state_class( + backend=SimpleNamespace(uses_causal_attention=lambda: False, page_size=1, infer_page_size=1), + _infer_state_weak_ref=weakref.ref(infer_state), + ) + + state.init_state() + + assert state.causal is False + + +@pytest.mark.parametrize("state_class", [Fa3DecodeAttState, MlaFa3DecodeAttState]) +def test_fa3_decode_state_owns_causality(state_class): + infer_state = InferStateInfo() + infer_state.b1_cu_q_seq_len = torch.tensor([0, 1, 2], dtype=torch.int32) + infer_state.b1_cu_kv_seq_len = torch.tensor([0, 3, 7], dtype=torch.int32) + infer_state.b_req_idx = torch.tensor([0, 1], dtype=torch.int32) + infer_state.b_seq_len = torch.tensor([3, 4], dtype=torch.int32) + model = SimpleNamespace( + is_mtp_draft_model=False, + mtp_manager=SimpleNamespace(get_decode_draft_step=lambda _: 0), + ) + state = state_class( + backend=SimpleNamespace( + model=model, + uses_causal_attention=lambda: False, + uses_dynamic_spec_verify_layout=lambda: False, + ), + _infer_state_weak_ref=weakref.ref(infer_state), + ) + state._init_page_table = lambda _: None + + state.init_state() + + assert state.causal is False + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_triton_mtp_decode_state_builds_group_markers(): + model = SimpleNamespace( + is_mtp_draft_model=False, + mtp_manager=SimpleNamespace(get_decode_draft_step=lambda _: 2), + req_manager=SimpleNamespace(HOLD_REQUEST_ID=-1), + ) + infer_state = InferStateInfo() + infer_state.b_req_idx = torch.tensor([7, 7, -1, -1], dtype=torch.int32, device="cuda") + infer_state.max_kv_seq_len = 4 + state = TritonDecodeAttState( + backend=SimpleNamespace(model=model), + _infer_state_weak_ref=weakref.ref(infer_state), + ) + + state.init_state() + + assert state.b_mark_mtp_shared_group.cpu().tolist() == [0, 2, 1, 1] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@pytest.mark.parametrize("state_class", [Fa3DecodeAttState, MlaFa3DecodeAttState]) +def test_fa3_dynamic_decode_state_builds_group_markers(state_class): + model = SimpleNamespace(req_manager=SimpleNamespace(HOLD_REQUEST_ID=-1)) + infer_state = InferStateInfo() + infer_state.b_req_idx = torch.tensor([7, 7, -1, -1], dtype=torch.int32, device="cuda") + infer_state.b_seq_len = torch.tensor([3, 4, 2, 2], dtype=torch.int32, device="cuda") + infer_state.batch_size = 4 + state = state_class( + backend=SimpleNamespace(model=model), + _infer_state_weak_ref=weakref.ref(infer_state), + ) + + b_att_req_idx = state._init_dynamic_spec_verify_state(draft_step=2) + torch.cuda.synchronize() + + assert b_att_req_idx.cpu().tolist() == [7, -1, -1, -1] + assert state.b_att_seq_len.cpu().tolist() == [4, 2, 2, 0] + assert state.cu_seqlens_q.cpu().tolist() == [0, 2, 3, 4, 4] + + +@pytest.mark.parametrize( + "model_type, spec_mode, expected_class_name", + [ + ("deepseek_v3", "vanilla_with_att", "Deepseek3MTPModel"), + ("deepseek_v3", "eagle_with_att", "Deepseek3MTPModel"), + ("glm4_moe_lite", "vanilla_with_att", "Glm4MoeLiteMTPModel"), + ("glm4_moe_lite", "eagle_with_att", "Glm4MoeLiteMTPModel"), + ("mistral", "vanilla_no_att", "MistralMTPModel"), + ("mistral", "eagle_no_att", "MistralMTPModel"), + ("qwen3_moe", "vanilla_no_att", "Qwen3MOEMTPModel"), + ("qwen3_moe", "eagle_no_att", "Qwen3MOEMTPModel"), + ("qwen3_5", "vanilla_with_att", "Qwen3_5MTPModel"), + ("qwen3_5_text", "eagle_with_att", "Qwen3_5MTPModel"), + ("qwen3_5_moe", "vanilla_with_att", "Qwen3_5MoeMTPModel"), + ("qwen3_5_moe_text", "eagle_with_att", "Qwen3_5MoeMTPModel"), + ("qwen3", "dflash", "Qwen3DFlashModel"), + ("qwen3", "dflash2", "Qwen3DFlash2Model"), + ("qwen3_5", "dflash", "Qwen3_5DFlashModel"), + ("qwen3_5_text", "dflash", "Qwen3_5DFlashModel"), + ("qwen3", "dspark", "Qwen3DSparkModel"), + ("qwen3_5", "dspark", "Qwen3_5DSparkModel"), + ("qwen3_5_text", "dspark", "Qwen3_5DSparkModel"), + ("qwen3", "eagle3", "Qwen3EagleModel"), + ], +) +def test_draft_model_registry(model_type, spec_mode, expected_class_name): + model_class = get_draft_model_class( + model_cfg={"model_type": model_type}, + spec_mode=spec_mode, + ) + + assert model_class.__name__ == expected_class_name + + +def test_draft_model_registry_rejects_unsupported_model_type(): + with pytest.raises(ValueError, match="Unsupported speculative draft model"): + get_draft_model_class( + model_cfg={"model_type": "gemma4"}, + spec_mode="dspark", + ) + + +@pytest.mark.parametrize( + "model_type, spec_mode", + [ + ("deepseek_v3", "eagle3"), + ("deepseek_v3", "dspark"), + ("deepseek_v3", "dflash"), + ("glm4_moe_lite", "eagle3"), + ("glm4_moe_lite", "dspark"), + ("glm4_moe_lite", "dflash"), + ("qwen3_5", "eagle3"), + ("qwen3_5_moe", "eagle3"), + ("qwen3_5_moe", "dspark"), + ("qwen3_5_moe", "dflash"), + ("qwen3_5", "dflash2"), + ("qwen3", "eagle_no_att"), + ], +) +def test_draft_model_registry_rejects_unsupported_mode(model_type, spec_mode): + with pytest.raises(ValueError, match="Unsupported speculative draft model"): + get_draft_model_class( + model_cfg={"model_type": model_type}, + spec_mode=spec_mode, + ) + + +def test_hidden_collector_reads_target_layer_ids(monkeypatch): + config_reads = [] + + def get_config_dict(path): + config_reads.append(path) + return {"target_layer_ids": [1, 20, 36]}, {} + + monkeypatch.setattr( + hidden_collector_module.PretrainedConfig, + "get_config_dict", + get_config_dict, + ) + monkeypatch.setattr( + hidden_collector_module, + "get_env_start_args", + lambda: SimpleNamespace(mtp_draft_model_dir=["/models/dspark"]), + ) + model = SimpleNamespace(is_mtp_draft_model=False, layers_num=40) + + collector = hidden_collector_module.LayerHiddenCollector(model=model) + new_collector = collector.new_instance() + + assert collector.layer_ids == frozenset((1, 20, 36)) + assert new_collector.layer_ids == collector.layer_ids + assert config_reads == ["/models/dspark"] + + +@pytest.mark.parametrize( + "mtp_mode, mtp_step, expected_layer_num", + [ + (None, 0, 0), + ("vanilla_no_att", 7, 0), + ("eagle_no_att", 7, 0), + ("vanilla_with_att", 7, 7), + ("eagle_with_att", 7, 1), + ], +) +def test_fixed_added_mtp_kv_layer_num_by_mode(monkeypatch, mtp_mode, mtp_step, expected_layer_num): + monkeypatch.setattr( + envs_utils, + "get_env_start_args", + lambda: SimpleNamespace(mtp_mode=mtp_mode, mtp_step=mtp_step), + ) + envs_utils.get_added_mtp_kv_layer_num.cache_clear() + + assert envs_utils.get_added_mtp_kv_layer_num() == expected_layer_num + + +@pytest.mark.parametrize("mtp_mode", ["dflash", "dflash2"]) +def test_dflash_added_kv_layers_come_from_draft_config(tmp_path, mtp_mode): + config_path = tmp_path / "config.json" + config_path.write_text(json.dumps({"num_hidden_layers": 5})) + + envs_utils.get_env_start_args.cache_clear() + envs_utils.get_added_mtp_kv_layer_num.cache_clear() + envs_utils.set_env_start_args( + { + "mtp_mode": mtp_mode, + "mtp_step": 7, + "mtp_dynamic_verify": False, + "mtp_draft_model_dir": [str(tmp_path)], + } + ) + + assert envs_utils.get_added_mtp_kv_layer_num() == 5 + + +def test_qwen35_dflash_added_kv_layers_come_from_nested_draft_config(tmp_path): + config_path = tmp_path / "config.json" + config_path.write_text( + json.dumps( + { + "num_hidden_layers": 48, + "dflash_config": {"num_hidden_layers": 5}, + } + ) + ) + + envs_utils.get_env_start_args.cache_clear() + envs_utils.get_added_mtp_kv_layer_num.cache_clear() + envs_utils.set_env_start_args( + { + "mtp_mode": "dflash", + "mtp_step": 7, + "mtp_dynamic_verify": False, + "mtp_draft_model_dir": [str(tmp_path)], + } + ) + + assert envs_utils.get_added_mtp_kv_layer_num() == 5 + + +@pytest.mark.parametrize( + "draft_config,expected_layer_num", + [ + ({"num_hidden_layers": 6}, 6), + ( + { + "model_type": "deepseek_v4", + "num_hidden_layers": 61, + "compress_ratios": [4] * 61 + [0] * 3, + "dspark_target_layer_ids": [15, 31, 47], + "dspark_block_size": 8, + }, + 3, + ), + ], +) +def test_dspark_added_kv_layers_come_from_draft_config(tmp_path, draft_config, expected_layer_num): + config_path = tmp_path / "config.json" + config_path.write_text(json.dumps(draft_config)) + + envs_utils.get_env_start_args.cache_clear() + envs_utils.get_added_mtp_kv_layer_num.cache_clear() + envs_utils.set_env_start_args( + { + "mtp_mode": "dspark", + "mtp_step": 7, + "mtp_dynamic_verify": True, + "mtp_draft_model_dir": [str(tmp_path)], + } + ) + + assert envs_utils.get_added_mtp_kv_layer_num() == expected_layer_num + + +def test_eagle3_added_kv_layers_come_from_draft_config(tmp_path): + config_path = tmp_path / "config.json" + config_path.write_text(json.dumps({"num_hidden_layers": 2})) + + envs_utils.get_env_start_args.cache_clear() + envs_utils.get_added_mtp_kv_layer_num.cache_clear() + envs_utils.set_env_start_args( + { + "mtp_mode": "eagle3", + "mtp_step": 7, + "mtp_dynamic_verify": False, + "mtp_draft_model_dir": [str(tmp_path)], + } + ) + + assert envs_utils.get_added_mtp_kv_layer_num() == 2 diff --git a/unit_tests/utils/test_start_utils.py b/unit_tests/utils/test_start_utils.py new file mode 100644 index 0000000000..816e4000f7 --- /dev/null +++ b/unit_tests/utils/test_start_utils.py @@ -0,0 +1,319 @@ +from types import SimpleNamespace + +import pytest + +from lightllm.utils import start_utils + + +class FakeHttpServerProcess: + def __init__(self, return_code=None): + self.return_code = return_code + self.pid = 4321 + self.sent_signals = [] + self.wait_timeouts = [] + + def poll(self): + return self.return_code + + def send_signal(self, sig): + self.sent_signals.append(sig) + + def wait(self, timeout=None): + self.wait_timeouts.append(timeout) + return self.return_code + + +class FakeProcess: + def __init__(self, pid, running=True, children=None, name="process", exitcode=None, wait_timeout=False): + self.pid = pid + self.running = running + self._children = children or [] + self._name = name + self.exitcode = exitcode + self.wait_timeout = wait_timeout + + def is_running(self): + return self.running + + def children(self, recursive=False): + return self._children + + def name(self): + return self._name + + def wait(self, timeout=None): + if self.wait_timeout: + raise start_utils.psutil.TimeoutExpired(timeout, pid=self.pid, name=self._name) + return self.exitcode + + +def test_start_submodule_processes_returns_and_manages_psutil_processes(monkeypatch): + class FakePipeReader: + def recv(self): + # 子进程应在等待初始化结果之前就进入 manager，保证此时 Ctrl-C 可以清理它们。 + assert len(process_manager.processes) == 2 + return "init ok" + + class FakeMpProcess: + next_pid = 1000 + + def __init__(self, target, args): + self.pid = self.next_pid + FakeMpProcess.next_pid += 1 + + def start(self): + pass + + def is_alive(self): + return True + + monkeypatch.setattr(start_utils.mp, "Pipe", lambda duplex: (FakePipeReader(), object())) + monkeypatch.setattr(start_utils.mp, "Process", FakeMpProcess) + monkeypatch.setattr( + start_utils.psutil, + "Process", + lambda pid: FakeProcess(pid=pid, name=f"process-{pid}"), + ) + process_manager = start_utils.SubmoduleManager() + + processes = process_manager.start_submodule_processes( + start_funcs=[lambda pipe_writer: None, lambda pipe_writer: None], + start_args=[(), ()], + ) + + assert processes == process_manager.processes + assert [process.pid for process in processes] == [1000, 1001] + assert process_manager.process_names == { + processes[0]: "process-1000", + processes[1]: "process-1001", + } + + +def test_register_process_tree_adds_recursive_descendants(monkeypatch): + descendants = [ + FakeProcess(pid=1001, name="lightllm::model_infer"), + FakeProcess(pid=1002, name="lightllm::pd_manager"), + FakeProcess(pid=1003, name="lightllm::pd_worker"), + ] + router_process = FakeProcess(pid=1000, children=descendants) + process_manager = start_utils.SubmoduleManager() + monkeypatch.setattr(start_utils, "is_process_active", lambda pid: True) + + process_manager.register_process_tree(router_process) + + assert process_manager.processes == descendants + assert process_manager.process_names == { + descendants[0]: "lightllm::model_infer", + descendants[1]: "lightllm::pd_manager", + descendants[2]: "lightllm::pd_worker", + } + + +def test_register_process_tree_filters_short_lived_helper_processes(monkeypatch): + model_process = FakeProcess(pid=1001, name="lightllm::model_infer") + compile_worker = FakeProcess(pid=1002, name="python") + pd_process = FakeProcess(pid=1003, name="lightllm::decode_trans") + router_process = FakeProcess(pid=1000, children=[model_process, compile_worker, pd_process]) + process_manager = start_utils.SubmoduleManager() + monkeypatch.setattr(start_utils, "is_process_active", lambda pid: True) + + process_manager.register_process_tree(router_process) + + assert process_manager.processes == [model_process, pd_process] + assert process_manager.process_names == { + model_process: "lightllm::model_infer", + pd_process: "lightllm::decode_trans", + } + + +def test_register_process_tree_ignores_processes_that_exit_during_scan(): + class ExitedProcess(FakeProcess): + def name(self): + raise start_utils.psutil.NoSuchProcess(self.pid) + + exited_process = ExitedProcess(pid=1001) + router_process = FakeProcess(pid=1000, children=[exited_process]) + process_manager = start_utils.SubmoduleManager() + + process_manager.register_process_tree(router_process) + + assert process_manager.processes == [] + assert process_manager.process_names == {} + + +def test_setup_signal_handlers_registers_and_handles_sigterm(monkeypatch): + http_server_process = FakeHttpServerProcess() + process_manager = start_utils.SubmoduleManager() + registered_handlers = {} + terminate_calls = [] + monkeypatch.setattr( + start_utils.signal, + "signal", + lambda sig, handler: registered_handlers.__setitem__(sig, handler), + ) + monkeypatch.setattr(process_manager, "terminate_all_processes", lambda: terminate_calls.append(True)) + + process_manager.setup_signal_handlers(http_server_process) + + assert set(registered_handlers) == { + start_utils.signal.SIGTERM, + start_utils.signal.SIGINT, + start_utils.signal.SIGHUP, + } + with pytest.raises(SystemExit) as exc_info: + registered_handlers[start_utils.signal.SIGTERM](start_utils.signal.SIGTERM, None) + + assert exc_info.value.code == 0 + assert http_server_process.sent_signals == [start_utils.signal.SIGTERM] + assert http_server_process.wait_timeouts == [60] + assert terminate_calls == [True] + + +def test_setup_signal_handlers_uses_configured_graceful_shutdown_signal(monkeypatch): + http_server_process = FakeHttpServerProcess() + process_manager = start_utils.SubmoduleManager() + registered_handlers = {} + monkeypatch.setattr( + start_utils.signal, + "signal", + lambda sig, handler: registered_handlers.__setitem__(sig, handler), + ) + monkeypatch.setattr(process_manager, "terminate_all_processes", lambda: None) + + process_manager.setup_signal_handlers( + http_server_process, + graceful_shutdown_signal=start_utils.signal.SIGUSR1, + ) + + with pytest.raises(SystemExit) as exc_info: + registered_handlers[start_utils.signal.SIGTERM](start_utils.signal.SIGTERM, None) + + assert exc_info.value.code == 0 + assert http_server_process.sent_signals == [start_utils.signal.SIGUSR1] + assert http_server_process.wait_timeouts == [60] + + +def test_setup_exit_controller_starts_cleanup_process_without_registering_signals(monkeypatch): + process_manager = start_utils.SubmoduleManager() + cleanup_process_calls = [] + registered_handlers = {} + monkeypatch.setattr(start_utils, "get_unique_server_name", lambda: "service_0") + monkeypatch.setattr( + start_utils, + "start_launcher_shm_cleanup_process", + lambda service_name: cleanup_process_calls.append(service_name), + ) + monkeypatch.setattr( + start_utils.signal, + "signal", + lambda sig, handler: registered_handlers.__setitem__(sig, handler), + ) + + process_manager.setup_exit_controller() + + assert cleanup_process_calls == ["service_0"] + assert registered_handlers == {} + + +def test_supervisor_fails_when_http_server_exits(monkeypatch): + http_server_process = FakeHttpServerProcess(return_code=0) + process_manager = start_utils.SubmoduleManager() + terminate_calls = [] + kill_calls = [] + monkeypatch.setattr(process_manager, "terminate_all_processes", lambda: terminate_calls.append(True)) + monkeypatch.setattr(start_utils, "kill_recursive", lambda process: kill_calls.append(process)) + + with pytest.raises(RuntimeError, match="HTTP server exited unexpectedly with return code 0"): + process_manager.supervise_processes(http_server_process) + + assert kill_calls == [] + assert terminate_calls == [True] + + +def test_supervisor_fails_and_cleans_up_when_submodule_exits(monkeypatch): + http_server_process = FakeHttpServerProcess() + dead_process = FakeProcess(pid=1234, running=False, name="router", exitcode=-9) + process_manager = start_utils.SubmoduleManager() + process_manager.processes = [dead_process] + process_manager.process_names = {dead_process: dead_process.name()} + terminate_calls = [] + kill_calls = [] + monkeypatch.setattr(process_manager, "terminate_all_processes", lambda: terminate_calls.append(True)) + monkeypatch.setattr(start_utils, "is_process_active", lambda pid: True) + monkeypatch.setattr(start_utils, "kill_recursive", lambda process: kill_calls.append(process)) + + with pytest.raises( + RuntimeError, + match="Critical LightLLM submodule exited unexpectedly: name=router pid=1234 exitcode=-9", + ): + process_manager.supervise_processes(http_server_process) + + assert kill_calls == [http_server_process] + assert terminate_calls == [True] + + +def test_supervisor_treats_zombie_submodule_as_dead(monkeypatch): + http_server_process = FakeHttpServerProcess() + zombie_process = FakeProcess(pid=1234, name="router", exitcode=-9) + process_manager = start_utils.SubmoduleManager() + process_manager.processes = [zombie_process] + process_manager.process_names = {zombie_process: zombie_process.name()} + terminate_calls = [] + kill_calls = [] + monkeypatch.setattr(process_manager, "terminate_all_processes", lambda: terminate_calls.append(True)) + monkeypatch.setattr(start_utils, "is_process_active", lambda pid: False) + monkeypatch.setattr(start_utils, "kill_recursive", lambda process: kill_calls.append(process)) + + with pytest.raises( + RuntimeError, + match="Critical LightLLM submodule exited unexpectedly: name=router pid=1234 exitcode=-9", + ): + process_manager.supervise_processes(http_server_process) + + assert kill_calls == [http_server_process] + assert terminate_calls == [True] + + +def test_supervisor_keeps_polling_while_all_processes_are_alive(monkeypatch): + http_server_process = FakeHttpServerProcess() + child_process = FakeProcess(pid=1234) + process_manager = start_utils.SubmoduleManager() + process_manager.processes = [child_process] + process_manager.process_names = {child_process: child_process.name()} + terminate_calls = [] + sleep_calls = [] + monkeypatch.setattr(process_manager, "terminate_all_processes", lambda: terminate_calls.append(True)) + + def stop_after_first_poll(interval): + sleep_calls.append(interval) + http_server_process.return_code = 2 + + monkeypatch.setattr(start_utils, "is_process_active", lambda pid: True) + monkeypatch.setattr(start_utils.time, "sleep", stop_after_first_poll) + + with pytest.raises(RuntimeError, match="HTTP server exited unexpectedly with return code 2"): + process_manager.supervise_processes(http_server_process) + + assert sleep_calls == [5.0] + assert terminate_calls == [True] + + +def test_supervisor_supports_submodule_only_processes(monkeypatch): + dead_process = FakeProcess(pid=1234, running=False, name="model_infer", wait_timeout=True) + process_manager = start_utils.SubmoduleManager() + process_manager.processes = [dead_process] + process_manager.process_names = {dead_process: dead_process.name()} + terminate_calls = [] + kill_calls = [] + monkeypatch.setattr(process_manager, "terminate_all_processes", lambda: terminate_calls.append(True)) + monkeypatch.setattr(start_utils, "is_process_active", lambda pid: False) + monkeypatch.setattr(start_utils, "kill_recursive", lambda process: kill_calls.append(process)) + + with pytest.raises( + RuntimeError, + match="Critical LightLLM submodule exited unexpectedly: name=model_infer pid=1234 exitcode=None", + ): + process_manager.supervise_processes() + + assert kill_calls == [] + assert terminate_calls == [True] From 917e8d114e2f2a34cb263dc06da73bd75eef29bc Mon Sep 17 00:00:00 2001 From: sufubao Date: Fri, 13 Nov 2026 17:38:14 +0800 Subject: [PATCH 3/6] fix(openai): preserve sampled scores and align chat logprobs --- .../triton_kernel/logprobs_capture.py | 37 +- lightllm/server/api_openai.py | 92 +-- lightllm/server/chat_logprob_buffer.py | 336 +++------- lightllm/server/core/objs/logprob_utils.py | 67 +- lightllm/server/function_call_parser.py | 21 + lightllm/server/reasoning_parser.py | 195 +++++- .../model_infer/mode_backend/base_backend.py | 4 +- .../mode_backend/generic_post_process.py | 25 +- lightllm/server/text_spans.py | 29 + .../structured_output/test_sampling.py | 44 +- .../server/test_chat_output_logprobs.py | 578 ++++++++++++++---- 11 files changed, 948 insertions(+), 480 deletions(-) create mode 100644 lightllm/server/text_spans.py diff --git a/lightllm/common/basemodel/triton_kernel/logprobs_capture.py b/lightllm/common/basemodel/triton_kernel/logprobs_capture.py index 504a1a6ffb..60dd9d0096 100644 --- a/lightllm/common/basemodel/triton_kernel/logprobs_capture.py +++ b/lightllm/common/basemodel/triton_kernel/logprobs_capture.py @@ -77,7 +77,7 @@ def _output_logprobs_partial( rows, candidates, sampled_ids, - higher_counts, + ranks, stride_row: tl.constexpr, vocab_size: tl.constexpr, tiles: tl.constexpr, @@ -91,10 +91,9 @@ def _output_logprobs_partial( row_logits = logits + row * stride_row ids = tile * BLOCK + tl.arange(0, BLOCK) scores = tl.load(row_logits + ids, ids < vocab_size, -float("inf")).to(tl.float32) - sampled = tl.load(row_logits + tl.load(sampled_ids + row)).to(tl.float32) - # 只数严格更高的 logit。同分不占名次，采样 token 仍算在前 20 里。 - higher = tl.sum(((ids < vocab_size) & (scores > sampled)).to(tl.int32), 0) - tl.store(higher_counts + index * tiles + tile, higher) + sampled = tl.load(row_logits + tl.load(sampled_ids + row)) + rank = tl.sum(((ids < vocab_size) & (scores > sampled)).to(tl.int32), 0) + tl.store(ranks + index * tiles + tile, rank) # tl.topk 只能比较整数。高 32 位放可排序的分数，低 32 位放反转后的 token id。 # 取最大的键时，分数高的在前；分数相同则 id 更小的在前。词表外的位置写成最小整数。 bits = scores.to(tl.int32, bitcast=True) @@ -115,7 +114,8 @@ def _output_logprobs_finish( candidates, token_ids, logprobs, - higher_counts, + partial_ranks, + sampled_ranks, stride_row: tl.constexpr, vocab_size: tl.constexpr, tiles: tl.constexpr, @@ -142,11 +142,10 @@ def _output_logprobs_finish( row_logprobs = scores - sampled + sampled_logprob tl.store(token_ids + offset + columns, stored_ids, columns < width) tl.store(logprobs + offset + columns, row_logprobs, columns < width) - # 候选平移用的是原始采样 logprob。至少 20 个 token 严格更高时，再把采样分数写成 -9999。 - higher_tiles = tl.arange(0, TILE_BLOCK) - higher = tl.sum(tl.load(higher_counts + index * tiles + higher_tiles, higher_tiles < tiles, 0), 0) - outside = (higher >= 20) | ~(tl.abs(sampled) < float("inf")) - tl.store(sampled_logprobs + row, -9999.0, mask=outside) + rank_tiles = tl.arange(0, TILE_BLOCK) + rank = 1 + tl.sum(tl.load(partial_ranks + index * tiles + rank_tiles, rank_tiles < tiles, 0), 0) + rank = tl.where(tl.abs(sampled) < float("inf"), rank, vocab_size + 1) + tl.store(sampled_ranks + index, rank) def capture_output_logprobs( @@ -156,12 +155,12 @@ def capture_output_logprobs( rows: torch.Tensor, max_topk: int, output_count: int, -) -> Tuple[torch.Tensor, torch.Tensor]: +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """返回每行的 top-k token id 和 logprob。 rows 每项是 (logits 行号, K, 写入起点)。写入起点按 token 个数累加， - 第 r 行的候选是两个数组的 [起点, 起点 + K)。至少 20 个 token 的 logit - 严格高于采样 token 时，把 sampled_logprobs 里对应位置写成 -9999。K=0 也判断。 + 第 r 行的候选是两个数组的 [起点, 起点 + K)。另返回每行采样 token 的 + 1-based rank（同分共享排名），不修改内部采样 logprob。K=0 仍计算排名。 """ block = 1024 # 先按词表分段留下局部 top-k，再逐行合并成最终候选。 @@ -169,13 +168,14 @@ def capture_output_logprobs( # Triton 的 topk 块长必须是 2 的幂。请求的 K 最大 20，这里向上取整；真正写入仍按每行的 K 截断。 topk = triton.next_power_of_2(max(2, max_topk)) candidates = logits.new_empty((rows.shape[0], tiles, topk), dtype=torch.int64) - higher_counts = logits.new_empty((rows.shape[0], tiles), dtype=torch.int32) + partial_ranks = logits.new_empty((rows.shape[0], tiles), dtype=torch.int32) + sampled_ranks = logits.new_empty((rows.shape[0],), dtype=torch.int32) _output_logprobs_partial[(rows.shape[0], tiles)]( logits=logits, rows=rows, candidates=candidates, sampled_ids=sampled_ids, - higher_counts=higher_counts, + ranks=partial_ranks, stride_row=logits.stride(0), vocab_size=logits.shape[1], tiles=tiles, @@ -193,7 +193,8 @@ def capture_output_logprobs( candidates=candidates, token_ids=token_ids, logprobs=logprobs, - higher_counts=higher_counts, + partial_ranks=partial_ranks, + sampled_ranks=sampled_ranks, stride_row=logits.stride(0), vocab_size=logits.shape[1], tiles=tiles, @@ -201,4 +202,4 @@ def capture_output_logprobs( TILE_BLOCK=triton.next_power_of_2(tiles), num_warps=4, ) - return token_ids, logprobs + return token_ids, logprobs, sampled_ranks diff --git a/lightllm/server/api_openai.py b/lightllm/server/api_openai.py index 2104a40da6..1afff9fa5f 100644 --- a/lightllm/server/api_openai.py +++ b/lightllm/server/api_openai.py @@ -12,6 +12,7 @@ from lightllm.server.reasoning_parser import ReasoningParser from .chat_logprob_buffer import ChatLogprobBuffer +from .text_spans import TextSpan from .function_call_parser import TOOLS_TAG_LIST, FunctionCallParser, ToolCallItem from .build_prompt import build_prompt, init_tokenizer @@ -228,7 +229,8 @@ def _process_reasoning_stream( reasoning_parser_dict: Dict[int, ReasoningParser], metadata: Dict[str, Any], request: ChatCompletionRequest, -) -> tuple[Optional[str], str]: + with_spans: bool = False, +): """Process reasoning content and update its token usage.""" if index not in reasoning_parser_dict: reasoning_parser_dict[index] = ReasoningParser( @@ -240,10 +242,12 @@ def _process_reasoning_stream( token_id = metadata.get("id") if token_id is not None: parser.update_reasoning_token_count(int(token_id)) - return parser.parse_stream_chunk(delta) + return parser.parse_stream_chunk_with_spans(delta) if with_spans else parser.parse_stream_chunk(delta) -def _process_tools_stream(index: int, delta: str, parser_dict: Dict, request: ChatCompletionRequest): +def _process_tools_stream( + index: int, delta: str, parser_dict: Dict, request: ChatCompletionRequest, with_spans: bool = False +): from .api_http import g_objs if index not in parser_dict: @@ -256,8 +260,7 @@ def _process_tools_stream(index: int, delta: str, parser_dict: Dict, request: Ch parser = parser_dict[index] # parse_increment => returns (normal_text, calls) - normal_text, calls = parser.parse_stream_chunk(delta) - return normal_text, calls + return parser.parse_stream_chunk_with_spans(delta) if with_spans else parser.parse_stream_chunk(delta) def _split_tool_argument_delta(arguments: Optional[str]) -> List[str]: @@ -424,31 +427,21 @@ async def chat_completions_impl(request: ChatCompletionRequest, raw_request: Req sampling_params.init(tokenizer=g_objs.httpserver_manager.tokenizer, **sampling_params_dict) await sampling_params.verify_cpu_async(g_objs.httpserver_manager.output_grammar_compiler) - logprob_model_type = get_env_start_args().reasoning_parser - logprob_force_reasoning = _is_force_thinking_mode(request) if logprob_model_type else False - - def _new_logprob_buffer() -> ChatLogprobBuffer: - return ChatLogprobBuffer( - tokenizer=g_objs.httpserver_manager.tokenizer, - model_type=logprob_model_type, - force_reasoning=logprob_force_reasoning, - separate_reasoning=bool(request.separate_reasoning), - tools=request.tools if request.tool_choice != "none" else None, - tool_call_parser=getattr(g_objs.args, "tool_call_parser", None) or "llama3", - ) - - logprob_buffers = None - if request.logprobs: - # 构造函数会检查思维链和工具调用的分隔符能不能编成稳定的 token 序列，不合法就抛 ValueError。 - # defaultdict 要等第一条输出 token 才创建 buffer。流式响应那时已经返回，这个错误进不了接口的 400 处理。 - # 这里先构造一次，只为在返回前完成检查；这个对象随即丢掉，真正使用的仍是下面按子请求创建的 buffer。 - # generate 放在这次检查之后，参数不合法时请求不会进入推理。 - _new_logprob_buffer() - logprob_buffers = collections.defaultdict(_new_logprob_buffer) - results_generator = g_objs.httpserver_manager.generate( prompt, sampling_params, multimodal_params, request=raw_request ) + logprob_buffers = ( + collections.defaultdict( + lambda: ChatLogprobBuffer( + tokenizer=g_objs.httpserver_manager.tokenizer, + separate_reasoning=bool(request.separate_reasoning), + skip_special_tokens=bool(sampling_params.skip_special_tokens), + spaces_between_special_tokens=bool(sampling_params.add_spaces_between_special_tokens), + ) + ) + if request.logprobs + else None + ) # Non-streaming case if not request.stream: @@ -499,13 +492,14 @@ def _new_logprob_buffer() -> ChatLogprobBuffer: sub_req_id = sub_ids[i] finish_reason = finish_reason_dict[sub_req_id] text = "".join(final_output_dict[sub_req_id]) + output_spans = [TextSpan(0, len(text))] # Handle reasoning content reasoning_text = None if reasoning_parser: try: parser = reasoning_parser_dict[sub_req_id] - reasoning_text, text = parser.parse_non_stream(text) + reasoning_text, text, output_spans = parser.parse_non_stream_with_spans(text) except Exception as e: logger.error(f"Reasoning parsing error: {e}") return create_error_response( @@ -516,8 +510,12 @@ def _new_logprob_buffer() -> ChatLogprobBuffer: text = (reasoning_text or "") + (text or "") reasoning_text = None + if logprob_buffers is not None: + logprob_buffers[sub_req_id].assign(output_spans, len(text)) + # Handle tool_calls parsing tool_calls = None + content_spans = None tool_choice = request.tool_choice tools = request.tools if tool_choice != "none" and any([i in text for i in TOOLS_TAG_LIST]): @@ -525,7 +523,7 @@ def _new_logprob_buffer() -> ChatLogprobBuffer: # 为 tool_call_parser 提供默认值 tool_parser = getattr(g_objs.args, "tool_call_parser", None) or "llama3" parser = FunctionCallParser(tools, tool_parser) - text, call_info_list = parser.parse_non_stream(text) + text, call_info_list, content_spans = parser.parse_non_stream_with_spans(text) tool_calls = [] history_tool_calls_cnt = _get_history_tool_calls_cnt(request) for call_info in call_info_list: @@ -548,7 +546,7 @@ def _new_logprob_buffer() -> ChatLogprobBuffer: finish_reason = "tool_calls" choice_logprobs = None if logprob_buffers is not None: - content_logprobs, reasoning_logprobs = logprob_buffers[sub_req_id].message_logprobs() + content_logprobs, reasoning_logprobs = logprob_buffers[sub_req_id].message_logprobs(content_spans) choice_logprobs = ChatLogprobs(content=content_logprobs, reasoning_content=reasoning_logprobs) chat_message = ChatMessage( role="assistant", @@ -612,7 +610,7 @@ async def stream_results() -> AsyncGenerator[bytes, None]: logprob_buffer = None if logprob_buffers is not None: logprob_buffer = logprob_buffers[sub_req_id] - logprob_buffer.record(text=delta, metadata=metadata, finish_status=finish_status) + raw_span = logprob_buffer.record(text=delta, metadata=metadata, finish_status=finish_status) current_finish_reason = finish_status.get_finish_reason() if current_finish_reason == "error" and completion_tokens == 1: raise GenerationError("Generation failed before producing output") @@ -636,9 +634,12 @@ async def stream_results() -> AsyncGenerator[bytes, None]: # Handle reasoning content if get_env_start_args().reasoning_parser: - reasoning_text, delta = _process_reasoning_stream( - choice_index, delta, reasoning_parser_dict, metadata, request + reasoning_text, delta, output_spans = _process_reasoning_stream( + choice_index, delta, reasoning_parser_dict, metadata, request, with_spans=True ) + if logprob_buffer is not None: + content_length = len(delta or "") + (0 if request.separate_reasoning else len(reasoning_text or "")) + logprob_buffer.assign(output_spans, content_length) if reasoning_text: if request.separate_reasoning: choice_logprobs = logprob_buffer.reasoning_logprobs() if logprob_buffer is not None else None @@ -656,11 +657,13 @@ async def stream_results() -> AsyncGenerator[bytes, None]: yield f"data: {_serialize_sse_chunk(chunk, _choice_nulls, logprobs=choice_logprobs)}\n\n" else: delta = reasoning_text + (delta or "") + elif logprob_buffer is not None: + logprob_buffer.assign([raw_span], len(delta)) if request.tool_choice != "none" and request.tools: # parse_increment => returns (normal_text, calls) - normal_text, calls = _process_tools_stream( - index=choice_index, delta=delta, parser_dict=parser_dict, request=request + normal_text, calls, content_spans = _process_tools_stream( + index=choice_index, delta=delta, parser_dict=parser_dict, request=request, with_spans=True ) # 1) if there's normal_text, output it as normal content @@ -676,8 +679,12 @@ async def stream_results() -> AsyncGenerator[bytes, None]: choices=[choice_data], model=request.model, ) - choice_logprobs = logprob_buffer.content_logprobs() if logprob_buffer is not None else None + choice_logprobs = ( + logprob_buffer.content_logprobs(content_spans) if logprob_buffer is not None else None + ) yield f"data: {_serialize_sse_chunk(chunk, _choice_nulls, logprobs=choice_logprobs)}\n\n" + elif logprob_buffer is not None: + logprob_buffer.discard_content(content_spans) # 2) if we found calls, we output them as separate chunk(s) history_tool_calls_cnt = _get_history_tool_calls_cnt(request) @@ -828,7 +835,12 @@ async def stream_results() -> AsyncGenerator[bytes, None]: if get_env_start_args().reasoning_parser: parser = reasoning_parser_dict.get(choice_index) if parser is not None: - flush_reasoning, flush_text = parser.flush() + flush_reasoning, flush_text, output_spans = parser.flush_with_spans() + if logprob_buffer is not None: + content_length = len(flush_text or "") + ( + 0 if request.separate_reasoning else len(flush_reasoning or "") + ) + flush_spans = logprob_buffer.assign(output_spans, content_length) if flush_reasoning: if request.separate_reasoning: flush_logprobs = ( @@ -841,7 +853,7 @@ async def stream_results() -> AsyncGenerator[bytes, None]: ) else: flush_logprobs = ( - logprob_buffer.content_logprobs() if logprob_buffer is not None else None + logprob_buffer.content_logprobs(flush_spans) if logprob_buffer is not None else None ) # vLLM compat: emit buffered thinking as content flush_choice = ChatCompletionStreamResponseChoice( @@ -869,7 +881,9 @@ async def stream_results() -> AsyncGenerator[bytes, None]: model=request.model, choices=[flush_choice], ) - flush_logprobs = logprob_buffer.content_logprobs() if logprob_buffer is not None else None + flush_logprobs = ( + logprob_buffer.content_logprobs(flush_spans) if logprob_buffer is not None else None + ) payload = _serialize_sse_chunk(flush_chunk, _choice_nulls, logprobs=flush_logprobs) yield f"data: {payload}\n\n" if has_emitted_tool_calls[sub_req_id] and current_finish_reason == "stop": diff --git a/lightllm/server/chat_logprob_buffer.py b/lightllm/server/chat_logprob_buffer.py index 2f41449b9c..fb77e139a2 100644 --- a/lightllm/server/chat_logprob_buffer.py +++ b/lightllm/server/chat_logprob_buffer.py @@ -1,265 +1,121 @@ -from __future__ import annotations - -from typing import Any, NamedTuple +from dataclasses import dataclass +from typing import Any, Dict, List, Optional, Tuple from lightllm.server.core.objs import FinishStatus -from lightllm.server.core.objs.logprob_utils import chat_token_logprob -from lightllm.server.function_call_parser import FunctionCallParser -from lightllm.server.reasoning_parser import ReasoningParser +from lightllm.server.core.objs.logprob_utils import ChatTokenDecoder +from lightllm.server.text_spans import TextSpan + + +@dataclass +class PendingLogprob: + start: int + end: int + token: Dict[str, Any] class ChatLogprobBuffer: - """按 token 序列识别分隔符，把可见文本的 logprob 放进正文或思维链。 + """Attribute generated tokens to the visible source ranges returned by parsers. - 分隔符可以是一个或多个 token。凑齐才切换通道，分隔符本身不输出。 - 空增量的分数先累加，加到后面第一个非空文本上。生成结束时还没拼上的空增量丢掉。 + A token spanning two channels belongs to the first visible channel. Empty + detokenization increments stay pending until the character they complete is + classified; recovering raw bytes is independent of retaining that token. """ def __init__( self, tokenizer: Any, - model_type: str | None, - force_reasoning: bool, separate_reasoning: bool, - tools: list | None = None, - tool_call_parser: str | None = None, + skip_special_tokens: bool = True, + spaces_between_special_tokens: bool = True, ): self.tokenizer = tokenizer self._separate_reasoning = separate_reasoning - self._channel = "content" # 当前输出通道：content、reasoning，或 commentary 的 drop。 - self._in_tool = False # 是否处在工具调用开始和结束分隔符之间。 - self._pending_logprob = 0.0 # 尚未拼成字符的空增量分数，等下一个非空文本一起输出。 - self._reasoning: list[dict[str, Any]] = [] # 分开思考时的思维链 logprob。 - self._content: list[dict[str, Any]] = [] # 正文 logprob；不分开思考时思维链也放这里。 - self._sequences: list[tuple[tuple[int, ...], str]] = [] # 分隔符 token id 序列和切换动作。 - self._held: list[_Held] = [] # 可能是分隔符前缀、尚未判定的 token。 - self._mode = "body" # gpt-oss 解析位置：body 正文、header 角色头、channel 通道名。 - self._channel_name = "" # <|channel|> 和 <|message|> 之间拼出来的通道名。 - if model_type: - self._init_reasoning(tokenizer=tokenizer, model_type=model_type, force_reasoning=force_reasoning) - if tools: - self._init_tools(tokenizer=tokenizer, tool_call_parser=tool_call_parser) - - def record(self, text: str, metadata: dict[str, Any], finish_status: FinishStatus) -> None: + self._hold: List[PendingLogprob] = [] + self._reasoning: List[PendingLogprob] = [] + self._content: List[PendingLogprob] = [] + self._source_length = 0 + self._content_length = 0 + self._decoder = ChatTokenDecoder(tokenizer, skip_special_tokens, spaces_between_special_tokens) + + def record(self, text: str, metadata: Dict[str, Any], finish_status: FinishStatus) -> TextSpan: + start = self._source_length + self._source_length += len(text) token_id, score = metadata.get("id"), metadata.get("logprob") - if token_id is not None and score is not None: - self._held.append(_Held(token_id=token_id, text=text, score=score, metadata=metadata)) - self._resolve_held(closing=False) - if finish_status.is_finished(): - self._resolve_held(closing=True) - self._pending_logprob = 0.0 - - def reasoning_logprobs(self) -> dict[str, Any] | None: - tokens = self._drain(items=self._reasoning) - return {"content": None, "reasoning_content": tokens} if tokens else None - - def content_logprobs(self) -> dict[str, Any] | None: - tokens = self._drain(items=self._content) - return {"content": tokens, "reasoning_content": None} if tokens else None - - def message_logprobs(self) -> tuple[list[dict[str, Any]] | None, list[dict[str, Any]] | None]: - return self._drain(items=self._content) or None, self._drain(items=self._reasoning) or None - - def _init_reasoning(self, tokenizer: Any, model_type: str, force_reasoning: bool) -> None: - key = model_type.lower() - detector_class = ReasoningParser.DetectorMap.get(key) - if detector_class is None: - raise ValueError(f"logprobs does not support reasoning parser {model_type!r}") - if key == "gpt-oss": - for text, action in ( - ("<|start|>", "harmony_start"), - ("<|channel|>", "harmony_channel"), - ("<|message|>", "harmony_message"), - ("<|end|>", "harmony_end"), - ("<|return|>", "harmony_end"), - ("<|call|>", "harmony_end"), - ): - # <|call|> 不是每个词表都有，缺了就跳过。 - self._add_sequence(tokenizer=tokenizer, text=text, action=action, required=text != "<|call|>") - if force_reasoning: - self._channel = "reasoning" - return - detector = detector_class(stream_reasoning=True, force_reasoning=force_reasoning) - self._add_sequence(tokenizer=tokenizer, text=detector.think_start_token, action="reasoning") - self._add_sequence(tokenizer=tokenizer, text=detector.think_end_token, action="content") - # 以解析器为准。Gemma 4 会忽略请求里的 force_reasoning，生成不是从已打开的思考通道开始。 - if detector._in_reasoning: - self._channel = "reasoning" - - def _init_tools(self, tokenizer: Any, tool_call_parser: str | None) -> None: - if not tool_call_parser: - raise ValueError("logprobs requires tool_call_parser when tools are set") - if tool_call_parser == "mistral": - raise ValueError("logprobs does not support mistral tool_call_parser") - parser = FunctionCallParser(tools=[], tool_call_parser=tool_call_parser) - self._add_sequence(tokenizer=tokenizer, text=parser.detector.bot_token, action="tool_on") - if parser.detector.eot_token: - self._add_sequence(tokenizer=tokenizer, text=parser.detector.eot_token, action="tool_off") - - def _add_sequence(self, tokenizer: Any, text: str, action: str, required: bool = True) -> None: - token_ids = _delimiter_ids(tokenizer=tokenizer, text=text) - if token_ids is None: - if required: - raise ValueError(f"logprobs requires {text!r} to encode to the same token sequence") - return - sequence = tuple(token_ids) - for existing, existing_action in self._sequences: - if existing == sequence and existing_action != action: - raise ValueError(f"logprobs delimiter {text!r} maps to conflicting actions") - if any(existing == sequence for existing, _ in self._sequences): - return - self._sequences.append((sequence, action)) - - def _resolve_held(self, closing: bool) -> None: - # 还是更长分隔符的前缀时先留着。生成结束则按已经凑齐的短分隔符结算，否则当正文。 - while self._held: - token_ids = tuple(item.token_id for item in self._held) - exact, waiting, completed = self._lookup(token_ids=token_ids) - if waiting and not closing: - return - if exact is not None: - self._apply(action=exact) - self._held.clear() - return - if completed is not None: - action, size = completed - self._apply(action=action) - del self._held[:size] - continue - self._release(item=self._held.pop(0)) - - def _lookup(self, token_ids: tuple[int, ...]) -> tuple[str | None, bool, tuple[str, int] | None]: - exact = None - waiting = False - completed = None - for sequence, action in self._sequences: - if sequence == token_ids: - exact = action - elif sequence[: len(token_ids)] == token_ids: - waiting = True - elif token_ids[: len(sequence)] == sequence and (completed is None or len(sequence) > completed[1]): - completed = (action, len(sequence)) - return exact, waiting, completed - - def _apply(self, action: str) -> None: - match action: - case "reasoning": - # 思维链开始。这组 token 是分隔符，不输出。 - self._channel = "reasoning" - case "content": - # 思维链结束，后面的文本回到正文。 - self._channel = "content" - case "tool_on": - # 工具调用开始，参数不进入正文。 - self._in_tool = True - case "tool_off": - # 工具调用结束，后面的文本回到当前通道。 - self._in_tool = False - case "harmony_start": - # <|start|> 到下一个结构标记之间是角色头。 - self._mode = "header" - case "harmony_channel": - # <|channel|> 后面开始收集通道名。 - self._mode = "channel" - self._channel_name = "" - case "harmony_message": - # <|message|>：analysis 进入思维链，commentary 整段丢弃，其余进入正文。 - name = self._channel_name.strip().lower() - if name.startswith("analysis"): - self._channel = "reasoning" - elif name.startswith("commentary"): - self._channel = "drop" - else: - self._channel = "content" - self._mode = "body" - self._channel_name = "" - case "harmony_end": - # 这段消息结束。通道保持不变，下一次 <|message|> 再切换。 - self._mode = "body" - - def _release(self, item: _Held) -> None: - token_id, text, score, metadata = item - match self._mode: - case "channel": - # 通道名本身，如 analysis，不输出。 - self._channel_name += text - case "header": - # 头部里的角色名，如 assistant，不输出。 - return - case _: - if self._in_tool or self._channel == "drop": - # 工具参数或 commentary 正文不输出。 - return - self._save(text=text, score=score, metadata=metadata, token_id=token_id) - - def _save(self, text: str, score: float, metadata: dict[str, Any], token_id: int) -> None: - if not text: - # 未完成的 UTF-8 字节并进后面的字符。特殊 token 的空文本直接丢掉。 - if not self._is_special(metadata=metadata, token_id=token_id): - self._pending_logprob += score - return - logprob = self._pending_logprob + score - self._pending_logprob = 0.0 - token = { - "token": text, - "logprob": logprob, - "bytes": None if "\ufffd" in text else list(text.encode("utf-8")), - "top_logprobs": [ - _candidate_logprob(tokenizer=self.tokenizer, token_id=candidate, logprob=value) + if "logprob_prompt_ids" in metadata: + self._decoder.set_prompt(metadata["logprob_prompt_ids"]) + special = metadata.get("special", token_id in getattr(self.tokenizer, "all_special_ids", ())) + if token_id is not None and score is not None and (text or not special): + rank = metadata.get("logprobs", {}).get(token_id, {}).get("rank") + token = self._decoder.token_logprob(token_id, -9999.0 if rank is not None and rank > 20 else score) + token["top_logprobs"] = [ + self._decoder.token_logprob(candidate, value) for candidate, value in metadata.get("output_top_logprobs", ()) - ], - } - # 分开思考时思维链单独列出；否则和正文共用 content。 - if self._channel == "reasoning" and self._separate_reasoning: - self._reasoning.append(token) - else: - self._content.append(token) + ] + self._hold.append(PendingLogprob(start, self._source_length, token)) + if token_id is not None: + self._decoder.advance(token_id) + return TextSpan(start, self._source_length) - def _is_special(self, metadata: dict[str, Any], token_id: int) -> bool: - if "special" in metadata: - return bool(metadata["special"]) - return token_id in getattr(self.tokenizer, "all_special_ids", ()) + @staticmethod + def _overlaps(item, span): + if item.start == item.end: + return span.start <= item.start < span.end + return item.start < span.end and item.end > span.start + + def assign(self, spans: List[TextSpan], content_length: int) -> List[TextSpan]: + spans = [span for span in spans if span.start < span.end] + visible = [TextSpan(self._content_length, self._content_length + content_length)] if content_length else [] + content_items = [] + content_offset = self._content_length + for span in spans: + selected = [item for item in self._hold if self._overlaps(item, span)] + self._hold = [item for item in self._hold if not self._overlaps(item, span)] + if span.channel == "reasoning" and self._separate_reasoning: + self._reasoning.extend(selected) + else: + for item in selected: + content_items.append( + PendingLogprob( + content_offset + max(item.start - span.start, 0), + content_offset + min(item.end, span.end) - span.start, + item.token, + ) + ) + content_offset += span.end - span.start + if spans: + consumed = max(span.end for span in spans) + self._hold = [item for item in self._hold if item.end > consumed] + self._content_length += content_length + self._content.extend(content_items) + return visible + + def _take_content(self, spans=None): + if spans is None: + return self._drain(self._content) + selected = [item for item in self._content if any(self._overlaps(item, span) for span in spans)] + if spans: + consumed = max(span.end for span in spans) + self._content = [item for item in self._content if item.end > consumed and item not in selected] + return [item.token for item in selected] + + def discard_content(self, spans: List[TextSpan]) -> None: + self._take_content(spans) @staticmethod - def _drain(items: list[dict[str, Any]]) -> list[dict[str, Any]]: - tokens = list(items) + def _drain(items): + tokens = [item.token for item in items] items.clear() return tokens + def reasoning_logprobs(self) -> Optional[Dict[str, Any]]: + tokens = self._drain(self._reasoning) + return {"content": None, "reasoning_content": tokens} if tokens else None -class _Held(NamedTuple): - token_id: int - text: str - score: float - metadata: dict[str, Any] - - -def _candidate_logprob(tokenizer: Any, token_id: int, logprob: float) -> dict[str, Any]: - return chat_token_logprob(tokenizer=tokenizer, token_id=token_id, logprob=logprob) - - -def _delimiter_ids(tokenizer: Any, text: str) -> list[int] | None: - if not text: - return None - encode = getattr(tokenizer, "encode", None) - if encode is not None: - token_ids = list(encode(text, add_special_tokens=False)) - if token_ids and tokenizer.decode(token_ids, skip_special_tokens=False) == text: - return token_ids - return _longest_token_ids(tokenizer=tokenizer, text=text) - + def content_logprobs(self, spans=None) -> Optional[Dict[str, Any]]: + tokens = self._take_content(spans) + return {"content": tokens, "reasoning_content": None} if tokens else None -def _longest_token_ids(tokenizer: Any, text: str) -> list[int] | None: - vocab = tokenizer.get_vocab() - pieces = sorted((piece for piece in vocab if piece), key=len, reverse=True) - token_ids = [] - index = 0 - while index < len(text): - piece = next((item for item in pieces if text.startswith(item, index)), None) - if piece is None: - return None - token_ids.append(vocab[piece]) - index += len(piece) - if tokenizer.decode(token_ids, skip_special_tokens=False) != text: - return None - return token_ids + def message_logprobs( + self, content_spans=None + ) -> Tuple[Optional[List[Dict[str, Any]]], Optional[List[Dict[str, Any]]]]: + return self._take_content(content_spans) or None, self._drain(self._reasoning) or None diff --git a/lightllm/server/core/objs/logprob_utils.py b/lightllm/server/core/objs/logprob_utils.py index 055fb9e213..75f0348fc1 100644 --- a/lightllm/server/core/objs/logprob_utils.py +++ b/lightllm/server/core/objs/logprob_utils.py @@ -1,3 +1,6 @@ +import json +import re + from typing import Any, Dict @@ -20,13 +23,69 @@ def logprob_info(tokenizer, token_id: int, logprob: float, rank: int): ) -def chat_token_logprob(tokenizer: Any, token_id: int, logprob: float) -> Dict[str, Any]: +def _decoder_types(decoder): + if decoder is None: + return set() + config = json.loads(decoder.__getstate__()) + + def types(node): + result = {node.get("type")} + for child in node.get("decoders", []): + result.update(types(child)) + return result + + return types(config) + + +def chat_token_logprob( + tokenizer: Any, token_id: int, logprob: float, prefix_ids=(), pending_ids=(), decode_options=None +) -> Dict[str, Any]: token = tokenizer.decode([token_id], skip_special_tokens=False) + standalone_token = token + special = token_id in getattr(tokenizer, "all_special_ids", ()) + if (prefix_ids or pending_ids) and "\ufffd" not in standalone_token and not special: + options = decode_options or {"skip_special_tokens": False} + prefix = tokenizer.decode([*prefix_ids, *pending_ids], **options) + combined = tokenizer.decode([*prefix_ids, *pending_ids, token_id], **options) + if combined.startswith(prefix): + token = combined[len(prefix) :] raw_bytes = list(token.encode("utf-8")) if "\ufffd" not in token else None decoder = getattr(getattr(tokenizer, "backend_tokenizer", None), "decoder", None) - added_tokens = getattr(tokenizer, "added_tokens_decoder", {}) - if type(decoder).__name__ == "ByteLevel" and token_id not in added_tokens: + decoder_types = _decoder_types(decoder) + if token_id not in getattr(tokenizer, "added_tokens_decoder", {}): piece = tokenizer.convert_ids_to_tokens(token_id) - if piece and all(char in _BYTE_LEVEL_BYTES for char in piece): + if "ByteLevel" in decoder_types and piece and all(char in _BYTE_LEVEL_BYTES for char in piece): raw_bytes = [_BYTE_LEVEL_BYTES[char] for char in piece] + token = standalone_token + elif "ByteFallback" in decoder_types and piece and re.fullmatch(r"<0x[0-9a-fA-F]{2}>", piece): + raw_bytes = [int(piece[3:5], 16)] + token = standalone_token return {"token": token, "logprob": logprob, "bytes": raw_bytes} + + +class ChatTokenDecoder: + def __init__(self, tokenizer, skip_special_tokens=True, spaces_between_special_tokens=True): + self.tokenizer = tokenizer + self.prefix_ids = [] + self.pending_ids = [] + self.options = { + "skip_special_tokens": skip_special_tokens, + "spaces_between_special_tokens": spaces_between_special_tokens, + } + + def set_prompt(self, token_ids): + from lightllm.server.embed_cache.utils import MULTIMODAL_TOKEN_ID_START + + self.prefix_ids = [token_id for token_id in token_ids if token_id < MULTIMODAL_TOKEN_ID_START] + self.pending_ids = [] + + def token_logprob(self, token_id, score): + return chat_token_logprob(self.tokenizer, token_id, score, self.prefix_ids, self.pending_ids, self.options) + + def advance(self, token_id): + prefix = self.tokenizer.decode(self.prefix_ids, **self.options) + self.pending_ids.append(token_id) + combined = self.tokenizer.decode(self.prefix_ids + self.pending_ids, **self.options) + if len(combined) > len(prefix) and not combined.endswith("\ufffd"): + self.prefix_ids = self.pending_ids + self.pending_ids = [] diff --git a/lightllm/server/function_call_parser.py b/lightllm/server/function_call_parser.py index 7140d3209c..d46268bd55 100644 --- a/lightllm/server/function_call_parser.py +++ b/lightllm/server/function_call_parser.py @@ -1995,6 +1995,27 @@ def __init__(self, tools: List[Tool], tool_call_parser: str): self.detector = detector self.tools = tools + def parse_non_stream_with_spans(self, full_text: str): + from lightllm.server.text_spans import select_text_spans + + normal, calls = self.parse_non_stream(full_text) + return normal, calls, select_text_spans(full_text, normal) + + def parse_stream_chunk_with_spans(self, chunk_text: str): + from lightllm.server.text_spans import TextSpan, select_text_spans + + source = self.detector._normal_text_buffer if isinstance(self.detector, Qwen25Detector) else "" + source += self.detector._buffer + chunk_text + end = getattr(self, "_source_length", 0) + len(chunk_text) + self._source_length = end + offset = end - len(source) + normal, calls = self.parse_stream_chunk(chunk_text) + marker = self.detector.eot_token + if marker: + source = source.replace(marker, "\0" * len(marker)) + spans = [TextSpan(s.start + offset, s.end + offset) for s in select_text_spans(source, normal)] + return normal, calls, spans + def has_tool_call(self, text: str) -> bool: """ Check if the given text contains a tool call in the format supported by this parser. diff --git a/lightllm/server/reasoning_parser.py b/lightllm/server/reasoning_parser.py index 9570ce0536..612d3aaa88 100644 --- a/lightllm/server/reasoning_parser.py +++ b/lightllm/server/reasoning_parser.py @@ -12,10 +12,11 @@ # limitations under the License. import re -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Iterator, List, Tuple, Dict, Optional, Type from lightllm.utils.config_utils import get_token_id +from lightllm.server.text_spans import TextSpan, select_text_spans @dataclass @@ -25,6 +26,7 @@ class Event: event_type: str content: str raw_text: str = None # Original text including structural markers + spans: List[TextSpan] = field(default_factory=list) @dataclass @@ -165,7 +167,7 @@ def parse(self, text: str) -> Tuple[List[Event], str]: if pos == len(tokens) - 1: # Last token emit, hold = prefix_hold(text[token.start : token.end], self.guard_tokens) if emit: - events.append(Event("normal", emit)) + events.append(Event("normal", emit, spans=[TextSpan(token.start, token.start + len(emit))])) return events, hold else: # Check if this might be commentary filler between blocks @@ -176,7 +178,7 @@ def parse(self, text: str) -> Tuple[List[Event], str]: content = text[token.start : token.end] # Skip standalone structural tokens that shouldn't be emitted as normal text if not self._is_standalone_structural_token(content): - events.append(Event("normal", content)) + events.append(Event("normal", content, spans=[TextSpan(token.start, token.end)])) pos += 1 elif token.type in ("START", "CHANNEL"): @@ -206,7 +208,7 @@ def parse(self, text: str) -> Tuple[List[Event], str]: # Unexpected token - only emit as text if it's not a standalone structural token content = text[token.start : token.end] if not self._is_standalone_structural_token(content): - events.append(Event("normal", content)) + events.append(Event("normal", content, spans=[TextSpan(token.start, token.end)])) pos += 1 return events, "" @@ -248,7 +250,7 @@ def _parse_partial_analysis(self, text: str, tokens: List[Token], start_pos: int # Return partial reasoning content and preserve the channel structure for next parse remaining_text = text[tokens[start_pos].start : content_start] - return Event("reasoning", content), remaining_text + return Event("reasoning", content, spans=[TextSpan(content_start, len(text), "reasoning")]), remaining_text def _extract_channel_type(self, header_text: str) -> Optional[str]: """Extract channel type from header, ignoring other attributes like to=... or <|constrain|>...""" @@ -298,7 +300,10 @@ def _parse_block(self, text: str, tokens: List[Token], start_pos: int) -> Option if end_token_pos is None: return None # Incomplete content = text[content_start : tokens[end_token_pos].start] - return Event("normal", content), end_token_pos + 1 + return ( + Event("normal", content, spans=[TextSpan(content_start, tokens[end_token_pos].start)]), + end_token_pos + 1, + ) # Standard channel block processing - message_pos is already found above pos = channel_pos + 1 # Skip CHANNEL token @@ -334,7 +339,7 @@ def _parse_block(self, text: str, tokens: List[Token], start_pos: int) -> Option if channel_type == "final": # Final blocks can end at end of input without requiring <|return|> content = text[content_start:] - return Event("normal", content), end_pos + return Event("normal", content, spans=[TextSpan(content_start, len(text))]), end_pos return None # Analysis and commentary need proper end tokens end_token = tokens[end_pos] @@ -345,15 +350,24 @@ def _parse_block(self, text: str, tokens: List[Token], start_pos: int) -> Option if end_token.type == "CALL": # Built-in tools (browser, python) use analysis channel with <|call|> raw_text = text[tokens[start_pos].start : end_token.end] - return Event("tool_call", content.strip(), raw_text), end_pos + 1 + return ( + Event("tool_call", content.strip(), raw_text, [TextSpan(tokens[start_pos].start, end_token.end)]), + end_pos + 1, + ) else: - return Event("reasoning", content), end_pos + 1 + return ( + Event("reasoning", content, spans=[TextSpan(content_start, end_token.start, "reasoning")]), + end_pos + 1, + ) elif channel_type == "commentary": if end_token.type == "CALL": raw_text = text[tokens[start_pos].start : end_token.end] - return Event("tool_call", content.strip(), raw_text), end_pos + 1 + return ( + Event("tool_call", content.strip(), raw_text, [TextSpan(tokens[start_pos].start, end_token.end)]), + end_pos + 1, + ) else: - return Event("normal", content), end_pos + 1 + return Event("normal", content, spans=[TextSpan(content_start, end_token.start)]), end_pos + 1 elif channel_type == "final": # For final blocks, include any trailing TEXT immediately after <|return|> final_content = content @@ -361,8 +375,18 @@ def _parse_block(self, text: str, tokens: List[Token], start_pos: int) -> Option next_token = tokens[end_pos + 1] if next_token.type == "TEXT": final_content += text[next_token.start : next_token.end] - return Event("normal", final_content), end_pos + 2 - return Event("normal", final_content), end_pos + 1 + return ( + Event( + "normal", + final_content, + spans=[ + TextSpan(content_start, end_token.start), + TextSpan(next_token.start, next_token.end), + ], + ), + end_pos + 2, + ) + return Event("normal", final_content, spans=[TextSpan(content_start, end_token.start)]), end_pos + 1 return None, end_pos + 1 @@ -439,11 +463,45 @@ def parse(self, text: str) -> Tuple[List[Event], str]: if m: channel, reasoning, final = m.groups() if channel.lower() == "analysis" and reasoning.strip(): - events.append(Event("reasoning", reasoning.strip())) + events.append( + Event( + "reasoning", + reasoning.strip(), + spans=[ + TextSpan( + m.start(2) + len(reasoning) - len(reasoning.lstrip()), + m.end(2) - len(reasoning) + len(reasoning.rstrip()), + "reasoning", + ) + ], + ) + ) elif channel.lower() == "commentary" and reasoning.strip(): - events.append(Event("normal", reasoning.strip())) + events.append( + Event( + "normal", + reasoning.strip(), + spans=[ + TextSpan( + m.start(2) + len(reasoning) - len(reasoning.lstrip()), + m.end(2) - len(reasoning) + len(reasoning.rstrip()), + ) + ], + ) + ) if final.strip(): - events.append(Event("normal", final.strip())) + events.append( + Event( + "normal", + final.strip(), + spans=[ + TextSpan( + m.start(m.lastindex) + len(final) - len(final.lstrip()), + m.end(m.lastindex) - len(final) + len(final.rstrip()), + ) + ], + ) + ) return events, "" # If assistantfinal appears to be incomplete (e.g., 'assistantfin'), hold entire buffer @@ -456,7 +514,18 @@ def parse(self, text: str) -> Tuple[List[Event], str]: if m: final = m.group(1) if final.strip(): - events.append(Event("normal", final.strip())) + events.append( + Event( + "normal", + final.strip(), + spans=[ + TextSpan( + m.start(m.lastindex) + len(final) - len(final.lstrip()), + m.end(m.lastindex) - len(final) + len(final.rstrip()), + ) + ], + ) + ) return events, "" m = self.patterns["analysis_only"].match(text) @@ -465,7 +534,9 @@ def parse(self, text: str) -> Tuple[List[Event], str]: emit, hold = prefix_hold(content, ["assistantfinal"]) if channel.lower() == "analysis" and emit: # Stream reasoning content as-is based on structural markers only. - events.append(Event("reasoning", emit)) + events.append( + Event("reasoning", emit, spans=[TextSpan(m.start(2), m.start(2) + len(emit), "reasoning")]) + ) # Keep the channel header in the remaining buffer to continue parsing # subsequent chunks in the text fallback format. Preserve any held # prefix that may complete into "assistantfinal". @@ -476,7 +547,7 @@ def parse(self, text: str) -> Tuple[List[Event], str]: elif channel.lower() == "commentary" and emit: # For commentary, stream as normal text. Preserve spaces unless holding. content_out = emit if hold else emit.strip() - events.append(Event("normal", content_out)) + events.append(Event("normal", content_out, spans=select_text_spans(text, content_out))) if hold: return events, text[: m.start(2)] + hold else: @@ -486,7 +557,7 @@ def parse(self, text: str) -> Tuple[List[Event], str]: emit, hold = prefix_hold(text, ["analysis", "commentary", "assistantfinal"]) if emit: - events.append(Event("normal", emit)) + events.append(Event("normal", emit, spans=[TextSpan(0, len(emit))])) return events, hold @@ -494,12 +565,20 @@ class HarmonyParser: """Facade for parsing Harmony format, switching between strategies.""" def __init__(self): + self.track_spans = False + self.last_spans = [] + self._source_offsets = [] + self._source_length = 0 self.strategy = None self._buffer = "" self._should_filter_commentary = False # Track if we should filter commentary in next chunks self._partial_commentary = "" # Track partial commentary being built across chunks def parse(self, chunk: str) -> List[Event]: + if self.track_spans: + self._source_offsets.extend(range(self._source_length, self._source_length + len(chunk))) + self._source_length += len(chunk) + self.last_spans = [] self._buffer += chunk if self.strategy is None: @@ -524,6 +603,18 @@ def parse(self, chunk: str) -> List[Event]: # Check if we should start filtering commentary (after <|call|> token or tool_call event) buffer_has_call_token = self._buffer.rstrip().endswith("<|call|>") + if self.track_spans: + source = self._buffer + offsets = self._source_offsets + start = source.rfind(remaining) + if start >= 0: + self._source_offsets = offsets[start : start + len(remaining)] + else: + prefix = 0 + while prefix < min(len(source), len(remaining)) and source[prefix] == remaining[prefix]: + prefix += 1 + suffix = len(remaining) - prefix + self._source_offsets = offsets[:prefix] + (offsets[-suffix:] if suffix else []) self._buffer = remaining # Filter events for streaming case @@ -565,6 +656,20 @@ def parse(self, chunk: str) -> List[Event]: elif buffer_has_call_token: self._should_filter_commentary = True # Filter commentary after <|call|> token + if self.track_spans: + spans = event.spans + channel = "reasoning" if event.event_type == "reasoning" else "content" + for span in spans: + for pos in range(span.start, span.end): + absolute = offsets[pos] + if ( + self.last_spans + and self.last_spans[-1].end == absolute + and self.last_spans[-1].channel == channel + ): + self.last_spans[-1].end += 1 + else: + self.last_spans.append(TextSpan(absolute, absolute + 1, channel)) filtered_events.append(event) return filtered_events @@ -986,6 +1091,56 @@ def parse_stream_chunk(self, chunk_text: str) -> Tuple[Optional[str], Optional[s ret = self.detector.parse_streaming_increment(chunk_text) return ret.reasoning_text, ret.normal_text + def parse_stream_chunk_with_spans(self, chunk_text: str): + if isinstance(self.detector, GptOssDetector): + self.detector.parser.track_spans = True + reasoning, content = self.parse_stream_chunk(chunk_text) + return reasoning, content, self.detector.parser.last_spans + source = self.detector._buffer + chunk_text + end = getattr(self, "_source_length", 0) + len(chunk_text) + self._source_length = end + offset = end - len(source) + markers = [] + if not self.detector.stripped_think_start and self.detector.think_start_token in source: + markers.append(self.detector.think_start_token) + if self.detector._in_reasoning or markers: + markers.append(self.detector.think_end_token) + reasoning, content = self.parse_stream_chunk(chunk_text) + spans = self._visible_spans(source, reasoning, content, offset, markers) + return reasoning, content, spans + + def _visible_spans(self, source, reasoning, content, offset, markers=()): + visible = source + positions = list(range(offset, offset + len(source))) + for marker in markers: + if not marker: + continue + start = visible.find(marker) + while start >= 0: + visible = visible[:start] + visible[start + len(marker) :] + del positions[start : start + len(marker)] + start = visible.find(marker) + spans = [] + cursor = 0 + for text, channel in ((reasoning, "reasoning"), (content, "content")): + next_cursor = cursor + for span in select_text_spans(visible[cursor:], text or "", channel): + start, end = span.start + cursor, span.end + cursor + next_cursor = end + for pos in positions[start:end]: + if spans and spans[-1].end == pos and spans[-1].channel == channel: + spans[-1].end += 1 + else: + spans.append(TextSpan(pos, pos + 1, channel)) + cursor = next_cursor + return spans + + def flush_with_spans(self): + source = self.detector._buffer + reasoning, content = self.flush() + offset = getattr(self, "_source_length", 0) - len(source) + return reasoning, content, self._visible_spans(source, reasoning, content, offset) + def flush(self) -> Tuple[Optional[str], Optional[str]]: """Flush remaining buffered content when generation ends prematurely.""" ret = self.detector.flush() 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 208fc46fbc..9a64e36af9 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -939,7 +939,9 @@ def _post_handle( pack: InferReqUpdatePack = pack output_top_logprobs = None if topk_row is not None: - output_top_logprobs = (topk_row.token_ids.tolist(), topk_row.logprobs.tolist()) + next_token_rank = int(topk_row.sampled_rank.item()) + if topk_row.token_ids.numel(): + output_top_logprobs = (topk_row.token_ids.tolist(), topk_row.logprobs.tolist()) pack.handle( next_token_id=next_token_id, next_token_logprob=next_token_logprob, diff --git a/lightllm/server/router/model_infer/mode_backend/generic_post_process.py b/lightllm/server/router/model_infer/mode_backend/generic_post_process.py index e3e01da612..d48e092e26 100644 --- a/lightllm/server/router/model_infer/mode_backend/generic_post_process.py +++ b/lightllm/server/router/model_infer/mode_backend/generic_post_process.py @@ -17,11 +17,13 @@ class TopkLogprobRow: """一行 top-k 候选。token_ids 与 logprobs 等长，按分数从高到低对齐。 batch_index 是该行在本次采样 logits 中的下标。返回列表与 logits 行对齐，未请求 logprobs 的位置是 None。 + K=0 的行只有 sampled_rank，候选数组为空。 """ batch_index: int token_ids: torch.Tensor logprobs: torch.Tensor + sampled_rank: torch.Tensor TopkLogprobRows = List[Optional[TopkLogprobRow]] @@ -143,14 +145,13 @@ def _collect_topk_logprobs_rows( ) -> Tuple[List[Tuple[int, int]], int]: """Rows that asked for candidate topk logprobs, as (batch_index, topk), plus the maximum topk. - topk == 0 still runs the fixed top-20 cutoff on the sampled logprob and writes no candidates. - batch_index is the position in the original reqs and logits. max_topk is 0 when the list is empty. + topk == 0 still needs the sampled token rank. batch_index is the position + in the original reqs and logits. max_topk is 0 when the list is empty. """ rows = [] max_topk = 0 for batch_index, req_obj in enumerate(reqs): topk = req_obj.sampling_param.shm_param.output_top_logprobs - # topk == 0 也要进来：没有候选列表，但仍要把落在前 20 名之外的采样 logprob 写成 -9999。 if topk >= 0 and (has_output is None or has_output[batch_index]): rows.append((batch_index, topk)) max_topk = max(max_topk, topk) @@ -172,7 +173,7 @@ def _capture_topk_logprobs( rows = g_pin_mem_manager.alloc_pin_tensor("output_logprob_rows", 3 * len(capture_rows), torch.int32).view(-1, 3) rows.numpy()[:] = capture_rows - token_ids, logprobs = capture_output_logprobs( + token_ids, logprobs, ranks = capture_output_logprobs( logits=logits, sampled_ids=next_token_ids, sampled_logprobs=next_token_logprobs, @@ -182,16 +183,16 @@ def _capture_topk_logprobs( ) token_ids_cpu = g_pin_mem_manager.async_copy_from_gpu_tensor(key="output_topk_token_ids", gpu_tensor=token_ids) logprobs_cpu = g_pin_mem_manager.async_copy_from_gpu_tensor(key="output_topk_logprobs", gpu_tensor=logprobs) + ranks_cpu = g_pin_mem_manager.async_copy_from_gpu_tensor(key="output_sampled_ranks", gpu_tensor=ranks) topk_rows: TopkLogprobRows = [None] * logits.shape[0] offset = 0 - for batch_index, topk in topk_logprobs_rows: - # topk == 0 只参与上面的采样分数改写，不返回候选。 - if topk > 0: - topk_rows[batch_index] = TopkLogprobRow( - batch_index=batch_index, - token_ids=token_ids_cpu[offset : offset + topk], - logprobs=logprobs_cpu[offset : offset + topk], - ) + for index, (batch_index, topk) in enumerate(topk_logprobs_rows): + topk_rows[batch_index] = TopkLogprobRow( + batch_index=batch_index, + token_ids=token_ids_cpu[offset : offset + topk], + logprobs=logprobs_cpu[offset : offset + topk], + sampled_rank=ranks_cpu[index : index + 1], + ) offset += topk return next_token_logprobs, topk_rows diff --git a/lightllm/server/text_spans.py b/lightllm/server/text_spans.py new file mode 100644 index 0000000000..2a58eb4f16 --- /dev/null +++ b/lightllm/server/text_spans.py @@ -0,0 +1,29 @@ +from dataclasses import dataclass +from typing import List + + +@dataclass +class TextSpan: + start: int + end: int + channel: str = "content" + + +def select_text_spans(source: str, text: str, channel: str = "content") -> List[TextSpan]: + spans = [] + cursor = 0 + while text: + start = source.find(text, cursor) + if start >= 0: + spans.append(TextSpan(start, start + len(text), channel)) + break + start = source.find(text[0], cursor) + if start < 0: + break + length = 1 + while length < len(text) and start + length < len(source) and source[start + length] == text[length]: + length += 1 + spans.append(TextSpan(start, start + length, channel)) + cursor = start + length + text = text[length:] + return spans diff --git a/unit_tests/server/router/model_infer/structured_output/test_sampling.py b/unit_tests/server/router/model_infer/structured_output/test_sampling.py index bad4430c71..4bf33aaec4 100644 --- a/unit_tests/server/router/model_infer/structured_output/test_sampling.py +++ b/unit_tests/server/router/model_infer/structured_output/test_sampling.py @@ -96,8 +96,6 @@ def test_sample_returns_requested_output_logprobs(compiler, sampling_manager, mo reference = logits.log_softmax(dim=-1) sampled_id = int(token_ids[0]) expected = reference[0, sampled_id] - if int((logits[0] > logits[0, sampled_id]).sum()) >= 20: - expected = expected.new_tensor(-9999.0) torch.testing.assert_close(logprobs[0], expected) if topk: assert len(output_topk_logprobs) == 1 @@ -106,7 +104,8 @@ def test_sample_returns_requested_output_logprobs(compiler, sampling_manager, mo assert row.token_ids.tolist() == list(range(req.vocab_size - 1, req.vocab_size - 1 - topk, -1)) torch.testing.assert_close(row.logprobs, reference[0, row.token_ids].cpu()) else: - assert output_topk_logprobs == [None] + assert output_topk_logprobs[0].token_ids.numel() == 0 + assert output_topk_logprobs[0].sampled_rank.item() == req.vocab_size - sampled_id def test_sample_candidate_logprobs_follow_mtp_row_selection(compiler, sampling_manager, monkeypatch): @@ -149,7 +148,9 @@ def test_sample_candidate_logprobs_follow_mtp_row_selection(compiler, sampling_m assert middle["next_token_logprob"] == pytest.approx(reference[2, 10].item()) assert first["output_top_logprobs"][1] == pytest.approx(reference[0, [256, 255, 254]].tolist()) assert last["output_top_logprobs"][1] == pytest.approx(reference[3, [256, 255]].tolist()) - assert logprobs_cpu[1].item() == -9999.0 + assert logprobs_cpu[1].item() == pytest.approx(reference[1, 0].item()) + assert first["next_token_rank"] == 1 + assert last["next_token_rank"] == 1 @pytest.mark.parametrize("topk", [0, 1, 3, 20]) @@ -170,8 +171,6 @@ def test_sample_reuses_sampled_normalization(compiler, sampling_manager, monkeyp topk_by_index = output_topk_logprobs for row in range(4): expected = reference[row, ids[row]] - if row in rows and int((logits[row] > logits[row, ids[row]]).sum()) >= 20: - expected = expected.new_tensor(-9999.0) torch.testing.assert_close(logprobs[row], expected) if topk and row in rows: top_ids, top_scores = topk_by_index[row].token_ids, topk_by_index[row].logprobs @@ -179,7 +178,7 @@ def test_sample_reuses_sampled_normalization(compiler, sampling_manager, monkeyp torch.testing.assert_close(top_scores, reference[row, top_ids].cpu()) -def test_sample_preserves_sampled_score_and_keeps_ties(compiler, sampling_manager, monkeypatch): +def test_sample_rank_keeps_ties_and_includes_zero_width_candidates(compiler, sampling_manager, monkeypatch): reqs = [sampling_req(compiler, sampling_manager, req_idx=i, regular_constraint=None) for i in range(4)] for req, topk in zip(reqs, [0, 0, 3, -1]): req.sampling_param.shm_param.top_k = req.vocab_size @@ -190,12 +189,10 @@ def test_sample_preserves_sampled_score_and_keeps_ties(compiler, sampling_manage ids = torch.full((4,), 256, dtype=torch.int64, device="cuda") monkeypatch.setattr(sampling, "_random_sample", lambda *args: ids) _, logprobs, output_topk_logprobs = sampling.sample(logits, reqs, [256]) - assert [None if row is None else row.batch_index for row in output_topk_logprobs] == [None, None, 2, None] + assert [None if row is None else row.batch_index for row in output_topk_logprobs] == [0, 1, 2, None] + assert [row.sampled_rank.item() for row in output_topk_logprobs[:3]] == [20, 21, 21] reference = logits.log_softmax(dim=-1) - expected = reference[:, 256].clone() - expected[1] = -9999.0 - expected[2] = -9999.0 - torch.testing.assert_close(logprobs, expected) + torch.testing.assert_close(logprobs, reference[:, 256]) candidate_ids, candidate_logprobs = output_topk_logprobs[2].token_ids, output_topk_logprobs[2].logprobs assert len(set(candidate_ids.tolist())) == 3 and (candidate_ids < 20).all() torch.testing.assert_close(candidate_logprobs, reference[2, candidate_ids].cpu()) @@ -442,16 +439,31 @@ def copy_snapshot_and_wait(manager): completed.synchronize() +@pytest.mark.parametrize("topk", [0, 1, 3, 20]) +@pytest.mark.parametrize("sampled_rank", [2, 20, 21]) +def test_sample_rank_is_independent_of_candidate_count(compiler, sampling_manager, monkeypatch, topk, sampled_rank): + req = sampling_req(compiler, sampling_manager, regular_constraint=None) + req.sampling_param.shm_param.top_k = req.vocab_size + req.sampling_param.shm_param.output_top_logprobs = topk + logits = torch.linspace(3.0, 0.1, req.vocab_size, device="cuda").view(1, -1) + ids = torch.tensor([sampled_rank - 1], device="cuda") + monkeypatch.setattr(sampling, "_random_sample", lambda *args: ids) + _, logprobs, rows = sampling.sample(logits, [req], [256]) + torch.testing.assert_close(logprobs[0], logits.log_softmax(-1)[0, ids[0]]) + assert rows[0].sampled_rank.item() == sampled_rank + assert rows[0].token_ids.numel() == topk + + @pytest.mark.parametrize("vocab_size", [30, 2051, 32769]) -def test_capture_across_vocab_tiles_preserves_sampled_scores(vocab_size): +def test_capture_rank_across_vocab_tiles_preserves_sampled_scores(vocab_size): logits = torch.linspace(3.0, 0.1, vocab_size, device="cuda").repeat(4, 1) ids = torch.tensor([1, 19, 20, vocab_size - 1], device="cuda") logprobs = logits.log_softmax(-1).gather(1, ids[:, None]).squeeze(1) original = logprobs.clone() rows = torch.tensor([[0, 0, 0], [1, 1, 0], [2, 3, 1], [3, 20, 4]], device="cuda", dtype=torch.int32) - candidate_ids, scores = sampling.capture_output_logprobs(logits, ids, logprobs, rows, 20, 24) - torch.testing.assert_close(logprobs[:2], original[:2], rtol=0, atol=0) - assert logprobs[2].item() == -9999.0 and logprobs[3].item() == -9999.0 + candidate_ids, scores, ranks = sampling.capture_output_logprobs(logits, ids, logprobs, rows, 20, 24) + torch.testing.assert_close(logprobs, original, rtol=0, atol=0) + assert ranks.tolist() == [2, 20, 21, vocab_size] for row, width, offset in rows.tolist(): assert candidate_ids[offset : offset + width].tolist() == list(range(width)) torch.testing.assert_close(scores[offset : offset + width], logits.log_softmax(-1)[row, :width]) diff --git a/unit_tests/server/test_chat_output_logprobs.py b/unit_tests/server/test_chat_output_logprobs.py index 18d4e64b11..9129dc7b59 100644 --- a/unit_tests/server/test_chat_output_logprobs.py +++ b/unit_tests/server/test_chat_output_logprobs.py @@ -7,7 +7,7 @@ import pytest from pydantic import ValidationError -from tokenizers import Tokenizer, decoders, models, pre_tokenizers +from tokenizers import Tokenizer, decoders, models from transformers import PreTrainedTokenizerFast from lightllm.server import api_http, api_openai, build_prompt as prompt_builder, reasoning_parser @@ -107,9 +107,8 @@ async def generate(prompt, params, *args, **kwargs): assert [entry["token"] for entry in reasoning] == (["think"] if separate_reasoning else []) expected_content = "你 yes" if separate_reasoning else "think你 yes" assert bytes(byte for entry in content for byte in entry["bytes"]).decode() == expected_content - assert [entry["token"] for entry in content] == (["你", " yes"] if separate_reasoning else ["think", "你", " yes"]) - assert [entry["logprob"] for entry in content] == ([-0.75, -0.25] if separate_reasoning else [-0.25, -0.75, -0.25]) - assert [entry["logprob"] for entry in reasoning] == ([-0.25] if separate_reasoning else []) + assert len(content) == (4 if separate_reasoning else 5) + assert all(entry["logprob"] == -0.25 for entry in content + reasoning) expected = [{"token": " yes", "logprob": -0.125, "bytes": [32, 121, 101, 115]}] if topk > 0 else [] assert all(entry["top_logprobs"] == expected for entry in content + reasoning) @@ -166,11 +165,21 @@ async def test_pd_first_token_uses_the_node_that_sampled_it(monkeypatch, needs_p params.output_top_logprobs = 2 params.group_request_id = 0 finish = FinishStatus() - d_first = (0, "first", {"count_output_tokens": 1, "node_mode": "decode", "output_top_logprobs": []}, finish) + d_first = ( + 0, + "first", + {"count_output_tokens": 1, "node_mode": "decode", "output_top_logprobs": [], "logprob_prompt_ids": prompt_ids}, + finish, + ) p_first = ( 0, "first", - {"count_output_tokens": 1, "node_mode": "prefill", "output_top_logprobs": [(1, -0.5)]}, + { + "count_output_tokens": 1, + "node_mode": "prefill", + "output_top_logprobs": [(1, -0.5)], + "logprob_prompt_ids": prompt_ids, + }, finish, ) d_next = (0, "second", {"count_output_tokens": 2}, FinishStatus(end_status)) @@ -198,6 +207,7 @@ async def test_pd_first_token_uses_the_node_that_sampled_it(monkeypatch, needs_p try: _, text, metadata, _ = await stream.__anext__() assert text == "first" + assert metadata["logprob_prompt_ids"] == prompt_ids assert metadata["output_top_logprobs"] == ([(1, -0.5)] if needs_prefill else []) _, text, _, status = await stream.__anext__() assert text == "second" and status.status == end_status @@ -207,7 +217,23 @@ async def test_pd_first_token_uses_the_node_that_sampled_it(monkeypatch, needs_p @pytest.fixture def chat_response(monkeypatch): - async def run(tokenizer, pieces, *, stream, separate_reasoning=True, model_type=None, tools=None, topk=0, rank=2): + async def run( + tokenizer, + pieces, + *, + stream, + separate_reasoning=True, + model_type=None, + tools=None, + topk=0, + rank=2, + tool_parser="qwen25", + scores=None, + first_metadata=None, + candidates=None, + logprobs=True, + return_choices=False, + ): async def build_prompt(*args): return "prompt" @@ -215,10 +241,11 @@ async def generate(prompt, params, *args, **kwargs): for index, (token_id, delta) in enumerate(pieces): yield 0, delta, { "id": token_id, - "logprob": -2.401099, + **(first_metadata or {} if index == 0 else {}), + "logprob": scores[index] if scores is not None else -2.401099, "logprobs": {token_id: {"rank": rank}}, "prompt_tokens": 10, - "output_top_logprobs": [(0, -0.125)] * topk, + "output_top_logprobs": candidates if candidates is not None else [(0, -0.125)] * topk, }, FinishStatus(FinishStatus.FINISHED_LENGTH if index == len(pieces) - 1 else FinishStatus.NO_FINISH) monkeypatch.setattr(api_openai, "build_prompt", build_prompt) @@ -235,13 +262,13 @@ async def generate(prompt, params, *args, **kwargs): ), raising=False, ) - monkeypatch.setattr(api_http.g_objs, "args", SimpleNamespace(tool_call_parser="qwen25"), raising=False) + monkeypatch.setattr(api_http.g_objs, "args", SimpleNamespace(tool_call_parser=tool_parser), raising=False) request = ChatCompletionRequest( messages=[], stream=stream, separate_reasoning=separate_reasoning, - logprobs=True, - top_logprobs=topk, + logprobs=logprobs, + top_logprobs=topk if logprobs else None, tools=tools, ) response = await api_openai.chat_completions_impl(request, SimpleNamespace()) @@ -255,6 +282,8 @@ async def generate(prompt, params, *args, **kwargs): choices.extend(json.loads(line[6:])["choices"]) else: choices = response.model_dump()["choices"] + if return_choices: + return choices messages = [choice["delta" if stream else "message"] for choice in choices] logprobs = [choice["logprobs"] for choice in choices if choice["logprobs"]] return ( @@ -278,22 +307,26 @@ def make_tokenizer(pieces, decoder=None): @pytest.mark.parametrize("stream", [False, True]) @pytest.mark.parametrize("topk", [0, 1, 3, 20]) @pytest.mark.parametrize("rank", [2, 20, 21]) -async def test_chat_keeps_the_sampled_logprob(chat_response, stream, topk, rank): +async def test_chat_sampled_score_uses_fixed_rank_cutoff(chat_response, stream, topk, rank): tokenizer = make_tokenizer(["answer"]) _, _, content, _ = await chat_response(tokenizer, [(0, "answer")], stream=stream, topk=topk, rank=rank) - assert content[0]["logprob"] == -2.401099 + assert content[0]["logprob"] == (-9999.0 if rank > 20 else -2.401099) assert len(content[0]["top_logprobs"]) == topk @pytest.mark.anyio @pytest.mark.parametrize("stream", [False, True]) -async def test_empty_increments_fuse_into_the_following_text(chat_response, stream): - tokenizer = make_tokenizer(["<0xE4>", "<0xBD>", "<0xA0>"]) +@pytest.mark.parametrize("sequence", [False, True]) +async def test_byte_fallback_keeps_every_generated_token(chat_response, stream, sequence): + decoder = decoders.ByteFallback() + if sequence: + decoder = decoders.Sequence([decoder, decoders.Fuse()]) + tokenizer = make_tokenizer(["<0xE4>", "<0xBD>", "<0xA0>"], decoder) text, _, content, _ = await chat_response(tokenizer, [(0, ""), (1, ""), (2, "你")], stream=stream) assert text == "你" - assert [entry["token"] for entry in content] == ["你"] - assert content[0]["logprob"] == pytest.approx(-2.401099 * 3) - assert content[0]["bytes"] == list("你".encode()) + assert len(content) == 3 + assert [entry["bytes"] for entry in content] == [[0xE4], [0xBD], [0xA0]] + assert bytes(byte for entry in content for byte in entry["bytes"]).decode() == text @pytest.mark.anyio @@ -371,10 +404,9 @@ async def test_harmony_headers_are_not_visible_logprobs(chat_response, stream, s if split_header: pieces[5:6] = ["<|sta", "rt|>"] tokenizer = make_tokenizer(list(dict.fromkeys(pieces)), decoders.Fuse()) - generated = [(tokenizer.convert_tokens_to_ids(piece), piece) for piece in pieces] text, reasoning_text, content, reasoning = await chat_response( tokenizer, - generated, + [(tokenizer.convert_tokens_to_ids(piece), piece) for piece in pieces], stream=stream, separate_reasoning=separate_reasoning, model_type="gpt-oss", @@ -416,153 +448,439 @@ async def test_harmony_repeated_header_words_are_real_content(chat_response, str @pytest.mark.anyio -async def test_tool_delimiter_fused_into_another_token_is_rejected(chat_response): - pieces = ["Checking.", '{"name":"weather","arguments":{"city":"Paris"}}', ""] +@pytest.mark.parametrize("valid_tool", [False, True]) +async def test_nonstream_tool_logprobs_follow_actual_extraction(chat_response, valid_tool): + pieces = ["Checking.", '{"name":"weather","arguments":{"city":"Paris"}}'] + if valid_tool: + pieces += [""] tokenizer = make_tokenizer(pieces, decoders.Fuse()) tools = [{"type": "function", "function": {"name": "weather", "parameters": {"type": "object"}}}] - with pytest.raises(ValueError, match="token sequence"): - await chat_response(tokenizer, list(enumerate(pieces)), stream=False, tools=tools) + text, _, content, _ = await chat_response(tokenizer, list(enumerate(pieces)), stream=False, tools=tools) + assert text == ("Checking." if valid_tool else "".join(pieces)) + assert [entry["token"] for entry in content] == (pieces[:1] if valid_tool else pieces) -def test_empty_increments_fuse_and_trailing_ones_are_dropped(): +def test_unknown_token_bytes_do_not_drop_empty_increments(): from lightllm.server.chat_logprob_buffer import ChatLogprobBuffer - tokenizer = SimpleNamespace(decode=lambda ids, **kwargs: "") - buffer = ChatLogprobBuffer(tokenizer, None, False, True) + tokenizer = SimpleNamespace( + decode=lambda ids, **kwargs: "\ufffd" if len(ids) == 1 else "你", + convert_ids_to_tokens=lambda token_id: "unknown-byte", + ) + buffer = ChatLogprobBuffer(tokenizer, separate_reasoning=True) for token_id, text in enumerate(["", "", "你"]): - buffer.record(text, {"id": token_id, "logprob": -0.25}, FinishStatus()) + span = buffer.record(text, {"id": token_id, "logprob": -0.25}, FinishStatus()) + buffer.assign([span], len(text)) content, _ = buffer.message_logprobs() - assert [token["token"] for token in content] == ["你"] - assert content[0]["logprob"] == -0.75 - assert content[0]["bytes"] == list("你".encode()) + assert len(content) == 3 + assert all(token["bytes"] is None for token in content) - buffer.record("", {"id": 3, "logprob": -0.5}, FinishStatus(FinishStatus.FINISHED_STOP)) - assert buffer.message_logprobs() == (None, None) - -def test_reasoning_delimiter_must_be_one_token(): +def test_combined_reasoning_and_content_token_is_only_returned_once(monkeypatch): from lightllm.server.chat_logprob_buffer import ChatLogprobBuffer - tokenizer = make_tokenizer(["", "ThinkingAnswer"], decoders.Fuse()) - with pytest.raises(ValueError, match="token sequence"): - ChatLogprobBuffer(tokenizer, "qwen3", False, True) + pieces = ["", "ThinkingAnswer"] + tokenizer = make_tokenizer(pieces, decoders.Fuse()) + monkeypatch.setattr(reasoning_parser, "get_token_id", tokenizer.convert_tokens_to_ids) + buffer = ChatLogprobBuffer(tokenizer, separate_reasoning=True) + parser = reasoning_parser.ReasoningParser("qwen3", force_reasoning=False) + for token_id, text in enumerate(pieces): + buffer.record(text, {"id": token_id, "logprob": -0.25}, FinishStatus()) + _, content, spans = parser.parse_stream_chunk_with_spans(text) + buffer.assign(spans, len(content)) + content, reasoning = buffer.message_logprobs() + assert content is None + assert [token["token"] for token in reasoning] == pieces[1:] + +def test_reasoning_span_api_preserves_literal_delimiters_and_flush(monkeypatch): + monkeypatch.setattr(reasoning_parser, "get_token_id", lambda token: 0) + parser = reasoning_parser.ReasoningParser("qwen3", force_reasoning=False) + text = "literal " + reasoning, content, spans = parser.parse_stream_chunk_with_spans(text) + assert not reasoning and content == text + assert "".join(text[span.start : span.end] for span in spans) == text + parser = reasoning_parser.ReasoningParser("qwen3", force_reasoning=True) + assert parser.parse_stream_chunk_with_spans("{"name":"weather","arguments":{"city":"', + "Paris", + '"}};', + "Paris", + ] + tokenizer = make_tokenizer(list(dict.fromkeys(pieces)), decoders.Fuse()) + scores = [-0.1 * (index + 1) for index in range(len(pieces))] + tools = [{"type": "function", "function": {"name": "weather", "parameters": {"type": "object"}}}] + text, _, content, _ = await chat_response( + tokenizer, + [(tokenizer.convert_tokens_to_ids(piece), piece) for piece in pieces], + stream=False, + tools=tools, + tool_parser="llama3", + scores=scores, + ) + assert text == prefix + "Paris" + assert [item["token"] for item in content] == ([prefix] if prefix else []) + ["Paris"] + assert [item["logprob"] for item in content] == ([scores[0]] if prefix else []) + [scores[-1]] -def test_multitoken_reasoning_delimiter_includes_newline(): - from lightllm.server.chat_logprob_buffer import ChatLogprobBuffer - pieces = ["<|channel>", "thought", "\n", "secret", "", "answer"] +@pytest.mark.anyio +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("separate_reasoning", [False, True]) +async def test_nonstream_reasoning_uses_message_trimming(chat_response, stream, separate_reasoning): + pieces = ["Thinking", "", "\n\n", "Answer"] tokenizer = make_tokenizer(pieces, decoders.Fuse()) - buffer = ChatLogprobBuffer(tokenizer, "gemma4", False, True) - for piece in pieces: - buffer.record( - piece, - {"id": tokenizer.convert_tokens_to_ids(piece), "logprob": -0.2}, - FinishStatus(FinishStatus.FINISHED_STOP if piece == pieces[-1] else FinishStatus.NO_FINISH), + text, reasoning_text, content, reasoning = await chat_response( + tokenizer, + list(enumerate(pieces)), + stream=stream, + separate_reasoning=separate_reasoning, + model_type="qwen3", + scores=[-0.1, -0.2, -0.3, -0.4], + ) + expected = (["Thinking"] if not separate_reasoning else []) + (["\n\n"] if stream else []) + ["Answer"] + assert text == "".join(expected) + assert [item["token"] for item in content] == expected + assert [item["logprob"] for item in content] == ( + ([-0.1] if not separate_reasoning else []) + ([-0.3] if stream else []) + [-0.4] + ) + assert bytes(byte for item in content for byte in item["bytes"]).decode() == text + assert [item["token"] for item in reasoning] == (["Thinking"] if separate_reasoning else []) + + +async def first_token_metadata(prompt_ids, output_index=1): + import asyncio + from unittest.mock import Mock + from lightllm.server.httpserver.manager import HttpServerManager + + manager = HttpServerManager.__new__(HttpServerManager) + manager.is_multinode_tp_slave = False + manager.first_time_costs = Mock() + manager.latest_success_infer_time_mark = Mock() + status = SimpleNamespace( + event=asyncio.Event(), + lock=asyncio.Lock(), + has_timed_out_waiting_for_inference=lambda: False, + out_token_info_list=[(0, "", {"count_output_tokens": output_index}, FinishStatus())], + ) + status.event.set() + params = SamplingParams() + params.best_of = 1 + params.output_top_logprobs = 1 + generator = manager._wait_to_token_package(0, prompt_ids, 0, params, status, None) + try: + _, _, metadata, _ = await generator.__anext__() + return metadata + finally: + await generator.aclose() + + +@pytest.mark.anyio +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("sequence", [False, True]) +async def test_prompt_generation_sentencepiece_boundary(chat_response, stream, sequence): + import numpy as np + from lightllm.server.detokenization.decode import decode_token + from lightllm.server.detokenization.decode_req import DecodeReq + from lightllm.server.embed_cache.utils import MULTIMODAL_TOKEN_ID_START + + decoder = ( + decoders.Sequence( + [ + decoders.Replace("▁", " "), + decoders.ByteFallback(), + decoders.Fuse(), + decoders.Strip(" ", 1, 0), + ] ) - content, reasoning = buffer.message_logprobs() - assert [token["token"] for token in reasoning] == ["secret"] - assert [token["token"] for token in content] == ["answer"] + if sequence + else decoders.Metaspace() + ) + tokenizer = make_tokenizer(["▁Hello", "▁world", "▁there"], decoder) + prompt_ids = [0] * 6 + [MULTIMODAL_TOKEN_ID_START] + req = SimpleNamespace( + request_id=0, + group_req_id=0, + input_len=len(prompt_ids), + shm_prompt_ids=SimpleNamespace(arr=np.array(prompt_ids + [1])), + sample_params=SimpleNamespace( + print_eos_token=False, skip_special_tokens=True, add_spaces_between_special_tokens=True + ), + ) + decode_req = DecodeReq(req) + decode_req.output_ids.append(1) + delta = decode_token(tokenizer, decode_req, 1, []) + metadata = await first_token_metadata(prompt_ids) + assert metadata["logprob_prompt_ids"] == prompt_ids[-5:] + assert "logprob_prompt_ids" not in await first_token_metadata(prompt_ids, output_index=2) + text, _, content, _ = await chat_response( + tokenizer, + [(1, delta)], + stream=stream, + first_metadata=metadata, + topk=1, + candidates=[(2, -0.2)], + ) + assert text == delta == " world" + assert content[0]["token"] == " world" + assert bytes(content[0]["bytes"]).decode() == text + candidate = content[0]["top_logprobs"][0] + assert candidate["token"] == " there" + assert bytes(candidate["bytes"]).decode() == " there" -def test_gemma4_answer_without_thought_header_stays_in_content(): +@pytest.mark.anyio +@pytest.mark.parametrize("parser_name", ["glm47", "qwen3_coder"]) +@pytest.mark.parametrize("chunk_size", [1, 7, 4096]) +async def test_stream_tool_body_uses_generated_position(parser_name, chunk_size): from lightllm.server.chat_logprob_buffer import ChatLogprobBuffer + from lightllm.server.function_call_parser import FunctionCallParser + from lightllm.server.api_models import Tool - pieces = ["<|channel>", "thought", "\n", "", "2"] + call = ( + "weathercityParis" + if parser_name == "glm47" + else "Paris" + ) + pieces = [call, "Paris"] tokenizer = make_tokenizer(pieces, decoders.Fuse()) - buffer = ChatLogprobBuffer(tokenizer, "gemma4", True, True) - buffer.record("2", {"id": tokenizer.convert_tokens_to_ids("2"), "logprob": -0.2}, FinishStatus()) - content, reasoning = buffer.message_logprobs() - assert reasoning is None - assert [token["token"] for token in content] == ["2"] + tools = [ + Tool( + type="function", + function={"name": "weather", "parameters": {"type": "object", "properties": {"city": {"type": "string"}}}}, + ) + ] + parser = FunctionCallParser(tools, parser_name) + raw = "".join(pieces) + normal_text, all_calls, positions = "", [], [] + for start in range(0, len(raw), chunk_size): + normal, calls, spans = parser.parse_stream_chunk_with_spans(raw[start : start + chunk_size]) + normal_text += normal + all_calls.extend(calls) + positions.extend(pos for span in spans for pos in range(span.start, span.end)) + assert normal_text == "Paris" and all_calls + assert positions == list(range(len(call), len(raw))) + buffer = ChatLogprobBuffer(tokenizer, separate_reasoning=True) + parser = FunctionCallParser(tools, parser_name) + for index, piece in enumerate(pieces): + raw_span = buffer.record(piece, {"id": index, "logprob": -0.1 * (index + 1)}, FinishStatus()) + buffer.assign([raw_span], len(piece)) + _, _, spans = parser.parse_stream_chunk_with_spans(piece) + content = buffer.content_logprobs(spans)["content"] + assert [(item["token"], item["logprob"]) for item in content] == [("Paris", -0.2)] -def test_incomplete_delimiter_prefix_is_visible_text(): - from lightllm.server.chat_logprob_buffer import ChatLogprobBuffer +@pytest.mark.anyio +@pytest.mark.parametrize("text", ['{"name":"weather","arguments":{"city":"Paris"}}', '{"name":123}']) +@pytest.mark.parametrize("tool_parser", [None, "llama3"]) +async def test_plain_json_with_tools_keeps_scores(chat_response, text, tool_parser): + tokenizer = make_tokenizer([text], decoders.Fuse()) + tools = [{"type": "function", "function": {"name": "weather", "parameters": {"type": "object"}}}] + for enabled in (False, True): + choices = await chat_response( + tokenizer, + [(0, text)], + stream=False, + tools=tools, + tool_parser=tool_parser, + scores=[-0.6], + logprobs=enabled, + return_choices=True, + ) + assert choices[0]["message"]["content"] == text + assert choices[0]["message"]["tool_calls"] is None + if enabled: + assert [(item["token"], item["logprob"]) for item in choices[0]["logprobs"]["content"]] == [(text, -0.6)] + else: + assert choices[0]["logprobs"] is None + - pieces = ["", "", "nope"] +@pytest.mark.anyio +@pytest.mark.parametrize("tool_parser", ["glm47", "qwen3_coder"]) +async def test_suppressed_tool_whitespace_has_no_scores(chat_response, tool_parser): + call = ( + "weathercityParis" + if tool_parser == "glm47" + else "Paris" + ) + pieces = [call, "\n", "Done"] tokenizer = make_tokenizer(pieces, decoders.Fuse()) - buffer = ChatLogprobBuffer(tokenizer, "qwen3", False, True) - for piece in ["", "ok"], decoders.Fuse()) - buffer = ChatLogprobBuffer( - tokenizer, - None, - False, - True, - tools=[{"type": "function", "function": {"name": "ok"}}], - tool_call_parser="llama3", + tokenizer = make_tokenizer( + ["▁Hello", "<0xE4>", "<0xBD>", "<0xA0>", "▁world"], + decoders.Sequence( + [ + decoders.Replace("▁", " "), + decoders.ByteFallback(), + decoders.Fuse(), + decoders.Strip(" ", 1, 0), + ] + ), + ) + req = SimpleNamespace( + request_id=0, + group_req_id=0, + input_len=1, + shm_prompt_ids=SimpleNamespace(arr=np.array([0, 1, 2, 3])), + sample_params=SimpleNamespace( + print_eos_token=False, skip_special_tokens=True, add_spaces_between_special_tokens=True + ), ) - assert buffer._sequences + decode_req = DecodeReq(req) + decoder = ChatTokenDecoder(tokenizer) + decoder.set_prompt([0]) + deltas = [] + for token_id in (1, 2, 3): + decode_req.output_ids.append(token_id) + deltas.append(decode_token(tokenizer, decode_req, token_id, [])) + assert decoder.token_logprob(token_id, -0.1)["bytes"] == [[0xE4], [0xBD], [0xA0]][token_id - 1] + decoder.advance(token_id) + candidate = decoder.token_logprob(4, -0.2) + assert candidate == {"token": " world", "logprob": -0.2, "bytes": list(b" world")} + assert deltas == ["", "", "你"] -def test_byte_level_candidate_preserves_raw_bytes(): - from lightllm.server.chat_logprob_buffer import _candidate_logprob +@pytest.mark.anyio +@pytest.mark.parametrize("stream", [False, True]) +async def test_special_candidates_keep_identity(chat_response, monkeypatch, stream): + from lightllm.server.core.objs import sampling_params as params_module - backend = Tokenizer(models.BPE({"Ã": 0}, [])) - backend.pre_tokenizer = pre_tokenizers.ByteLevel() - backend.decoder = decoders.ByteLevel() - tokenizer = PreTrainedTokenizerFast(tokenizer_object=backend) - assert _candidate_logprob(tokenizer, 0, -0.5)["bytes"] == [195] + tokenizer = make_tokenizer(["prefix", "answer", "", ""], decoders.Fuse()) + tokenizer.add_special_tokens({"eos_token": "", "pad_token": ""}) + monkeypatch.setattr(params_module, "SKIP_SPECIAL_TOKENS", True) + _, _, content, _ = await chat_response( + tokenizer, + [(1, "answer")], + stream=stream, + topk=2, + first_metadata={"logprob_prompt_ids": [0]}, + candidates=[(2, -0.2), (3, -0.3)], + ) + assert content[0]["top_logprobs"] == [ + {"token": "", "logprob": -0.2, "bytes": list(b"")}, + {"token": "", "logprob": -0.3, "bytes": list(b"")}, + ] -def test_harmony_commentary_attributes_are_hidden_logprobs(): - from lightllm.server.chat_logprob_buffer import ChatLogprobBuffer +@pytest.mark.anyio +@pytest.mark.parametrize("stream", [False, True]) +async def test_pending_byte_candidates_in_chat(chat_response, stream): + tokenizer = make_tokenizer( + ["▁Hello", "<0xE4>", "<0xBD>", "<0xA0>", "▁world"], + decoders.Sequence( + [ + decoders.Replace("▁", " "), + decoders.ByteFallback(), + decoders.Fuse(), + decoders.Strip(" ", 1, 0), + ] + ), + ) + text, _, content, _ = await chat_response( + tokenizer, + [(1, ""), (2, ""), (3, "你"), (4, " world")], + stream=stream, + first_metadata={"logprob_prompt_ids": [0]}, + topk=1, + candidates=[(4, -0.2)], + ) + assert text == "你 world" + assert len(content) == 4 + assert bytes(byte for item in content for byte in item["bytes"]).decode() == text + assert all( + item["top_logprobs"] == [{"token": " world", "logprob": -0.2, "bytes": list(b" world")}] for item in content + ) + +@pytest.mark.anyio +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("separate_reasoning", [False, True]) +async def test_reasoning_then_tool_uses_api_spans(chat_response, stream, separate_reasoning): pieces = [ - "<|start|>", - "assistant", - "<|channel|>", - "commentary to=functions.weather", - "<|message|>", - '{"city":"Paris"}', - "<|end|>", - "<|start|>", - "assistant", - "<|channel|>", - "final", - "<|message|>", - "Answer", - "<|return|>", + "Thinking", + "", + "Checking.", + "", + '{"name":"weather","arguments":{"city":"Paris"}}', + "", ] - unique = list(dict.fromkeys(pieces)) - ids = {piece: index for index, piece in enumerate(unique)} - tokenizer = make_tokenizer(unique, decoders.Fuse()) - buffer = ChatLogprobBuffer(tokenizer, "gpt-oss", False, True) - for index, piece in enumerate(pieces): - buffer.record( - piece, - {"id": ids[piece], "logprob": -0.1}, - FinishStatus(FinishStatus.FINISHED_STOP if index == len(pieces) - 1 else FinishStatus.NO_FINISH), - ) - content, reasoning = buffer.message_logprobs() - assert [token["token"] for token in content] == ["Answer"] - assert reasoning is None - + tokenizer = make_tokenizer(pieces, decoders.Fuse()) + tools = [{"type": "function", "function": {"name": "weather", "parameters": {"type": "object"}}}] + text, _, content, reasoning = await chat_response( + tokenizer, + list(enumerate(pieces)), + stream=stream, + separate_reasoning=separate_reasoning, + model_type="qwen3", + tools=tools, + ) + expected = ["Checking."] if separate_reasoning else ["Thinking", "Checking."] + assert text == "".join(expected) + assert [item["token"] for item in content] == expected + assert [item["token"] for item in reasoning] == (["Thinking"] if separate_reasoning else []) -def test_mistral_tool_parser_is_rejected(): - from lightllm.server.chat_logprob_buffer import ChatLogprobBuffer - tokenizer = make_tokenizer(["[TOOL_CALLS] [", "]"], decoders.Fuse()) - with pytest.raises(ValueError, match="mistral"): - ChatLogprobBuffer( - tokenizer, - None, - False, - True, - tools=[{"type": "function", "function": {"name": "ok"}}], - tool_call_parser="mistral", - ) +@pytest.mark.anyio +async def test_reasoning_flush_does_not_include_hidden_tool_scores(chat_response): + call = "weathercityParis" + pieces = [call, " Date: Fri, 13 Nov 2026 18:21:25 +0800 Subject: [PATCH 4/6] fix(openai): track logprob source spans and prompt context --- lightllm/server/function_call_parser.py | 213 +++++++++++++----------- lightllm/server/httpserver/manager.py | 6 + lightllm/server/reasoning_parser.py | 131 +++++++-------- lightllm/server/text_spans.py | 47 ++++-- 4 files changed, 213 insertions(+), 184 deletions(-) diff --git a/lightllm/server/function_call_parser.py b/lightllm/server/function_call_parser.py index d46268bd55..ba0fa2bf8f 100644 --- a/lightllm/server/function_call_parser.py +++ b/lightllm/server/function_call_parser.py @@ -28,6 +28,7 @@ from pydantic import BaseModel, Field from .api_models import Tool +from .text_spans import TextSpan, extract_text, slice_spans logger = logging.getLogger(__name__) ENABLE_TOOL_NAME_CHECK = os.getenv("LIGHTLLM_ENABLE_TOOL_NAME_CHECK", "False").upper() in ["ON", "TRUE", "1"] @@ -55,8 +56,9 @@ class ToolCallItem(BaseModel): class StreamingParseResult: """Result of streaming incremental parsing.""" - def __init__(self, normal_text: str = "", calls: Optional[List[ToolCallItem]] = None): + def __init__(self, normal_text: str = "", calls: Optional[List[ToolCallItem]] = None, spans=None): self.normal_text = normal_text + self.spans = spans or [] self.calls = calls or [] @@ -117,6 +119,7 @@ def __init__(self): # Streaming state management # Buffer for accumulating incomplete patterns that arrive across multiple streaming chunks self._buffer = "" + self._source_length = 0 # Stores complete tool call info (name and arguments) for each tool being parsed. # Used by serving layer for completion handling when streaming ends. # Format: [{"name": str, "arguments": dict}, ...] @@ -227,11 +230,11 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami ): # Only clear buffer if we're sure no tool call is starting if not self._ends_with_partial_token(self._buffer, self.bot_token): - normal_text = self._buffer + normal_text, normal_spans = extract_text( + self._buffer, self._source_length - len(self._buffer), remove=(self.eot_token,) + ) self._buffer = "" - if self.eot_token in normal_text: - normal_text = normal_text.replace(self.eot_token, "") - return StreamingParseResult(normal_text=normal_text) + return StreamingParseResult(normal_text=normal_text, spans=normal_spans) else: # Might be partial bot_token, keep buffering return StreamingParseResult() @@ -405,6 +408,7 @@ def __init__(self): self.eot_token = "" self.tool_call_separator = "\n" self._normal_text_buffer = "" # Buffer for handling partial end tokens + self._normal_text_spans = [] def has_tool_call(self, text: str) -> bool: """Check if the text contains a Qwen 2.5 format tool call.""" @@ -419,9 +423,9 @@ def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult :return: ParseResult indicating success or failure, consumed text, leftover text, and parsed calls. """ idx = text.find(self.bot_token) - normal_text = text[:idx].strip() if idx != -1 else text + normal_text, normal_spans = extract_text(text[:idx], strip=True) if idx != -1 else extract_text(text) if self.bot_token not in text: - return StreamingParseResult(normal_text=normal_text, calls=[]) + return StreamingParseResult(normal_text=normal_text, spans=normal_spans, calls=[]) # Find all \n...\n blocks pattern = rf"{re.escape(self.bot_token)}(.*?){re.escape(self.eot_token)}" @@ -434,7 +438,7 @@ def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult except json.JSONDecodeError as e: logger.warning(f"Failed to parse JSON part: {match_result}, JSON parse error: {str(e)}") continue - return StreamingParseResult(normal_text=normal_text, calls=calls) + return StreamingParseResult(normal_text=normal_text, spans=normal_spans, calls=calls) def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> StreamingParseResult: """ @@ -446,11 +450,19 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami # Handle partial end tokens that are streamed character by character if result.normal_text: self._normal_text_buffer += result.normal_text + self._normal_text_spans.extend(result.spans) # Check if buffer contains complete end token (without leading newline) end_token_without_newline = self.eot_token # "" if end_token_without_newline in self._normal_text_buffer: - cleaned_text = self._normal_text_buffer.replace(end_token_without_newline, "") + source = self._normal_text_buffer + cleaned_text, relative_spans = extract_text(source, remove=(end_token_without_newline,)) + result.spans = [ + part + for span in relative_spans + for part in slice_spans(self._normal_text_spans, span.start, span.end) + ] + self._normal_text_spans = [] self._normal_text_buffer = "" result.normal_text = cleaned_text else: @@ -459,12 +471,17 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami if partial_match_len: # Keep potential partial match in buffer, return the rest + end = len(self._normal_text_buffer) - partial_match_len + result.spans = slice_spans(self._normal_text_spans, 0, end) + self._normal_text_spans = slice_spans(self._normal_text_spans, end, len(self._normal_text_buffer)) result.normal_text = self._normal_text_buffer[:-partial_match_len] self._normal_text_buffer = self._normal_text_buffer[-partial_match_len:] else: # No partial match, return all buffered text result.normal_text = self._normal_text_buffer + result.spans = self._normal_text_spans self._normal_text_buffer = "" + self._normal_text_spans = [] return result @@ -507,16 +524,16 @@ def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult :return: ParseResult indicating success or failure, consumed text, leftover text, and parsed calls. """ idx = text.find(self.bot_token) - normal_text = text[:idx].strip() if idx != -1 else text + normal_text, normal_spans = extract_text(text[:idx], strip=True) if idx != -1 else extract_text(text) if self.bot_token not in text: - return StreamingParseResult(normal_text=normal_text, calls=[]) + return StreamingParseResult(normal_text=normal_text, spans=normal_spans, calls=[]) # Extract the JSON array part from [TOOL_CALLS] [...] # Use bracket counting to properly handle nested brackets in JSON content json_array_str = self._extract_json_array(text) if not json_array_str: - return StreamingParseResult(normal_text=normal_text, calls=[]) + return StreamingParseResult(normal_text=normal_text, spans=normal_spans, calls=[]) calls = [] try: @@ -528,7 +545,7 @@ def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult except json.JSONDecodeError as e: logger.warning(f"Failed to parse JSON part: {json_array_str}, JSON parse error: {str(e)}") - return StreamingParseResult(normal_text=normal_text, calls=calls) + return StreamingParseResult(normal_text=normal_text, spans=normal_spans, calls=calls) def _extract_json_array(self, text: str) -> str: """ @@ -601,7 +618,7 @@ def has_tool_call(self, text: str) -> bool: def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult: """Parse function calls from text, handling multiple JSON objects.""" if "<|python_tag|>" not in text and not text.startswith("{"): - return StreamingParseResult(normal_text=text, calls=[]) + return StreamingParseResult(normal_text=text, spans=[TextSpan(0, len(text))], calls=[]) if "<|python_tag|>" in text: normal_text, action_text = text.split("<|python_tag|>", maxsplit=1) @@ -631,8 +648,10 @@ def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult # Only process if we found valid JSON objects calls = self.parse_base_json(all_actions, tools) if all_actions else [] # Use safe_idx to avoid idx containing the last part of an invalid JSON object - trailing_text = action_text[safe_idx:].strip() if safe_idx < action_text_len else "" - return StreamingParseResult(normal_text=normal_text + trailing_text, calls=calls) + action_offset = len(text) - len(action_text) + trailing_text, trailing_spans = extract_text(action_text[safe_idx:], action_offset + safe_idx, strip=True) + spans = ([TextSpan(0, len(normal_text))] if normal_text else []) + trailing_spans + return StreamingParseResult(normal_text=normal_text + trailing_text, spans=spans, calls=calls) class KimiK2Detector(BaseFormatDetector): @@ -684,7 +703,7 @@ def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult :return: ParseResult indicating success or failure, consumed text, leftover text, and parsed calls. """ if self.bot_token not in text: - return StreamingParseResult(normal_text=text, calls=[]) + return StreamingParseResult(normal_text=text, spans=[TextSpan(0, len(text))], calls=[]) try: # there are two possible captures - between tags, or between a # tag and end-of-string so the result of @@ -715,12 +734,12 @@ def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult ) content = text[: text.find(self.bot_token)] - return StreamingParseResult(normal_text=content, calls=tool_calls) + return StreamingParseResult(normal_text=content, spans=[TextSpan(0, len(content))], calls=tool_calls) except Exception as e: logger.error(f"Error in detect_and_parse: {e}") # return the normal text if parsing fails - return StreamingParseResult(normal_text=text) + return StreamingParseResult(normal_text=text, spans=[TextSpan(0, len(text))]) def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> StreamingParseResult: """ @@ -734,10 +753,10 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami if not has_tool_call: self._buffer = "" - for e_token in [self.eot_token, self.tool_call_end_token]: - if e_token in new_text: - new_text = new_text.replace(e_token, "") - return StreamingParseResult(normal_text=new_text) + new_text, spans = extract_text( + new_text, self._source_length - len(new_text), remove=[self.eot_token, self.tool_call_end_token] + ) + return StreamingParseResult(normal_text=new_text, spans=spans) if not hasattr(self, "_tool_indices"): self._tool_indices = self._get_tool_indices(tools) @@ -829,7 +848,9 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami except Exception as e: logger.error(f"Error in parse_streaming_increment: {e}") - return StreamingParseResult(normal_text=current_text) + return StreamingParseResult( + normal_text=current_text, spans=[TextSpan(self._source_length - len(current_text), self._source_length)] + ) class DeepSeekV31Detector(BaseFormatDetector): @@ -883,9 +904,9 @@ def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult :return: ParseResult indicating success or failure, consumed text, leftover text, and parsed calls. """ idx = text.find(self.bot_token) - normal_text = text[:idx].strip() if idx != -1 else text + normal_text, normal_spans = extract_text(text[:idx], strip=True) if idx != -1 else extract_text(text) if self.bot_token not in text: - return StreamingParseResult(normal_text=normal_text, calls=[]) + return StreamingParseResult(normal_text=normal_text, spans=normal_spans, calls=[]) match_result_list = re.findall(self.func_call_regex, text, re.DOTALL) calls = [] try: @@ -898,11 +919,11 @@ def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult # construct match_result for parse_base_json match_result = {"name": func_name, "parameters": func_args} calls.extend(self.parse_base_json(match_result, tools)) - return StreamingParseResult(normal_text=normal_text, calls=calls) + return StreamingParseResult(normal_text=normal_text, spans=normal_spans, calls=calls) except Exception as e: logger.error(f"Error in detect_and_parse: {e}") # return the normal text if parsing fails - return StreamingParseResult(normal_text=text) + return StreamingParseResult(normal_text=text, spans=[TextSpan(0, len(text))]) def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> StreamingParseResult: """ @@ -916,10 +937,10 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami if not has_tool_call: self._buffer = "" - for e_token in [self.eot_token, "<｜tool▁call▁end｜>"]: - if e_token in new_text: - new_text = new_text.replace(e_token, "") - return StreamingParseResult(normal_text=new_text) + new_text, spans = extract_text( + new_text, self._source_length - len(new_text), remove=[self.eot_token, "<｜tool▁call▁end｜>"] + ) + return StreamingParseResult(normal_text=new_text, spans=spans) if not hasattr(self, "_tool_indices"): self._tool_indices = self._get_tool_indices(tools) @@ -1009,7 +1030,9 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami except Exception as e: logger.error(f"Error in parse_streaming_increment: {e}") - return StreamingParseResult(normal_text=current_text) + return StreamingParseResult( + normal_text=current_text, spans=[TextSpan(self._source_length - len(current_text), self._source_length)] + ) class DeepSeekV3Detector(BaseFormatDetector): @@ -1064,9 +1087,9 @@ def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult :return: ParseResult indicating success or failure, consumed text, leftover text, and parsed calls. """ idx = text.find(self.bot_token) - normal_text = text[:idx].strip() if idx != -1 else text + normal_text, normal_spans = extract_text(text[:idx], strip=True) if idx != -1 else extract_text(text) if self.bot_token not in text: - return StreamingParseResult(normal_text=normal_text, calls=[]) + return StreamingParseResult(normal_text=normal_text, spans=normal_spans, calls=[]) match_result_list = re.findall(self.func_call_regex, text, re.DOTALL) calls = [] try: @@ -1079,11 +1102,11 @@ def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult # construct match_result for parse_base_json match_result = {"name": func_name, "parameters": func_args} calls.extend(self.parse_base_json(match_result, tools)) - return StreamingParseResult(normal_text=normal_text, calls=calls) + return StreamingParseResult(normal_text=normal_text, spans=normal_spans, calls=calls) except Exception as e: logger.error(f"Error in detect_and_parse: {e}") # return the normal text if parsing fails - return StreamingParseResult(normal_text=text) + return StreamingParseResult(normal_text=text, spans=[TextSpan(0, len(text))]) def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> StreamingParseResult: """ @@ -1097,10 +1120,10 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami if not has_tool_call: self._buffer = "" - for e_token in [self.eot_token, "```", "<｜tool▁call▁end｜>"]: - if e_token in new_text: - new_text = new_text.replace(e_token, "") - return StreamingParseResult(normal_text=new_text) + new_text, spans = extract_text( + new_text, self._source_length - len(new_text), remove=[self.eot_token, "```", "<｜tool▁call▁end｜>"] + ) + return StreamingParseResult(normal_text=new_text, spans=spans) if not hasattr(self, "_tool_indices"): self._tool_indices = self._get_tool_indices(tools) @@ -1192,7 +1215,9 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami except Exception as e: logger.error(f"Error in parse_streaming_increment: {e}") - return StreamingParseResult(normal_text=current_text) + return StreamingParseResult( + normal_text=current_text, spans=[TextSpan(self._source_length - len(current_text), self._source_length)] + ) class Glm47Detector(BaseFormatDetector): @@ -1295,10 +1320,10 @@ def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult :return: StreamingParseResult with normal_text and parsed calls. """ idx = text.find(self.bot_token) - normal_text = text[:idx].strip() if idx != -1 else text + normal_text, normal_spans = extract_text(text[:idx], strip=True) if idx != -1 else extract_text(text) if self.bot_token not in text: - return StreamingParseResult(normal_text=normal_text, calls=[]) + return StreamingParseResult(normal_text=normal_text, spans=normal_spans, calls=[]) tool_indices = self._get_tool_indices(tools) calls = [] @@ -1336,22 +1361,27 @@ def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult logger.warning(f"Failed to parse GLM-4.7 tool call: {match_result}, error: {str(e)}") continue - return StreamingParseResult(normal_text=normal_text, calls=calls) + return StreamingParseResult(normal_text=normal_text, spans=normal_spans, calls=calls) def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> StreamingParseResult: """Buffer XML arguments while emitting tool deltas to keep the stream alive.""" self._buffer += new_text normal_text = "" + normal_spans = [] calls: List[ToolCallItem] = [] while self._buffer: start = self._buffer.find(self.bot_token) if start == -1: keep = self._ends_with_partial_token(self._buffer, self.bot_token) end = len(self._buffer) - keep + offset = self._source_length - len(self._buffer) + normal_spans.extend([TextSpan(offset, offset + end)] if end else []) normal_text += self._buffer[:end] self._buffer = self._buffer[end:] break + offset = self._source_length - len(self._buffer) + normal_spans.extend([TextSpan(offset, offset + start)] if start else []) normal_text += self._buffer[:start] self._buffer = self._buffer[start:] end = self._buffer.find(self.eot_token) @@ -1384,7 +1414,7 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami self._streaming_tool_name = None self._buffer = self._buffer[end:] - return StreamingParseResult(normal_text=normal_text, calls=calls) + return StreamingParseResult(normal_text=normal_text, spans=normal_spans, calls=calls) class DeepSeekV32Detector(BaseFormatDetector): @@ -1462,9 +1492,9 @@ def _dsml_params_to_json(self, params: List[tuple]) -> str: def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult: """One-time parsing for DSML format tool calls.""" idx = text.find(self.bot_token) - normal_text = text[:idx].strip() if idx != -1 else text + normal_text, normal_spans = extract_text(text[:idx], strip=True) if idx != -1 else extract_text(text) if self.bot_token not in text: - return StreamingParseResult(normal_text=normal_text, calls=[]) + return StreamingParseResult(normal_text=normal_text, spans=normal_spans, calls=[]) tool_indices = self._get_tool_indices(tools) calls = [] @@ -1486,7 +1516,7 @@ def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult ) ) - return StreamingParseResult(normal_text=normal_text, calls=calls) + return StreamingParseResult(normal_text=normal_text, spans=normal_spans, calls=calls) def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> StreamingParseResult: """Streaming incremental parsing for DSML format tool calls.""" @@ -1502,10 +1532,10 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami return StreamingParseResult() self._buffer = "" - for e_token in [self.eot_token, self.invoke_end_token]: - if e_token in new_text: - new_text = new_text.replace(e_token, "") - return StreamingParseResult(normal_text=new_text) + new_text, spans = extract_text( + new_text, self._source_length - len(new_text), remove=[self.eot_token, self.invoke_end_token] + ) + return StreamingParseResult(normal_text=new_text, spans=spans) # Mark that we're inside a function_calls block if self.has_tool_call(current_text): @@ -1849,10 +1879,10 @@ def _build_partial_arguments_json(self, func_name: str, partial_body: str, tools def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult: idx = text.find(self.bot_token) - normal_text = text[:idx].strip() if idx != -1 else text + normal_text, normal_spans = extract_text(text[:idx], strip=True) if idx != -1 else extract_text(text) if " StreamingParseResult item.tool_index = len(calls) calls.append(item) - return StreamingParseResult(normal_text=normal_text, calls=calls) + return StreamingParseResult(normal_text=normal_text, spans=normal_spans, calls=calls) def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> StreamingParseResult: """Streaming incremental parsing for Qwen3-Coder XML tool calls.""" @@ -1878,6 +1908,7 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami self._tool_indices = self._get_tool_indices(tools) normal_text = "" + normal_spans = [] calls: List[ToolCallItem] = [] while True: @@ -1887,13 +1918,19 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami if tool_call_start == -1: partial_len = self._ends_with_partial_token(current_text, self.bot_token) if partial_len: - return StreamingParseResult(normal_text=normal_text, calls=calls) + return StreamingParseResult(normal_text=normal_text, spans=normal_spans, calls=calls) if current_text: - normal_text += current_text.replace(self.eot_token, "") + visible, spans = extract_text( + current_text, self._source_length - len(current_text), remove=(self.eot_token,) + ) + normal_text += visible + normal_spans.extend(spans) self._buffer = "" - return StreamingParseResult(normal_text=normal_text, calls=calls) + return StreamingParseResult(normal_text=normal_text, spans=normal_spans, calls=calls) if tool_call_start > 0: + offset = self._source_length - len(current_text) + normal_spans.append(TextSpan(offset, offset + tool_call_start)) normal_text += current_text[:tool_call_start] self._buffer = current_text[tool_call_start:] current_text = self._buffer @@ -1903,7 +1940,7 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami if not self.current_tool_name_sent: function_name = self._get_incomplete_function_name(current_text) if function_name is None: - return StreamingParseResult(normal_text=normal_text, calls=calls) + return StreamingParseResult(normal_text=normal_text, spans=normal_spans, calls=calls) if self.current_tool_id == -1: self.current_tool_id = 0 @@ -1925,7 +1962,7 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami parameters="", ) ) - return StreamingParseResult(normal_text=normal_text, calls=calls) + return StreamingParseResult(normal_text=normal_text, spans=normal_spans, calls=calls) complete_block = current_text[: eot_pos + len(self.eot_token)] func_matches = self.function_regex.findall(complete_block) @@ -1996,25 +2033,20 @@ def __init__(self, tools: List[Tool], tool_call_parser: str): self.tools = tools def parse_non_stream_with_spans(self, full_text: str): - from lightllm.server.text_spans import select_text_spans - - normal, calls = self.parse_non_stream(full_text) - return normal, calls, select_text_spans(full_text, normal) + if not self.tools: + return full_text, [], [TextSpan(0, len(full_text))] + result = self.detector.detect_and_parse(full_text, self.tools) + if not result.calls: + return full_text, [], [TextSpan(0, len(full_text))] + return result.normal_text, result.calls, result.spans def parse_stream_chunk_with_spans(self, chunk_text: str): - from lightllm.server.text_spans import TextSpan, select_text_spans - - source = self.detector._normal_text_buffer if isinstance(self.detector, Qwen25Detector) else "" - source += self.detector._buffer + chunk_text - end = getattr(self, "_source_length", 0) + len(chunk_text) - self._source_length = end - offset = end - len(source) - normal, calls = self.parse_stream_chunk(chunk_text) - marker = self.detector.eot_token - if marker: - source = source.replace(marker, "\0" * len(marker)) - spans = [TextSpan(s.start + offset, s.end + offset) for s in select_text_spans(source, normal)] - return normal, calls, spans + self.detector._source_length += len(chunk_text) + if not self.tools: + end = self.detector._source_length + return chunk_text, [], [TextSpan(end - len(chunk_text), end)] + result = self.detector.parse_streaming_increment(chunk_text, self.tools) + return result.normal_text, result.calls, result.spans def has_tool_call(self, text: str) -> bool: """ @@ -2043,14 +2075,8 @@ def parse_non_stream(self, full_text: str) -> Tuple[str, list[ToolCallItem]]: - The remaining text after parsing that was not consumed by the detector (can be treated as normal text) - A list of tool calls parsed from the text """ - if not self.tools: - return full_text, [] - parsed_result = self.detector.detect_and_parse(full_text, self.tools) - tool_call_list = parsed_result.calls - if tool_call_list: - return parsed_result.normal_text, tool_call_list - else: - return full_text, [] + normal, calls, _ = self.parse_non_stream_with_spans(full_text) + return normal, calls def parse_stream_chunk(self, chunk_text: str) -> Tuple[str, list[ToolCallItem]]: """ @@ -2064,16 +2090,5 @@ def parse_stream_chunk(self, chunk_text: str) -> Tuple[str, list[ToolCallItem]]: - The normal text that should be displayed to the user - A list of tool calls parsed from the chunk """ - if not self.tools: - return chunk_text, [] - final_normal_text = "" - final_calls = [] - - sp_result = self.detector.parse_streaming_increment(chunk_text, self.tools) - if sp_result.normal_text: - final_normal_text = sp_result.normal_text - if sp_result.calls: - final_calls.extend(sp_result.calls) - final_normal_text = sp_result.normal_text - - return final_normal_text, final_calls + normal, calls, _ = self.parse_stream_chunk_with_spans(chunk_text) + return normal, calls diff --git a/lightllm/server/httpserver/manager.py b/lightllm/server/httpserver/manager.py index 95dd057212..e89bffd515 100644 --- a/lightllm/server/httpserver/manager.py +++ b/lightllm/server/httpserver/manager.py @@ -800,6 +800,8 @@ async def _wait_to_token_package( request: Request, ): + from lightllm.server.detokenization.decode_req import LIGHTLLM_DECODE_PREFIX_LENGTH + event = req_status.event unfinished_count = sampling_params.best_of out_token_counter = 0 @@ -838,6 +840,10 @@ async def _wait_to_token_package( for sub_req_id, out_str, metadata, finish_status in req_status.out_token_info_list: # pd master 节点需要这个做统计信息， 所以放在元数据中返回给 pd master 节点 metadata["prompt_tokens"] = prompt_tokens + if sampling_params.output_top_logprobs >= 0 and metadata.get("count_output_tokens") == 1: + metadata["logprob_prompt_ids"] = prompt_ids[ + max(len(prompt_ids) - LIGHTLLM_DECODE_PREFIX_LENGTH, 0) : + ] gpu_prompt_cache_len = metadata.pop("prompt_cache_len", 0) cpu_prompt_cache_len = metadata.pop("cpu_prompt_cache_len", 0) diff --git a/lightllm/server/reasoning_parser.py b/lightllm/server/reasoning_parser.py index 612d3aaa88..e6c1535211 100644 --- a/lightllm/server/reasoning_parser.py +++ b/lightllm/server/reasoning_parser.py @@ -16,7 +16,7 @@ from typing import Iterator, List, Tuple, Dict, Optional, Type from lightllm.utils.config_utils import get_token_id -from lightllm.server.text_spans import TextSpan, select_text_spans +from lightllm.server.text_spans import TextSpan, extract_text, slice_spans @dataclass @@ -547,7 +547,7 @@ def parse(self, text: str) -> Tuple[List[Event], str]: elif channel.lower() == "commentary" and emit: # For commentary, stream as normal text. Preserve spaces unless holding. content_out = emit if hold else emit.strip() - events.append(Event("normal", content_out, spans=select_text_spans(text, content_out))) + events.append(Event("normal", content_out, spans=extract_text(emit, m.start(2), strip=not hold)[1])) if hold: return events, text[: m.start(2)] + hold else: @@ -682,7 +682,9 @@ def __init__( self, normal_text: Optional[str] = None, reasoning_text: Optional[str] = None, + spans=None, ): + self.spans = spans or [] self.normal_text = normal_text or "" self.reasoning_text = reasoning_text or "" @@ -703,6 +705,7 @@ def __init__( self.stream_reasoning = stream_reasoning self._buffer = "" + self._source_length = 0 self.stripped_think_start = False @property @@ -718,21 +721,29 @@ def detect_and_parse(self, text: str) -> StreamingParseResult: in_reasoning = self._in_reasoning or self.think_start_token in text if not in_reasoning: - return StreamingParseResult(normal_text=text) + return StreamingParseResult(normal_text=text, spans=[TextSpan(0, len(text))]) # The text is considered to be in a reasoning block. - processed_text = text.replace(self.think_start_token, "").strip() + processed_text, spans = extract_text(text, strip=True, remove=(self.think_start_token,)) if self.think_end_token not in processed_text: # Assume reasoning was truncated before `` token - return StreamingParseResult(reasoning_text=processed_text) + for span in spans: + span.channel = "reasoning" + return StreamingParseResult(reasoning_text=processed_text, spans=spans) # Extract reasoning content splits = processed_text.split(self.think_end_token, maxsplit=1) reasoning_text = splits[0] normal_text = splits[1].strip() - - return StreamingParseResult(normal_text=normal_text, reasoning_text=reasoning_text) + content_start = len(reasoning_text) + len(self.think_end_token) + len(splits[1]) - len(splits[1].lstrip()) + reasoning_spans = slice_spans(spans, 0, len(reasoning_text)) + for span in reasoning_spans: + span.channel = "reasoning" + content_spans = slice_spans(spans, content_start, content_start + len(normal_text)) + return StreamingParseResult( + normal_text=normal_text, reasoning_text=reasoning_text, spans=reasoning_spans + content_spans + ) def flush(self) -> StreamingParseResult: """ @@ -745,10 +756,12 @@ def flush(self) -> StreamingParseResult: return StreamingParseResult() remaining = self._buffer self._buffer = "" + span = TextSpan(self._source_length - len(remaining), self._source_length) if self._in_reasoning: - return StreamingParseResult(reasoning_text=remaining) + span.channel = "reasoning" + return StreamingParseResult(reasoning_text=remaining, spans=[span]) else: - return StreamingParseResult(normal_text=remaining) + return StreamingParseResult(normal_text=remaining, spans=[span]) def parse_streaming_increment(self, new_text: str) -> StreamingParseResult: """ @@ -763,6 +776,8 @@ def parse_streaming_increment(self, new_text: str) -> StreamingParseResult: """ self._buffer += new_text current_text = self._buffer + offset = self._source_length - len(current_text) + spans = [TextSpan(offset, self._source_length)] if current_text else [] # If the current text is a prefix of the think token, keep buffering # until we can confirm or refute the tag. @@ -774,7 +789,7 @@ def parse_streaming_increment(self, new_text: str) -> StreamingParseResult: # Strip `` token if present if not self.stripped_think_start and self.think_start_token in current_text: - current_text = current_text.replace(self.think_start_token, "") + current_text, spans = extract_text(current_text, offset, remove=(self.think_start_token,)) self.stripped_think_start = True self._in_reasoning = True @@ -788,18 +803,27 @@ def parse_streaming_increment(self, new_text: str) -> StreamingParseResult: self._in_reasoning = False normal_text = current_text[end_idx + len(self.think_end_token) :] - return StreamingParseResult(normal_text=normal_text, reasoning_text=reasoning_text.rstrip()) + reasoning_text = reasoning_text.rstrip() + reasoning_spans = slice_spans(spans, 0, len(reasoning_text)) + for span in reasoning_spans: + span.channel = "reasoning" + content_spans = slice_spans(spans, end_idx + len(self.think_end_token), len(current_text)) + return StreamingParseResult( + normal_text=normal_text, reasoning_text=reasoning_text, spans=reasoning_spans + content_spans + ) # Always stream reasoning content immediately. # stream_reasoning flag is ignored for streaming responses. if self._in_reasoning: self._buffer = "" - return StreamingParseResult(reasoning_text=current_text) + for span in spans: + span.channel = "reasoning" + return StreamingParseResult(reasoning_text=current_text, spans=spans) # If we're not in a reasoning block return as normal text if not self._in_reasoning: self._buffer = "" - return StreamingParseResult(normal_text=current_text) + return StreamingParseResult(normal_text=current_text, spans=spans) return StreamingParseResult() @@ -899,9 +923,12 @@ def grammar_start_token(self) -> str: return "<|channel|>final<|message|>" def detect_and_parse(self, text: str) -> StreamingParseResult: + self.parser.track_spans = True events = self.parser.parse(text) + spans = list(self.parser.last_spans) # Flush the buffer for one-shot parsing events += self.parser.parse("") + spans.extend(self.parser.last_spans) reasoning_text = "".join([e.content for e in events if e.event_type == "reasoning"]) normal_parts = [] @@ -917,10 +944,13 @@ def detect_and_parse(self, text: str) -> StreamingParseResult: return StreamingParseResult( normal_text=normal_text, reasoning_text=reasoning_text, + spans=spans, ) def parse_streaming_increment(self, new_text: str) -> StreamingParseResult: + self.parser.track_spans = True events = self.parser.parse(new_text) + spans = self.parser.last_spans reasoning_text = "".join([e.content for e in events if e.event_type == "reasoning"]) normal_parts = [] @@ -935,6 +965,7 @@ def parse_streaming_increment(self, new_text: str) -> StreamingParseResult: return StreamingParseResult( normal_text=normal_text, reasoning_text=reasoning_text, + spans=spans, ) @@ -954,13 +985,14 @@ def __init__(self, stream_reasoning: bool = True, force_reasoning: bool = False) self.is_first_chunk = False def parse_streaming_increment(self, new_text: str) -> StreamingParseResult: + spans = [TextSpan(self._source_length - len(new_text), self._source_length)] if new_text else [] if not self.is_first_chunk: self.is_first_chunk = True new_text = self.think_start_token + new_text - return StreamingParseResult(normal_text=new_text) + return StreamingParseResult(normal_text=new_text, spans=spans) def detect_and_parse(self, text: str) -> StreamingParseResult: - return StreamingParseResult(normal_text=self.think_start_token + text) + return StreamingParseResult(normal_text=self.think_start_token + text, spans=[TextSpan(0, len(text))]) class NanoV3Detector(BaseReasoningFormatDetector): @@ -1083,65 +1115,30 @@ def update_reasoning_token_count(self, token_id: int) -> None: def parse_non_stream(self, full_text: str) -> Tuple[Optional[str], Optional[str]]: """Non-streaming call: one-time parsing""" - ret = self.detector.detect_and_parse(full_text) - return ret.reasoning_text, ret.normal_text + reasoning, content, _ = self.parse_non_stream_with_spans(full_text) + return reasoning, content + + def parse_non_stream_with_spans(self, full_text: str): + result = self.detector.detect_and_parse(full_text) + spans = sorted(result.spans, key=lambda span: span.channel != "reasoning") + return result.reasoning_text, result.normal_text, spans def parse_stream_chunk(self, chunk_text: str) -> Tuple[Optional[str], Optional[str]]: """Streaming call: incremental parsing""" - ret = self.detector.parse_streaming_increment(chunk_text) - return ret.reasoning_text, ret.normal_text + reasoning, content, _ = self.parse_stream_chunk_with_spans(chunk_text) + return reasoning, content def parse_stream_chunk_with_spans(self, chunk_text: str): - if isinstance(self.detector, GptOssDetector): - self.detector.parser.track_spans = True - reasoning, content = self.parse_stream_chunk(chunk_text) - return reasoning, content, self.detector.parser.last_spans - source = self.detector._buffer + chunk_text - end = getattr(self, "_source_length", 0) + len(chunk_text) - self._source_length = end - offset = end - len(source) - markers = [] - if not self.detector.stripped_think_start and self.detector.think_start_token in source: - markers.append(self.detector.think_start_token) - if self.detector._in_reasoning or markers: - markers.append(self.detector.think_end_token) - reasoning, content = self.parse_stream_chunk(chunk_text) - spans = self._visible_spans(source, reasoning, content, offset, markers) - return reasoning, content, spans - - def _visible_spans(self, source, reasoning, content, offset, markers=()): - visible = source - positions = list(range(offset, offset + len(source))) - for marker in markers: - if not marker: - continue - start = visible.find(marker) - while start >= 0: - visible = visible[:start] + visible[start + len(marker) :] - del positions[start : start + len(marker)] - start = visible.find(marker) - spans = [] - cursor = 0 - for text, channel in ((reasoning, "reasoning"), (content, "content")): - next_cursor = cursor - for span in select_text_spans(visible[cursor:], text or "", channel): - start, end = span.start + cursor, span.end + cursor - next_cursor = end - for pos in positions[start:end]: - if spans and spans[-1].end == pos and spans[-1].channel == channel: - spans[-1].end += 1 - else: - spans.append(TextSpan(pos, pos + 1, channel)) - cursor = next_cursor - return spans + self.detector._source_length += len(chunk_text) + result = self.detector.parse_streaming_increment(chunk_text) + spans = sorted(result.spans, key=lambda span: span.channel != "reasoning") + return result.reasoning_text, result.normal_text, spans def flush_with_spans(self): - source = self.detector._buffer - reasoning, content = self.flush() - offset = getattr(self, "_source_length", 0) - len(source) - return reasoning, content, self._visible_spans(source, reasoning, content, offset) + result = self.detector.flush() + return result.reasoning_text, result.normal_text, result.spans def flush(self) -> Tuple[Optional[str], Optional[str]]: """Flush remaining buffered content when generation ends prematurely.""" - ret = self.detector.flush() - return ret.reasoning_text, ret.normal_text + reasoning, content, _ = self.flush_with_spans() + return reasoning, content diff --git a/lightllm/server/text_spans.py b/lightllm/server/text_spans.py index 2a58eb4f16..212fbf160c 100644 --- a/lightllm/server/text_spans.py +++ b/lightllm/server/text_spans.py @@ -9,21 +9,32 @@ class TextSpan: channel: str = "content" -def select_text_spans(source: str, text: str, channel: str = "content") -> List[TextSpan]: - spans = [] - cursor = 0 - while text: - start = source.find(text, cursor) - if start >= 0: - spans.append(TextSpan(start, start + len(text), channel)) - break - start = source.find(text[0], cursor) - if start < 0: - break - length = 1 - while length < len(text) and start + length < len(source) and source[start + length] == text[length]: - length += 1 - spans.append(TextSpan(start, start + length, channel)) - cursor = start + length - text = text[length:] - return spans +def slice_spans(spans: List[TextSpan], start: int, end: int) -> List[TextSpan]: + selected = [] + offset = 0 + for span in spans: + size = span.end - span.start + left, right = max(start - offset, 0), min(end - offset, size) + if left < right: + selected.append(TextSpan(span.start + left, span.start + right, span.channel)) + offset += size + return selected + + +def extract_text(text: str, offset=0, channel="content", strip=False, remove=()): + spans = [TextSpan(offset, offset + len(text), channel)] if text else [] + for marker in remove: + if not marker: + continue + start = text.find(marker) + while start >= 0: + end = start + len(marker) + spans = slice_spans(spans, 0, start) + slice_spans(spans, end, len(text)) + text = text[:start] + text[end:] + start = text.find(marker, start) + if strip: + start = len(text) - len(text.lstrip()) + end = len(text.rstrip()) + spans = slice_spans(spans, start, end) + text = text.strip() + return text, spans From 4c6ffb3b3fdf148de2e18c797b3286a07d979511 Mon Sep 17 00:00:00 2001 From: sufubao Date: Fri, 13 Nov 2026 21:06:01 +0800 Subject: [PATCH 5/6] fix(openai): flush fallback content with aligned logprobs --- docs/CN/source/tutorial/function_calling.rst | 2 - docs/EN/source/tutorial/function_calling.rst | 2 - lightllm/server/api_openai.py | 22 ++++++- lightllm/server/function_call_parser.py | 23 ++++++- lightllm/server/httpserver/manager.py | 3 - unit_tests/server/test_application_stop.py | 48 ++++++++++++++ .../server/test_chat_output_logprobs.py | 64 +++++++++++++++++-- 7 files changed, 151 insertions(+), 13 deletions(-) diff --git a/docs/CN/source/tutorial/function_calling.rst b/docs/CN/source/tutorial/function_calling.rst index e6eab454ec..afa588572f 100644 --- a/docs/CN/source/tutorial/function_calling.rst +++ b/docs/CN/source/tutorial/function_calling.rst @@ -53,8 +53,6 @@ Mistral **格式**: ``[TOOL_CALLS] [{"name": "func", "arguments": {...}}, ...]`` -启用 chat logprobs 时暂不支持 ``mistral`` 工具调用解析器。 - DeepSeek-V3 ~~~~~~~~~~~ diff --git a/docs/EN/source/tutorial/function_calling.rst b/docs/EN/source/tutorial/function_calling.rst index 99db50b0a8..d038eb9022 100644 --- a/docs/EN/source/tutorial/function_calling.rst +++ b/docs/EN/source/tutorial/function_calling.rst @@ -53,8 +53,6 @@ Mistral **Format**: ``[TOOL_CALLS] [{"name": "func", "arguments": {...}}, ...]`` -The ``mistral`` tool-call parser is not supported when chat logprobs are enabled. - DeepSeek-V3 ~~~~~~~~~~~ diff --git a/lightllm/server/api_openai.py b/lightllm/server/api_openai.py index 1afff9fa5f..b7dde16480 100644 --- a/lightllm/server/api_openai.py +++ b/lightllm/server/api_openai.py @@ -684,7 +684,8 @@ async def stream_results() -> AsyncGenerator[bytes, None]: ) yield f"data: {_serialize_sse_chunk(chunk, _choice_nulls, logprobs=choice_logprobs)}\n\n" elif logprob_buffer is not None: - logprob_buffer.discard_content(content_spans) + if current_finish_reason is None or not parser_dict[choice_index].detector._buffer: + logprob_buffer.discard_content(content_spans) # 2) if we found calls, we output them as separate chunk(s) history_tool_calls_cnt = _get_history_tool_calls_cnt(request) @@ -886,6 +887,25 @@ async def stream_results() -> AsyncGenerator[bytes, None]: ) payload = _serialize_sse_chunk(flush_chunk, _choice_nulls, logprobs=flush_logprobs) yield f"data: {payload}\n\n" + if request.tool_choice != "none" and request.tools and choice_index in parser_dict: + flush_text, _, output_spans = parser_dict[choice_index].flush_with_spans() + if flush_text: + flush_choice = ChatCompletionStreamResponseChoice( + index=choice_index, + delta=DeltaMessage(content=flush_text), + finish_reason=None, + ) + flush_chunk = ChatCompletionStreamResponse( + id=chat_completion_id, + created=created_time, + model=request.model, + choices=[flush_choice], + ) + flush_logprobs = ( + logprob_buffer.content_logprobs(output_spans) if logprob_buffer is not None else None + ) + payload = _serialize_sse_chunk(flush_chunk, _choice_nulls, logprobs=flush_logprobs) + yield f"data: {payload}\n\n" if has_emitted_tool_calls[sub_req_id] and current_finish_reason == "stop": current_finish_reason = "tool_calls" final_choice = ChatCompletionStreamResponseChoice( diff --git a/lightllm/server/function_call_parser.py b/lightllm/server/function_call_parser.py index ba0fa2bf8f..c4959ae03f 100644 --- a/lightllm/server/function_call_parser.py +++ b/lightllm/server/function_call_parser.py @@ -810,7 +810,6 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami parsed_args_diff = argument_diff.split("<|tool_call_end|>", 1)[0] if parsed_args_diff: - calls.append( ToolCallItem( tool_index=self.current_tool_id, @@ -2031,6 +2030,7 @@ def __init__(self, tools: List[Tool], tool_call_parser: str): self.detector = detector self.tools = tools + self._emitted_calls = False def parse_non_stream_with_spans(self, full_text: str): if not self.tools: @@ -2046,8 +2046,29 @@ def parse_stream_chunk_with_spans(self, chunk_text: str): end = self.detector._source_length return chunk_text, [], [TextSpan(end - len(chunk_text), end)] result = self.detector.parse_streaming_increment(chunk_text, self.tools) + if result.calls: + self._emitted_calls = True return result.normal_text, result.calls, result.spans + def flush_with_spans(self): + """Return buffered input as visible content when generation ends early.""" + if not self.tools: + return "", [], [] + buffered = self.detector._buffer + if self._emitted_calls: + if self.detector.eot_token: + buffered = buffered.replace(self.detector.eot_token, "") + if self.detector.tool_call_separator and buffered.startswith(self.detector.tool_call_separator): + buffered = buffered[len(self.detector.tool_call_separator) :] + else: + buffered = getattr(self.detector, "_normal_text_buffer", "") + buffered + self.detector._normal_text_buffer = "" + self.detector._buffer = "" + if not buffered: + return "", [], [] + end = self.detector._source_length + return buffered, [], [TextSpan(end - len(buffered), end)] + def has_tool_call(self, text: str) -> bool: """ Check if the given text contains a tool call in the format supported by this parser. diff --git a/lightllm/server/httpserver/manager.py b/lightllm/server/httpserver/manager.py index e89bffd515..677b236acb 100644 --- a/lightllm/server/httpserver/manager.py +++ b/lightllm/server/httpserver/manager.py @@ -344,7 +344,6 @@ async def _generate( # 用于等待 pd_master 下发的交换信息 pd_event: asyncio.Event = None, ) -> AsyncGenerator[Tuple[int, str, dict, FinishStatus], None]: - start_time = time.time() # 在请求登记、编码和 PD 握手前获取语法；缓存未命中时先完成编译，避免后续流程等待或校验失败。 compiled_grammar = await sampling_params.get_compiled_grammar(self.output_grammar_compiler) @@ -756,7 +755,6 @@ async def transfer_to_next_module( self, group_req_objs: Optional[GroupReqObjs] = None, ): - if self.pd_mode.is_P_or_NORMAL(): if not self.args.disable_vision: self.send_to_visual.send_pyobj(group_req_objs.to_group_req_index(), protocol=pickle.HIGHEST_PROTOCOL) @@ -799,7 +797,6 @@ async def _wait_to_token_package( req_status: "ReqStatus", request: Request, ): - from lightllm.server.detokenization.decode_req import LIGHTLLM_DECODE_PREFIX_LENGTH event = req_status.event diff --git a/unit_tests/server/test_application_stop.py b/unit_tests/server/test_application_stop.py index 0603c29c67..6825028434 100644 --- a/unit_tests/server/test_application_stop.py +++ b/unit_tests/server/test_application_stop.py @@ -113,6 +113,54 @@ async def raw_results(*args, **kwargs): asyncio.run(run()) +@pytest.mark.parametrize("output_top_logprobs", [-1, 0, 2]) +def test_tgi_num_beam_keeps_true_scores(output_top_logprobs): + async def run(): + async def raw_results( + prompt, + sampling_params, + multimodal, + request, + pd_upload_websocket, + pd_event, + ): + assert sampling_params.output_top_logprobs == output_top_logprobs + for sub_id, scores in enumerate(((-2.302585, -3.912023), (-2.995732, -6.907755))): + text = "A" if sub_id == 0 else "B" + for index, score in enumerate(scores): + cumlogprob = sum(scores[: index + 1]) / (index + 1) + yield ( + sub_id, + text, + { + "prompt_tokens": 1, + "logprobs": {}, + "logprob": score, + "cumlogprob": cumlogprob, + }, + FinishStatus(FinishStatus.FINISHED_LENGTH if index else FinishStatus.NO_FINISH), + ) + + manager = _manager( + SimpleNamespace(encode=lambda text, **kwargs: [ord(c) for c in text]), + SimpleNamespace(reasoning_parser=None, enable_rl=False, enable_return_routed_experts=False), + raw_results, + ) + request = SimpleNamespace( + json=AsyncMock( + return_value={ + "inputs": "prompt", + "num_beam": 2, + "parameters": {"output_top_logprobs": output_top_logprobs}, + } + ) + ) + response = await tgi_generate_impl(request, manager) + assert json.loads(response.body)["generated_text"] == "AA" + + asyncio.run(run()) + + @pytest.mark.parametrize("mode", [NodeRole.NORMAL, NodeRole.P, NodeRole.D]) @pytest.mark.parametrize("slave", [False, True]) def test_http_generate_only_wraps_client_output(mode, slave): diff --git a/unit_tests/server/test_chat_output_logprobs.py b/unit_tests/server/test_chat_output_logprobs.py index 9129dc7b59..d072f15c0c 100644 --- a/unit_tests/server/test_chat_output_logprobs.py +++ b/unit_tests/server/test_chat_output_logprobs.py @@ -382,6 +382,28 @@ async def test_tool_logprobs_only_contain_message_content(chat_response, stream, assert reasoning == [] +@pytest.mark.anyio +@pytest.mark.parametrize("stream", [False, True]) +async def test_mistral_tool_parser_keeps_message_scores(chat_response, stream): + pieces = [ + "Checking.", + "[TOOL_CALLS] [", + '{"name":"weather","arguments":{"city":"Paris"}}', + "]", + ] + tokenizer = make_tokenizer(pieces, decoders.Fuse()) + tools = [{"type": "function", "function": {"name": "weather", "parameters": {"type": "object"}}}] + text, _, content, _ = await chat_response( + tokenizer, + list(enumerate(pieces)), + stream=stream, + tools=tools, + tool_parser="mistral", + ) + assert text == "Checking." + assert [entry["token"] for entry in content] == ["Checking."] + + @pytest.mark.anyio @pytest.mark.parametrize("stream", [False, True]) @pytest.mark.parametrize("separate_reasoning", [False, True]) @@ -447,15 +469,48 @@ async def test_harmony_repeated_header_words_are_real_content(chat_response, str assert [entry["token"] for entry in reasoning] == (["analysis"] if separate_reasoning else []) +@pytest.mark.anyio +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("separate_reasoning", [False, True]) +async def test_harmony_commentary_content_keeps_its_scores(chat_response, stream, separate_reasoning): + pieces = [ + "<|channel|>", + "commentary to=functions.weather", + "<|message|>", + "Checking.", + "<|end|>", + "<|start|>", + "assistant", + "<|channel|>", + "final", + "<|message|>", + "Answer", + "<|return|>", + ] + ids = {piece: index for index, piece in enumerate(dict.fromkeys(pieces))} + tokenizer = make_tokenizer(list(ids), decoders.Fuse()) + text, _, content, reasoning = await chat_response( + tokenizer, + [(ids[piece], piece) for piece in pieces], + stream=stream, + separate_reasoning=separate_reasoning, + model_type="gpt-oss", + ) + assert text == "Checking.Answer" + assert [entry["token"] for entry in content] == ["Checking.", "Answer"] + assert reasoning == [] + + @pytest.mark.anyio @pytest.mark.parametrize("valid_tool", [False, True]) -async def test_nonstream_tool_logprobs_follow_actual_extraction(chat_response, valid_tool): +@pytest.mark.parametrize("stream", [False, True]) +async def test_nonstream_tool_logprobs_follow_actual_extraction(chat_response, valid_tool, stream): pieces = ["Checking.", '{"name":"weather","arguments":{"city":"Paris"}}'] if valid_tool: pieces += [""] tokenizer = make_tokenizer(pieces, decoders.Fuse()) tools = [{"type": "function", "function": {"name": "weather", "parameters": {"type": "object"}}}] - text, _, content, _ = await chat_response(tokenizer, list(enumerate(pieces)), stream=False, tools=tools) + text, _, content, _ = await chat_response(tokenizer, list(enumerate(pieces)), stream=stream, tools=tools) assert text == ("Checking." if valid_tool else "".join(pieces)) assert [entry["token"] for entry in content] == (pieces[:1] if valid_tool else pieces) @@ -509,7 +564,8 @@ def test_reasoning_span_api_preserves_literal_delimiters_and_flush(monkeypatch): @pytest.mark.anyio @pytest.mark.parametrize("prefix", ["", "Paris"]) -async def test_tool_body_uses_generated_position(chat_response, prefix): +@pytest.mark.parametrize("stream", [False, True]) +async def test_tool_body_uses_generated_position(chat_response, prefix, stream): pieces = ([prefix] if prefix else []) + [ '<|python_tag|>{"name":"weather","arguments":{"city":"', "Paris", @@ -522,7 +578,7 @@ async def test_tool_body_uses_generated_position(chat_response, prefix): text, _, content, _ = await chat_response( tokenizer, [(tokenizer.convert_tokens_to_ids(piece), piece) for piece in pieces], - stream=False, + stream=stream, tools=tools, tool_parser="llama3", scores=scores, From b183ec080a1b83122983c64b108da7af9b117aa7 Mon Sep 17 00:00:00 2001 From: sufubao Date: Fri, 13 Nov 2026 21:16:13 +0800 Subject: [PATCH 6/6] fix(openai): preserve logprob metadata across numeric and PD edges --- .../triton_kernel/logprobs_capture.py | 45 ++++++++++++++++--- lightllm/server/core/objs/req.py | 5 ++- lightllm/server/pd_io_struct.py | 11 +++-- .../server/router/model_infer/infer_batch.py | 11 ++++- .../model_infer/mode_backend/base_backend.py | 5 ++- .../decode_node_impl/decode_trans_process.py | 4 ++ .../pd/prefill_node_impl/prefill_impl.py | 12 ++++- .../prefill_node_impl/prefill_impl_for_dp.py | 16 ++++++- .../prefill_trans_process.py | 2 - .../att_resources/test_network_transport.py | 2 + .../mode_backend/test_pd_notify.py | 13 ++++++ .../mode_backend/test_pd_page_io.py | 17 ++++++- .../structured_output/test_sampling.py | 30 ++++++++++++- .../structured_output/test_state.py | 11 +++++ .../server/test_chat_output_logprobs.py | 12 +++-- 15 files changed, 171 insertions(+), 25 deletions(-) diff --git a/lightllm/common/basemodel/triton_kernel/logprobs_capture.py b/lightllm/common/basemodel/triton_kernel/logprobs_capture.py index 60dd9d0096..b7ed312469 100644 --- a/lightllm/common/basemodel/triton_kernel/logprobs_capture.py +++ b/lightllm/common/basemodel/triton_kernel/logprobs_capture.py @@ -78,6 +78,9 @@ def _output_logprobs_partial( candidates, sampled_ids, ranks, + partial_maxes, + partial_sums, + partial_posinf_counts, stride_row: tl.constexpr, vocab_size: tl.constexpr, tiles: tl.constexpr, @@ -94,6 +97,14 @@ def _output_logprobs_partial( sampled = tl.load(row_logits + tl.load(sampled_ids + row)) rank = tl.sum(((ids < vocab_size) & (scores > sampled)).to(tl.int32), 0) tl.store(ranks + index * tiles + tile, rank) + posinf_count = tl.sum(((ids < vocab_size) & (scores == float("inf"))).to(tl.int32), 0) + finite_scores = tl.where((scores != float("inf")) & (scores != -float("inf")), scores, -float("inf")) + tile_max = tl.max(finite_scores, 0) + tile_sum = tl.sum(tl.exp(finite_scores - tile_max), 0) + tile_sum = tl.where(tile_max > -float("inf"), tile_sum, 0.0) + tl.store(partial_maxes + index * tiles + tile, tile_max) + tl.store(partial_sums + index * tiles + tile, tile_sum) + tl.store(partial_posinf_counts + index * tiles + tile, posinf_count) # tl.topk 只能比较整数。高 32 位放可排序的分数，低 32 位放反转后的 token id。 # 取最大的键时，分数高的在前；分数相同则 id 更小的在前。词表外的位置写成最小整数。 bits = scores.to(tl.int32, bitcast=True) @@ -116,6 +127,9 @@ def _output_logprobs_finish( logprobs, partial_ranks, sampled_ranks, + partial_maxes, + partial_sums, + partial_posinf_counts, stride_row: tl.constexpr, vocab_size: tl.constexpr, tiles: tl.constexpr, @@ -127,7 +141,6 @@ def _output_logprobs_finish( row = tl.load(rows + index * 3).to(tl.int64) row_logits = logits + row * stride_row sampled = tl.load(row_logits + tl.load(sampled_ids + row)) - sampled_logprob = tl.load(sampled_logprobs + row) width = tl.load(rows + index * 3 + 1) offset = tl.load(rows + index * 3 + 2) columns = tl.arange(0, TOPK) @@ -137,14 +150,25 @@ def _output_logprobs_finish( # 键的低 32 位是反转后的 token id，这里还原。结果仍按分数从高到低排列。 ids = 0xFFFFFFFF - (best & 0xFFFFFFFF) scores = tl.load(row_logits + ids, (columns < width) & (ids < vocab_size), -float("inf")) - stored_ids = tl.where(tl.abs(scores) < float("inf"), ids, -1).to(tl.int64) - # 候选和采样 token 的 logsumexp 相同，用采样 logprob 平移即可，不必再做整词表 softmax。 - row_logprobs = scores - sampled + sampled_logprob + stored_ids = tl.where(scores == scores, ids, -1).to(tl.int64) + rank_tiles = tl.arange(0, TILE_BLOCK) + tile_maxes = tl.load(partial_maxes + index * tiles + rank_tiles, rank_tiles < tiles, -float("inf")) + tile_sums = tl.load(partial_sums + index * tiles + rank_tiles, rank_tiles < tiles, 0.0) + posinf_counts = tl.load(partial_posinf_counts + index * tiles + rank_tiles, rank_tiles < tiles, 0) + row_max = tl.max(tile_maxes, 0) + row_sum = tl.sum(tl.where(tile_maxes > -float("inf"), tile_sums * tl.exp(tile_maxes - row_max), 0.0), 0) + posinf_count = tl.sum(posinf_counts, 0) + finite_logprobs = scores - (row_max + tl.log(row_sum)) + uniform_logprob = -tl.log(float(vocab_size)) + row_logprobs = tl.where( + posinf_count > 0, + tl.where(scores == float("inf"), -tl.log(posinf_count.to(tl.float32)), -float("inf")), + tl.where(row_max > -float("inf"), finite_logprobs, uniform_logprob), + ) tl.store(token_ids + offset + columns, stored_ids, columns < width) tl.store(logprobs + offset + columns, row_logprobs, columns < width) - rank_tiles = tl.arange(0, TILE_BLOCK) rank = 1 + tl.sum(tl.load(partial_ranks + index * tiles + rank_tiles, rank_tiles < tiles, 0), 0) - rank = tl.where(tl.abs(sampled) < float("inf"), rank, vocab_size + 1) + rank = tl.where(sampled > -float("inf"), rank, vocab_size + 1) tl.store(sampled_ranks + index, rank) @@ -170,12 +194,18 @@ def capture_output_logprobs( candidates = logits.new_empty((rows.shape[0], tiles, topk), dtype=torch.int64) partial_ranks = logits.new_empty((rows.shape[0], tiles), dtype=torch.int32) sampled_ranks = logits.new_empty((rows.shape[0],), dtype=torch.int32) + partial_maxes = logits.new_empty((rows.shape[0], tiles), dtype=torch.float32) + partial_sums = logits.new_empty((rows.shape[0], tiles), dtype=torch.float32) + partial_posinf_counts = logits.new_empty((rows.shape[0], tiles), dtype=torch.int32) _output_logprobs_partial[(rows.shape[0], tiles)]( logits=logits, rows=rows, candidates=candidates, sampled_ids=sampled_ids, ranks=partial_ranks, + partial_maxes=partial_maxes, + partial_sums=partial_sums, + partial_posinf_counts=partial_posinf_counts, stride_row=logits.stride(0), vocab_size=logits.shape[1], tiles=tiles, @@ -195,6 +225,9 @@ def capture_output_logprobs( logprobs=logprobs, partial_ranks=partial_ranks, sampled_ranks=sampled_ranks, + partial_maxes=partial_maxes, + partial_sums=partial_sums, + partial_posinf_counts=partial_posinf_counts, stride_row=logits.stride(0), vocab_size=logits.shape[1], tiles=tiles, diff --git a/lightllm/server/core/objs/req.py b/lightllm/server/core/objs/req.py index 29f5e29d98..dc074256fb 100644 --- a/lightllm/server/core/objs/req.py +++ b/lightllm/server/core/objs/req.py @@ -349,7 +349,8 @@ def _create_output_topk_logprobs_shm_array(self, topk: int) -> None: self.shm_output_topk_logprobs = ShmArray( name=f"shm_output_topk_logprobs_{self.index_in_shm_mem}", shape=(self.sample_params.max_new_tokens + 128, 2, topk), - dtype=np.float32, + # Token IDs must stay exact beyond float32's 2^24 integer limit. + dtype=np.float64, ) self.shm_output_topk_logprobs.create_shm() # PD decode 可能直接复用 prefill 已采样的第一个 token，先把 token id 标成无效。 @@ -360,7 +361,7 @@ def _link_output_topk_logprobs_shm_array(self, topk: int) -> None: self.shm_output_topk_logprobs = ShmArray( name=f"shm_output_topk_logprobs_{self.index_in_shm_mem}", shape=(self.sample_params.max_new_tokens + 128, 2, topk), - dtype=np.float32, + dtype=np.float64, ) self.shm_output_topk_logprobs.link_shm() diff --git a/lightllm/server/pd_io_struct.py b/lightllm/server/pd_io_struct.py index 52e27992ec..fd787c6eea 100644 --- a/lightllm/server/pd_io_struct.py +++ b/lightllm/server/pd_io_struct.py @@ -104,7 +104,6 @@ class PDUpKVStatus: pd_kv_trans_params: bytes # pd kv 传输建立连接所使用的元数据对象 def __post_init__(self): - if not isinstance(self.group_request_id, int): error_info = "group_request_id only can be int" logger.error(error_info) @@ -173,8 +172,10 @@ class PDChunckedTransTask: decode_num_pages: Optional[int] decode_page_reg_desc: Optional[bytes] - first_gen_token_id: Optional[int] - first_gen_token_logprob: Optional[float] + first_gen_token_id: Optional[int] = None + first_gen_token_logprob: Optional[float] = None + first_gen_token_rank: Optional[int] = None + first_gen_token_output_top_logprobs: Optional[tuple[list[int], list[float]]] = None write_stage: Optional[str] = None @@ -273,6 +274,8 @@ def createRetObj(self) -> "PDChunckedTransTaskRet": error_info=self.error_info, first_gen_token_id=self.first_gen_token_id, first_gen_token_logprob=self.first_gen_token_logprob, + first_gen_token_rank=self.first_gen_token_rank, + first_gen_token_output_top_logprobs=self.first_gen_token_output_top_logprobs, ) return ret @@ -302,6 +305,8 @@ class PDChunckedTransTaskRet: error_info: str = None first_gen_token_id: Optional[int] = None first_gen_token_logprob: Optional[float] = None + first_gen_token_rank: Optional[int] = None + first_gen_token_output_top_logprobs: Optional[tuple[list[int], list[float]]] = None def get_key(self) -> str: return f"{self.request_id}_{self.start_kv_index}_{self.end_kv_index}" diff --git a/lightllm/server/router/model_infer/infer_batch.py b/lightllm/server/router/model_infer/infer_batch.py index 0026c7cbf0..aea1e7ffb5 100644 --- a/lightllm/server/router/model_infer/infer_batch.py +++ b/lightllm/server/router/model_infer/infer_batch.py @@ -1009,13 +1009,20 @@ def handle( eos_ids: List[int], is_master_in_dp: bool, extra_post_req_handle_func: Optional[Callable[[InferReq, int, float], None]] = None, - pd_prefill_chunked_handle_func: Optional[Callable[[InferReq, int, float, int], None]] = None, + pd_prefill_chunked_handle_func: Optional[Callable[[InferReq, int, float, int, int, object], None]] = None, output_top_logprobs=None, ): # pd_prefill_chunked_handle_func 主要是为了处理 pd prefill 模式下 # 分块 prefill 后，形成对应的pd 分块传输处理。 if pd_prefill_chunked_handle_func is not None: - pd_prefill_chunked_handle_func(self.req_obj, next_token_id, next_token_logprob, self.output_len) + pd_prefill_chunked_handle_func( + self.req_obj, + next_token_id, + next_token_logprob, + self.output_len, + next_token_rank, + output_top_logprobs, + ) if self.output_len <= 0: return 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 9a64e36af9..6ab4920fc4 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -603,9 +603,12 @@ def _read_pd_trans_io_buffer_and_update_req_status(self): InferReqUpdatePack(req_obj=req, output_len=req.cur_output_len).handle( next_token_id=obj.first_gen_token_id, next_token_logprob=obj.first_gen_token_logprob, - next_token_rank=-1, + next_token_rank=( + obj.first_gen_token_rank if obj.first_gen_token_rank is not None else -1 + ), eos_ids=self.eos_id, is_master_in_dp=self.is_master_in_dp, + output_top_logprobs=obj.first_gen_token_output_top_logprobs, ) return 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 ca4b51e9b5..480e28fba9 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 @@ -333,6 +333,10 @@ def accept_peer_task_loop( local_trans_task.shared_time_out_mark.update() local_trans_task.first_gen_token_id = remote_trans_task.first_gen_token_id local_trans_task.first_gen_token_logprob = remote_trans_task.first_gen_token_logprob + local_trans_task.first_gen_token_rank = remote_trans_task.first_gen_token_rank + local_trans_task.first_gen_token_output_top_logprobs = ( + remote_trans_task.first_gen_token_output_top_logprobs + ) self.ready_page_task_queue.put(local_trans_task) if self.args.pd_detail_log: logger.debug(f"recv WRITE done from prefill: {remote_trans_task.to_str()}") diff --git a/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_impl.py b/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_impl.py index da4ce0d073..1d7ec7d9a1 100644 --- a/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_impl.py +++ b/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_impl.py @@ -66,7 +66,13 @@ def _filter_not_ready_reqs(self, req_ids: List[int]) -> List[InferReq]: return ans_list def _prefill_chuncked_handle_func( - self, req_obj: InferReq, next_token_id: int, next_token_prob: float, output_len: int + self, + req_obj: InferReq, + next_token_id: int, + next_token_prob: float, + output_len: int, + next_token_rank: int, + output_top_logprobs=None, ): """ 在每一步chuncked prefill 后，尝试生成chuncked 传输任务，发个 kv_move_manager 进行处理。 @@ -105,6 +111,8 @@ def _prefill_chuncked_handle_func( ) trans_task_list[-1].first_gen_token_id = next_token_id trans_task_list[-1].first_gen_token_logprob = next_token_prob + trans_task_list[-1].first_gen_token_rank = next_token_rank + trans_task_list[-1].first_gen_token_output_top_logprobs = output_top_logprobs if self.is_master_in_dp: for trans_task in trans_task_list: @@ -155,6 +163,8 @@ def _create_pd_trans_task( decode_page_reg_desc=pd_decode_node_info.page_reg_desc, first_gen_token_id=None, first_gen_token_logprob=None, + first_gen_token_rank=None, + first_gen_token_output_top_logprobs=None, page_kind=page_kind, req_idx=req_idx, ) diff --git a/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_impl_for_dp.py b/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_impl_for_dp.py index 5f4465a730..673a17365e 100644 --- a/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_impl_for_dp.py +++ b/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_impl_for_dp.py @@ -24,10 +24,22 @@ def _filter_not_ready_reqs(self, req_ids: List[int]) -> List[InferReq]: return PDChunkedPrefillForPrefillNode._filter_not_ready_reqs(self, req_ids) def _prefill_chuncked_handle_func( - self, req_obj: InferReq, next_token_id: int, next_token_prob: float, output_len: int + self, + req_obj: InferReq, + next_token_id: int, + next_token_prob: float, + output_len: int, + next_token_rank: int, + output_top_logprobs=None, ): return PDChunkedPrefillForPrefillNode._prefill_chuncked_handle_func( - self, req_obj=req_obj, next_token_id=next_token_id, next_token_prob=next_token_prob, output_len=output_len + self, + req_obj=req_obj, + next_token_id=next_token_id, + next_token_prob=next_token_prob, + output_len=output_len, + next_token_rank=next_token_rank, + output_top_logprobs=output_top_logprobs, ) def _create_pd_trans_task( 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 b3c1dcc609..dfeb322d51 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 @@ -446,8 +446,6 @@ def success_loop(self): self.transporter.release_xfer_handle(trans_task.xfer_handle) ret = trans_task.createRetObj() - ret.first_gen_token_id = None - ret.first_gen_token_logprob = None self.task_out_queue.put(ret) if self.args.pd_detail_log: diff --git a/unit_tests/common/att_resources/test_network_transport.py b/unit_tests/common/att_resources/test_network_transport.py index f66767b9f4..c416d12cdc 100644 --- a/unit_tests/common/att_resources/test_network_transport.py +++ b/unit_tests/common/att_resources/test_network_transport.py @@ -122,6 +122,8 @@ def _task(peer: PDAgentMetadata, kind: str): decode_page_reg_desc=peer.page_reg_desc, first_gen_token_id=42, first_gen_token_logprob=-0.5, + first_gen_token_rank=21, + first_gen_token_output_top_logprobs=([42, 7], [-0.5, -0.6]), src_page_index=1 if kind == "kv" else 2, page_kind=kind, req_idx=1, 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 8f99c6add3..98042570c1 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 @@ -40,6 +40,8 @@ def make_task(page_kind): decode_page_reg_desc=b"decode pages", first_gen_token_id=42, first_gen_token_logprob=-0.5, + first_gen_token_rank=21, + first_gen_token_output_top_logprobs=([42, 7], [-0.5, -0.6]), page_kind=page_kind, req_idx=17, src_page_index=3, @@ -77,6 +79,17 @@ def test_copy_for_notify_preserves_local_task(page_kind, with_mark): assert task.shared_time_out_mark.time_mark == time_mark +def test_task_return_preserves_first_token_metadata(): + task = make_task("kv") + + result = task.createRetObj() + + assert result.first_gen_token_id == 42 + assert result.first_gen_token_logprob == -0.5 + assert result.first_gen_token_rank == 21 + assert result.first_gen_token_output_top_logprobs == ([42, 7], [-0.5, -0.6]) + + _UNPICKLE_WITHOUT_DECODE = """ import importlib.abc import pickle diff --git a/unit_tests/server/router/model_infer/mode_backend/test_pd_page_io.py b/unit_tests/server/router/model_infer/mode_backend/test_pd_page_io.py index d2adf388f4..3c7ad49f4a 100644 --- a/unit_tests/server/router/model_infer/mode_backend/test_pd_page_io.py +++ b/unit_tests/server/router/model_infer/mode_backend/test_pd_page_io.py @@ -29,7 +29,11 @@ def test_hybrid_pd_uses_the_shared_att_state_task_without_model_flags(monkeypatc def create(**kwargs): task = SimpleNamespace( - page_kind=kwargs.get("page_kind", "kv"), first_gen_token_id=None, first_gen_token_logprob=None + page_kind=kwargs.get("page_kind", "kv"), + first_gen_token_id=None, + first_gen_token_logprob=None, + first_gen_token_rank=None, + first_gen_token_output_top_logprobs=None, ) tasks.append(task) if "group" in kwargs: @@ -38,8 +42,17 @@ def create(**kwargs): backend._create_pd_trans_task = create if role == "prefill": - backend._prefill_chuncked_handle_func(req, next_token_id=17, next_token_prob=-0.2, output_len=1) + backend._prefill_chuncked_handle_func( + req, + next_token_id=17, + next_token_prob=-0.2, + output_len=1, + next_token_rank=21, + output_top_logprobs=([17, 16], [-0.2, -0.3]), + ) assert tasks[-1].first_gen_token_id == 17 + assert tasks[-1].first_gen_token_rank == 21 + assert tasks[-1].first_gen_token_output_top_logprobs == ([17, 16], [-0.2, -0.3]) else: backend.radix_cache = None backend.model = SimpleNamespace( diff --git a/unit_tests/server/router/model_infer/structured_output/test_sampling.py b/unit_tests/server/router/model_infer/structured_output/test_sampling.py index 4bf33aaec4..e26c60d9a2 100644 --- a/unit_tests/server/router/model_infer/structured_output/test_sampling.py +++ b/unit_tests/server/router/model_infer/structured_output/test_sampling.py @@ -147,7 +147,7 @@ def test_sample_candidate_logprobs_follow_mtp_row_selection(compiler, sampling_m reference = logits.log_softmax(dim=-1).cpu() assert middle["next_token_logprob"] == pytest.approx(reference[2, 10].item()) assert first["output_top_logprobs"][1] == pytest.approx(reference[0, [256, 255, 254]].tolist()) - assert last["output_top_logprobs"][1] == pytest.approx(reference[3, [256, 255]].tolist()) + assert last["output_top_logprobs"][1] == pytest.approx(reference[3, [256, 255]].tolist(), rel=1e-5) assert logprobs_cpu[1].item() == pytest.approx(reference[1, 0].item()) assert first["next_token_rank"] == 1 assert last["next_token_rank"] == 1 @@ -467,3 +467,31 @@ def test_capture_rank_across_vocab_tiles_preserves_sampled_scores(vocab_size): for row, width, offset in rows.tolist(): assert candidate_ids[offset : offset + width].tolist() == list(range(width)) torch.testing.assert_close(scores[offset : offset + width], logits.log_softmax(-1)[row, :width]) + + +@pytest.mark.parametrize( + "logits,sampled_id,expected_ids,expected_rank", + [ + ([0.0, -torch.inf, 1.0], 1, [2, 0, 1], 4), + ([0.0, torch.inf, 1.0], 1, [1, 2, 0], 1), + ([-torch.inf, -torch.inf, -torch.inf], 0, [0, 1, 2], 4), + ], +) +def test_capture_preserves_nonfinite_candidates(logits, sampled_id, expected_ids, expected_rank): + logits = torch.tensor([logits], device="cuda") + ids = torch.tensor([sampled_id], device="cuda") + logprobs = torch.tensor([[0.0]], device="cuda") + rows = torch.tensor([[0, 3, 0]], device="cuda", dtype=torch.int32) + + candidate_ids, scores, ranks = sampling.capture_output_logprobs(logits, ids, logprobs, rows, 3, 3) + + assert candidate_ids.tolist() == expected_ids + assert ranks.tolist() == [expected_rank] + assert not torch.isnan(scores).any() + if logits[0, sampled_id] == -torch.inf and not torch.isneginf(logits).all(): + torch.testing.assert_close(scores[:2], logits.log_softmax(-1)[0, [2, 0]]) + assert scores[2].item() == -torch.inf + elif logits[0, sampled_id] == torch.inf: + assert scores.tolist() == [0.0, -torch.inf, -torch.inf] + else: + torch.testing.assert_close(scores, torch.full_like(scores, -torch.log(torch.tensor(3.0)))) diff --git a/unit_tests/server/router/model_infer/structured_output/test_state.py b/unit_tests/server/router/model_infer/structured_output/test_state.py index d76269b0d6..9a76597ce3 100644 --- a/unit_tests/server/router/model_infer/structured_output/test_state.py +++ b/unit_tests/server/router/model_infer/structured_output/test_state.py @@ -5,6 +5,7 @@ from unittest.mock import patch import pytest +import numpy as np import torch import xgrammar as xgr @@ -308,6 +309,9 @@ def test_pd_first_token_updates_constraint_before_sampling( req.pd_task_success_num = 0 req.shm_req.shm_prompt_ids.arr.append(0) req.shm_req.shm_logprobs = SimpleNamespace(arr=[None] * (req.shm_req.input_len + 1)) + req.sampling_param.shm_param.output_top_logprobs = 2 + req.shm_req.shm_output_topk_logprobs = SimpleNamespace(arr=np.full((2, 2, 2), -1, dtype=np.float64)) + req.shm_req.output_topk_logprobs_row = lambda output_index: req.shm_req.shm_output_topk_logprobs.arr[output_index] req.set_next_gen_token_id = MethodType(infer_batch.InferReq.set_next_gen_token_id, req) req.update_finish_status = lambda **kwargs: None sampling_params = make_mask_buffers([req, good]) @@ -316,6 +320,8 @@ def test_pd_first_token_updates_constraint_before_sampling( has_error=False, first_gen_token_id=first_token, first_gen_token_logprob=0.0, + first_gen_token_rank=21, + first_gen_token_output_top_logprobs=([first_token, 7], [-0.1, -0.2]), ) backend = SimpleNamespace( shm_pd_trans_io_buffer=SimpleNamespace(read_obj=lambda: [transfer], sub_state=lambda: None), @@ -333,9 +339,14 @@ def test_pd_first_token_updates_constraint_before_sampling( assert req.cur_output_len == 1 and req.pd_task_success_num == 1 assert req.shm_req.shm_prompt_ids.arr == [ord("P")] + list(history) + [first_token] + assert req.shm_req.shm_logprobs.arr[req.shm_req.input_len] == (0.0, 21) assert sampling_params.req_to_next_token_ids[req.req_idx, 0].item() == first_token if is_master_in_dp: assert req.shm_req.shm_cur_output_len == 1 + assert req.shm_req.shm_output_topk_logprobs.arr[0].tolist() == [ + [first_token, 7], + [-0.1, -0.2], + ] state = req.output_constraint if first_token == ord("!"): assert req.finish_status.status == FinishStatus.FINISHED_ERROR diff --git a/unit_tests/server/test_chat_output_logprobs.py b/unit_tests/server/test_chat_output_logprobs.py index d072f15c0c..f7d4704921 100644 --- a/unit_tests/server/test_chat_output_logprobs.py +++ b/unit_tests/server/test_chat_output_logprobs.py @@ -128,13 +128,14 @@ def test_candidate_shm_is_reused_and_resized(monkeypatch): assert worker.shm_output_topk_logprobs.shm.size == worker.shm_output_topk_logprobs.arr.nbytes assert worker.output_topk_logprobs_row(output_index=0)[0].tolist() == [-1, -1] worker.shm_logprobs.arr[3] = (-0.25, 1) - worker.output_topk_logprobs_row(output_index=0)[:] = ([9, 8], [-0.5, -1.0]) + worker.output_topk_logprobs_row(output_index=0)[:] = ([16_777_217, 8], [-0.5, -1.0]) worker.output_topk_logprobs_row(output_index=1)[:] = ([7, 6], [-0.25, -1.5]) worker.detach_shm_arrays() owner.detach_shm_arrays() owner.link_logprobs_shm_array() assert owner.shm_logprobs.arr[3].tolist() == (-0.25, 1) + assert owner.output_topk_logprobs_row(output_index=0)[0].tolist() == [16_777_217, 8] assert owner.output_topk_logprobs_row(output_index=1).tolist() == [[7, 6], [-0.25, -1.5]] owner.detach_shm_arrays() @@ -168,7 +169,12 @@ async def test_pd_first_token_uses_the_node_that_sampled_it(monkeypatch, needs_p d_first = ( 0, "first", - {"count_output_tokens": 1, "node_mode": "decode", "output_top_logprobs": [], "logprob_prompt_ids": prompt_ids}, + { + "count_output_tokens": 1, + "node_mode": "decode", + "output_top_logprobs": [(1, -0.5)], + "logprob_prompt_ids": prompt_ids, + }, finish, ) p_first = ( @@ -208,7 +214,7 @@ async def test_pd_first_token_uses_the_node_that_sampled_it(monkeypatch, needs_p _, text, metadata, _ = await stream.__anext__() assert text == "first" assert metadata["logprob_prompt_ids"] == prompt_ids - assert metadata["output_top_logprobs"] == ([(1, -0.5)] if needs_prefill else []) + assert metadata["output_top_logprobs"] == [(1, -0.5)] _, text, _, status = await stream.__anext__() assert text == "second" and status.status == end_status finally: