From a9c30b491bfedbb949992efdcd38b49b4a38bb9d Mon Sep 17 00:00:00 2001 From: sufubao Date: Thu, 12 Nov 2026 00:58:41 +0800 Subject: [PATCH 1/6] feat: support Qwen3.8 Flash Next with ReplaySSM --- docs/qwen38_adaptation.rst | 48 + docs/qwen38_h100_performance.rst | 76 ++ .../bounded_att/gdn/create_gdn_state.py | 18 +- .../bounded_att/gdn/replay_gdn_state.py | 835 ++++++++++++++++++ .../bounded_att/kda/create_kda_state.py | 18 +- .../bounded_att/kda/replay_kda_state.py | 834 +++++++++++++++++ .../attention_state/bounded_att/utils.py | 64 +- .../attention_state/full_att/gqa/gqa_state.py | 27 +- .../attention/create_linear_utils.py | 26 +- .../common/basemodel/attention/linear/gdn.py | 53 +- .../basemodel/attention/linear/replay_gdn.py | 352 ++++++++ .../attention/linear/replay_gdn_decode.py | 86 ++ .../basemodel/attention/linear/replay_kda.py | 241 +++++ .../attention/linear/replay_kda_decode.py | 86 ++ .../meta_weights/pinned_embedding_weight.py | 66 ++ .../row_sharded_embedding_weight.py | 64 ++ .../common/basemodel/multimodal_tokenizer.py | 3 + lightllm/common/basemodel/tp_reduce.py | 58 ++ .../triton_kernel/destindex_copy_kv.py | 74 ++ .../common/basemodel/triton_kernel/kv_move.py | 97 ++ .../linear_att/fla/ops/__init__.py | 3 +- .../linear_att/fla/ops/fused_recurrent.py | 165 ++++ .../linear_att/mtp_state_params.py | 2 +- .../triton_kernel/linear_att/replayssm.py | 29 + .../linear_att/replayssm_cursor.py | 24 + .../linear_att/replayssm_decode.py | 792 +++++++++++++++++ .../linear_att/replayssm_delta.py | 21 + .../linear_att/replayssm_fold.py | 104 +++ .../linear_att/ssm_acceptance.py | 55 ++ .../triton_kernel/linear_att/ssm_autotune.py | 180 ++++ .../basemodel/triton_kernel/mtp_utils.py | 6 +- .../basemodel/triton_kernel/scalar_linear.py | 42 + lightllm/common/gated_residual/__init__.py | 85 ++ lightllm/common/gated_residual/kernels.py | 103 +++ lightllm/common/gated_residual/projection.py | 156 ++++ lightllm/common/ngram_embedding/__init__.py | 132 +++ lightllm/common/ngram_embedding/kernels.py | 94 ++ lightllm/common/ngram_embedding/ple.py | 198 +++++ .../common/ngram_embedding/ple_kernels.py | 100 +++ lightllm/common/qsa/__init__.py | 1 + lightllm/common/qsa/attention.py | 120 +++ lightllm/common/qsa/attention_kernels.py | 137 +++ lightllm/common/qsa/indexer.py | 279 ++++++ lightllm/common/qsa/indexer_kernels.py | 216 +++++ lightllm/common/qsa/topk.py | 127 +++ .../common/quantization/aligned_projection.py | 51 ++ .../common/quantization/scalar_projection.py | 21 + ...fp8_w8a8=false}_NVIDIA_H100_80GB_HBM3.json | 92 ++ ...fp8_w8a8=false}_NVIDIA_H100_80GB_HBM3.json | 92 ++ ...num=1,use_fp8_w8a8=false}_NVIDIA_H200.json | 47 + ...um=11,use_fp8_w8a8=false}_NVIDIA_H200.json | 65 ++ lightllm/distributed/symm_mem_all_reduce.py | 19 + lightllm/models/builtin.py | 8 + lightllm/models/glm5_next/config.py | 5 +- lightllm/models/glm5_next/model.py | 84 +- .../qwen2_vl/triton_kernel/rotary_pos_emb.py | 9 +- .../layer_infer/transformer_layer_infer.py | 12 +- lightllm/models/qwen3next/model.py | 30 +- lightllm/models/qwen4_exp/__init__.py | 1 + lightllm/models/qwen4_exp/attention_state.py | 100 +++ lightllm/models/qwen4_exp/batch.py | 119 +++ lightllm/models/qwen4_exp/checkpoint.py | 38 + lightllm/models/qwen4_exp/configuration.py | 24 + lightllm/models/qwen4_exp/image_processing.py | 9 + lightllm/models/qwen4_exp/infer_struct.py | 76 ++ .../models/qwen4_exp/layer_infer/__init__.py | 1 + .../qwen4_exp/layer_infer/post_layer_infer.py | 6 + .../qwen4_exp/layer_infer/pre_layer_infer.py | 23 + .../layer_infer/transformer_layer_infer.py | 93 ++ .../qwen4_exp/layer_weights/__init__.py | 1 + .../qwen4_exp/layer_weights/components.py | 129 +++ .../layer_weights/embedding_weight.py | 14 + .../qwen4_exp/layer_weights/lm_head_weight.py | 25 + .../qwen4_exp/layer_weights/norm_weight.py | 30 + .../pre_and_post_layer_weight.py | 26 + .../layer_weights/transformer_layer_weight.py | 85 ++ lightllm/models/qwen4_exp/lexical_cache.py | 89 ++ lightllm/models/qwen4_exp/lexical_state.py | 257 ++++++ lightllm/models/qwen4_exp/model.py | 114 +++ lightllm/models/qwen4_exp/mtp/__init__.py | 1 + lightllm/models/qwen4_exp/mtp/infer.py | 29 + lightllm/models/qwen4_exp/mtp/model.py | 51 ++ lightllm/models/qwen4_exp/mtp/weights.py | 55 ++ lightllm/models/qwen4_exp/multimodal.py | 43 + lightllm/models/qwen4_exp/pinned_memory.py | 72 ++ lightllm/models/qwen4_exp/position.py | 18 + .../models/qwen4_exp/speculative_cache.py | 108 +++ lightllm/models/qwen4_exp/state.py | 136 +++ lightllm/models/qwen4_exp/tokenizer.py | 6 + .../qwen4_exp/triton_kernel/__init__.py | 1 + .../qwen4_exp/triton_kernel/cache_pages.py | 71 ++ .../qwen4_exp/triton_kernel/gated_norm.py | 92 ++ lightllm/models/qwen4_exp/vision.py | 36 + lightllm/server/api_cli.py | 21 + lightllm/server/api_lightllm.py | 13 +- lightllm/server/api_start.py | 55 ++ lightllm/server/core/objs/start_args_type.py | 3 + lightllm/server/httpserver/manager.py | 19 +- .../server/httpserver/slot_availability.py | 38 + .../model_infer/mode_backend/base_backend.py | 4 + .../mtp_speculative/planner/lightspec.py | 4 + lightllm/server/router/req_queue/__init__.py | 8 + .../server/router/req_queue/cached_prefix.py | 71 ++ .../chunked_prefill/cached_prefix_queue.py | 42 + .../router/req_queue/chunked_prefill/impl.py | 35 +- .../router/req_queue/peak_reservation.py | 52 ++ lightllm/server/tokenizer.py | 9 + .../visualserver/model_infer/model_rpc.py | 4 + lightllm/utils/config_utils.py | 18 +- test/benchmark/kernels/benchmark_kda_ssm.py | 162 ++++ .../kernels/benchmark_replayssm_service.py | 119 +++ .../qwen38/benchmark_gated_residual.py | 93 ++ test/benchmark/qwen38/benchmark_lm_head.py | 111 +++ test/benchmark/qwen38/benchmark_ngram_hash.py | 129 +++ test/benchmark/qwen38/benchmark_ple.py | 120 +++ test/benchmark/qwen38/benchmark_qsa.py | 165 ++++ test/benchmark/qwen38/check_aiperf_timing.py | 60 ++ test/benchmark/qwen38/check_baseline.py | 102 +++ test/benchmark/qwen38/check_cold_streams.py | 159 ++++ test/benchmark/qwen38/check_cpu_cache.py | 146 +++ test/benchmark/qwen38/check_generation.py | 228 +++++ test/benchmark/qwen38/check_math_quality.py | 129 +++ test/benchmark/qwen38/check_multimodal.py | 291 ++++++ .../benchmark/qwen38/check_prefix_eviction.py | 174 ++++ test/benchmark/qwen38/check_retrieval.py | 114 +++ test/benchmark/qwen38/check_scheduler.py | 215 +++++ test/benchmark/qwen38/check_streaming.py | 76 ++ .../qwen38/generation_observations.py | 48 + test/benchmark/qwen38/math_quality.py | 127 +++ test/benchmark/qwen38/measure_aiperf.py | 310 +++++++ test/benchmark/qwen38/measure_fixed_work.py | 152 ++++ test/benchmark/qwen38/prepare_aiperf_trace.py | 72 ++ test/benchmark/qwen38/retrieval_quality.py | 51 ++ test/benchmark/qwen38/score_reference.py | 208 +++++ test/benchmark/service/cache_capacity.py | 50 ++ test/benchmark/service/completion.py | 61 ++ test/benchmark/service/completion_runner.py | 63 ++ test/benchmark/service/stream_cohort.py | 166 ++++ test/benchmark/service/stream_timing.py | 80 ++ test/benchmark/service/token_stream.py | 237 +++++ .../benchmark/qwen38/test_aiperf_replay.py | 139 +++ .../qwen38/test_check_cold_streams.py | 61 ++ .../benchmark/qwen38/test_check_cpu_cache.py | 66 ++ .../benchmark/qwen38/test_check_generation.py | 95 ++ .../qwen38/test_check_math_quality.py | 133 +++ .../benchmark/qwen38/test_check_multimodal.py | 98 ++ .../qwen38/test_check_prefix_eviction.py | 86 ++ .../benchmark/qwen38/test_check_scheduler.py | 121 +++ .../qwen38/test_generation_observations.py | 62 ++ .../benchmark/qwen38/test_math_quality.py | 145 +++ .../qwen38/test_measure_fixed_work.py | 55 ++ .../qwen38/test_retrieval_quality.py | 62 ++ .../benchmark/qwen38/test_score_reference.py | 49 + .../benchmark/service/test_cache_capacity.py | 62 ++ .../benchmark/service/test_completion.py | 88 ++ .../service/test_completion_runner.py | 111 +++ .../benchmark/service/test_stream_cohort.py | 185 ++++ .../benchmark/service/test_stream_timing.py | 79 ++ .../benchmark/service/test_token_stream.py | 287 ++++++ .../common/att_resources/model_helpers.py | 5 + .../test_hybrid_layer_mapping.py | 15 +- .../att_resources/test_page_pool_storage.py | 92 +- .../basemodel/attention/linear/test_gdn.py | 114 +++ .../linear_att/test_acceptance_fusion.py | 243 +++++ .../linear_att/test_causal_conv1d_mtp.py | 25 + .../test_fused_recurrent_packed_decode.py | 62 ++ .../test_mtp_fused_recurrent_equiv.py | 52 ++ .../linear_att/test_mtp_state_params.py | 65 ++ .../linear_att/test_replayssm.py | 776 ++++++++++++++++ .../linear_att/test_replayssm_execution.py | 473 ++++++++++ .../linear_att/test_ssm_autotune_memory.py | 82 ++ .../test_gated_rmsnorm_strided_gate.py | 41 + .../triton_kernel/test_scalar_linear.py | 70 ++ .../gated_residual/test_gated_residual.py | 178 ++++ .../common/ngram_embedding/test_hash.py | 206 +++++ .../ngram_embedding/test_pinned_embedding.py | 88 ++ unit_tests/common/ngram_embedding/test_ple.py | 259 ++++++ .../test_row_sharded_embedding.py | 161 ++++ unit_tests/common/qsa/test_attention.py | 133 +++ unit_tests/common/qsa/test_flashinfer_topk.py | 84 ++ unit_tests/common/qsa/test_indexer.py | 419 +++++++++ .../common/qsa/test_topk_compatibility.py | 142 +++ unit_tests/common/qsa/test_topk_cuda.py | 62 ++ .../quantization/test_aligned_projection.py | 97 ++ .../test_hybrid_req_manager_pd_sharing.py | 144 +++ unit_tests/common/test_linear_state.py | 268 ++++++ .../common/test_linear_state_contract.py | 201 +++++ unit_tests/conftest.py | 13 + unit_tests/distributed/test_ordered_tp_sum.py | 110 +++ unit_tests/models/glm5_next/test_cache.py | 25 +- unit_tests/models/glm5_next/test_mtp.py | 103 ++- unit_tests/models/glm5_next/test_pd_cache.py | 35 +- .../models/qwen2_vl/test_rotary_pos_emb.py | 17 +- ...st_transformer_layer_infer_strided_gate.py | 33 + unit_tests/models/qwen4_exp/__init__.py | 1 + unit_tests/models/qwen4_exp/conftest.py | 84 ++ unit_tests/models/qwen4_exp/helpers.py | 133 +++ .../models/qwen4_exp/lexical_helpers.py | 32 + unit_tests/models/qwen4_exp/mtp_reference.py | 27 + unit_tests/models/qwen4_exp/reference.py | 235 +++++ .../models/qwen4_exp/resource_helpers.py | 73 ++ unit_tests/models/qwen4_exp/runtime.py | 133 +++ .../models/qwen4_exp/test_attention_state.py | 205 +++++ unit_tests/models/qwen4_exp/test_batch.py | 97 ++ unit_tests/models/qwen4_exp/test_decoder.py | 214 +++++ .../qwen4_exp/test_embedding_residency.py | 45 + unit_tests/models/qwen4_exp/test_eos.py | 20 + .../models/qwen4_exp/test_lexical_cache.py | 100 +++ .../models/qwen4_exp/test_linear_decoder.py | 183 ++++ unit_tests/models/qwen4_exp/test_lm_head.py | 45 + unit_tests/models/qwen4_exp/test_mtp.py | 148 ++++ .../models/qwen4_exp/test_multimodal.py | 213 +++++ unit_tests/models/qwen4_exp/test_norm.py | 48 + .../models/qwen4_exp/test_pinned_memory.py | 112 +++ unit_tests/models/qwen4_exp/test_reduction.py | 104 +++ unit_tests/models/qwen4_exp/test_runtime.py | 228 +++++ .../qwen4_exp/test_sequence_parallel.py | 87 ++ .../qwen4_exp/test_speculative_cache.py | 159 ++++ unit_tests/models/qwen4_exp/test_vision.py | 49 + unit_tests/models/qwen4_exp/test_weights.py | 99 +++ .../models/qwen4_exp/vision_reference.py | 93 ++ .../models/sensenova_loop/test_gdn_state.py | 150 ++++ .../httpserver/test_slot_availability.py | 274 ++++++ .../mtp_speculative/test_planner.py | 83 ++ .../req_queue/test_cached_prefix_queue.py | 224 +++++ .../router/req_queue/test_peak_reservation.py | 96 ++ unit_tests/server/test_hybrid_start_args.py | 34 + .../server/test_native_stream_details.py | 69 ++ unit_tests/server/test_ssm_state_mode.py | 167 ++++ 229 files changed, 24167 insertions(+), 199 deletions(-) create mode 100644 docs/qwen38_adaptation.rst create mode 100644 docs/qwen38_h100_performance.rst create mode 100644 lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py create mode 100644 lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py create mode 100644 lightllm/common/basemodel/attention/linear/replay_gdn.py create mode 100644 lightllm/common/basemodel/attention/linear/replay_gdn_decode.py create mode 100644 lightllm/common/basemodel/attention/linear/replay_kda.py create mode 100644 lightllm/common/basemodel/attention/linear/replay_kda_decode.py create mode 100644 lightllm/common/basemodel/layer_weights/meta_weights/pinned_embedding_weight.py create mode 100644 lightllm/common/basemodel/layer_weights/meta_weights/row_sharded_embedding_weight.py create mode 100644 lightllm/common/basemodel/tp_reduce.py create mode 100644 lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py create mode 100644 lightllm/common/basemodel/triton_kernel/linear_att/replayssm_cursor.py create mode 100644 lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py create mode 100644 lightllm/common/basemodel/triton_kernel/linear_att/replayssm_delta.py create mode 100644 lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py create mode 100644 lightllm/common/basemodel/triton_kernel/linear_att/ssm_acceptance.py create mode 100644 lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py create mode 100644 lightllm/common/basemodel/triton_kernel/scalar_linear.py create mode 100644 lightllm/common/gated_residual/__init__.py create mode 100644 lightllm/common/gated_residual/kernels.py create mode 100644 lightllm/common/gated_residual/projection.py create mode 100644 lightllm/common/ngram_embedding/__init__.py create mode 100644 lightllm/common/ngram_embedding/kernels.py create mode 100644 lightllm/common/ngram_embedding/ple.py create mode 100644 lightllm/common/ngram_embedding/ple_kernels.py create mode 100644 lightllm/common/qsa/__init__.py create mode 100644 lightllm/common/qsa/attention.py create mode 100644 lightllm/common/qsa/attention_kernels.py create mode 100644 lightllm/common/qsa/indexer.py create mode 100644 lightllm/common/qsa/indexer_kernels.py create mode 100644 lightllm/common/qsa/topk.py create mode 100644 lightllm/common/quantization/aligned_projection.py create mode 100644 lightllm/common/quantization/scalar_projection.py create mode 100644 lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_H100_80GB_HBM3/grouped_matmul:v1/{K=2560,N=160,expert_num=513,mul_routed_weight=false,out_dtype=torch.bfloat16,topk_num=11,use_fp8_w8a8=false}_NVIDIA_H100_80GB_HBM3.json create mode 100644 lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_H100_80GB_HBM3/grouped_matmul:v1/{K=80,N=2560,expert_num=513,mul_routed_weight=true,out_dtype=torch.bfloat16,topk_num=1,use_fp8_w8a8=false}_NVIDIA_H100_80GB_HBM3.json create mode 100644 lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_H200/grouped_matmul:v1/{K=160,N=2560,expert_num=513,mul_routed_weight=true,out_dtype=torch.bfloat16,topk_num=1,use_fp8_w8a8=false}_NVIDIA_H200.json create mode 100644 lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_H200/grouped_matmul:v1/{K=2560,N=320,expert_num=513,mul_routed_weight=false,out_dtype=torch.bfloat16,topk_num=11,use_fp8_w8a8=false}_NVIDIA_H200.json create mode 100644 lightllm/models/qwen4_exp/__init__.py create mode 100644 lightllm/models/qwen4_exp/attention_state.py create mode 100644 lightllm/models/qwen4_exp/batch.py create mode 100644 lightllm/models/qwen4_exp/checkpoint.py create mode 100644 lightllm/models/qwen4_exp/configuration.py create mode 100644 lightllm/models/qwen4_exp/image_processing.py create mode 100644 lightllm/models/qwen4_exp/infer_struct.py create mode 100644 lightllm/models/qwen4_exp/layer_infer/__init__.py create mode 100644 lightllm/models/qwen4_exp/layer_infer/post_layer_infer.py create mode 100644 lightllm/models/qwen4_exp/layer_infer/pre_layer_infer.py create mode 100644 lightllm/models/qwen4_exp/layer_infer/transformer_layer_infer.py create mode 100644 lightllm/models/qwen4_exp/layer_weights/__init__.py create mode 100644 lightllm/models/qwen4_exp/layer_weights/components.py create mode 100644 lightllm/models/qwen4_exp/layer_weights/embedding_weight.py create mode 100644 lightllm/models/qwen4_exp/layer_weights/lm_head_weight.py create mode 100644 lightllm/models/qwen4_exp/layer_weights/norm_weight.py create mode 100644 lightllm/models/qwen4_exp/layer_weights/pre_and_post_layer_weight.py create mode 100644 lightllm/models/qwen4_exp/layer_weights/transformer_layer_weight.py create mode 100644 lightllm/models/qwen4_exp/lexical_cache.py create mode 100644 lightllm/models/qwen4_exp/lexical_state.py create mode 100644 lightllm/models/qwen4_exp/model.py create mode 100644 lightllm/models/qwen4_exp/mtp/__init__.py create mode 100644 lightllm/models/qwen4_exp/mtp/infer.py create mode 100644 lightllm/models/qwen4_exp/mtp/model.py create mode 100644 lightllm/models/qwen4_exp/mtp/weights.py create mode 100644 lightllm/models/qwen4_exp/multimodal.py create mode 100644 lightllm/models/qwen4_exp/pinned_memory.py create mode 100644 lightllm/models/qwen4_exp/position.py create mode 100644 lightllm/models/qwen4_exp/speculative_cache.py create mode 100644 lightllm/models/qwen4_exp/state.py create mode 100644 lightllm/models/qwen4_exp/tokenizer.py create mode 100644 lightllm/models/qwen4_exp/triton_kernel/__init__.py create mode 100644 lightllm/models/qwen4_exp/triton_kernel/cache_pages.py create mode 100644 lightllm/models/qwen4_exp/triton_kernel/gated_norm.py create mode 100644 lightllm/models/qwen4_exp/vision.py create mode 100644 lightllm/server/httpserver/slot_availability.py create mode 100644 lightllm/server/router/req_queue/cached_prefix.py create mode 100644 lightllm/server/router/req_queue/chunked_prefill/cached_prefix_queue.py create mode 100644 lightllm/server/router/req_queue/peak_reservation.py create mode 100644 test/benchmark/kernels/benchmark_kda_ssm.py create mode 100644 test/benchmark/kernels/benchmark_replayssm_service.py create mode 100644 test/benchmark/qwen38/benchmark_gated_residual.py create mode 100644 test/benchmark/qwen38/benchmark_lm_head.py create mode 100644 test/benchmark/qwen38/benchmark_ngram_hash.py create mode 100644 test/benchmark/qwen38/benchmark_ple.py create mode 100644 test/benchmark/qwen38/benchmark_qsa.py create mode 100644 test/benchmark/qwen38/check_aiperf_timing.py create mode 100644 test/benchmark/qwen38/check_baseline.py create mode 100644 test/benchmark/qwen38/check_cold_streams.py create mode 100644 test/benchmark/qwen38/check_cpu_cache.py create mode 100644 test/benchmark/qwen38/check_generation.py create mode 100644 test/benchmark/qwen38/check_math_quality.py create mode 100644 test/benchmark/qwen38/check_multimodal.py create mode 100644 test/benchmark/qwen38/check_prefix_eviction.py create mode 100644 test/benchmark/qwen38/check_retrieval.py create mode 100644 test/benchmark/qwen38/check_scheduler.py create mode 100644 test/benchmark/qwen38/check_streaming.py create mode 100644 test/benchmark/qwen38/generation_observations.py create mode 100644 test/benchmark/qwen38/math_quality.py create mode 100644 test/benchmark/qwen38/measure_aiperf.py create mode 100644 test/benchmark/qwen38/measure_fixed_work.py create mode 100644 test/benchmark/qwen38/prepare_aiperf_trace.py create mode 100644 test/benchmark/qwen38/retrieval_quality.py create mode 100644 test/benchmark/qwen38/score_reference.py create mode 100644 test/benchmark/service/cache_capacity.py create mode 100644 test/benchmark/service/completion.py create mode 100644 test/benchmark/service/completion_runner.py create mode 100644 test/benchmark/service/stream_cohort.py create mode 100644 test/benchmark/service/stream_timing.py create mode 100644 test/benchmark/service/token_stream.py create mode 100644 unit_tests/benchmark/qwen38/test_aiperf_replay.py create mode 100644 unit_tests/benchmark/qwen38/test_check_cold_streams.py create mode 100644 unit_tests/benchmark/qwen38/test_check_cpu_cache.py create mode 100644 unit_tests/benchmark/qwen38/test_check_generation.py create mode 100644 unit_tests/benchmark/qwen38/test_check_math_quality.py create mode 100644 unit_tests/benchmark/qwen38/test_check_multimodal.py create mode 100644 unit_tests/benchmark/qwen38/test_check_prefix_eviction.py create mode 100644 unit_tests/benchmark/qwen38/test_check_scheduler.py create mode 100644 unit_tests/benchmark/qwen38/test_generation_observations.py create mode 100644 unit_tests/benchmark/qwen38/test_math_quality.py create mode 100644 unit_tests/benchmark/qwen38/test_measure_fixed_work.py create mode 100644 unit_tests/benchmark/qwen38/test_retrieval_quality.py create mode 100644 unit_tests/benchmark/qwen38/test_score_reference.py create mode 100644 unit_tests/benchmark/service/test_cache_capacity.py create mode 100644 unit_tests/benchmark/service/test_completion.py create mode 100644 unit_tests/benchmark/service/test_completion_runner.py create mode 100644 unit_tests/benchmark/service/test_stream_cohort.py create mode 100644 unit_tests/benchmark/service/test_stream_timing.py create mode 100644 unit_tests/benchmark/service/test_token_stream.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_packed_decode.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.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/linear_att/test_ssm_autotune_memory.py create mode 100644 unit_tests/common/basemodel/triton_kernel/test_gated_rmsnorm_strided_gate.py create mode 100644 unit_tests/common/basemodel/triton_kernel/test_scalar_linear.py create mode 100644 unit_tests/common/gated_residual/test_gated_residual.py create mode 100644 unit_tests/common/ngram_embedding/test_hash.py create mode 100644 unit_tests/common/ngram_embedding/test_pinned_embedding.py create mode 100644 unit_tests/common/ngram_embedding/test_ple.py create mode 100644 unit_tests/common/ngram_embedding/test_row_sharded_embedding.py create mode 100644 unit_tests/common/qsa/test_attention.py create mode 100644 unit_tests/common/qsa/test_flashinfer_topk.py create mode 100644 unit_tests/common/qsa/test_indexer.py create mode 100644 unit_tests/common/qsa/test_topk_compatibility.py create mode 100644 unit_tests/common/qsa/test_topk_cuda.py create mode 100644 unit_tests/common/quantization/test_aligned_projection.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/distributed/test_ordered_tp_sum.py create mode 100644 unit_tests/models/qwen3next/test_transformer_layer_infer_strided_gate.py create mode 100644 unit_tests/models/qwen4_exp/__init__.py create mode 100644 unit_tests/models/qwen4_exp/conftest.py create mode 100644 unit_tests/models/qwen4_exp/helpers.py create mode 100644 unit_tests/models/qwen4_exp/lexical_helpers.py create mode 100644 unit_tests/models/qwen4_exp/mtp_reference.py create mode 100644 unit_tests/models/qwen4_exp/reference.py create mode 100644 unit_tests/models/qwen4_exp/resource_helpers.py create mode 100644 unit_tests/models/qwen4_exp/runtime.py create mode 100644 unit_tests/models/qwen4_exp/test_attention_state.py create mode 100644 unit_tests/models/qwen4_exp/test_batch.py create mode 100644 unit_tests/models/qwen4_exp/test_decoder.py create mode 100644 unit_tests/models/qwen4_exp/test_embedding_residency.py create mode 100644 unit_tests/models/qwen4_exp/test_eos.py create mode 100644 unit_tests/models/qwen4_exp/test_lexical_cache.py create mode 100644 unit_tests/models/qwen4_exp/test_linear_decoder.py create mode 100644 unit_tests/models/qwen4_exp/test_lm_head.py create mode 100644 unit_tests/models/qwen4_exp/test_mtp.py create mode 100644 unit_tests/models/qwen4_exp/test_multimodal.py create mode 100644 unit_tests/models/qwen4_exp/test_norm.py create mode 100644 unit_tests/models/qwen4_exp/test_pinned_memory.py create mode 100644 unit_tests/models/qwen4_exp/test_reduction.py create mode 100644 unit_tests/models/qwen4_exp/test_runtime.py create mode 100644 unit_tests/models/qwen4_exp/test_sequence_parallel.py create mode 100644 unit_tests/models/qwen4_exp/test_speculative_cache.py create mode 100644 unit_tests/models/qwen4_exp/test_vision.py create mode 100644 unit_tests/models/qwen4_exp/test_weights.py create mode 100644 unit_tests/models/qwen4_exp/vision_reference.py create mode 100644 unit_tests/models/sensenova_loop/test_gdn_state.py create mode 100644 unit_tests/server/httpserver/test_slot_availability.py create mode 100644 unit_tests/server/router/req_queue/test_cached_prefix_queue.py create mode 100644 unit_tests/server/router/req_queue/test_peak_reservation.py create mode 100644 unit_tests/server/test_native_stream_details.py create mode 100644 unit_tests/server/test_ssm_state_mode.py diff --git a/docs/qwen38_adaptation.rst b/docs/qwen38_adaptation.rst new file mode 100644 index 0000000000..d4bc2d4f9b --- /dev/null +++ b/docs/qwen38_adaptation.rst @@ -0,0 +1,48 @@ +Qwen3.8-Flash-Next +================== + +Qwen3.8-Flash-Next checkpoints use the ``qwen4_exp`` model type. The model +implementation combines gated residuals, n-gram embeddings, QSA and GDN +attention, with dedicated multimodal, MTP and cache-state integration. +Reusable operations live in ``lightllm/common/gated_residual``, +``lightllm/common/ngram_embedding`` and ``lightllm/common/qsa``. + +CPU embedding offload is disabled by default. The optional QSA TopK selector +is enabled by the model only for its qualified H100/Torch/Triton geometry; +other configurations use Torch. Shared-prefix admission is limited to the +model's single-node, DP1 normal serving mode with one-page state blocks. + +Validation +---------- + +Historical H100 measurements and their configuration are recorded in +``qwen38_h100_performance.rst``. They cover TP8 BF16, FP32 SSM states, +fixed MTP2 and one-page linear-attention state blocks. These results do not +qualify PD/cross-DP serving or other cache block configurations. + +The four-page-block pressure-recovery configuration has unresolved text +differences. Low-load latency still trails SGLang in some metrics. Archived +HC and sampling experiments are not part of the runtime implementation. + +Rebase to upstream, 2026-11-03 +------------------------------ + +The implementation now uses AttentionStateManager. QSA keys and positions are +owned by the model's GQA State and participate in token copies and CPU/PD/DP +transfer layouts. Lexical histories use a separate bounded State for reset, +prefix snapshots, CPU/PD/DP transfers and accepted MTP candidate commits. +The removed request/cache manager classes are not restored. + +Model and draft registration use the upstream lazy registry. HTTP slot waiting +preserves atomic group allocation and current PD timeouts. The GDN MTP path uses +the upstream recurrent kernel and its explicit run configuration. + +The historical H100 results above predate this interface migration. They do not +validate the rebased runtime. GPU generation, complete PD transfer and performance +must be rechecked before deploying a new image of this revision. + +CPU validation after migration: 1022 checks passed and 930 GPU cases were +skipped. The DeepSeek V4 integration module could not be collected because +FlashMLA is absent. Two upstream State identity cases require CUDA-backed +pinned memory; both failures reproduce on unmodified upstream/main with CUDA +hidden. No GPU serving, PD end-to-end or performance result is claimed here. diff --git a/docs/qwen38_h100_performance.rst b/docs/qwen38_h100_performance.rst new file mode 100644 index 0000000000..bb9bd4be39 --- /dev/null +++ b/docs/qwen38_h100_performance.rst @@ -0,0 +1,76 @@ +Qwen3.8-Flash-Next H100 性能与精度 +================================= + +当前结论：4× 的两轮吞吐、TTFT P95、完整 HTTP P95 和末次可见文本 P95 均优于冻结基线。 +低负载下仍有延迟差距，尚未完成“超过 vLLM 和 SGLang”的整体目标。 + +测试口径 +-------- + +* 8 × NVIDIA H100 80GB HBM3，TP8，BF16，SSM 状态 FP32；MTP 固定 2。 +* 模型版本 de4b8e4d43b917e7706784d8bb445c9af86a3540；4096 上下文、最多 8 个并发请求、512-token prefill chunk、16384 KV 容量。 +* 当前已验证的线性注意力状态页为 512 tokens，每块一页。 +* AIPerf 0.9 统一回放公开 Mooncake 轨迹中的 256 个请求，倍率为 1×、2×、4×。 +* 公开轨迹提供匿名长度、时间戳和 hash；请求正文由 AIPerf 合成，并非原始私有正文。 +* 每个 LightLLM 样本使用全新服务；vLLM、SGLang 只使用此前冻结的原始结果，未重新推理。 +* 原生 Torch 2.9.1+cu128 / Triton 3.5.1；vLLM 0.31；SGLang 0.5.22。 +* 1×/2× 吞吐受请求到达速率限制，不能当作最大处理容量；保留所有重复样本。 + +已完成的性能样本 +---------------- + +.. csv-table:: 延迟单位 ms；吞吐单位成功输出 tok/s + :header: "倍率", "框架", "样本", "吞吐 ↑", "TTFT P95 ↓", "完整 HTTP P95 ↓", "末次文本 P95 ↓", "空文本错误/256" + + "1×", "LightLLM", "1", "276.14", "516.09", "3832.07", "3615.36", "2" + "1×", "LightLLM", "2", "276.16", "515.16", "3818.23", "3622.51", "2" + "1×", "SGLang", "冻结", "276.99", "610.40", "3864.63", "3329.87", "2" + "1×", "vLLM", "冻结", "277.72", "804.93", "4255.93", "3548.29", "1" + "2×", "LightLLM", "1", "546.33", "475.34", "4466.54", "4187.74", "2" + "2×", "SGLang", "冻结", "545.80", "600.51", "5302.52", "4178.47", "2" + "2×", "vLLM", "冻结", "513.20", "2933.67", "8695.77", "7227.73", "2" + "4×", "LightLLM", "1", "887.26", "376.17", "5794.25", "5188.47", "2" + "4×", "LightLLM", "2", "886.14", "390.13", "5770.75", "5151.58", "2" + "4×", "SGLang", "冻结", "835.28", "864.93", "5924.62", "5204.95", "2" + "4×", "vLLM", "冻结", "523.95", "3367.76", "9085.75", "7930.38", "1" + +完整 HTTP 统计包含全部 256 个请求；TTFT 和末次文本仅包含有可见文本的请求。 +空文本仍计为错误，其输出 token 不计入成功吞吐。AIPerf 导出的 request_latency +在这些记录中对应末次可见文本，不能替代完整 HTTP/DONE 边界。 + +精度与兼容性 +------------ + +* 数学题正确 1272/1319，准确率 96.44%。 +* 全部 1319 个数学样例的生成 token 序列与原实现一致；分别保留 19 例截断、19 例无有效最终答案。 +* 480 个冷缓存请求和 48 个调度、缓存、EOS 请求核对通过；2788 个训练模型中间张量完全一致。 +* H100 共享组件/默认行为回归 134 项通过；本地 CPU 兼容性检查 124 项通过，零失败、错误、跳过。 +* 已完成的 5 轮回放共 1280 个输出文本与原实现一致，空文本错误单独保留。 +* TopK 为独立组件，共享选项默认关闭，Qwen4 在工厂启用，select 接口保持兼容。 + +剩余工作与限制 +-------------- + +* 1× 首轮末次文本 P95 为 3615.36 ms，高于 SGLang 3329.87 ms；中位请求延迟也仍落后。 +* 2× 末次文本 P95 为 4187.74 ms，比 SGLang 4178.47 ms 高约 9.28 ms。 +* 历史四页块压力恢复的文本差异仍未解决，不能将当前一页块精度结论推广到该配置。 +* 五轮回放及低负载剖析已完成；按用户要求收尾，优化目标已暂停，未标记为全面达成。 + +原始结果与证据 +-------------- + +* 五轮独立汇总：``artifacts/h100/261019/audits/topk-short-five-sample-summary.json``。 +* 性能 CSV：``artifacts/h100/261019/audits/topk-short-performance-current.csv``。 +* 完整精度：``artifacts/h100/261019/audits/topk-short-full-math-token-parity.json``。 +* 源码与本地验证：``artifacts/h100/261019/audits/topk-short-local-adoption.json``。 +* 重复回放状态：``artifacts/h100/261019/tuning/topk-short-canonical-candidate/progress.json``。 + +收尾状态 +-------- + +* 最终保留完整精度和五轮回放验证通过的 TopK 版本；H100 隔离源码已恢复到同一版本。 +* HC 小批量上投影候选已归档：30 项 CPU、203 项 GPU 检查通过，2788 个模型中间张量一致。 + 完整模型轨迹实际覆盖 8 行 BD8；其余选择形状只有组件检查。HC 尚未完成服务精度与公开回放验证，未纳入最终实现或上述性能结果。 +* 首次 HC 模型轨迹因测试插桩作用域断言失败而回滚；修复插桩并验证后重跑通过，失败记录一并保留。 +* H100 恢复后核对 949 项源码和 42 项测试/基准文件哈希；GPU 无遗留计算进程。未启动新的推理测试。 +* 收尾记录：``artifacts/h100/261019/audits/user-wrapup-qualified-topk-restored.json``。 diff --git a/lightllm/common/attention_state/bounded_att/gdn/create_gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/create_gdn_state.py index 3fe6f2317e..82a07fc486 100644 --- a/lightllm/common/attention_state/bounded_att/gdn/create_gdn_state.py +++ b/lightllm/common/attention_state/bounded_att/gdn/create_gdn_state.py @@ -13,8 +13,24 @@ def create_gdn_state( request_capacity: int, mtp_step: int = 0, device: str | torch.device = "cuda", + ssm_state_mode: str = "native", + ssm_replay_config: dict | None = None, + gate_lower_bound: float = -5.0, ) -> GDNState: - """创建 GDN 请求运行态，大小页池由后续容量分配流程初始化。""" + """创建 GDN 请求运行态。replay 模式使用独立的 ReplayGDNState。""" + if ssm_state_mode == "replay": + from .replay_gdn_state import ReplayGDNState + + return ReplayGDNState( + config=config, + request_capacity=request_capacity, + mtp_step=mtp_step, + device=device, + ssm_replay_config=ssm_replay_config, + gate_lower_bound=gate_lower_bound, + ) + if ssm_state_mode != "native": + raise ValueError(f"unknown GDN state mode: {ssm_state_mode}") return GDNState( config=config, request_capacity=request_capacity, diff --git a/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py new file mode 100644 index 0000000000..d927891ada --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/gdn/replay_gdn_state.py @@ -0,0 +1,835 @@ +"""GDN replay 的 bounded attention 状态。 + +Conv、canonical SSM、replay history、页布局和已接受快照都在本类。 +一层 decode 重放在 replay_gdn_decode。 +""" + +from __future__ import annotations + +from functools import lru_cache +from math import prod +from typing import NamedTuple + +import torch +import triton + +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ( + _advance_cursor, + _fold_prefix, +) +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_fold import _reset_folded_cursors +from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix +from lightllm.utils.log_utils import init_logger + +from ...types import AxisMapping, PageKind, StateTransferLayout, TensorField, TransferRegion +from ...utils import deepcopy_result, dtype_name +from ..base_state import BoundedAttentionState +from ..conv_ssm_types import ConvSSMStateConfig +from ..copy_kernel import copy_bounded_att_tensors +from ..transfer_utils import bounded_att_buffer_copies +from .gdn_pages_pool import GDNPagesPool + +logger = init_logger(__name__) + + +class _AcceptedTensors(NamedTuple): + conv: torch.Tensor + ssm: torch.Tensor + + +class ReplayGDNState(BoundedAttentionState): + """GDN replay 运行态。Conv、单槽 SSM 和 history 都在本类，不沿用 native 的 MTP 槽。""" + + # ======================================================================== + # 一、资源定义与初始化 + # ======================================================================== + + def __init__( + self, + config: ConvSSMStateConfig, + request_capacity: int, + mtp_step: int = 0, + device: str | torch.device = "cuda", + ssm_state_mode: str = "replay", + ssm_replay_config: dict | None = None, + gate_lower_bound: float = -5.0, + ) -> None: + """创建单槽 canonical SSM，并挂上只属于当前请求的 replay history。""" + if ssm_state_mode != "replay": + raise ValueError(f"ReplayGDNState only implements replay mode, got {ssm_state_mode}") + assert isinstance(config, ConvSSMStateConfig), "config must be ConvSSMStateConfig" + self.config = config + self.state_name = None + self.request_capacity = request_capacity + self.mtp_step = mtp_step + self.gate_lower_bound = gate_lower_bound + self._algorithm = "gdn" + self._logical_to_storage_layer = {layer_id: index for index, layer_id in enumerate(config.logical_layer_ids)} + layer_count = len(config.logical_layer_ids) + self.conv_state = torch.zeros( + ( + layer_count, + request_capacity, + config.conv_channels, + config.conv_width + mtp_step, + ), + dtype=config.conv_dtype, + device=device, + ) + self.ssm_state = torch.zeros( + ( + layer_count, + request_capacity, + config.num_v_heads, + config.key_dim, + config.value_dim, + ), + dtype=config.ssm_dtype, + device=device, + ) + self.mtp_state_index = torch.zeros(request_capacity, dtype=torch.int32, device=device) if mtp_step else None + self.big_pages_pool = None + self.small_pages_pool = None + replay_config = dict(ssm_replay_config or {}) + capacity = int(replay_config.get("cache_len", 16)) + compute_mode = replay_config.get("compute_mode", "precompute" if mtp_step > 0 else "inline") + force_compress = bool(replay_config.get("force_compress", False)) + _bind_history( + self, + self.ssm_state, + capacity, + mtp_step + 1, + config.conv_dtype, + config.num_k_heads, + gate_lower_bound, + compute_mode, + None, + force_compress, + ) + + def get_state_type(self) -> str: + """返回算法名和 Conv/SSM 存储精度。""" + return ( + f"{self._algorithm}_conv_{dtype_name(dtype=self.config.conv_dtype)}" + f"_ssm_{dtype_name(dtype=self.config.ssm_dtype)}" + ) + + def initialize_page_pools( + self, + big_page_capacity: int, + small_page_capacity: int = 0, + device: str | torch.device | int = "cpu", + pin_memory: bool = True, + ) -> None: + """容量规划完成后创建本算法的大小页池；容量为零的页池保持 None。""" + assert self.big_pages_pool is None and self.small_pages_pool is None, "page pools already initialized" + assert big_page_capacity >= 0 and small_page_capacity >= 0, "page capacities must be nonnegative" + self.big_pages_pool = ( + GDNPagesPool( + size=big_page_capacity, + config=self.config, + device=device, + pin_memory=pin_memory, + ) + if big_page_capacity + else None + ) + self.small_pages_pool = ( + GDNPagesPool( + size=small_page_capacity, + config=self.config, + device=device, + pin_memory=pin_memory, + ) + if small_page_capacity + else None + ) + + # ======================================================================== + # 二、计算视图、运行态重置、已接受快照与 MTP 提交 + # ======================================================================== + + def contains_logical_layer(self, logical_layer_id: int) -> bool: + """判断指定逻辑层是否属于当前 State 的层映射。""" + return logical_layer_id in self._logical_to_storage_layer + + def storage_layer(self, logical_layer_id: int) -> int: + """返回逻辑层在 history 和 canonical state 里的存储下标。""" + + return self._logical_to_storage_layer[logical_layer_id] + + def get_layer_view(self, logical_layer_id: int): + index = self.storage_layer(logical_layer_id) + return self.conv_state[index], self.ssm_state[index] + + def reset_bounded_att_state(self, req_indices: list[int]) -> None: + """清零请求的 Conv、单槽 SSM、history 和 MTP 游标。""" + assert isinstance(req_indices, list), "request indices must be a list" + for req_idx in req_indices: + self.conv_state[:, req_idx].zero_() + self.ssm_state[:, req_idx].zero_() + self.clear_history(req_idx) + if self.mtp_state_index is not None: + self.mtp_state_index[req_idx].fill_(0) + + def save_bounded_att_runtime_to_pages( + self, + req_indices: list[int], + state_token_lens: list[int], + page_kind, + page_slot_ids: list[int], + ) -> None: + """把已接受 Conv 窗口和物化 SSM 写入页，不改 canonical state。""" + assert isinstance(req_indices, list) and isinstance(page_slot_ids, list), "indices must be lists" + assert len(req_indices) == len(state_token_lens) == len(page_slot_ids), "request, length and page counts differ" + pool = self._pool(page_kind) + copy_bounded_att_tensors( + copies=[ + (self._accepted_conv(req_idx), pool.conv_state[slot_id]) + for req_idx, slot_id in zip(req_indices, page_slot_ids) + ] + ) + self.export_ssm(req_indices, [pool.ssm_state[slot_id] for slot_id in page_slot_ids]) + + def restore_bounded_att_runtime_from_pages( + self, + req_indices: list[int], + state_token_lens: list[int], + page_kind, + page_slot_ids: list[int], + ) -> None: + """从页恢复基准 Conv/SSM，并清空目标请求的 history。""" + assert isinstance(req_indices, list) and isinstance(page_slot_ids, list), "indices must be lists" + assert len(req_indices) == len(state_token_lens) == len(page_slot_ids), "request, length and page counts differ" + pool = self._pool(page_kind) + width = self.config.conv_width + conv_copies, ssm_copies = [], [] + for req_idx, slot_id in zip(req_indices, page_slot_ids): + self.clear_history(req_idx) + conv_copies.append((pool.conv_state[slot_id], self.conv_state[:, req_idx, :, :width])) + ssm_copies.append((pool.ssm_state[slot_id], self.ssm_state[:, req_idx])) + copy_bounded_att_tensors(copies=conv_copies) + copy_bounded_att_tensors(copies=ssm_copies) + if self.mtp_state_index is not None: + for req_idx in req_indices: + self.mtp_state_index[req_idx].fill_(0) + + def commit_bounded_att_state(self, req_indices: torch.Tensor, accepted_steps: torch.Tensor) -> None: + """提交已接受长度，再按容量或 force_compress 决定是否整段压缩。""" + if self.mtp_state_index is None: + return + self.commit_accepted(self.mtp_state_index, req_indices, accepted_steps) + self.maybe_compress_prefix(req_indices) + + def get_accepted_state(self, req_idx: int) -> _AcceptedTensors: + """返回已接受 Conv 窗口和物化 SSM，不改写 canonical state。""" + return _AcceptedTensors(self._accepted_conv(req_idx), self.accepted_ssm(req_idx)) + + # ======================================================================== + # 三、CPU Cache 数据交换 + # ======================================================================== + + @deepcopy_result + @lru_cache(maxsize=None) + def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: + """返回已接受 Conv/SSM 快照的字段区间。不含 history 或 MTP 候选区。""" + return self._transfer_regions() + + def write_state_pages_to_cpu( + self, + page_kind: PageKind, + page_slot_ids: list[int], + cpu_page_buffers: torch.Tensor, + cpu_page_indices: list[int], + layout: StateTransferLayout, + cpu_page_ready_flags: list[bool] | None = None, + ) -> None: + """将指定页池槽位的快照写入 CPU 页。""" + assert layout.state_name == self.state_name, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" + pool = self._pool(page_kind) + assert isinstance(page_slot_ids, list), "slot indices must be a list" + assert isinstance(cpu_page_indices, list), "CPU page indices must be a list" + assert len(page_slot_ids) == len(cpu_page_indices), "CPU page and slot counts differ" + if cpu_page_ready_flags is not None: + assert len(cpu_page_ready_flags) == len(cpu_page_indices), "CPU page and ready counts differ" + page_fields, page_indices = [], [] + for index, (slot, page) in enumerate(zip(page_slot_ids, cpu_page_indices)): + # ready 页不参与复制；同一传输区域的所有页面组成一个同形状批次。 + if cpu_page_ready_flags is not None and cpu_page_ready_flags[index]: + continue + page_fields.append(self._tensor_fields(pool.conv_state[slot], pool.ssm_state[slot])) + page_indices.append(page) + for copies in bounded_att_buffer_copies( + page_fields=page_fields, + page_buffers=cpu_page_buffers, + layout=layout, + page_indices=page_indices, + write=True, + ): + copy_bounded_att_tensors(copies=copies) + + def read_state_pages_from_cpu( + self, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: list[int], + big_page_slot_ids: list[int], + layout: StateTransferLayout, + ) -> None: + """读入指定大页槽位。请求运行态由 restore_bounded_att_runtime_from_pages 另行恢复。""" + assert layout.state_name == self.state_name, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" + pool = self.big_pages_pool + assert isinstance(big_page_slot_ids, list), "slot indices must be a list" + assert isinstance(cpu_page_indices, list), "CPU page indices must be a list" + assert len(cpu_page_indices) == len(big_page_slot_ids), "CPU page and slot counts differ" + # 本地目标保留本 TP rank 的分片形状。layout.regions 描述 CPU 全局字段。 + pool_page_fields = [ + self._tensor_fields(pool.conv_state[slot], pool.ssm_state[slot]) for slot in big_page_slot_ids + ] + for copies in bounded_att_buffer_copies( + page_fields=pool_page_fields, + page_buffers=cpu_page_buffers, + layout=layout, + page_indices=cpu_page_indices, + write=False, + ): + copy_bounded_att_tensors(copies=copies) + + # ======================================================================== + # 四、DP 迁移：大页快照与基准运行态 + # ======================================================================== + + def copy_big_pages_from_dp(self, src_page_ids: list[int], dst_page_ids: list[int], remote_dp_bounded_state) -> None: + """复制同一 TP 分片的大页 Conv/SSM 快照，保留源页和双方请求运行态。""" + assert isinstance(src_page_ids, list) and isinstance(dst_page_ids, list), "page indices must be lists" + assert len(src_page_ids) == len(dst_page_ids), "source and destination page counts differ" + assert ( + type(remote_dp_bounded_state) is type(self) and remote_dp_bounded_state.config == self.config + ), "DP copy requires the same algorithm and equivalent TP state configurations" + src_pool = remote_dp_bounded_state.big_pages_pool + dst_pool = self.big_pages_pool + assert src_pool is not None and dst_pool is not None, "DP big page copy requires both page pools" + # 大页是基准快照，与请求槽位无关。页池须在本进程注册，当前流访问完才能归还页面。 + conv_copies, ssm_copies = [], [] + for src_page, dst_page in zip(src_page_ids, dst_page_ids): + conv_copies.append((src_pool.conv_state[src_page], dst_pool.conv_state[dst_page])) + ssm_copies.append((src_pool.ssm_state[src_page], dst_pool.ssm_state[dst_page])) + copy_bounded_att_tensors(copies=conv_copies) + copy_bounded_att_tensors(copies=ssm_copies) + + def copy_runtime_state_from_dp( + self, + src_req_indices: list[int], + dst_req_indices: list[int], + remote_dp_bounded_state: "ReplayGDNState", + ) -> None: + """复制已接受运行态到目标单槽，并清空目标 history。""" + assert isinstance(src_req_indices, list) and isinstance(dst_req_indices, list), "request indices must be lists" + assert len(src_req_indices) == len(dst_req_indices), "source and destination request counts differ" + assert ( + type(remote_dp_bounded_state) is type(self) and remote_dp_bounded_state.config == self.config + ), "DP copy requires the same algorithm and equivalent TP state configurations" + width = self.config.conv_width + conv_copies = [] + for src_req, dst_req in zip(src_req_indices, dst_req_indices): + self.clear_history(dst_req) + conv_copies.append( + ( + remote_dp_bounded_state._accepted_conv(src_req), + self.conv_state[:, dst_req, :, :width], + ) + ) + copy_bounded_att_tensors(copies=conv_copies) + remote_dp_bounded_state.export_ssm(src_req_indices, [self.ssm_state[:, dst_req] for dst_req in dst_req_indices]) + if self.mtp_state_index is not None: + for dst_req in dst_req_indices: + self.mtp_state_index[dst_req].fill_(0) + + # ======================================================================== + # 五、PD 分离传输 + # ======================================================================== + + @deepcopy_result + @lru_cache(maxsize=None) + def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: + """返回 PD 传输字段与切片，交由管理器安排其在联合页中的位置。""" + return self._transfer_regions() + + def write_to_pd_page( + self, + req_idx: int, + pd_page_buffers: torch.Tensor, + pd_page_index: int, + state_token_len: int, + layout: StateTransferLayout, + ) -> None: + """写出已接受 Conv 窗口和物化 SSM，不改写源 history。""" + assert layout.state_name == self.state_name, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + accepted = self.get_accepted_state(req_idx) + fields = self._tensor_fields(accepted.conv, accepted.ssm) + for copies in bounded_att_buffer_copies( + page_fields=[fields], + page_buffers=pd_page_buffers, + layout=layout, + page_indices=[pd_page_index], + write=True, + ): + copy_bounded_att_tensors(copies=copies) + + def read_from_pd_page( + self, + pd_page_buffers: torch.Tensor, + req_idx: int, + pd_page_index: int, + state_token_len: int, + layout: StateTransferLayout, + ) -> None: + """读入基准 Conv/SSM，并清空该请求的 history。""" + assert layout.state_name == self.state_name, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + self.clear_history(req_idx) + fields = self._tensor_fields( + self.conv_state[:, req_idx, :, : self.config.conv_width], + self.ssm_state[:, req_idx], + ) + for copies in bounded_att_buffer_copies( + page_fields=[fields], + page_buffers=pd_page_buffers, + layout=layout, + page_indices=[pd_page_index], + write=False, + ): + copy_bounded_att_tensors(copies=copies) + if self.mtp_state_index is not None: + self.mtp_state_index[req_idx].fill_(0) + + # ======================================================================== + # 六、Replay history + # ======================================================================== + + def finish_prefill(self, logical_layer_id: int, reqs: torch.Tensor) -> None: + """最后一层把已接受前缀整段压进 SSM。""" + + layers = () if self.config is None else self.config.logical_layer_ids + if layers and logical_layer_id == layers[-1]: + self.compress_prefix(reqs) + + def finish_decode_token(self, logical_layer_id: int, reqs: torch.Tensor) -> None: + """最后一层记下这 1 个 token，再按容量决定是否整段压缩。""" + + layers = () if self.config is None else self.config.logical_layer_ids + if not layers or logical_layer_id != layers[-1]: + return + self.accept_single_token(reqs) + self.maybe_compress_prefix(reqs) + + def clear_history(self, req: int | torch.Tensor) -> None: + """清空复用或恢复请求的临时 history，保留 canonical SSM。""" + + if type(req) is int: + # 标量赋值会把 CPU tensor 复制到 CUDA slot。 + self.cursors[req].zero_() + else: + self.cursors[req] = 0 + + def prepare_decode(self, reqs: torch.Tensor, cu_seqlens: torch.Tensor | None = None) -> torch.Tensor: + """返回本步写入前的前缀长度。decode 不在这里修改 SSM。""" + + del cu_seqlens + return self.cursors.index_select(0, reqs.long()) + + def accept_updates( + self, reqs: torch.Tensor, accepted_state_indices: torch.Tensor | None = None, width: int | None = None + ) -> None: + """更新 cursor。width 缺省时按 verify_width 取已接受长度，传 1 时固定加一个 token。""" + + _advance_cursor[(triton.cdiv(reqs.numel(), 256),)]( + self.cursors, + reqs, + accepted_state_indices, + reqs.numel(), + self.hold_req_id, + self.verify_width if width is None else width, + 256, + ) + + def accept_single_token(self, reqs: torch.Tensor) -> None: + """把一个已接受 token 记入前缀长度。""" + + if reqs.numel() == 0: + return + # width=1 不读 index。仍传 reqs，kernel 不能收到空指针。 + self.accept_updates(reqs, reqs, width=1) + + def merge_accepted_updates(self, reqs: torch.Tensor) -> None: + """把 cursor 覆盖的已接受前缀整段折回 canonical SSM，并清空 cursor。""" + + if reqs.numel() == 0: + return + self._materialize_accepted(reqs, self.ssm_state, snapshot=False) + self.cursors.index_fill_(0, reqs.long(), 0) + + def compress_prefix(self, reqs: torch.Tensor) -> None: + """把 cursor 覆盖的已接受前缀整段压进 SSM。长度为 0 时 kernel 直接返回。""" + + if reqs.numel() == 0: + return + if self.ssm_state.is_cuda: + self.merge_accepted_updates(reqs) + else: + self.clear_history(reqs) + + def maybe_compress_prefix(self, reqs: torch.Tensor) -> None: + """长度加上一次最大 verify_width 放不下，或 force_compress 打开时，整段压缩。""" + + if reqs.numel() == 0 or not self.ssm_state.is_cuda: + return + if self.force_compress: + self.compress_prefix(reqs) + return + self._materialize_accepted(reqs, self.ssm_state, snapshot=False, when_full=True) + + def accepted_ssm(self, req_idx: int) -> torch.Tensor: + """返回不修改 canonical SSM 和 history 的单请求快照。""" + + layers, _, hv, k, v = self.ssm_state.shape + output = torch.empty((layers, hv, k, v), dtype=self.ssm_state.dtype, device=self.ssm_state.device) + reqs = torch.tensor([req_idx], dtype=torch.int32, device=self.ssm_state.device) + self._materialize_accepted(reqs, output, snapshot=True) + return output + + def export_accepted_state( + self, + req_indices: list[int], + targets: list[torch.Tensor], + flush: bool = False, + ) -> None: + """将多个 accepted 快照直接写入调用方目标页。""" + + if not req_indices: + return + shape = (self.ssm_state.shape[0], *self.ssm_state.shape[2:]) + strides = targets[0].stride() + assert len(req_indices) == len(targets) + assert all(t.shape == shape and t.stride() == strides and t.dtype == self.ssm_state.dtype for t in targets) + device = self.ssm_state.device + self._materialize_accepted( + torch.tensor(req_indices, dtype=torch.int32, pin_memory=True).to(device, non_blocking=True), + torch.tensor([t.data_ptr() for t in targets], dtype=torch.int64, pin_memory=True).to( + device, non_blocking=True + ), + snapshot=not flush, + flush_export=flush, + output_strides=strides, + ) + + def commit_accepted(self, indices: torch.Tensor, reqs: torch.Tensor, lengths: torch.Tensor) -> None: + """把接受长度写入 MTP 槽位索引，并推进本请求的 replay cursor。""" + + if self.verify_width > 1: + commit_ssm_prefix(indices, reqs, lengths, self.cursors) + else: + commit_ssm_prefix(indices, reqs, lengths) + self.accept_updates(reqs, indices) + + def export_ssm(self, req_indices: list[int], targets: list[torch.Tensor], flush: bool = False) -> None: + """把已接受 SSM 写入调用方页面。CUDA 目标直接写指针，其他设备先物化再拷贝。""" + + if not req_indices: + return + target = targets[0] + if target.is_cuda and target.device == self.ssm_state.device: + self.export_accepted_state(req_indices, targets, flush=flush) + return + for req, dst in zip(req_indices, targets): + snap = self.accepted_ssm(req) + dst.copy_(snap) + if flush: + self.ssm_state[:, req].copy_(snap) + self.clear_history(req) + + _run_config: dict[str, int] | None = None + + @property + def run_config(self) -> dict[str, int] | None: + """返回物化和 autotune 使用的 Triton 配置，不是 SSM 布局。""" + + return self._run_config + + @run_config.setter + def run_config(self, config: dict | None) -> None: + """保存一份与调用方脱开的 kernel 配置。""" + + if config is None: + self._run_config = None + return + self._run_config = { + "BV": int(config["BV"]), + "num_warps": int(config["num_warps"]), + "num_stages": int(config.get("num_stages", 3)), + } + + # ======================================================================== + # 七、私有辅助实现 + # ======================================================================== + + def _transfer_regions(self) -> tuple[TransferRegion, ...]: + """依次描述全局 Conv 和 SSM 区域，并标明当前 TP rank 的读写范围。""" + # 只传有效历史窗口，不包含 MTP 候选区和游标。 + config = self.config + storage_layer_count = len(config.logical_layer_ids) + conv_channels = 2 * config.global_k_heads * config.key_dim + config.global_v_heads * config.value_dim + conv_shape = (storage_layer_count, conv_channels, config.conv_width) + conv_byte_end = prod(conv_shape) * config.conv_dtype.itemsize + regions = [] + logger.debug( + "%s building transfer regions: TP=%d/%d, Conv shape=%s", + self.state_name, + config.tp_rank, + config.tp_world_size, + conv_shape, + ) + + # Conv 按 Q/K/V 拼接。三个映射共用同一块全局 Conv 区域，不能重复累计字节数。 + local_channel_start = 0 + global_channel_start = 0 + for name, local_head_count, global_head_count, head_dim in ( + ("Q", config.num_k_heads, config.global_k_heads, config.key_dim), + ("K", config.num_k_heads, config.global_k_heads, config.key_dim), + ("V", config.num_v_heads, config.global_v_heads, config.value_dim), + ): + # 每组连续 rank 保存相同 heads，仅组内首 rank 写出。 + replica_count = config.tp_world_size * local_head_count // global_head_count + head_start = config.tp_rank // replica_count * local_head_count + write_enabled = config.tp_rank % replica_count == 0 + local_channel_count = local_head_count * head_dim + target_channel_start = global_channel_start + head_start * head_dim + regions.append( + TransferRegion( + field="conv", + byte_start=0, + byte_end=conv_byte_end, + dtype=config.conv_dtype, + axes=("layer", "channel", "window"), + shape=conv_shape, + axis_mappings=( + AxisMapping( + axis="channel", + local_start=local_channel_start, + external_start=target_channel_start, + length=local_channel_count, + ), + AxisMapping(axis="window", local_start=0, external_start=0, length=config.conv_width), + ), + write_enabled=write_enabled, + ) + ) + logger.debug( + "%s Conv %s: local_channel_start=%d, global_channel_start=%d, " + "channels=%d, replicas=%d, write_enabled=%s", + self.state_name, + name, + local_channel_start, + target_channel_start, + local_channel_count, + replica_count, + write_enabled, + ) + local_channel_start += local_channel_count + global_channel_start += global_head_count * head_dim + + # SSM 沿 V head 切分，接在 Conv 之后，并按 SSM dtype 对齐字节起点。 + replica_count = config.tp_world_size * config.num_v_heads // config.global_v_heads + head_start = config.tp_rank // replica_count * config.num_v_heads + write_enabled = config.tp_rank % replica_count == 0 + ssm_shape = (storage_layer_count, config.global_v_heads, config.key_dim, config.value_dim) + element_bytes = config.ssm_dtype.itemsize + ssm_byte_start = (conv_byte_end + element_bytes - 1) // element_bytes * element_bytes + ssm_region = TransferRegion( + field="ssm", + byte_start=ssm_byte_start, + byte_end=ssm_byte_start + prod(ssm_shape) * element_bytes, + dtype=config.ssm_dtype, + axes=("layer", "head", "key_dim", "value_dim"), + shape=ssm_shape, + axis_mappings=( + AxisMapping(axis="head", local_start=0, external_start=head_start, length=config.num_v_heads), + AxisMapping(axis="key_dim", local_start=0, external_start=0, length=config.key_dim), + AxisMapping(axis="value_dim", local_start=0, external_start=0, length=config.value_dim), + ), + write_enabled=write_enabled, + ) + regions.append(ssm_region) + logger.debug( + "%s transfer regions built: SSM head_start=%d, heads=%d, write_enabled=%s, SSM bytes=[%d, %d)", + self.state_name, + head_start, + config.num_v_heads, + write_enabled, + ssm_region.byte_start, + ssm_region.byte_end, + ) + return tuple(regions) + + def _pool(self, page_kind: PageKind): + """根据大小页类型取得已启用的页池，不访问 allocator。""" + assert page_kind in ("big", "small"), f"unknown page kind {page_kind}" + pool = self.big_pages_pool if page_kind == "big" else self.small_pages_pool + assert pool is not None, f"{page_kind} pages are disabled for {self.state_name}" + return pool + + @staticmethod + def _tensor_fields(conv: torch.Tensor, ssm: torch.Tensor): + """描述一页里的 Conv 和 SSM 字段。""" + return { + "conv": TensorField(tensor=conv, axes=("layer", "channel", "window")), + "ssm": TensorField(tensor=ssm, axes=("layer", "head", "key_dim", "value_dim")), + } + + def _accepted_conv(self, req_idx: int) -> torch.Tensor: + """取出该请求已接受的 Conv 窗口，不包含 MTP 候选区。""" + width = self.config.conv_width + conv = self.conv_state[:, req_idx] + if self.mtp_state_index is None: + return conv[:, :, :width] + offset = self.mtp_state_index[req_idx] + return conv.index_select(-1, torch.arange(width, device=conv.device) + offset) + + def _materialize_accepted( + self, + reqs: torch.Tensor, + output: torch.Tensor, + snapshot: bool, + output_strides=None, + flush_export: bool = False, + when_full: bool = False, + ) -> None: + """将请求的 accepted prefix 写入 canonical SSM 或调用方提供的输出页。""" + + assert not (snapshot and flush_export), "flush export must update the canonical state" + layers, slots, hv, k, v = self.ssm_state.shape + config = self.run_config or {"BV": 32, "num_warps": 4, "num_stages": 3} + bv = config["BV"] + _fold_prefix[(triton.cdiv(v, bv), reqs.numel(), layers * hv)]( + self.ssm_state, + output, + self.raw_keys, + self.raw_values, + self.gates, + self.betas, + self.cursors, + reqs, + self.ssm_state.stride(), + slots, + hv, + self.num_key_heads, + k, + v, + self.capacity, + self.hold_req_id, + triton.next_power_of_2(k), + bv, + snapshot, + False, + FLUSH_EXPORT=flush_export, + OUTPUT_STRIDES=output_strides, + WHEN_FULL=when_full, + NEXT=self.verify_width, + num_warps=config["num_warps"], + num_stages=config["num_stages"], + ) + if flush_export or when_full: + _reset_folded_cursors[(triton.cdiv(reqs.numel(), 256),)]( + self.cursors, + reqs, + reqs.numel(), + self.hold_req_id, + 0 if flush_export else self.capacity, + self.verify_width, + 256, + ) + + +def _history_owner( + ssm: torch.Tensor, + capacity: int, + verify_width: int, + activation_dtype: torch.dtype = torch.bfloat16, + num_key_heads: int | None = None, + gate_lower_bound: float = -5.0, + compute_mode: str = "inline", + run_config: dict[str, int] | None = None, +) -> ReplayGDNState: + """把 history 挂到调用方已有的 SSM 上，不再分配第二份状态。供保持原 stride 的 kernel 测试使用。""" + + owner = ReplayGDNState.__new__(ReplayGDNState) + owner.config = None + owner.state_name = None + owner.request_capacity = ssm.shape[1] + owner.mtp_step = verify_width - 1 + owner._algorithm = "gdn" + owner._logical_to_storage_layer = {} + owner.conv_state = None + owner.mtp_state_index = None + owner.big_pages_pool = None + owner.small_pages_pool = None + _bind_history( + owner, + ssm, + capacity, + verify_width, + activation_dtype, + num_key_heads, + gate_lower_bound, + compute_mode, + run_config, + ) + return owner + + +def _bind_history( + owner, + ssm: torch.Tensor, + capacity: int, + verify_width: int, + activation_dtype: torch.dtype, + num_key_heads: int | None, + gate_lower_bound: float, + compute_mode: str, + run_config: dict[str, int] | None, + force_compress: bool = False, +) -> None: + """在 owner 上创建 GDN replay 的连续前缀 history 和 cursor。""" + + assert ssm.dtype in (torch.float32, torch.bfloat16), "GDN replay requires FP32 or BF16 SSM state" + assert verify_width > 0 and capacity >= max(4, verify_width) and capacity & (capacity - 1) == 0 + assert compute_mode in ("inline", "precompute") + owner.ssm_state = ssm + owner.kda = False + owner.gate_lower_bound = gate_lower_bound + owner.capacity = capacity + owner.verify_width = verify_width + owner.force_compress = bool(force_compress) + owner.compute_mode = compute_mode + owner.run_config = run_config + owner._config_is_fixed = run_config is not None + owner.batch_configs = {} + layers, slots, hv, k, v = ssm.shape + owner.num_key_heads = hv if num_key_heads is None else num_key_heads + assert owner.num_key_heads > 0 and hv % owner.num_key_heads == 0 + owner.hold_req_id = slots - 1 + owner.cursors = torch.zeros(slots, dtype=torch.int32, device=ssm.device) + # 归一化 k 可用 fp16；delta 保留 fp32，避免有限 BF16 输入的修正量溢出 fp16。 + # 原始 k/v 保持激活精度。g 和 beta 保持 fp32。压缩在 fp32 里从原始值折回。 + reconstruction_dtype = torch.float16 if activation_dtype == torch.bfloat16 else torch.float32 + key_shape = (layers, slots, owner.num_key_heads, capacity, k) + owner.keys = torch.empty(key_shape, dtype=reconstruction_dtype, device=ssm.device) + owner.deltas = torch.empty((layers, slots, hv, capacity, v), dtype=torch.float32, device=ssm.device) + owner.gates = torch.empty((layers, slots, hv, capacity), dtype=torch.float32, device=ssm.device) + owner.raw_keys = torch.empty(key_shape, dtype=activation_dtype, device=ssm.device) + owner.raw_values = torch.empty((layers, slots, hv, capacity, v), dtype=activation_dtype, device=ssm.device) + owner.betas = torch.empty((layers, slots, hv, capacity), dtype=torch.float32, device=ssm.device) diff --git a/lightllm/common/attention_state/bounded_att/kda/create_kda_state.py b/lightllm/common/attention_state/bounded_att/kda/create_kda_state.py index 49c9ad0ce5..1742441389 100644 --- a/lightllm/common/attention_state/bounded_att/kda/create_kda_state.py +++ b/lightllm/common/attention_state/bounded_att/kda/create_kda_state.py @@ -13,8 +13,24 @@ def create_kda_state( request_capacity: int, mtp_step: int = 0, device: str | torch.device = "cuda", + ssm_state_mode: str = "native", + ssm_replay_config: dict | None = None, + gate_lower_bound: float = -5.0, ) -> KDAState: - """创建 KDA 请求运行态，大小页池由后续容量分配流程初始化。""" + """创建 KDA 请求运行态。replay 模式使用独立的 ReplayKDAState。""" + if ssm_state_mode == "replay": + from .replay_kda_state import ReplayKDAState + + return ReplayKDAState( + config=config, + request_capacity=request_capacity, + mtp_step=mtp_step, + device=device, + ssm_replay_config=ssm_replay_config, + gate_lower_bound=gate_lower_bound, + ) + if ssm_state_mode != "native": + raise ValueError(f"unknown KDA state mode: {ssm_state_mode}") return KDAState( config=config, request_capacity=request_capacity, diff --git a/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py new file mode 100644 index 0000000000..c2507e03e9 --- /dev/null +++ b/lightllm/common/attention_state/bounded_att/kda/replay_kda_state.py @@ -0,0 +1,834 @@ +"""KDA replay 的 bounded attention 状态。 + +Conv、canonical SSM、replay history、页布局和已接受快照都在本类。 +一层 decode 重放在 replay_kda_decode。 +""" + +from __future__ import annotations + +from functools import lru_cache +from math import prod +from typing import NamedTuple + +import torch +import triton + +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import ( + _advance_cursor, + _fold_prefix, +) +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm_fold import _reset_folded_cursors +from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix +from lightllm.utils.log_utils import init_logger + +from ...types import AxisMapping, PageKind, StateTransferLayout, TensorField, TransferRegion +from ...utils import deepcopy_result, dtype_name +from ..base_state import BoundedAttentionState +from ..conv_ssm_types import ConvSSMStateConfig +from ..copy_kernel import copy_bounded_att_tensors +from ..transfer_utils import bounded_att_buffer_copies +from .kda_pages_pool import KDAPagesPool + +logger = init_logger(__name__) + + +class _AcceptedTensors(NamedTuple): + conv: torch.Tensor + ssm: torch.Tensor + + +class ReplayKDAState(BoundedAttentionState): + """KDA replay 运行态。Conv、单槽 SSM 和 history 都在本类，不沿用 native 的 MTP 槽。""" + + # ======================================================================== + # 一、资源定义与初始化 + # ======================================================================== + + def __init__( + self, + config: ConvSSMStateConfig, + request_capacity: int, + mtp_step: int = 0, + device: str | torch.device = "cuda", + ssm_state_mode: str = "replay", + ssm_replay_config: dict | None = None, + gate_lower_bound: float = -5.0, + ) -> None: + """创建单槽 canonical SSM，并挂上只属于当前请求的 replay history。""" + if ssm_state_mode != "replay": + raise ValueError(f"ReplayKDAState only implements replay mode, got {ssm_state_mode}") + assert isinstance(config, ConvSSMStateConfig), "config must be ConvSSMStateConfig" + self.config = config + self.state_name = None + self.request_capacity = request_capacity + self.mtp_step = mtp_step + self.gate_lower_bound = gate_lower_bound + self._algorithm = "kda" + self._logical_to_storage_layer = {layer_id: index for index, layer_id in enumerate(config.logical_layer_ids)} + layer_count = len(config.logical_layer_ids) + self.conv_state = torch.zeros( + ( + layer_count, + request_capacity, + config.conv_channels, + config.conv_width + mtp_step, + ), + dtype=config.conv_dtype, + device=device, + ) + self.ssm_state = torch.zeros( + ( + layer_count, + request_capacity, + config.num_v_heads, + config.key_dim, + config.value_dim, + ), + dtype=config.ssm_dtype, + device=device, + ) + self.mtp_state_index = torch.zeros(request_capacity, dtype=torch.int32, device=device) if mtp_step else None + self.big_pages_pool = None + self.small_pages_pool = None + replay_config = dict(ssm_replay_config or {}) + capacity = int(replay_config.get("cache_len", 16)) + compute_mode = replay_config.get("compute_mode", "precompute" if mtp_step > 0 else "inline") + force_compress = bool(replay_config.get("force_compress", False)) + _bind_history( + self, + self.ssm_state, + capacity, + mtp_step + 1, + config.conv_dtype, + config.num_k_heads, + gate_lower_bound, + compute_mode, + None, + force_compress, + ) + + def get_state_type(self) -> str: + """返回算法名和 Conv/SSM 存储精度。""" + return ( + f"{self._algorithm}_conv_{dtype_name(dtype=self.config.conv_dtype)}" + f"_ssm_{dtype_name(dtype=self.config.ssm_dtype)}" + ) + + def initialize_page_pools( + self, + big_page_capacity: int, + small_page_capacity: int = 0, + device: str | torch.device | int = "cpu", + pin_memory: bool = True, + ) -> None: + """容量规划完成后创建本算法的大小页池；容量为零的页池保持 None。""" + assert self.big_pages_pool is None and self.small_pages_pool is None, "page pools already initialized" + assert big_page_capacity >= 0 and small_page_capacity >= 0, "page capacities must be nonnegative" + self.big_pages_pool = ( + KDAPagesPool( + size=big_page_capacity, + config=self.config, + device=device, + pin_memory=pin_memory, + ) + if big_page_capacity + else None + ) + self.small_pages_pool = ( + KDAPagesPool( + size=small_page_capacity, + config=self.config, + device=device, + pin_memory=pin_memory, + ) + if small_page_capacity + else None + ) + + # ======================================================================== + # 二、计算视图、运行态重置、已接受快照与 MTP 提交 + # ======================================================================== + + def contains_logical_layer(self, logical_layer_id: int) -> bool: + """判断指定逻辑层是否属于当前 State 的层映射。""" + return logical_layer_id in self._logical_to_storage_layer + + def storage_layer(self, logical_layer_id: int) -> int: + """返回逻辑层在 history 和 canonical state 里的存储下标。""" + + return self._logical_to_storage_layer[logical_layer_id] + + def get_layer_view(self, logical_layer_id: int): + index = self.storage_layer(logical_layer_id) + return self.conv_state[index], self.ssm_state[index] + + def reset_bounded_att_state(self, req_indices: list[int]) -> None: + """清零请求的 Conv、单槽 SSM、history 和 MTP 游标。""" + assert isinstance(req_indices, list), "request indices must be a list" + for req_idx in req_indices: + self.conv_state[:, req_idx].zero_() + self.ssm_state[:, req_idx].zero_() + self.clear_history(req_idx) + if self.mtp_state_index is not None: + self.mtp_state_index[req_idx].fill_(0) + + def save_bounded_att_runtime_to_pages( + self, + req_indices: list[int], + state_token_lens: list[int], + page_kind, + page_slot_ids: list[int], + ) -> None: + """把已接受 Conv 窗口和物化 SSM 写入页，不改 canonical state。""" + assert isinstance(req_indices, list) and isinstance(page_slot_ids, list), "indices must be lists" + assert len(req_indices) == len(state_token_lens) == len(page_slot_ids), "request, length and page counts differ" + pool = self._pool(page_kind) + copy_bounded_att_tensors( + copies=[ + (self._accepted_conv(req_idx), pool.conv_state[slot_id]) + for req_idx, slot_id in zip(req_indices, page_slot_ids) + ] + ) + self.export_ssm(req_indices, [pool.ssm_state[slot_id] for slot_id in page_slot_ids]) + + def restore_bounded_att_runtime_from_pages( + self, + req_indices: list[int], + state_token_lens: list[int], + page_kind, + page_slot_ids: list[int], + ) -> None: + """从页恢复基准 Conv/SSM，并清空目标请求的 history。""" + assert isinstance(req_indices, list) and isinstance(page_slot_ids, list), "indices must be lists" + assert len(req_indices) == len(state_token_lens) == len(page_slot_ids), "request, length and page counts differ" + pool = self._pool(page_kind) + width = self.config.conv_width + conv_copies, ssm_copies = [], [] + for req_idx, slot_id in zip(req_indices, page_slot_ids): + self.clear_history(req_idx) + conv_copies.append((pool.conv_state[slot_id], self.conv_state[:, req_idx, :, :width])) + ssm_copies.append((pool.ssm_state[slot_id], self.ssm_state[:, req_idx])) + copy_bounded_att_tensors(copies=conv_copies) + copy_bounded_att_tensors(copies=ssm_copies) + if self.mtp_state_index is not None: + for req_idx in req_indices: + self.mtp_state_index[req_idx].fill_(0) + + def commit_bounded_att_state(self, req_indices: torch.Tensor, accepted_steps: torch.Tensor) -> None: + """提交已接受长度，再按容量或 force_compress 决定是否整段压缩。""" + if self.mtp_state_index is None: + return + self.commit_accepted(self.mtp_state_index, req_indices, accepted_steps) + self.maybe_compress_prefix(req_indices) + + def get_accepted_state(self, req_idx: int) -> _AcceptedTensors: + """返回已接受 Conv 窗口和物化 SSM，不改写 canonical state。""" + return _AcceptedTensors(self._accepted_conv(req_idx), self.accepted_ssm(req_idx)) + + # ======================================================================== + # 三、CPU Cache 数据交换 + # ======================================================================== + + @deepcopy_result + @lru_cache(maxsize=None) + def get_cpu_transfer_regions(self) -> tuple[TransferRegion, ...]: + """返回已接受 Conv/SSM 快照的字段区间。不含 history 或 MTP 候选区。""" + return self._transfer_regions() + + def write_state_pages_to_cpu( + self, + page_kind: PageKind, + page_slot_ids: list[int], + cpu_page_buffers: torch.Tensor, + cpu_page_indices: list[int], + layout: StateTransferLayout, + cpu_page_ready_flags: list[bool] | None = None, + ) -> None: + """将指定页池槽位的快照写入 CPU 页。""" + assert layout.state_name == self.state_name, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" + pool = self._pool(page_kind) + assert isinstance(page_slot_ids, list), "slot indices must be a list" + assert isinstance(cpu_page_indices, list), "CPU page indices must be a list" + assert len(page_slot_ids) == len(cpu_page_indices), "CPU page and slot counts differ" + if cpu_page_ready_flags is not None: + assert len(cpu_page_ready_flags) == len(cpu_page_indices), "CPU page and ready counts differ" + page_fields, page_indices = [], [] + for index, (slot, page) in enumerate(zip(page_slot_ids, cpu_page_indices)): + # ready 页不参与复制；同一传输区域的所有页面组成一个同形状批次。 + if cpu_page_ready_flags is not None and cpu_page_ready_flags[index]: + continue + page_fields.append(self._tensor_fields(pool.conv_state[slot], pool.ssm_state[slot])) + page_indices.append(page) + for copies in bounded_att_buffer_copies( + page_fields=page_fields, + page_buffers=cpu_page_buffers, + layout=layout, + page_indices=page_indices, + write=True, + ): + copy_bounded_att_tensors(copies=copies) + + def read_state_pages_from_cpu( + self, + cpu_page_buffers: torch.Tensor, + cpu_page_indices: list[int], + big_page_slot_ids: list[int], + layout: StateTransferLayout, + ) -> None: + """读入指定大页槽位。请求运行态由 restore_bounded_att_runtime_from_pages 另行恢复。""" + assert layout.state_name == self.state_name, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + assert cpu_page_buffers.device.type == "cpu", "CPU cache buffer must reside on CPU" + pool = self.big_pages_pool + assert isinstance(big_page_slot_ids, list), "slot indices must be a list" + assert isinstance(cpu_page_indices, list), "CPU page indices must be a list" + assert len(cpu_page_indices) == len(big_page_slot_ids), "CPU page and slot counts differ" + # 本地目标保留本 TP rank 的分片形状。layout.regions 描述 CPU 全局字段。 + pool_page_fields = [ + self._tensor_fields(pool.conv_state[slot], pool.ssm_state[slot]) for slot in big_page_slot_ids + ] + for copies in bounded_att_buffer_copies( + page_fields=pool_page_fields, + page_buffers=cpu_page_buffers, + layout=layout, + page_indices=cpu_page_indices, + write=False, + ): + copy_bounded_att_tensors(copies=copies) + + # ======================================================================== + # 四、DP 迁移：大页快照与基准运行态 + # ======================================================================== + + def copy_big_pages_from_dp(self, src_page_ids: list[int], dst_page_ids: list[int], remote_dp_bounded_state) -> None: + """复制同一 TP 分片的大页 Conv/SSM 快照，保留源页和双方请求运行态。""" + assert isinstance(src_page_ids, list) and isinstance(dst_page_ids, list), "page indices must be lists" + assert len(src_page_ids) == len(dst_page_ids), "source and destination page counts differ" + assert ( + type(remote_dp_bounded_state) is type(self) and remote_dp_bounded_state.config == self.config + ), "DP copy requires the same algorithm and equivalent TP state configurations" + src_pool = remote_dp_bounded_state.big_pages_pool + dst_pool = self.big_pages_pool + assert src_pool is not None and dst_pool is not None, "DP big page copy requires both page pools" + # 大页是基准快照，与请求槽位无关。页池须在本进程注册，当前流访问完才能归还页面。 + conv_copies, ssm_copies = [], [] + for src_page, dst_page in zip(src_page_ids, dst_page_ids): + conv_copies.append((src_pool.conv_state[src_page], dst_pool.conv_state[dst_page])) + ssm_copies.append((src_pool.ssm_state[src_page], dst_pool.ssm_state[dst_page])) + copy_bounded_att_tensors(copies=conv_copies) + copy_bounded_att_tensors(copies=ssm_copies) + + def copy_runtime_state_from_dp( + self, + src_req_indices: list[int], + dst_req_indices: list[int], + remote_dp_bounded_state: "ReplayKDAState", + ) -> None: + """复制已接受运行态到目标单槽，并清空目标 history。""" + assert isinstance(src_req_indices, list) and isinstance(dst_req_indices, list), "request indices must be lists" + assert len(src_req_indices) == len(dst_req_indices), "source and destination request counts differ" + assert ( + type(remote_dp_bounded_state) is type(self) and remote_dp_bounded_state.config == self.config + ), "DP copy requires the same algorithm and equivalent TP state configurations" + width = self.config.conv_width + conv_copies = [] + for src_req, dst_req in zip(src_req_indices, dst_req_indices): + self.clear_history(dst_req) + conv_copies.append( + ( + remote_dp_bounded_state._accepted_conv(src_req), + self.conv_state[:, dst_req, :, :width], + ) + ) + copy_bounded_att_tensors(copies=conv_copies) + remote_dp_bounded_state.export_ssm(src_req_indices, [self.ssm_state[:, dst_req] for dst_req in dst_req_indices]) + if self.mtp_state_index is not None: + for dst_req in dst_req_indices: + self.mtp_state_index[dst_req].fill_(0) + + # ======================================================================== + # 五、PD 分离传输 + # ======================================================================== + + @deepcopy_result + @lru_cache(maxsize=None) + def get_pd_transfer_regions(self) -> tuple[TransferRegion, ...]: + """返回 PD 传输字段与切片，交由管理器安排其在联合页中的位置。""" + return self._transfer_regions() + + def write_to_pd_page( + self, + req_idx: int, + pd_page_buffers: torch.Tensor, + pd_page_index: int, + state_token_len: int, + layout: StateTransferLayout, + ) -> None: + """写出已接受 Conv 窗口和物化 SSM，不改写源 history。""" + assert layout.state_name == self.state_name, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + accepted = self.get_accepted_state(req_idx) + fields = self._tensor_fields(accepted.conv, accepted.ssm) + for copies in bounded_att_buffer_copies( + page_fields=[fields], + page_buffers=pd_page_buffers, + layout=layout, + page_indices=[pd_page_index], + write=True, + ): + copy_bounded_att_tensors(copies=copies) + + def read_from_pd_page( + self, + pd_page_buffers: torch.Tensor, + req_idx: int, + pd_page_index: int, + state_token_len: int, + layout: StateTransferLayout, + ) -> None: + """读入基准 Conv/SSM，并清空该请求的 history。""" + assert layout.state_name == self.state_name, "transfer layout does not describe this resource" + assert 0 <= layout.byte_start < layout.byte_end, "invalid State byte interval" + self.clear_history(req_idx) + fields = self._tensor_fields( + self.conv_state[:, req_idx, :, : self.config.conv_width], + self.ssm_state[:, req_idx], + ) + for copies in bounded_att_buffer_copies( + page_fields=[fields], + page_buffers=pd_page_buffers, + layout=layout, + page_indices=[pd_page_index], + write=False, + ): + copy_bounded_att_tensors(copies=copies) + if self.mtp_state_index is not None: + self.mtp_state_index[req_idx].fill_(0) + + # ======================================================================== + # 六、Replay history + # ======================================================================== + def finish_prefill(self, logical_layer_id: int, reqs: torch.Tensor) -> None: + """最后一层把已接受前缀整段压进 SSM。""" + + layers = () if self.config is None else self.config.logical_layer_ids + if layers and logical_layer_id == layers[-1]: + self.compress_prefix(reqs) + + def finish_decode_token(self, logical_layer_id: int, reqs: torch.Tensor) -> None: + """最后一层记下这 1 个 token，再按容量决定是否整段压缩。""" + + layers = () if self.config is None else self.config.logical_layer_ids + if not layers or logical_layer_id != layers[-1]: + return + self.accept_single_token(reqs) + self.maybe_compress_prefix(reqs) + + def clear_history(self, req: int | torch.Tensor) -> None: + """清空复用或恢复请求的临时 history，保留 canonical SSM。""" + + if type(req) is int: + # 标量赋值会把 CPU tensor 复制到 CUDA slot。 + self.cursors[req].zero_() + else: + self.cursors[req] = 0 + + def prepare_decode(self, reqs: torch.Tensor, cu_seqlens: torch.Tensor | None = None) -> torch.Tensor: + """返回本步写入前的前缀长度。decode 不在这里修改 SSM。""" + + del cu_seqlens + return self.cursors.index_select(0, reqs.long()) + + def accept_updates( + self, reqs: torch.Tensor, accepted_state_indices: torch.Tensor | None = None, width: int | None = None + ) -> None: + """更新 cursor。width 缺省时按 verify_width 取已接受长度，传 1 时固定加一个 token。""" + + _advance_cursor[(triton.cdiv(reqs.numel(), 256),)]( + self.cursors, + reqs, + accepted_state_indices, + reqs.numel(), + self.hold_req_id, + self.verify_width if width is None else width, + 256, + ) + + def accept_single_token(self, reqs: torch.Tensor) -> None: + """把一个已接受 token 记入前缀长度。""" + + if reqs.numel() == 0: + return + # width=1 不读 index。仍传 reqs，kernel 不能收到空指针。 + self.accept_updates(reqs, reqs, width=1) + + def merge_accepted_updates(self, reqs: torch.Tensor) -> None: + """把 cursor 覆盖的已接受前缀整段折回 canonical SSM，并清空 cursor。""" + + if reqs.numel() == 0: + return + self._materialize_accepted(reqs, self.ssm_state, snapshot=False) + self.cursors.index_fill_(0, reqs.long(), 0) + + def compress_prefix(self, reqs: torch.Tensor) -> None: + """把 cursor 覆盖的已接受前缀整段压进 SSM。长度为 0 时 kernel 直接返回。""" + + if reqs.numel() == 0: + return + if self.ssm_state.is_cuda: + self.merge_accepted_updates(reqs) + else: + self.clear_history(reqs) + + def maybe_compress_prefix(self, reqs: torch.Tensor) -> None: + """长度加上一次最大 verify_width 放不下，或 force_compress 打开时，整段压缩。""" + + if reqs.numel() == 0 or not self.ssm_state.is_cuda: + return + if self.force_compress: + self.compress_prefix(reqs) + return + self._materialize_accepted(reqs, self.ssm_state, snapshot=False, when_full=True) + + def accepted_ssm(self, req_idx: int) -> torch.Tensor: + """返回不修改 canonical SSM 和 history 的单请求快照。""" + + layers, _, hv, k, v = self.ssm_state.shape + output = torch.empty((layers, hv, k, v), dtype=self.ssm_state.dtype, device=self.ssm_state.device) + reqs = torch.tensor([req_idx], dtype=torch.int32, device=self.ssm_state.device) + self._materialize_accepted(reqs, output, snapshot=True) + return output + + def export_accepted_state( + self, + req_indices: list[int], + targets: list[torch.Tensor], + flush: bool = False, + ) -> None: + """将多个 accepted 快照直接写入调用方目标页。""" + + if not req_indices: + return + shape = (self.ssm_state.shape[0], *self.ssm_state.shape[2:]) + strides = targets[0].stride() + assert len(req_indices) == len(targets) + assert all(t.shape == shape and t.stride() == strides and t.dtype == self.ssm_state.dtype for t in targets) + device = self.ssm_state.device + self._materialize_accepted( + torch.tensor(req_indices, dtype=torch.int32, pin_memory=True).to(device, non_blocking=True), + torch.tensor([t.data_ptr() for t in targets], dtype=torch.int64, pin_memory=True).to( + device, non_blocking=True + ), + snapshot=not flush, + flush_export=flush, + output_strides=strides, + ) + + def commit_accepted(self, indices: torch.Tensor, reqs: torch.Tensor, lengths: torch.Tensor) -> None: + """把接受长度写入 MTP 槽位索引，并推进本请求的 replay cursor。""" + + if self.verify_width > 1: + commit_ssm_prefix(indices, reqs, lengths, self.cursors) + else: + commit_ssm_prefix(indices, reqs, lengths) + self.accept_updates(reqs, indices) + + def export_ssm(self, req_indices: list[int], targets: list[torch.Tensor], flush: bool = False) -> None: + """把已接受 SSM 写入调用方页面。CUDA 目标直接写指针，其他设备先物化再拷贝。""" + + if not req_indices: + return + target = targets[0] + if target.is_cuda and target.device == self.ssm_state.device: + self.export_accepted_state(req_indices, targets, flush=flush) + return + for req, dst in zip(req_indices, targets): + snap = self.accepted_ssm(req) + dst.copy_(snap) + if flush: + self.ssm_state[:, req].copy_(snap) + self.clear_history(req) + + _run_config: dict[str, int] | None = None + + @property + def run_config(self) -> dict[str, int] | None: + """返回物化和 autotune 使用的 Triton 配置，不是 SSM 布局。""" + + return self._run_config + + @run_config.setter + def run_config(self, config: dict | None) -> None: + """保存一份与调用方脱开的 kernel 配置。""" + + if config is None: + self._run_config = None + return + self._run_config = { + "BV": int(config["BV"]), + "num_warps": int(config["num_warps"]), + "num_stages": int(config.get("num_stages", 3)), + } + + # ======================================================================== + # 七、私有辅助实现 + # ======================================================================== + + def _transfer_regions(self) -> tuple[TransferRegion, ...]: + """依次描述全局 Conv 和 SSM 区域，并标明当前 TP rank 的读写范围。""" + # 只传有效历史窗口，不包含 MTP 候选区和游标。 + config = self.config + storage_layer_count = len(config.logical_layer_ids) + conv_channels = 2 * config.global_k_heads * config.key_dim + config.global_v_heads * config.value_dim + conv_shape = (storage_layer_count, conv_channels, config.conv_width) + conv_byte_end = prod(conv_shape) * config.conv_dtype.itemsize + regions = [] + logger.debug( + "%s building transfer regions: TP=%d/%d, Conv shape=%s", + self.state_name, + config.tp_rank, + config.tp_world_size, + conv_shape, + ) + + # Conv 按 Q/K/V 拼接。三个映射共用同一块全局 Conv 区域，不能重复累计字节数。 + local_channel_start = 0 + global_channel_start = 0 + for name, local_head_count, global_head_count, head_dim in ( + ("Q", config.num_k_heads, config.global_k_heads, config.key_dim), + ("K", config.num_k_heads, config.global_k_heads, config.key_dim), + ("V", config.num_v_heads, config.global_v_heads, config.value_dim), + ): + # 每组连续 rank 保存相同 heads，仅组内首 rank 写出。 + replica_count = config.tp_world_size * local_head_count // global_head_count + head_start = config.tp_rank // replica_count * local_head_count + write_enabled = config.tp_rank % replica_count == 0 + local_channel_count = local_head_count * head_dim + target_channel_start = global_channel_start + head_start * head_dim + regions.append( + TransferRegion( + field="conv", + byte_start=0, + byte_end=conv_byte_end, + dtype=config.conv_dtype, + axes=("layer", "channel", "window"), + shape=conv_shape, + axis_mappings=( + AxisMapping( + axis="channel", + local_start=local_channel_start, + external_start=target_channel_start, + length=local_channel_count, + ), + AxisMapping(axis="window", local_start=0, external_start=0, length=config.conv_width), + ), + write_enabled=write_enabled, + ) + ) + logger.debug( + "%s Conv %s: local_channel_start=%d, global_channel_start=%d, " + "channels=%d, replicas=%d, write_enabled=%s", + self.state_name, + name, + local_channel_start, + target_channel_start, + local_channel_count, + replica_count, + write_enabled, + ) + local_channel_start += local_channel_count + global_channel_start += global_head_count * head_dim + + # SSM 沿 V head 切分，接在 Conv 之后，并按 SSM dtype 对齐字节起点。 + replica_count = config.tp_world_size * config.num_v_heads // config.global_v_heads + head_start = config.tp_rank // replica_count * config.num_v_heads + write_enabled = config.tp_rank % replica_count == 0 + ssm_shape = (storage_layer_count, config.global_v_heads, config.key_dim, config.value_dim) + element_bytes = config.ssm_dtype.itemsize + ssm_byte_start = (conv_byte_end + element_bytes - 1) // element_bytes * element_bytes + ssm_region = TransferRegion( + field="ssm", + byte_start=ssm_byte_start, + byte_end=ssm_byte_start + prod(ssm_shape) * element_bytes, + dtype=config.ssm_dtype, + axes=("layer", "head", "key_dim", "value_dim"), + shape=ssm_shape, + axis_mappings=( + AxisMapping(axis="head", local_start=0, external_start=head_start, length=config.num_v_heads), + AxisMapping(axis="key_dim", local_start=0, external_start=0, length=config.key_dim), + AxisMapping(axis="value_dim", local_start=0, external_start=0, length=config.value_dim), + ), + write_enabled=write_enabled, + ) + regions.append(ssm_region) + logger.debug( + "%s transfer regions built: SSM head_start=%d, heads=%d, write_enabled=%s, SSM bytes=[%d, %d)", + self.state_name, + head_start, + config.num_v_heads, + write_enabled, + ssm_region.byte_start, + ssm_region.byte_end, + ) + return tuple(regions) + + def _pool(self, page_kind: PageKind): + """根据大小页类型取得已启用的页池，不访问 allocator。""" + assert page_kind in ("big", "small"), f"unknown page kind {page_kind}" + pool = self.big_pages_pool if page_kind == "big" else self.small_pages_pool + assert pool is not None, f"{page_kind} pages are disabled for {self.state_name}" + return pool + + @staticmethod + def _tensor_fields(conv: torch.Tensor, ssm: torch.Tensor): + """描述一页里的 Conv 和 SSM 字段。""" + return { + "conv": TensorField(tensor=conv, axes=("layer", "channel", "window")), + "ssm": TensorField(tensor=ssm, axes=("layer", "head", "key_dim", "value_dim")), + } + + def _accepted_conv(self, req_idx: int) -> torch.Tensor: + """取出该请求已接受的 Conv 窗口，不包含 MTP 候选区。""" + width = self.config.conv_width + conv = self.conv_state[:, req_idx] + if self.mtp_state_index is None: + return conv[:, :, :width] + offset = self.mtp_state_index[req_idx] + return conv.index_select(-1, torch.arange(width, device=conv.device) + offset) + + def _materialize_accepted( + self, + reqs: torch.Tensor, + output: torch.Tensor, + snapshot: bool, + output_strides=None, + flush_export: bool = False, + when_full: bool = False, + ) -> None: + """将请求的 accepted prefix 写入 canonical SSM 或调用方提供的输出页。""" + + assert not (snapshot and flush_export), "flush export must update the canonical state" + layers, slots, hv, k, v = self.ssm_state.shape + config = self.run_config or {"BV": 32, "num_warps": 4, "num_stages": 3} + bv = config["BV"] + _fold_prefix[(triton.cdiv(v, bv), reqs.numel(), layers * hv)]( + self.ssm_state, + output, + self.raw_keys, + self.raw_values, + self.gates, + self.betas, + self.cursors, + reqs, + self.ssm_state.stride(), + slots, + hv, + self.num_key_heads, + k, + v, + self.capacity, + self.hold_req_id, + triton.next_power_of_2(k), + bv, + snapshot, + True, + FLUSH_EXPORT=flush_export, + OUTPUT_STRIDES=output_strides, + WHEN_FULL=when_full, + NEXT=self.verify_width, + num_warps=config["num_warps"], + num_stages=config["num_stages"], + ) + if flush_export or when_full: + _reset_folded_cursors[(triton.cdiv(reqs.numel(), 256),)]( + self.cursors, + reqs, + reqs.numel(), + self.hold_req_id, + 0 if flush_export else self.capacity, + self.verify_width, + 256, + ) + + +def _history_owner( + ssm: torch.Tensor, + capacity: int, + verify_width: int, + activation_dtype: torch.dtype = torch.bfloat16, + num_key_heads: int | None = None, + gate_lower_bound: float = -5.0, + compute_mode: str = "inline", + run_config: dict[str, int] | None = None, +) -> ReplayKDAState: + """把 history 挂到调用方已有的 SSM 上，不再分配第二份状态。供保持原 stride 的 kernel 测试使用。""" + + owner = ReplayKDAState.__new__(ReplayKDAState) + owner.config = None + owner.state_name = None + owner.request_capacity = ssm.shape[1] + owner.mtp_step = verify_width - 1 + owner._algorithm = "kda" + owner._logical_to_storage_layer = {} + owner.conv_state = None + owner.mtp_state_index = None + owner.big_pages_pool = None + owner.small_pages_pool = None + _bind_history( + owner, + ssm, + capacity, + verify_width, + activation_dtype, + num_key_heads, + gate_lower_bound, + compute_mode, + run_config, + ) + return owner + + +def _bind_history( + owner, + ssm: torch.Tensor, + capacity: int, + verify_width: int, + activation_dtype: torch.dtype, + num_key_heads: int | None, + gate_lower_bound: float, + compute_mode: str, + run_config: dict[str, int] | None, + force_compress: bool = False, +) -> None: + """在 owner 上创建 KDA replay 的连续前缀 history 和 cursor。""" + + assert ssm.dtype in (torch.float32, torch.bfloat16), "KDA replay requires FP32 or BF16 SSM state" + assert verify_width > 0 and capacity >= max(4, verify_width) and capacity & (capacity - 1) == 0 + assert compute_mode in ("inline", "precompute") + owner.ssm_state = ssm + owner.kda = True + owner.gate_lower_bound = gate_lower_bound + owner.capacity = capacity + owner.verify_width = verify_width + owner.force_compress = bool(force_compress) + owner.compute_mode = compute_mode + owner.run_config = run_config + owner._config_is_fixed = run_config is not None + owner.batch_configs = {} + layers, slots, hv, k, v = ssm.shape + owner.num_key_heads = hv if num_key_heads is None else num_key_heads + assert owner.num_key_heads > 0 and hv % owner.num_key_heads == 0 + owner.hold_req_id = slots - 1 + owner.cursors = torch.zeros(slots, dtype=torch.int32, device=ssm.device) + # 归一化 k 可用 fp16；delta 保留 fp32，避免有限 BF16 输入的修正量溢出 fp16。 + # 原始 k/v 保持激活精度。g 和 beta 保持 fp32。压缩在 fp32 里从原始值折回。 + reconstruction_dtype = torch.float16 if activation_dtype == torch.bfloat16 else torch.float32 + key_shape = (layers, slots, owner.num_key_heads, capacity, k) + owner.keys = torch.empty(key_shape, dtype=reconstruction_dtype, device=ssm.device) + owner.deltas = torch.empty((layers, slots, hv, capacity, v), dtype=torch.float32, device=ssm.device) + owner.gates = torch.empty((layers, slots, hv, capacity, k), dtype=torch.float32, device=ssm.device) + owner.raw_keys = torch.empty(key_shape, dtype=activation_dtype, device=ssm.device) + owner.raw_values = torch.empty((layers, slots, hv, capacity, v), dtype=activation_dtype, device=ssm.device) + owner.betas = torch.empty((layers, slots, hv, capacity), dtype=torch.float32, device=ssm.device) diff --git a/lightllm/common/attention_state/bounded_att/utils.py b/lightllm/common/attention_state/bounded_att/utils.py index 93f5b2ebe3..be4be7e5ed 100644 --- a/lightllm/common/attention_state/bounded_att/utils.py +++ b/lightllm/common/attention_state/bounded_att/utils.py @@ -2,7 +2,8 @@ from __future__ import annotations -from math import prod +from math import lcm, prod +from mmap import PAGESIZE from weakref import finalize import torch @@ -17,6 +18,10 @@ CUDA_HOST_REGISTER_PORTABLE = 1 # Mapped：建立设备可访问映射，供 CUDA/Triton kernel 直接访问主机存储。 CUDA_HOST_REGISTER_MAPPED = 2 +# Keep DP-replicated shared page pools below the CUDA host-registration range +# pressure observed with eight ranks. Each chunk is still rounded to a whole +# snapshot, so DMA copies never cross registration boundaries. +_HOST_REGISTRATION_BYTES = 1024 * 1024 * 1024 # ======================================================================== @@ -43,16 +48,25 @@ def create_page_pool_tensor( # 阶段一：未启用 DP fetch 时沿用 torch 分配；GPU 页池也保持原设备分配方式。 if not shared: - tensor = torch.empty( - size=size, - dtype=dtype, - device=device, - pin_memory=pin_memory, - ) + try: + tensor = torch.empty( + size=size, + dtype=dtype, + device=device, + pin_memory=pin_memory, + ) + except RuntimeError: + if device.type != "cpu" or not pin_memory or not torch.cuda.is_available(): + raise + tensor = torch.empty(size=size, dtype=dtype, device=device, pin_memory=False) + register_page_pool_tensor(tensor=tensor) + allocation = "torch.empty + cudaHostRegister" + else: + allocation = "torch.empty" reason = "GPU device allocation" if device.type != "cpu" else "DP prompt cache fetch disabled" logger.info( f"Allocated bounded attention page pool: shape={size}, dtype={dtype}, bytes={nbytes}, " - f"device={device}, shared=False, pin_memory={pin_memory}, allocator=torch.empty, reason={reason}" + f"device={device}, shared=False, pin_memory={pin_memory}, allocator={allocation}, reason={reason}" ) return tensor @@ -91,24 +105,43 @@ def register_page_pool_tensor(tensor: torch.Tensor) -> None: """ if tensor.device.type != "cpu" or tensor.numel() == 0 or tensor.is_pinned(): return - assert tensor.is_shared(), "page pool registration requires shared CPU storage" storage = tensor.untyped_storage() ptr = storage.data_ptr() nbytes = storage.nbytes() device_id = torch.cuda.current_device() flags = CUDA_HOST_REGISTER_PORTABLE | CUDA_HOST_REGISTER_MAPPED - torch.cuda.check_error(torch.cuda.cudart().cudaHostRegister(ptr, nbytes, flags)) + # 大页池整块注册在多 GPU 初始化时可能失败；边界对齐完整快照，避免一次 DMA 跨注册区间。 + registration_unit = lcm(prod(tensor.shape[1:]) * tensor.element_size(), PAGESIZE) + chunk_bytes = max(registration_unit, _HOST_REGISTRATION_BYTES // registration_unit * registration_unit) + if not tensor.is_shared(): + chunk_bytes = nbytes + addresses = [] + runtime = torch.cuda.cudart() + try: + for offset in range(0, nbytes, chunk_bytes): + address = ptr + offset + torch.cuda.check_error(runtime.cudaHostRegister(address, min(chunk_bytes, nbytes - offset), flags)) + addresses.append(address) + except BaseException: + logger.error( + f"Failed to register bounded attention page pool: bytes={nbytes}, cuda_device={device_id}, " + f"chunk_bytes={chunk_bytes}, registered_chunks={len(addresses)}" + ) + if addresses: + _unregister_page_pool_storage(tuple(addresses), device_id) + raise # finalizer 绑定 Storage，所有本地 Tensor 及视图释放后才注销，不能绑在单个 Tensor 上。 # 正常释放时等待 GPU 完成使用；解释器退出时交给进程清理，避免 CUDA 已关闭后再调用。 cleanup = finalize( storage, _unregister_page_pool_storage, - ptr, + tuple(addresses), device_id, ) cleanup.atexit = False logger.info( - f"Registered shared bounded attention page pool: bytes={nbytes}, cuda_device={device_id}, " + f"Registered shared bounded attention page pool: bytes={nbytes}, chunks={len(addresses)}, " + f"cuda_device={device_id}, " "flags=Portable|Mapped; registration applies to this process's mapping" ) @@ -118,8 +151,9 @@ def register_page_pool_tensor(tensor: torch.Tensor) -> None: # ======================================================================== -def _unregister_page_pool_storage(ptr: int, device_id: int) -> None: +def _unregister_page_pool_storage(addresses: tuple[int, ...], device_id: int) -> None: """等待当前 GPU 使用结束，再注销本进程的映射；共享内存本身由 Storage 管理。""" torch.cuda.synchronize(device=device_id) - torch.cuda.check_error(torch.cuda.cudart().cudaHostUnregister(ptr)) - logger.debug(f"Unregistered shared bounded attention page pool: cuda_device={device_id}, ptr={ptr}") + for ptr in reversed(addresses): + torch.cuda.check_error(torch.cuda.cudart().cudaHostUnregister(ptr)) + logger.debug(f"Unregistered shared bounded attention page pool: cuda_device={device_id}, chunks={len(addresses)}") diff --git a/lightllm/common/attention_state/full_att/gqa/gqa_state.py b/lightllm/common/attention_state/full_att/gqa/gqa_state.py index fe184b4110..c658533abc 100644 --- a/lightllm/common/attention_state/full_att/gqa/gqa_state.py +++ b/lightllm/common/attention_state/full_att/gqa/gqa_state.py @@ -146,14 +146,13 @@ def write_full_att_kv(self, logical_layer_id: int, token_indices: torch.Tensor, if not len(token_indices): return assert self.k_storage.is_cuda, "KV writes require CUDA storage" - from lightllm.common.basemodel.triton_kernel.destindex_copy_kv import destindex_copy_kv + from lightllm.common.basemodel.triton_kernel.destindex_copy_kv import destindex_copy_kv_fused - # TODO：将同一层 K/V 的两次索引写入融合到一个 kernel，直接写入独立 K/V 存储。 - destindex_copy_kv( - K=kv[:, : self.config.head_num], DestLoc=token_indices, Out=self.k_storage[storage_layer_index] - ) - destindex_copy_kv( - K=kv[:, self.config.head_num :], DestLoc=token_indices, Out=self.v_storage[storage_layer_index] + destindex_copy_kv_fused( + KV=kv, + DestLoc=token_indices, + OutK=self.k_storage[storage_layer_index], + OutV=self.v_storage[storage_layer_index], ) def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Tensor) -> None: @@ -169,22 +168,16 @@ def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Ten assert self.k_storage.is_cuda, "token copies require CUDA storage" if not src_indices.numel(): return - from lightllm.common.basemodel.triton_kernel.kv_move import copy_kv_buffer_to_kv_buffer + from lightllm.common.basemodel.triton_kernel.kv_move import copy_kv_buffer_pair_to_kv_buffer_pair # 直接使用调用方提供的 CUDA 索引，保留原 dtype 和 stride。 # 按字节直接从源槽位写到目标槽位，覆盖量化 scale 和 packed padding，不创建中间 KV。 # 调用方保证源、目标集合不相交，且目标索引不重复，避免并行读写竞争。 - # TODO：将 K/V 的同索引 token 复制融合为一次 kernel 发射， - # 保持原始字节编码及各字段 stride，仍要求源目标不重叠、目标索引不重复。 - copy_kv_buffer_to_kv_buffer( - src_mem_index=src_indices, - dst_mem_index=dst_indices, - kv_buffer=self.k_storage, - ) - copy_kv_buffer_to_kv_buffer( + copy_kv_buffer_pair_to_kv_buffer_pair( src_mem_index=src_indices, dst_mem_index=dst_indices, - kv_buffer=self.v_storage, + k_buffer=self.k_storage, + v_buffer=self.v_storage, ) # ======================================================================== diff --git a/lightllm/common/basemodel/attention/create_linear_utils.py b/lightllm/common/basemodel/attention/create_linear_utils.py index 6a40a05183..fce8e8d78f 100644 --- a/lightllm/common/basemodel/attention/create_linear_utils.py +++ b/lightllm/common/basemodel/attention/create_linear_utils.py @@ -19,38 +19,48 @@ def get_qwen35_linear_prefill_att_backend_class(index=1, priority_list=("flashql args = get_env_start_args() backend_str = _get_backend_str(args.llm_prefill_att_backend, index) if backend_str != "auto": - return load_backend_class(linear_prefill_att_backend_classes[backend_str]) + return _linear_backend_class(linear_prefill_att_backend_classes[backend_str], backend_str) for backend_name in priority_list: backend_class = linear_prefill_att_backend_classes[backend_name] if backend_name == "triton": logger.info("Linear prefill attention backend: triton.") - return load_backend_class(backend_class) + return _linear_backend_class(backend_class, backend_name) if validate(backend_name): logger.info(f"Linear prefill attention backend: {backend_name} (validated).") - return load_backend_class(backend_class) + return _linear_backend_class(backend_class, backend_name) logger.warning("No linear prefill attention backend validation succeeded, falling back to Triton.") - return load_backend_class("linear.triton:TritonLinearAttBackend") + return _linear_backend_class("linear.triton:TritonLinearAttBackend", "triton") def get_qwen35_linear_decode_att_backend_class(index=1, priority_list=("triton",)): args = get_env_start_args() backend_str = _get_backend_str(args.llm_decode_att_backend, index) if backend_str != "auto": - return load_backend_class(linear_decode_att_backend_classes[backend_str]) + return _linear_backend_class(linear_decode_att_backend_classes[backend_str], backend_str) for backend_name in priority_list: backend_class = linear_decode_att_backend_classes[backend_name] if backend_name == "triton": logger.info("Linear decode attention backend: triton.") - return load_backend_class(backend_class) + return _linear_backend_class(backend_class, backend_name) if validate(backend_name): logger.info(f"Linear decode attention backend: {backend_name} (validated).") - return load_backend_class(backend_class) + return _linear_backend_class(backend_class, backend_name) logger.warning("No linear decode attention backend validation succeeded, falling back to Triton.") - return load_backend_class("linear.triton:TritonLinearAttBackend") + return _linear_backend_class("linear.triton:TritonLinearAttBackend", "triton") + + +def _linear_backend_class(native_path: str, backend_name: str): + """Native 仍按路径加载。replay 把 kernel 名字交给独立的 GDN replay 入口。""" + + if getattr(get_env_start_args(), "ssm_state_mode", "native") != "replay": + return load_backend_class(native_path) + from lightllm.common.basemodel.attention.linear.replay_gdn import replay_gdn_backend_class + + return replay_gdn_backend_class(backend_name) def _get_backend_str(backend_args, index): diff --git a/lightllm/common/basemodel/attention/linear/gdn.py b/lightllm/common/basemodel/attention/linear/gdn.py index 986214568f..b715b77128 100644 --- a/lightllm/common/basemodel/attention/linear/gdn.py +++ b/lightllm/common/basemodel/attention/linear/gdn.py @@ -8,7 +8,6 @@ from lightllm.utils.envs_utils import get_env_start_args from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import causal_conv1d_fn from lightllm.common.basemodel.triton_kernel.linear_att.fused_gdn_gating import fused_gdn_gating -from lightllm.common.basemodel.triton_kernel.linear_att.gdn_decode_pack import conv_pack_gdn_decode_inputs from lightllm.common.basemodel.triton_kernel.linear_att.mtp_fused_recurrent import ( mtp_fused_recurrent_gated_delta_rule, ) @@ -17,6 +16,7 @@ ) from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops import ( fused_recurrent_gated_delta_rule, + fused_recurrent_gated_delta_rule_packed_decode, ) if TYPE_CHECKING: @@ -37,7 +37,6 @@ def get_prefill_kernel(self): pass def _init_linear_layer_metadata(self, network_config, tp_world_size): - self.mtp_step = get_env_start_args().mtp_step # Linear attention specific dimensions @@ -109,7 +108,6 @@ def create_att_decode_state(self, infer_state: "InferStateInfo") -> "LinearAttDe @dataclasses.dataclass class LinearAttPrefillAttState(BasePrefillAttState): - b_conv_buffer_idx: torch.Tensor = None b_ssm_buffer_idx: torch.Tensor = None @@ -207,11 +205,11 @@ def _gdn_prefill_kernel( @dataclasses.dataclass class LinearAttDecodeAttState(BaseDecodeAttState): - b_conv_buffer_idx: torch.Tensor = None b_ssm_buffer_idx: torch.Tensor = None b1_mtp_cu_q_seq_len: torch.Tensor = None b_num_accepted_tokens: torch.Tensor = None + b1_mtp_cu_q_seq_len_int64: torch.Tensor = None def init_state(self): draft_step = self.backend.model.mtp_manager.get_decode_draft_step(self.backend.model.is_mtp_draft_model) @@ -232,6 +230,7 @@ def init_state(self): self._init_dynamic_mtp_decode_state(draft_step + 1, req_to_mtp_state_index=state.mtp_state_index) else: self._init_fixed_mtp_decode_state(draft_step, req_to_mtp_state_index=state.mtp_state_index) + self.b1_mtp_cu_q_seq_len_int64 = self.b1_mtp_cu_q_seq_len.to(torch.int64) def _init_normal_decode_state(self): self.b_conv_buffer_idx = self.infer_state.b_req_idx @@ -345,37 +344,30 @@ def _gdn_decode_kernel( ): backend: LinearAttBackend = self.backend - # Recurrent processing with fused gating. Decode uses a specialized - # conv+pack kernel to avoid materializing the post-conv qkv tensor - # before immediately splitting it into q/k/v. - query, key, value, z, a, b = conv_pack_gdn_decode_inputs( + from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import causal_conv1d_update + + # Keep the post-convolution q/k/v layout packed. The recurrent kernel + # performs the GQA rearrangement and recurrent update in one launch, + # matching the packed decode path used by vLLM and avoiding six + # short-lived q/k/v/z/a/b tensors on every linear-attention layer. + mixed_qkv = causal_conv1d_update( mixed_qkv, - z, - a, - b, conv_states, layer_weight.linear_conv1d.mm_param.weight, layer_weight.linear_conv1d.bias, - self.b_conv_buffer_idx, - backend.activation, - backend.conv_kernel_dim, - backend.tp_num_k_heads, - backend.head_k_dim, - backend.tp_num_v_heads, - backend.head_v_dim, + activation=backend.activation, + conv_state_indices=self.b_conv_buffer_idx, ) - core_attn_out, _ = fused_recurrent_gated_delta_rule( - q=query, - k=key, - v=value, + core_attn_out, _ = fused_recurrent_gated_delta_rule_packed_decode( + mixed_qkv=mixed_qkv, + a=a, + b=b, + A_log=layer_weight.linear_A_log.weight, + dt_bias=layer_weight.linear_dt_bias.weight, + scale=backend.head_k_dim ** -0.5, initial_state=ssm_states, - inplace_final_state=True, ssm_state_indices=self.b_ssm_buffer_idx, use_qk_l2norm_in_kernel=True, - A_log=layer_weight.linear_A_log.weight, - dt_bias=layer_weight.linear_dt_bias.weight, - a_raw=a, - b_raw=b, ) return core_attn_out, z @@ -418,7 +410,7 @@ def _gdn_mtp_kernel( k=key, v=value, initial_state=ssm_states, - cu_seqlens=cu_seqlens_q.to(torch.long), + cu_seqlens=self.b1_mtp_cu_q_seq_len_int64, ssm_state_indices=self.b_ssm_buffer_idx, ssm_state_write_indices=self.b_ssm_buffer_idx, num_accepted_tokens=self.b_num_accepted_tokens, @@ -426,5 +418,10 @@ def _gdn_mtp_kernel( dt_bias=layer_weight.linear_dt_bias.weight, a_raw=a, b_raw=b, + run_config=( + {"BV": 32, "num_warps": 4, "num_stages": 1} + if getattr(backend.model, "linear_att_mtp_match_decode", False) + else None + ), ) return core_attn_out diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn.py b/lightllm/common/basemodel/attention/linear/replay_gdn.py new file mode 100644 index 0000000000..d08a714d27 --- /dev/null +++ b/lightllm/common/basemodel/attention/linear/replay_gdn.py @@ -0,0 +1,352 @@ +"""GDN replay 的推理入口。 + +Triton 和 FlashQla 各有一个 backend。它们只差 prefill chunk kernel， +decode 都走 ReplayGDNState。两个类都直接实现 backend 接口，不继承 native GDN。 +""" + +from __future__ import annotations + +import dataclasses +import weakref +from typing import TYPE_CHECKING + +import torch + +from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ( + ReplayGDNState, +) +from lightllm.common.basemodel.attention.base_att import ( + AttControl, + BaseAttBackend, + BaseDecodeAttState, + BasePrefillAttState, +) +from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import ( + causal_conv1d_fn, +) +from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_mtp import ( + causal_conv1d_update as causal_conv1d_update_mtp, +) +from lightllm.common.basemodel.triton_kernel.linear_att.fused_gdn_gating import ( + fused_gdn_gating, +) +from lightllm.common.basemodel.triton_kernel.linear_att.gdn_decode_pack import ( + conv_pack_gdn_decode_inputs, +) +from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( + build_dynamic_mtp_linear_att_state_params, +) +from lightllm.utils.envs_utils import get_env_start_args + +from .replay_gdn_decode import replay + +if TYPE_CHECKING: + from lightllm.common.basemodel.infer_struct import InferStateInfo + + +def replay_gdn_backend_class(name: str): + """按已选中的 linear kernel 名字返回对应的 replay backend。""" + + classes = { + "triton": ReplayGDNTritonAttBackend, + "flashqla": ReplayGDNFlashQlaAttBackend, + } + try: + return classes[name] + except KeyError: + raise ValueError(f"replay GDN has no backend named {name!r}") from None + + +class ReplayGDNTritonAttBackend(BaseAttBackend): + """GDN replay 的 Triton 入口。Prefill 使用仓内 chunk kernel。""" + + def __init__(self, model): + super().__init__(model=model) + from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops import chunk_gated_delta_rule + + _configure_gdn(self, model, chunk_gated_delta_rule) + + def create_att_prefill_state(self, infer_state: "InferStateInfo"): + return ReplayGDNPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) + + def create_att_decode_state(self, infer_state: "InferStateInfo"): + return ReplayGDNDecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) + + +class ReplayGDNFlashQlaAttBackend(BaseAttBackend): + """GDN replay 的 FlashQla 入口。Prefill 使用 flash_qla chunk kernel。""" + + def __init__(self, model): + super().__init__(model=model) + from flash_qla import chunk_gated_delta_rule + + _configure_gdn(self, model, chunk_gated_delta_rule) + + def create_att_prefill_state(self, infer_state: "InferStateInfo"): + return ReplayGDNPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) + + def create_att_decode_state(self, infer_state: "InferStateInfo"): + return ReplayGDNDecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) + + +@dataclasses.dataclass +class ReplayGDNPrefillAttState(BasePrefillAttState): + b_conv_buffer_idx: torch.Tensor = None + b_ssm_buffer_idx: torch.Tensor = None + + def init_state(self): + reqs = self.infer_state.b_req_idx + self.b_conv_buffer_idx = reqs + self.b_ssm_buffer_idx = reqs + + def prefill_att( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + att_control: AttControl = AttControl(), + alloc_func=torch.empty, + ): + assert att_control.linear_att_prefill and att_control.linear_att_prefill_dict is not None + params = att_control.linear_att_prefill_dict + layer_num = params["layer_num"] + layer_weight = params["layer_weight"] + backend = self.backend + state = _replay_state(self, layer_num) + conv_states, ssm_states = state.get_layer_view(logical_layer_id=layer_num) + if backend.mtp_step > 0: + conv_states = conv_states[:, :, : -backend.mtp_step] + mixed_qkv, z, b, a = _split_qkvzba(backend, params["mixed_qkvzba"]) + core_attn_out = self._prefill(mixed_qkv, conv_states, ssm_states, a, b, layer_weight) + state.finish_prefill(layer_num, self.b_conv_buffer_idx) + return core_attn_out, z + + def _prefill(self, mixed_qkv, conv_states, ssm_states, a, b, layer_weight): + backend = self.backend + g, beta = fused_gdn_gating(layer_weight.linear_A_log.weight, a, b, layer_weight.linear_dt_bias.weight) + mixed_qkv = causal_conv1d_fn( + mixed_qkv.transpose(0, 1), + layer_weight.linear_conv1d.mm_param.weight, + bias=layer_weight.linear_conv1d.bias, + query_start_loc=self.infer_state.b1_cu_q_seq_len, + cache_indices=self.b_conv_buffer_idx, + has_initial_state=self.infer_state.b_ready_cache_len > 0, + conv_states=conv_states, + activation=backend.activation, + ).transpose(0, 1) + query, key, value = _rearrange_mixed_qkv(backend, mixed_qkv) + core_attn_out, last_state = backend.prefill_kernel( + q=query, + k=key, + v=value, + g=g.unsqueeze(0), + beta=beta.unsqueeze(0), + initial_state=ssm_states[self.b_ssm_buffer_idx], + output_final_state=True, + cu_seqlens=self.infer_state.b1_cu_q_seq_len, + use_qk_l2norm_in_kernel=True, + ) + ssm_states[self.b_ssm_buffer_idx] = last_state.to(ssm_states.dtype, copy=False) + return core_attn_out + + +@dataclasses.dataclass +class ReplayGDNDecodeAttState(BaseDecodeAttState): + b_conv_buffer_idx: torch.Tensor = None + b_ssm_buffer_idx: torch.Tensor = None + b1_mtp_cu_q_seq_len: torch.Tensor = None + b_num_accepted_tokens: torch.Tensor = None + + def init_state(self): + role = self.infer_state.att_state_role + matches = [ + state + for state in self.infer_state.att_state_manager.get_bounded_attention_states(role) + if isinstance(state, ReplayGDNState) + ] + if len(matches) != 1: + raise ValueError(f"expected one ReplayGDNState for {role}") + draft_step = self.backend.model.mtp_manager.get_decode_draft_step(self.backend.model.is_mtp_draft_model) + self._init_indices(draft_step, matches[0].mtp_state_index) + self.b_ssm_buffer_idx = self.b_conv_buffer_idx + + def decode_att( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + att_control: AttControl = AttControl(), + alloc_func=torch.empty, + ): + assert att_control.linear_att_decode and att_control.linear_att_decode_dict is not None + params = att_control.linear_att_decode_dict + layer_num = params["layer_num"] + layer_weight = params["layer_weight"] + backend = self.backend + mixed_qkv, z, b, a = _split_qkvzba(backend, params["mixed_qkvzba"]) + state = _replay_state(self, layer_num) + conv_states, _ = state.get_layer_view(logical_layer_id=layer_num) + draft_step = backend.model.mtp_manager.get_decode_draft_step(backend.model.is_mtp_draft_model) + if draft_step > 0: + core = self._mtp(mixed_qkv, conv_states, a, b, layer_weight, layer_num, draft_step, state) + else: + core, z = self._decode(mixed_qkv, z, conv_states, a, b, layer_weight, layer_num, state) + state.finish_decode_token(layer_num, self.b_conv_buffer_idx) + return core, z + + def _init_indices(self, draft_step: int, mtp_index: torch.Tensor | None) -> None: + infer = self.infer_state + if draft_step == 0: + self.b_conv_buffer_idx = infer.b_req_idx + self.b1_mtp_cu_q_seq_len = None + self.b_num_accepted_tokens = None + return + if self.backend.uses_dynamic_spec_verify_layout(): + cu, conv_idx, accepted = build_dynamic_mtp_linear_att_state_params( + b_req_idx=infer.b_req_idx, + b_mtp_index=infer.b_mtp_index, + req_to_mtp_state_index=mtp_index, + hold_req_id=infer.req_manager.HOLD_REQUEST_ID, + ) + self.b1_mtp_cu_q_seq_len = cu + self.b_conv_buffer_idx = conv_idx + self.b_num_accepted_tokens = accepted + return + width = draft_step + 1 + batch_size = infer.batch_size + if batch_size % width: + raise ValueError( + "GDN replay decode requires batch_size divisible by draft_step + 1, " + f"got batch_size={batch_size}, draft_step={draft_step}." + ) + self.b1_mtp_cu_q_seq_len = torch.arange( + 0, batch_size + 1, width, dtype=torch.int32, device=infer.b_req_idx.device + ) + self.b_conv_buffer_idx = infer.b_req_idx.view(-1, width)[:, 0].contiguous() + self.b_num_accepted_tokens = mtp_index[self.b_conv_buffer_idx] + 1 + + def _decode(self, mixed_qkv, z, conv_states, a, b, layer_weight, layer_num, state): + backend = self.backend + query, key, value, z, a, b = conv_pack_gdn_decode_inputs( + mixed_qkv, + z, + a, + b, + conv_states, + layer_weight.linear_conv1d.mm_param.weight, + layer_weight.linear_conv1d.bias, + self.b_conv_buffer_idx, + backend.activation, + backend.conv_kernel_dim, + backend.tp_num_k_heads, + backend.head_k_dim, + backend.tp_num_v_heads, + backend.head_v_dim, + ) + output, _ = replay( + state, + state.storage_layer(layer_num), + query, + key, + value, + a, + b, + layer_weight.linear_A_log.weight, + layer_weight.linear_dt_bias.weight, + self.b_conv_buffer_idx, + self.b1_mtp_cu_q_seq_len, + ) + return output, z + + def _mtp(self, mixed_qkv, conv_states, a, b, layer_weight, layer_num, draft_step, state): + backend = self.backend + mixed_qkv = causal_conv1d_update_mtp( + mixed_qkv, + conv_states, + layer_weight.linear_conv1d.mm_param.weight, + mtp_step=draft_step, + bias=layer_weight.linear_conv1d.bias, + activation=backend.activation, + conv_state_indices=self.b_conv_buffer_idx, + num_accepted_tokens=self.b_num_accepted_tokens, + query_start_loc=self.b1_mtp_cu_q_seq_len, + ) + query, key, value = _rearrange_mixed_qkv(backend, mixed_qkv, decode=False) + output, _ = replay( + state, + state.storage_layer(layer_num), + query, + key, + value, + a, + b, + layer_weight.linear_A_log.weight, + layer_weight.linear_dt_bias.weight, + self.b_conv_buffer_idx, + self.b1_mtp_cu_q_seq_len, + ) + return output + + +def _replay_state(att_state, layer_num: int) -> ReplayGDNState: + return att_state.infer_state.att_state_manager.get_state( + model_role=att_state.infer_state.att_state_role, + state_type_filter="gdn", + logical_layer_id=layer_num, + state_class=ReplayGDNState, + ) + + +def _configure_gdn(backend, model, prefill_kernel) -> None: + """写入 GDN 的张量并行尺寸，并绑定这个 backend 自己的 prefill kernel。""" + + config = model.config + tp_world_size = model.tp_world_size_ + backend.mtp_step = get_env_start_args().mtp_step + backend.num_v_heads = config["linear_num_value_heads"] + backend.num_k_heads = config["linear_num_key_heads"] + backend.head_k_dim = config["linear_key_head_dim"] + backend.head_v_dim = config["linear_value_head_dim"] + backend.key_dim = backend.head_k_dim * backend.num_k_heads + backend.value_dim = backend.head_v_dim * backend.num_v_heads + backend.conv_kernel_dim = config["linear_conv_kernel_dim"] + backend.activation = config["hidden_act"] + backend.tp_qkvz_dim = (backend.key_dim * 2 + backend.value_dim * 2) // tp_world_size + backend.tp_ba_dim = (backend.num_v_heads * 2) // tp_world_size + backend.tp_num_k_heads = backend.num_k_heads // tp_world_size + backend.tp_num_v_heads = backend.num_v_heads // tp_world_size + backend.tp_key_dim = backend.key_dim // tp_world_size + backend.tp_value_dim = backend.value_dim // tp_world_size + assert backend.num_v_heads % backend.num_k_heads == 0, "num_v_heads must be divisible by num_k_heads" + backend.num_v_heads_per_k_head = backend.num_v_heads // backend.num_k_heads + ssm_dtypes = {"bfloat16": torch.bfloat16, "float32": torch.float32} + backend.ssm_state_dtype = ssm_dtypes.get(get_env_start_args().linear_att_ssm_data_type, torch.bfloat16) + backend.prefill_kernel = prefill_kernel + + +def _split_qkvzba(backend, mixed_qkvzba): + qkv_dim = backend.tp_key_dim * 2 + backend.tp_value_dim + z_end = qkv_dim + backend.tp_value_dim + b_end = z_end + backend.tp_num_v_heads + mixed_qkv = mixed_qkvzba[:, :qkv_dim] + z = mixed_qkvzba[:, qkv_dim:z_end].view(-1, backend.tp_num_v_heads, backend.head_v_dim) + b = mixed_qkvzba[:, z_end:b_end] + a = mixed_qkvzba[:, b_end:] + return mixed_qkv, z, b, a + + +def _rearrange_mixed_qkv(backend, mixed_qkv, decode=False): + query, key, value = torch.split( + mixed_qkv, + [backend.tp_key_dim, backend.tp_key_dim, backend.tp_value_dim], + dim=-1, + ) + if decode: + batch_size = mixed_qkv.shape[0] + leading = (batch_size, 1) + else: + leading = (1, query.shape[0]) + query = query.view(*leading, backend.tp_num_k_heads, backend.head_k_dim) + key = key.view(*leading, backend.tp_num_k_heads, backend.head_k_dim) + value = value.view(*leading, backend.tp_num_v_heads, backend.head_v_dim) + return query, key, value diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py b/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py new file mode 100644 index 0000000000..71fcf8e1aa --- /dev/null +++ b/lightllm/common/basemodel/attention/linear/replay_gdn_decode.py @@ -0,0 +1,86 @@ +"""GDN replay decode 的 kernel 辅助。 + +History 张量和 cursor 在 ReplayGDNState。这里只选择 kernel 配置，并发起一层重放。 +""" + +from __future__ import annotations + +import torch +import triton + +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import _decode_replay +from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import resolve_run_config + + +def replay( + history, + linear_layer_idx: int, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + a_log: torch.Tensor, + bias: torch.Tensor, + reqs: torch.Tensor, + cu_seqlens: torch.Tensor | None = None, + run_config: dict[str, int] | None = None, +) -> tuple[torch.Tensor, dict[str, int]]: + """执行一层 GDN replay decode。Kernel 只接收这一层的 SSM 和 history；state 只提供共用 cursor 与配置。""" + + if history.kda: + raise ValueError("GDN replay step does not use the KDA gate") + config = run_config or resolve_run_config(history, "replay", q, k, v, a, b, a_log, bias, cu_seqlens) + hv, kd, vd = history.ssm_state.shape[-3:] + assert history.num_key_heads in (q.shape[-2], hv) + axis = 1 if cu_seqlens is not None else 0 + out = torch.empty_like(v) + bv = config["BV"] + # mtp_step == 0 时 verify_width 为 1。compute_mode=precompute 仍然投影；BV < 16 才退回逐 token。 + precompute_projections = history.compute_mode == "precompute" and bv >= 16 + # 在 Python 中解析 tile 维度；Triton JIT 中直接调用 max() 会降级为 tensor 运算，无法定义静态 shape。 + bw = triton.next_power_of_2(history.verify_width) + _decode_replay[(triton.cdiv(vd, bv), reqs.numel(), hv)]( + q, + k, + v, + a, + b, + a_log, + bias, + history.ssm_state[linear_layer_idx], + history.keys[linear_layer_idx], + history.deltas[linear_layer_idx], + history.gates[linear_layer_idx], + history.raw_keys[linear_layer_idx], + history.raw_values[linear_layer_idx], + history.betas[linear_layer_idx], + reqs, + history.cursors, + cu_seqlens, + out, + q.stride(axis), + k.stride(axis), + v.stride(axis), + a.stride(0), + b.stride(0), + q.shape[-2], + hv, + history.num_key_heads, + kd, + vd, + history.capacity, + history.hold_req_id, + triton.next_power_of_2(kd), + bv, + cu_seqlens is not None, + precompute_projections, + bw, + False, + history.gate_lower_bound, + BQK=max(16, 2 * bw), + BR=triton.next_power_of_2(max(16, history.capacity + bw)), + num_warps=config["num_warps"], + num_stages=config.get("num_stages", 3), + ) + return out, config diff --git a/lightllm/common/basemodel/attention/linear/replay_kda.py b/lightllm/common/basemodel/attention/linear/replay_kda.py new file mode 100644 index 0000000000..3090faa9a8 --- /dev/null +++ b/lightllm/common/basemodel/attention/linear/replay_kda.py @@ -0,0 +1,241 @@ +"""KDA replay 的推理入口。 + +这个 backend 只服务 ReplayKDAState。维度在这里自己读取，不继承 native KDA backend。 +""" + +from __future__ import annotations + +import dataclasses +import weakref +from typing import TYPE_CHECKING + +import torch + +from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import ( + ReplayKDAState, +) +from lightllm.common.basemodel.attention.base_att import ( + AttControl, + BaseAttBackend, + BaseDecodeAttState, + BasePrefillAttState, +) +from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import ( + causal_conv1d_fn, + causal_conv1d_update, +) +from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_mtp import ( + causal_conv1d_update as causal_conv1d_update_mtp, +) +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.mtp_state_params import ( + build_dynamic_mtp_linear_att_state_params, +) + +from .replay_kda_decode import replay + +if TYPE_CHECKING: + from lightllm.common.basemodel.infer_struct import InferStateInfo + + +class ReplayKDALinearAttBackend(BaseAttBackend): + """KDA replay 推理 backend。GLM-5-Next 在 replay 模式下只构造这一个类。""" + + def __init__(self, model): + super().__init__(model=model) + config = model.config["linear_attn_config"] + self.num_heads = config["num_heads"] + self.head_dim = config["head_dim"] + assert self.num_heads % model.tp_world_size_ == 0 + self.tp_num_heads = self.num_heads // model.tp_world_size_ + self.tp_hidden_size = self.tp_num_heads * self.head_dim + self.conv_kernel_size = config["short_conv_kernel_size"] + self.lower_bound = config.get("gate_lower_bound", -5.0) + self.mtp_step = model.args.mtp_step + + def split_qkv(self, mixed_qkv: torch.Tensor): + return mixed_qkv.split(self.tp_hidden_size, dim=-1) + + def create_att_prefill_state(self, infer_state: "InferStateInfo"): + return ReplayKDAPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) + + def create_att_decode_state(self, infer_state: "InferStateInfo"): + return ReplayKDADecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) + + +@dataclasses.dataclass +class ReplayKDAPrefillAttState(BasePrefillAttState): + b_conv_buffer_idx: torch.Tensor = None + b_ssm_buffer_idx: torch.Tensor = None + + def init_state(self): + reqs = self.infer_state.b_req_idx + self.b_conv_buffer_idx = reqs + self.b_ssm_buffer_idx = reqs + + def prefill_att( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + att_control: AttControl = AttControl(), + alloc_func=torch.empty, + ): + assert att_control.linear_att_prefill + params = att_control.linear_att_prefill_dict + layer_num = params["layer_num"] + layer_weight = params["layer_weight"] + backend: ReplayKDALinearAttBackend = self.backend + state = _replay_state(self, layer_num) + conv_states, ssm_states = state.get_layer_view(logical_layer_id=layer_num) + conv_states = conv_states[..., : backend.conv_kernel_size - 1] + mixed_qkv = causal_conv1d_fn( + params["mixed_qkv"].transpose(0, 1), + layer_weight.get_merged_kda_conv_weight(), + bias=None, + query_start_loc=self.infer_state.b1_cu_q_seq_len, + cache_indices=self.b_conv_buffer_idx, + has_initial_state=self.infer_state.b_ready_cache_len > 0, + conv_states=conv_states, + activation="silu", + ).transpose(0, 1) + q, k, v = backend.split_qkv(mixed_qkv) + q = q.view(1, -1, backend.tp_num_heads, backend.head_dim) + k = k.view(1, -1, backend.tp_num_heads, backend.head_dim) + v = v.view(1, -1, backend.tp_num_heads, backend.head_dim) + raw_gate = params["raw_gate"].view(1, -1, backend.tp_hidden_size) + raw_beta = params["raw_beta"].view(1, -1, backend.tp_num_heads) + output, final_state = chunk_kda_with_fused_gate( + q=q, + k=k, + v=v, + raw_g=raw_gate.view(1, -1, backend.tp_num_heads, backend.head_dim), + beta=raw_beta.float().sigmoid(), + A_log=layer_weight.linear_A_log.weight, + g_bias=layer_weight.linear_dt_bias.weight, + initial_state=ssm_states[self.b_ssm_buffer_idx].contiguous(), + output_final_state=True, + use_qk_l2norm_in_kernel=True, + cu_seqlens=self.infer_state.b1_cu_q_seq_len, + safe_gate=True, + lower_bound=backend.lower_bound, + ) + ssm_states[self.b_ssm_buffer_idx] = final_state.to(ssm_states.dtype, copy=False) + state.finish_prefill(layer_num, self.b_conv_buffer_idx) + return output + + +@dataclasses.dataclass +class ReplayKDADecodeAttState(BaseDecodeAttState): + b_conv_buffer_idx: torch.Tensor = None + b_ssm_buffer_idx: torch.Tensor = None + b1_mtp_cu_q_seq_len: torch.Tensor = None + b_num_accepted_tokens: torch.Tensor = None + + def init_state(self): + role = self.infer_state.att_state_role + matches = [ + state + for state in self.infer_state.att_state_manager.get_bounded_attention_states(role) + if isinstance(state, ReplayKDAState) + ] + if len(matches) != 1: + raise ValueError(f"expected one ReplayKDAState for {role}") + self._init_indices(self.backend.mtp_step, matches[0].mtp_state_index) + self.b_ssm_buffer_idx = self.b_conv_buffer_idx + + def decode_att( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + att_control: AttControl = AttControl(), + alloc_func=torch.empty, + ): + assert att_control.linear_att_decode + params = att_control.linear_att_decode_dict + layer_num = params["layer_num"] + layer_weight = params["layer_weight"] + backend: ReplayKDALinearAttBackend = self.backend + state = _replay_state(self, layer_num) + conv_states, _ = state.get_layer_view(logical_layer_id=layer_num) + conv_kwargs = dict(bias=None, activation="silu", conv_state_indices=self.b_conv_buffer_idx) + conv_update = causal_conv1d_update + if backend.mtp_step > 0: + conv_update = causal_conv1d_update_mtp + conv_kwargs.update( + mtp_step=backend.mtp_step, + num_accepted_tokens=self.b_num_accepted_tokens, + query_start_loc=self.b1_mtp_cu_q_seq_len, + ) + mixed_qkv = conv_update( + params["mixed_qkv"], + conv_states, + layer_weight.get_merged_kda_conv_weight(), + **conv_kwargs, + ) + q, k, v = backend.split_qkv(mixed_qkv) + shape = (1, -1) if backend.mtp_step > 0 else (-1, 1) + q = q.view(*shape, backend.tp_num_heads, backend.head_dim) + k = k.view(*shape, backend.tp_num_heads, backend.head_dim) + v = v.view(*shape, backend.tp_num_heads, backend.head_dim) + raw_gate = params["raw_gate"].view(*shape, backend.tp_hidden_size) + raw_beta = params["raw_beta"].view(*shape, backend.tp_num_heads) + output, _ = replay( + state, + state.storage_layer(layer_num), + q, + k, + v, + raw_gate.reshape(-1, backend.tp_hidden_size), + raw_beta.reshape(-1, backend.tp_num_heads), + layer_weight.linear_A_log.weight, + layer_weight.linear_dt_bias.weight, + self.b_conv_buffer_idx, + self.b1_mtp_cu_q_seq_len, + ) + if backend.mtp_step == 0: + state.finish_decode_token(layer_num, self.b_conv_buffer_idx) + return output + + def _init_indices(self, mtp_step: int, mtp_index: torch.Tensor | None) -> None: + infer = self.infer_state + if mtp_step == 0: + self.b_conv_buffer_idx = infer.b_req_idx + self.b1_mtp_cu_q_seq_len = None + self.b_num_accepted_tokens = None + return + if self.backend.uses_dynamic_spec_verify_layout(): + cu, conv_idx, accepted = build_dynamic_mtp_linear_att_state_params( + b_req_idx=infer.b_req_idx, + b_mtp_index=infer.b_mtp_index, + req_to_mtp_state_index=mtp_index, + hold_req_id=infer.req_manager.HOLD_REQUEST_ID, + ) + self.b1_mtp_cu_q_seq_len = cu + self.b_conv_buffer_idx = conv_idx + self.b_num_accepted_tokens = accepted + return + width = mtp_step + 1 + batch_size = infer.batch_size + if batch_size % width: + raise ValueError( + "KDA replay decode requires batch_size divisible by mtp_step + 1, " + f"got batch_size={batch_size}, mtp_step={mtp_step}." + ) + self.b1_mtp_cu_q_seq_len = torch.arange( + 0, batch_size + 1, width, dtype=torch.int32, device=infer.b_req_idx.device + ) + self.b_conv_buffer_idx = infer.b_req_idx.view(-1, width)[:, 0].contiguous() + self.b_num_accepted_tokens = mtp_index[self.b_conv_buffer_idx] + 1 + + +def _replay_state(att_state, layer_num: int) -> ReplayKDAState: + return att_state.infer_state.att_state_manager.get_state( + model_role=att_state.infer_state.att_state_role, + state_type_filter="kda", + logical_layer_id=layer_num, + state_class=ReplayKDAState, + ) diff --git a/lightllm/common/basemodel/attention/linear/replay_kda_decode.py b/lightllm/common/basemodel/attention/linear/replay_kda_decode.py new file mode 100644 index 0000000000..55514aeb57 --- /dev/null +++ b/lightllm/common/basemodel/attention/linear/replay_kda_decode.py @@ -0,0 +1,86 @@ +"""KDA replay decode 的 kernel 辅助。 + +History 张量和 cursor 在 ReplayKDAState。这里只选择 kernel 配置，并发起一层重放。 +""" + +from __future__ import annotations + +import torch +import triton + +from lightllm.common.basemodel.triton_kernel.linear_att.replayssm import _decode_replay +from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import resolve_run_config + + +def replay( + history, + linear_layer_idx: int, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + a_log: torch.Tensor, + bias: torch.Tensor, + reqs: torch.Tensor, + cu_seqlens: torch.Tensor | None = None, + run_config: dict[str, int] | None = None, +) -> tuple[torch.Tensor, dict[str, int]]: + """执行一层 KDA replay decode。Kernel 只接收这一层的 SSM 和 history；state 只提供共用 cursor 与配置。""" + + if not history.kda: + raise ValueError("KDA replay step requires the KDA gate") + config = run_config or resolve_run_config(history, "replay", q, k, v, a, b, a_log, bias, cu_seqlens) + hv, kd, vd = history.ssm_state.shape[-3:] + assert history.num_key_heads in (q.shape[-2], hv) + axis = 1 if cu_seqlens is not None else 0 + out = torch.empty_like(v) + bv = config["BV"] + # mtp_step == 0 时 verify_width 为 1。compute_mode=precompute 仍然投影；BV < 16 才退回逐 token。 + precompute_projections = history.compute_mode == "precompute" and bv >= 16 + # 在 Python 中解析 tile 维度；Triton JIT 中直接调用 max() 会降级为 tensor 运算，无法定义静态 shape。 + bw = triton.next_power_of_2(history.verify_width) + _decode_replay[(triton.cdiv(vd, bv), reqs.numel(), hv)]( + q, + k, + v, + a, + b, + a_log, + bias, + history.ssm_state[linear_layer_idx], + history.keys[linear_layer_idx], + history.deltas[linear_layer_idx], + history.gates[linear_layer_idx], + history.raw_keys[linear_layer_idx], + history.raw_values[linear_layer_idx], + history.betas[linear_layer_idx], + reqs, + history.cursors, + cu_seqlens, + out, + q.stride(axis), + k.stride(axis), + v.stride(axis), + a.stride(0), + b.stride(0), + q.shape[-2], + hv, + history.num_key_heads, + kd, + vd, + history.capacity, + history.hold_req_id, + triton.next_power_of_2(kd), + bv, + cu_seqlens is not None, + precompute_projections, + bw, + True, + history.gate_lower_bound, + BQK=max(16, 2 * bw), + BR=triton.next_power_of_2(max(16, history.capacity + bw)), + num_warps=config["num_warps"], + num_stages=config.get("num_stages", 3), + ) + return out, config diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/pinned_embedding_weight.py b/lightllm/common/basemodel/layer_weights/meta_weights/pinned_embedding_weight.py new file mode 100644 index 0000000000..3cb376594a --- /dev/null +++ b/lightllm/common/basemodel/layer_weights/meta_weights/pinned_embedding_weight.py @@ -0,0 +1,66 @@ +import torch + +from .row_sharded_embedding_weight import RowShardedEmbeddingWeight +from lightllm.utils.log_utils import init_logger + +logger = init_logger(__name__) + + +class PinnedRowShardedEmbeddingWeight(RowShardedEmbeddingWeight): + def __init__( + self, dim, vocab_size, weight_name, data_type, shard_count, *, cpu_allocator, tp_rank=None, tp_world_size=None + ): + if not callable(cpu_allocator): + raise ValueError("Pinned embeddings require a CPU allocation policy") + self.cpu_allocator = cpu_allocator + super().__init__(dim, vocab_size, weight_name, data_type, shard_count, tp_rank, tp_world_size) + + def _create_weight(self): + shape = (self.tp_vocab_end_id - self.tp_vocab_start_id, self.dim) + with torch.cuda.device(self.device_id_): + self.weight = self.cpu_allocator(shape, dtype=self.data_type_, device="cpu", pin_memory=True) + if ( + self.weight.device.type != "cpu" + or self.weight.dtype != self.data_type_ + or tuple(self.weight.shape) != shape + or not self.weight.is_contiguous() + or (self.weight.numel() and not self.weight.is_pinned()) + ): + raise ValueError( + "Embedding allocator must return contiguous pinned CPU storage with the requested shape/dtype" + ) + self.weight.load_ok = False + logger.info( + f"Pinned embedding table: rank={self.tp_rank_} rows={shape[0]} dim={self.dim} " + f"bytes={self.weight.numel() * self.weight.element_size()} dtype={self.weight.dtype} residency=cpu" + ) + + def load_hf_weights(self, weights): + names = {self.weight_name} | {f"{self._shard_prefix}.shard_{index}.weight" for index in self._required_shards} + host_weights = {name: value.cpu() if name in names else value for name, value in weights.items()} + super().load_hf_weights(host_weights) + + def _cuda_forward(self, input_ids, out=None, alloc_func=torch.empty): + if ( + not input_ids.is_cuda + or input_ids.ndim != 1 + or not input_ids.is_contiguous() + or input_ids.dtype not in (torch.int32, torch.int64) + ): + raise ValueError("Pinned embedding lookup requires a CUDA vector of integer IDs") + shape = (input_ids.shape[0], self.dim) + if out is None: + out = alloc_func(shape, dtype=self.weight.dtype, device=input_ids.device) + if ( + tuple(out.shape) != shape + or out.dtype != self.weight.dtype + or out.device != input_ids.device + or out.stride(1) != 1 + ): + raise ValueError("Pinned embedding output must match CUDA IDs, table dtype and lookup shape") + if not input_ids.numel(): + return out + return super()._triton_forward(input_ids, out=out, alloc_func=alloc_func) + + def _native_forward(self, *args, **kwargs): + raise ValueError("Pinned host embedding lookup requires the CUDA Triton implementation") diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/row_sharded_embedding_weight.py b/lightllm/common/basemodel/layer_weights/meta_weights/row_sharded_embedding_weight.py new file mode 100644 index 0000000000..ca2390c736 --- /dev/null +++ b/lightllm/common/basemodel/layer_weights/meta_weights/row_sharded_embedding_weight.py @@ -0,0 +1,64 @@ +from typing import Dict, Optional + +import torch + +from .embedding_weight import EmbeddingWeight + + +class RowShardedEmbeddingWeight(EmbeddingWeight): + def __init__( + self, + dim: int, + vocab_size: int, + weight_name: str, + data_type: torch.dtype, + shard_count: int, + tp_rank: Optional[int] = None, + tp_world_size: Optional[int] = None, + ): + if shard_count < 1 or vocab_size < 1 or vocab_size % shard_count: + raise ValueError("Row-sharded embeddings require a positive vocabulary divisible by shard count") + if not weight_name.endswith(".weight"): + raise ValueError("Embedding weight name must end with .weight") + self.shard_count = shard_count + self.shard_rows = vocab_size // shard_count + self._shard_prefix = weight_name[: -len(".weight")] + self._loaded_shards = set() + self._full_loaded = False + super().__init__(dim, vocab_size, weight_name, data_type, tp_rank, tp_world_size) + self._required_shards = { + index + for index in range(shard_count) + if max(index * self.shard_rows, self.tp_vocab_start_id) + < min((index + 1) * self.shard_rows, self.tp_vocab_end_id) + } + self.weight.load_ok = not self._required_shards + + def load_hf_weights(self, weights: Dict[str, torch.Tensor]): + shard_keys = [(index, f"{self._shard_prefix}.shard_{index}.weight") for index in range(self.shard_count)] + present = [(index, key) for index, key in shard_keys if key in weights] + if self.weight_name in weights: + if present or self._loaded_shards: + raise ValueError("Cannot mix full and row-sharded embedding checkpoint tensors") + if weights[self.weight_name].shape != (self.vocab_size, self.dim): + raise ValueError("Full embedding checkpoint dimensions do not match") + super().load_hf_weights(weights) + self._full_loaded = True + return + if self._full_loaded: + if present: + raise ValueError("Cannot mix full and row-sharded embedding checkpoint tensors") + return + for index, key in present: + source = weights[key] + if source.shape != (self.shard_rows, self.dim): + raise ValueError(f"Unexpected dimensions for embedding shard {index}: {tuple(source.shape)}") + shard_start = index * self.shard_rows + begin = max(shard_start, self.tp_vocab_start_id) + end = min(shard_start + self.shard_rows, self.tp_vocab_end_id) + if begin < end: + self.weight[begin - self.tp_vocab_start_id : end - self.tp_vocab_start_id].copy_( + source[begin - shard_start : end - shard_start].to(self.data_type_) + ) + self._loaded_shards.add(index) + self.weight.load_ok = self._loaded_shards >= self._required_shards diff --git a/lightllm/common/basemodel/multimodal_tokenizer.py b/lightllm/common/basemodel/multimodal_tokenizer.py index cdcbd7f089..6cd0f8f26c 100644 --- a/lightllm/common/basemodel/multimodal_tokenizer.py +++ b/lightllm/common/basemodel/multimodal_tokenizer.py @@ -40,6 +40,9 @@ def __getattr__(self, name): return obj_dict[name] return getattr(self.tokenizer, name) + def count_prompt_tokens(self, text_tokens, image_tokens, image_count, audio_tokens, audio_count): + return text_tokens + image_tokens + image_count + audio_tokens + audio_count + @abstractmethod def encode(self, prompt, multimodal_params: MultimodalParams = None, **kwargs): pass diff --git a/lightllm/common/basemodel/tp_reduce.py b/lightllm/common/basemodel/tp_reduce.py new file mode 100644 index 0000000000..f50c80b989 --- /dev/null +++ b/lightllm/common/basemodel/tp_reduce.py @@ -0,0 +1,58 @@ +import torch +import triton +import triton.language as tl + +from lightllm.distributed import all_gather_into_tensor + + +@triton.jit +def _sum_shards( + X, + Y, + N: tl.constexpr, + START: tl.constexpr, + SIZE: tl.constexpr, + WORLD: tl.constexpr, + BLOCK: tl.constexpr, +): + index = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + accumulator = tl.full((BLOCK,), 0, tl.float64) + for rank in range(WORLD): + value = tl.load(X + rank * N + START + index, index < SIZE, other=0) + accumulator += value.to(tl.float64) + tl.store(Y + index, accumulator, index < SIZE) + + +def ordered_tp_sum(input: torch.Tensor, group, world_size: int, *, scatter_rank=None) -> torch.Tensor: + if world_size == 1: + return input + if scatter_rank is not None: + if not 0 <= scatter_rank < world_size or input.shape[0] % world_size: + raise ValueError("Sequence-parallel input must divide evenly across valid ranks") + shape = (input.shape[0] // world_size, *input.shape[1:]) + output = torch.empty(shape, device=input.device, dtype=input.dtype) + start = scatter_rank * output.numel() + else: + output = input if input.is_contiguous() else torch.empty(input.shape, device=input.device, dtype=input.dtype) + start = 0 + if input.numel(): + source = input.contiguous() + symmetric = getattr(group, "symm_mem_reduce", None) if input.is_cuda else None + gathered = symmetric.ordered_all_gather(source) if symmetric is not None else None + if gathered is None: + gathered = torch.empty((world_size, *input.shape), device=input.device, dtype=input.dtype) + all_gather_into_tensor(gathered.view(-1, *input.shape[1:]), source, group=group) + if input.is_cuda: + _sum_shards[(triton.cdiv(output.numel(), 256),)]( + gathered, output, input.numel(), start, output.numel(), world_size, 256 + ) + else: + accumulator = torch.zeros(output.numel(), dtype=torch.float64, device=input.device) + shards = gathered.reshape(world_size, -1) + for rank in range(world_size): + accumulator.add_(shards[rank, start : start + output.numel()]) + output.copy_(accumulator.reshape(output.shape)) + if scatter_rank is None and output is not input: + input.copy_(output) + return input + return output diff --git a/lightllm/common/basemodel/triton_kernel/destindex_copy_kv.py b/lightllm/common/basemodel/triton_kernel/destindex_copy_kv.py index 060a92bf7c..f301a03a1e 100644 --- a/lightllm/common/basemodel/triton_kernel/destindex_copy_kv.py +++ b/lightllm/common/basemodel/triton_kernel/destindex_copy_kv.py @@ -34,6 +34,43 @@ def _fwd_kernel_destindex_copy_kv( return +@triton.jit +def _fwd_kernel_destindex_copy_kv_fused( + KV, + Dest_loc, + OutK, + OutV, + stride_kv_bs, + stride_kv_h, + stride_kv_d, + stride_ok_bs, + stride_ok_h, + stride_ok_d, + stride_ov_bs, + stride_ov_h, + stride_ov_d, + head_num, + head_dim, + BLOCK_DMODEL: tl.constexpr, + BLOCK_HEAD: tl.constexpr, +): + cur_index = tl.program_id(0) + offs_h = tl.arange(0, BLOCK_HEAD) + offs_d = (tl.arange(0, BLOCK_DMODEL)) % head_dim + + dest_index = tl.load(Dest_loc + cur_index).to(tl.int64) + mask = offs_h[:, None] < head_num + + kv_base = KV + cur_index * stride_kv_bs + offs_h[:, None] * stride_kv_h + offs_d[None, :] * stride_kv_d + v_base = KV + cur_index * stride_kv_bs + (head_num + offs_h[:, None]) * stride_kv_h + offs_d[None, :] * stride_kv_d + k_out = OutK + dest_index * stride_ok_bs + offs_h[:, None] * stride_ok_h + offs_d[None, :] * stride_ok_d + v_out = OutV + dest_index * stride_ov_bs + offs_h[:, None] * stride_ov_h + offs_d[None, :] * stride_ov_d + + tl.store(k_out, tl.load(kv_base, mask=mask, other=0.0), mask=mask) + tl.store(v_out, tl.load(v_base, mask=mask, other=0.0), mask=mask) + return + + @torch.no_grad() def destindex_copy_kv(K, DestLoc, Out): seq_len = DestLoc.shape[0] @@ -62,3 +99,40 @@ def destindex_copy_kv(K, DestLoc, Out): num_stages=1, ) return + + +@torch.no_grad() +def destindex_copy_kv_fused(KV, DestLoc, OutK, OutV): + """Copy packed K/V rows to separate storage buffers with one kernel launch.""" + seq_len = DestLoc.shape[0] + head_num = OutK.shape[1] + head_dim = OutK.shape[2] + assert KV.ndim == 3 and KV.shape[0] == seq_len + assert OutK.shape == OutV.shape + assert KV.shape[1] == 2 * head_num and KV.shape[2] == head_dim + + BLOCK_HEAD = triton.next_power_of_2(head_num) + grid = (seq_len,) + + _fwd_kernel_destindex_copy_kv_fused[grid]( + KV, + DestLoc, + OutK, + OutV, + KV.stride(0), + KV.stride(1), + KV.stride(2), + OutK.stride(0), + OutK.stride(1), + OutK.stride(2), + OutV.stride(0), + OutV.stride(1), + OutV.stride(2), + head_num, + head_dim, + BLOCK_DMODEL=triton.next_power_of_2(head_dim), + BLOCK_HEAD=BLOCK_HEAD, + num_warps=1, + num_stages=1, + ) + return diff --git a/lightllm/common/basemodel/triton_kernel/kv_move.py b/lightllm/common/basemodel/triton_kernel/kv_move.py index ce99a007bd..34b812327b 100644 --- a/lightllm/common/basemodel/triton_kernel/kv_move.py +++ b/lightllm/common/basemodel/triton_kernel/kv_move.py @@ -42,6 +42,64 @@ def _copy_kv_buffer_to_kv_buffer( return +@triton.jit +def _copy_kv_buffer_pair_to_kv_buffer_pair( + mem_num, + src_mem_index, + dst_mem_index, + src_index_stride: tl.constexpr, + dst_index_stride: tl.constexpr, + k_buffer, + v_buffer, + k_buffer_stride_l, + k_buffer_stride_s, + k_buffer_stride_d, + v_buffer_stride_l, + v_buffer_stride_s, + v_buffer_stride_d, + k_buffer_tail_dim, + v_buffer_tail_dim, + BLOCK: tl.constexpr, +): + layer_index = tl.program_id(0).to(tl.int64) + start_index = tl.program_id(1).to(tl.int64) + grid_num = tl.num_programs(1).to(tl.int64) + + k_buffer_stride_l = tl.cast(k_buffer_stride_l, dtype=tl.int64) + k_buffer_stride_s = tl.cast(k_buffer_stride_s, dtype=tl.int64) + k_buffer_stride_d = tl.cast(k_buffer_stride_d, dtype=tl.int64) + v_buffer_stride_l = tl.cast(v_buffer_stride_l, dtype=tl.int64) + v_buffer_stride_s = tl.cast(v_buffer_stride_s, dtype=tl.int64) + v_buffer_stride_d = tl.cast(v_buffer_stride_d, dtype=tl.int64) + + for i in range(start_index, mem_num, grid_num): + src_mem = tl.load(src_mem_index + i * src_index_stride) + dst_mem = tl.load(dst_mem_index + i * dst_index_stride) + for j in range(tl.cdiv(k_buffer_tail_dim, BLOCK)): + offs = j * BLOCK + tl.arange(0, BLOCK) + mask = offs < k_buffer_tail_dim + k_buffer_data = tl.load( + k_buffer + layer_index * k_buffer_stride_l + src_mem * k_buffer_stride_s + offs, mask=mask + ) + tl.store( + k_buffer + layer_index * k_buffer_stride_l + dst_mem * k_buffer_stride_s + offs, + k_buffer_data, + mask=mask, + ) + for j in range(tl.cdiv(v_buffer_tail_dim, BLOCK)): + offs = j * BLOCK + tl.arange(0, BLOCK) + mask = offs < v_buffer_tail_dim + v_buffer_data = tl.load( + v_buffer + layer_index * v_buffer_stride_l + src_mem * v_buffer_stride_s + offs, mask=mask + ) + tl.store( + v_buffer + layer_index * v_buffer_stride_l + dst_mem * v_buffer_stride_s + offs, + v_buffer_data, + mask=mask, + ) + return + + def copy_kv_buffer_to_kv_buffer( src_mem_index: torch.Tensor, dst_mem_index: torch.Tensor, @@ -70,3 +128,42 @@ def copy_kv_buffer_to_kv_buffer( BLOCK=BLOCK, ) return + + +def copy_kv_buffer_pair_to_kv_buffer_pair( + src_mem_index: torch.Tensor, + dst_mem_index: torch.Tensor, + k_buffer: torch.Tensor, + v_buffer: torch.Tensor, +): + assert len(src_mem_index) == len(dst_mem_index) + assert src_mem_index.is_cuda and dst_mem_index.is_cuda and k_buffer.is_cuda and v_buffer.is_cuda + assert k_buffer.shape == v_buffer.shape + k_buffer = k_buffer.view(k_buffer.shape[0], k_buffer.shape[1], -1).view(dtype=torch.uint8) + v_buffer = v_buffer.view(v_buffer.shape[0], v_buffer.shape[1], -1).view(dtype=torch.uint8) + assert k_buffer.shape == v_buffer.shape + BLOCK = 4096 + layer_num = k_buffer.shape[0] + grid = ( + layer_num, + 1024, + ) + _copy_kv_buffer_pair_to_kv_buffer_pair[grid]( + mem_num=len(src_mem_index), + src_mem_index=src_mem_index, + dst_mem_index=dst_mem_index, + src_index_stride=src_mem_index.stride(0), + dst_index_stride=dst_mem_index.stride(0), + k_buffer=k_buffer, + v_buffer=v_buffer, + k_buffer_stride_l=k_buffer.stride(0), + k_buffer_stride_s=k_buffer.stride(1), + k_buffer_stride_d=k_buffer.stride(2), + v_buffer_stride_l=v_buffer.stride(0), + v_buffer_stride_s=v_buffer.stride(1), + v_buffer_stride_d=v_buffer.stride(2), + k_buffer_tail_dim=k_buffer.shape[-1], + v_buffer_tail_dim=v_buffer.shape[-1], + BLOCK=BLOCK, + ) + return diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/__init__.py b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/__init__.py index cd3b0962a3..1fd7fabe3f 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/__init__.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/__init__.py @@ -7,9 +7,10 @@ # the following copyright notice: # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang from .chunk import chunk_gated_delta_rule -from .fused_recurrent import fused_recurrent_gated_delta_rule +from .fused_recurrent import fused_recurrent_gated_delta_rule, fused_recurrent_gated_delta_rule_packed_decode __all__ = [ "chunk_gated_delta_rule", "fused_recurrent_gated_delta_rule", + "fused_recurrent_gated_delta_rule_packed_decode", ] diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/fused_recurrent.py b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/fused_recurrent.py index 5dfbd6e4ab..53b29d0437 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/fused_recurrent.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/fused_recurrent.py @@ -380,6 +380,171 @@ def fused_recurrent_gated_delta_rule_fwd( return o, final_state +@triton.jit +def _fused_recurrent_gated_delta_rule_packed_decode_kernel( + mixed_qkv, + a, + b, + A_log, + dt_bias, + o, + h0, + ssm_state_indices, + scale, + stride_mixed_qkv_tok: tl.constexpr, + stride_a_tok: tl.constexpr, + stride_b_tok: tl.constexpr, + stride_state_slot: tl.constexpr, + stride_indices: tl.constexpr, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + SOFTPLUS_THRESHOLD: tl.constexpr, + USE_QK_L2NORM_IN_KERNEL: tl.constexpr, +): + """Packed one-token GDN decode for the non-MTP path. + + ``mixed_qkv`` stays in its post-convolution ``[q, k, v]`` layout. Each + program owns one value-head tile, so the kernel avoids materializing + reshaped q/k/v tensors and writes the recurrent state in-place. LightLLM + stores the recurrent state as ``[slot, head, key, value]``; this differs + from the ``[slot, head, value, key]`` layout used by the upstream FLA + kernel, so the state offsets intentionally keep key as the first matrix + dimension. + """ + i_v, i_nh = tl.program_id(0), tl.program_id(1) + i_n, i_hv = i_nh // HV, i_nh % HV + i_h = i_hv // (HV // H) + + o_k = tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + mask_k = o_k < K + mask_v = o_v < V + mask_h = mask_k[:, None] & mask_v[None, :] + + state_idx = tl.load(ssm_state_indices + i_n * stride_indices).to(tl.int64) + p_h0 = h0 + state_idx * stride_state_slot + p_h0 += i_hv * K * V + o_k[:, None] * V + o_v[None, :] + b_h = tl.load(p_h0, mask=mask_h, other=0).to(tl.float32) + + p_mixed = mixed_qkv + i_n * stride_mixed_qkv_tok + q_off = i_h * K + o_k + k_off = H * K + i_h * K + o_k + v_off = 2 * H * K + i_hv * V + o_v + b_q = tl.load(p_mixed + q_off, mask=mask_k, other=0).to(tl.float32) + b_k = tl.load(p_mixed + k_off, mask=mask_k, other=0).to(tl.float32) + b_v = tl.load(p_mixed + v_off, mask=mask_v, other=0).to(tl.float32) + + if USE_QK_L2NORM_IN_KERNEL: + b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6) + b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6) + b_q = b_q * scale + + a_val = tl.load(a + i_n * stride_a_tok + i_hv).to(tl.float32) + b_val = tl.load(b + i_n * stride_b_tok + i_hv).to(tl.float32) + A_log_val = tl.load(A_log + i_hv).to(tl.float32) + dt_bias_val = tl.load(dt_bias + i_hv).to(tl.float32) + x = a_val + dt_bias_val + softplus_x = tl.where(x <= SOFTPLUS_THRESHOLD, tl.log(1.0 + tl.exp(x)), x) + g_val = -tl.exp(A_log_val) * softplus_x + beta_val = tl.sigmoid(b_val) + + b_h *= exp(g_val) + b_v -= tl.sum(b_h * b_k[:, None], 0) + b_v *= beta_val + b_h += b_k[:, None] * b_v[None, :] + b_o = tl.sum(b_h * b_q[:, None], 0) + + p_o = o + (i_n * HV + i_hv) * V + o_v + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v) + tl.store(p_h0, b_h.to(p_h0.dtype.element_ty), mask=mask_h) + + +def fused_recurrent_gated_delta_rule_packed_decode( + mixed_qkv: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + ssm_state_indices: torch.Tensor, + use_qk_l2norm_in_kernel: bool = False, + out: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Run packed one-token GDN decode without q/k/v materialization.""" + if mixed_qkv.ndim != 2 or mixed_qkv.stride(-1) != 1: + raise ValueError("mixed_qkv must be a 2D tensor contiguous in its last dimension") + if a.ndim != 2 or b.ndim != 2 or a.stride(-1) != 1 or b.stride(-1) != 1: + raise ValueError("a and b must be 2D tensors contiguous in their last dimension") + if A_log.ndim != 1 or dt_bias.ndim != 1 or A_log.stride(0) != 1 or dt_bias.stride(0) != 1: + raise ValueError("A_log and dt_bias must be contiguous 1D tensors") + if ssm_state_indices.ndim != 1 or not ssm_state_indices.is_contiguous(): + raise ValueError("ssm_state_indices must be a contiguous 1D tensor") + if initial_state.ndim != 4 or initial_state.stride(-1) != 1: + raise ValueError("initial_state must be a 4D tensor contiguous in its last dimension") + + batch = mixed_qkv.shape[0] + num_v_heads, key_dim, value_dim = initial_state.shape[-3:] + if a.shape != (batch, num_v_heads) or b.shape != (batch, num_v_heads): + raise ValueError("a and b must have shape [batch, num_v_heads]") + if ssm_state_indices.shape[0] != batch: + raise ValueError("ssm_state_indices must have one entry per batch row") + if A_log.numel() != num_v_heads or dt_bias.numel() != num_v_heads: + raise ValueError("A_log and dt_bias must match num_v_heads") + + qk_dim = mixed_qkv.shape[1] - num_v_heads * value_dim + if qk_dim <= 0 or qk_dim % 2 != 0: + raise ValueError("mixed_qkv does not contain a valid q/k/v layout") + query_dim = qk_dim // 2 + if query_dim % key_dim != 0: + raise ValueError("query dimension must be divisible by key head dimension") + num_k_heads = query_dim // key_dim + if num_k_heads <= 0 or num_v_heads % num_k_heads != 0: + raise ValueError("invalid GQA head configuration") + + block_k = triton.next_power_of_2(key_dim) + if triton.cdiv(key_dim, block_k) != 1: + raise ValueError("packed decode currently supports one key tile") + block_v = min(triton.next_power_of_2(value_dim), 32) + if out is None: + out = mixed_qkv.new_empty((batch, num_v_heads, value_dim)) + elif out.shape != (batch, num_v_heads, value_dim) or not out.is_contiguous(): + raise ValueError("out must be contiguous with shape [batch, num_v_heads, value_dim]") + + grid = (triton.cdiv(value_dim, block_v), batch * num_v_heads) + _fused_recurrent_gated_delta_rule_packed_decode_kernel[grid]( + mixed_qkv=mixed_qkv, + a=a, + b=b, + A_log=A_log, + dt_bias=dt_bias, + o=out, + h0=initial_state, + ssm_state_indices=ssm_state_indices, + scale=scale, + stride_mixed_qkv_tok=mixed_qkv.stride(0), + stride_a_tok=a.stride(0), + stride_b_tok=b.stride(0), + stride_state_slot=initial_state.stride(0), + stride_indices=ssm_state_indices.stride(0), + H=num_k_heads, + HV=num_v_heads, + K=key_dim, + V=value_dim, + BK=block_k, + BV=block_v, + SOFTPLUS_THRESHOLD=20.0, + USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, + num_warps=1, + num_stages=3, + ) + return out.unsqueeze(1), initial_state + + class FusedRecurrentFunction(torch.autograd.Function): @staticmethod def forward( diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/mtp_state_params.py b/lightllm/common/basemodel/triton_kernel/linear_att/mtp_state_params.py index b81bf9ec16..0413aea898 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/mtp_state_params.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/mtp_state_params.py @@ -3,7 +3,7 @@ import triton.language as tl -@triton.jit +@triton.jit(do_not_specialize=["batch_size"]) def _build_dynamic_mtp_linear_att_state_params_kernel( b_req_idx, b_mtp_index, diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py new file mode 100644 index 0000000000..e54186eae4 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm.py @@ -0,0 +1,29 @@ +# SPDX-License-Identifier: Apache-2.0 +"""ReplaySSM kernel exports, in lifecycle order. + +decode writes history for the current window. +cursor records how many of those tokens were accepted. +fold applies that accepted prefix onto the canonical SSM. +""" + +from .replayssm_decode import ( + _decode_replay, + _gdn_precomputed_window, + _kda_precomputed_window, + _project_qk, + _replay_each_token, +) +from .replayssm_delta import _delta_rule_step +from .replayssm_cursor import _advance_cursor +from .replayssm_fold import _fold_prefix + +__all__ = [ + "_decode_replay", + "_project_qk", + "_gdn_precomputed_window", + "_kda_precomputed_window", + "_replay_each_token", + "_delta_rule_step", + "_advance_cursor", + "_fold_prefix", +] diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_cursor.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_cursor.py new file mode 100644 index 0000000000..5354c71260 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_cursor.py @@ -0,0 +1,24 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Advance the accepted-prefix cursor. This kernel does not fold state.""" + +import triton +import triton.language as tl + + +@triton.jit(do_not_specialize=["N"]) +def _advance_cursor( + cursors, + reqs, + accepted_state_indices, + N, + HOLD: tl.constexpr, + WIDTH: tl.constexpr, + BLOCK: tl.constexpr, +): + # N 是运行时请求数；BLOCK/WIDTH 仍是编译期参数，避免 eager batch 变化触发重新编译。 + i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + r = tl.load(reqs + i, i < N, HOLD) + valid = (i < N) & (r != HOLD) + cursor = tl.load(cursors + r, valid, 0) + count = tl.load(accepted_state_indices + r, valid, 0) + 1 if WIDTH > 1 else 1 + tl.store(cursors + r, cursor + count, valid) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py new file mode 100644 index 0000000000..5619869b64 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_decode.py @@ -0,0 +1,792 @@ +# SPDX-License-Identifier: Apache-2.0 +"""ReplaySSM decode kernels. + +`_decode_replay` is the entry. It projects the verify window when asked, then +solves a GDN or KDA window, or walks tokens one by one. Neither path writes the +canonical SSM. +""" + +import triton +import triton.language as tl + + +@triton.jit +def _history_coeff(rows, vec, BK: tl.constexpr): + # rows [R, BK] · vec [BK] -> [R]. tf32x3 是 tensor core 上最接近 fp32 的输出重建。 + if BK >= 16: + return tl.sum(tl.dot(rows, vec[:, None], input_precision="tf32x3"), 1) + return tl.sum(rows * vec[None, :], 1) + + +@triton.jit +def _apply_deltas(coeff, deltas, R: tl.constexpr, BV: tl.constexpr): + # coeff [R] · deltas [R, BV] -> [BV]. R 小于 16 时 tensor core 不能做这次点积。 + if R >= 16 and BV >= 16: + return tl.sum(tl.dot(coeff[None, :], deltas, input_precision="tf32x3"), 0) + return tl.sum(deltas * coeff[:, None], 0) + + +@triton.jit +def _decode_replay( + Q, + Kp, + Vp, + A, + B, + Alog, + Bias, + State, + Keys, + Deltas, + Gates, + RawKeys, + RawValues, + Betas, + Reqs, + Cursors, + Cu, + Out, + SQ: tl.constexpr, + SK: tl.constexpr, + SV: tl.constexpr, + SA: tl.constexpr, + SB: tl.constexpr, + H: tl.constexpr, + HV: tl.constexpr, + KH: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + L: tl.constexpr, + HOLD: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + VARLEN: tl.constexpr, + PRECOMPUTE_PROJECTIONS: tl.constexpr, + BW: tl.constexpr, + KDA: tl.constexpr, + LOWER: tl.constexpr, + BQK: tl.constexpr, + BR: tl.constexpr, +): + """执行一个请求窗口的 replay decode，并写入 history。 + + 只把本步 token 追加到长度为 L 的连续前缀，不写 canonical SSM。 + cursor 是已有前缀长度，本步从那里接着写。请求数属于运行时布局， + 不应因为 eager batch 改变而形成新的 kernel 变体。 + """ + iv, seq, hv = tl.program_id(0), tl.program_id(1), tl.program_id(2) + req = tl.load(Reqs + seq).to(tl.int64) + if VARLEN: + start, end = tl.load(Cu + seq), tl.load(Cu + seq + 1) + else: + start, end = seq, seq + 1 + vv = iv * BV + tl.arange(0, BV) + if req == HOLD: + for t in range(start, end): + tl.store(Out + (t * HV + hv) * V + vv, 0, vv < V) + return + if start == end: + return + + kk = tl.arange(0, BK) + ll = tl.arange(0, L) + slot = req * HV + hv + key_slot = req * KH + hv // (HV // KH) + n = tl.load(Cursors + req) + sp = State + slot * K * V + kk[:, None] * V + vv[None, :] + state = tl.load(sp, (kk[:, None] < K) & (vv[None, :] < V), 0).to(tl.float32) + + if PRECOMPUTE_PROJECTIONS: + # 输出用 tf32x3 tensor core 从冻结的 checkpoint 和 history 重建。 + # 状态不在这里更新；压缩仍用原始 k/v 做 fp32 串行折回。 + projection_dtype: tl.constexpr = tl.float32 + projection_precision: tl.constexpr = "tf32x3" + qk, projections, tt, tokens, valid, h = _project_qk( + Q, + Kp, + A, + Alog, + Bias, + Gates, + SQ, + SK, + SA, + H, + HV, + K, + L, + BW, + KDA, + LOWER, + BQK, + end, + hv, + kk, + ll, + n, + projection_dtype, + projection_precision, + slot, + start, + state, + ) + else: + projections, tt = None, None + + if PRECOMPUTE_PROJECTIONS and KDA: + _kda_precomputed_window( + Kp, + Vp, + A, + B, + Alog, + Bias, + Keys, + Deltas, + Gates, + RawKeys, + RawValues, + Betas, + Out, + SK, + SV, + SA, + SB, + HV, + KH, + K, + V, + L, + BV, + BW, + BR, + BK, + LOWER, + end, + h, + hv, + iv, + key_slot, + kk, + n, + projections, + qk, + slot, + start, + tt, + vv, + ) + elif PRECOMPUTE_PROJECTIONS: + _gdn_precomputed_window( + Kp, + Vp, + A, + B, + Alog, + Bias, + Keys, + Deltas, + Gates, + RawKeys, + RawValues, + Betas, + Out, + SK, + SV, + SA, + SB, + HV, + KH, + K, + V, + L, + BV, + BW, + BR, + end, + h, + hv, + iv, + key_slot, + kk, + n, + projection_dtype, + projection_precision, + projections, + qk, + slot, + start, + tokens, + tt, + valid, + vv, + ) + else: + _replay_each_token( + Q, + Kp, + Vp, + A, + B, + Alog, + Bias, + Keys, + Deltas, + Gates, + RawKeys, + RawValues, + Betas, + Out, + SQ, + SK, + SV, + SA, + SB, + H, + HV, + KH, + K, + V, + L, + PRECOMPUTE_PROJECTIONS, + BW, + KDA, + LOWER, + BK, + BV, + end, + hv, + iv, + key_slot, + kk, + ll, + n, + projections, + slot, + start, + state, + tt, + vv, + ) + + +@triton.jit +def _project_qk( + Q, + Kp, + A, + Alog, + Bias, + Gates, + SQ: tl.constexpr, + SK: tl.constexpr, + SA: tl.constexpr, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + L: tl.constexpr, + BW: tl.constexpr, + KDA: tl.constexpr, + LOWER: tl.constexpr, + BQK: tl.constexpr, + end, + hv, + kk, + ll, + n, + projection_dtype: tl.constexpr, + projection_precision: tl.constexpr, + slot, + start, + state, +): + # 预计算 verify window 的 q/k 投影。 + tt = tl.arange(0, BQK) + tokens = start + tt % BW + valid = (tt < 2 * BW) & (tokens < end) + safe_tokens = tl.where(valid, tokens, start) + h = hv // (HV // H) + q_rows = tl.load( + Q + safe_tokens[:, None] * SQ + h * K + kk[None, :], + (valid & (tt < BW))[:, None] & (kk[None, :] < K), + 0, + ).to(tl.float32) + k_rows = tl.load( + Kp + safe_tokens[:, None] * SK + h * K + kk[None, :], + (valid & (tt >= BW))[:, None] & (kk[None, :] < K), + 0, + ).to(tl.float32) + # Padded query/key rows share the MMA tile. Keep them at zero so they + # cannot change the valid rows. + row_valid = valid[:, None] + qk = tl.where(row_valid, q_rows + k_rows, 0) + qk *= tl.rsqrt(tl.sum(qk * qk, axis=1) + 1.0e-6)[:, None] + qk *= tl.where(tt < BW, K ** -0.5, 1.0)[:, None] + qk = tl.where(row_valid, qk, 0) + if KDA: + history = ll < n + gates = tl.load( + Gates + (slot * L + ll[:, None]) * K + kk[None, :], + history[:, None] & (kk[None, :] < K), + 0, + ).to(tl.float32) + total_g = tl.sum(gates, axis=0) + wi = tl.arange(0, BW) + gate_ok = start + wi < end + safe_wi = tl.where(gate_ok, wi, 0) + gate_inputs = tl.load( + A + (start + safe_wi[:, None]) * SA + hv * K + kk[None, :], + gate_ok[:, None] & (kk[None, :] < K), + 0, + ).to(tl.float32) + gate_inputs += tl.load(Bias + hv * K + kk, kk < K, 0).to(tl.float32)[None, :] + fresh_gates = LOWER * tl.sigmoid(tl.exp(tl.load(Alog + hv)) * gate_inputs) + fresh_gates = tl.where((start + wi < end)[:, None], fresh_gates, 0) + query_gates = total_g[None, :] + tl.gather( + tl.cumsum(fresh_gates, axis=0), + tl.broadcast_to((tt % BW)[:, None], qk.shape), + 0, + ) + # All exponents are forward decay. No inverse prefix scaling can + # overflow when a long history contains strongly negative gates. + # Padded rows stay zero: 0 * exp(garbage) is NaN and would enter the dot. + scaled = tl.where(row_valid, qk * tl.exp(query_gates), 0) + projections = tl.dot( + scaled.to(projection_dtype), + state.to(projection_dtype), + input_precision=projection_precision, + ) + else: + projections = tl.dot(qk.to(projection_dtype), state.to(projection_dtype), input_precision=projection_precision) + return qk, projections, tt, tokens, valid, h + + +@triton.jit +def _gdn_precomputed_window( + Kp, + Vp, + A, + B, + Alog, + Bias, + Keys, + Deltas, + Gates, + RawKeys, + RawValues, + Betas, + Out, + SK: tl.constexpr, + SV: tl.constexpr, + SA: tl.constexpr, + SB: tl.constexpr, + HV: tl.constexpr, + KH: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + L: tl.constexpr, + BV: tl.constexpr, + BW: tl.constexpr, + BR: tl.constexpr, + end, + h, + hv, + iv, + key_slot, + kk, + n, + projection_dtype: tl.constexpr, + projection_precision: tl.constexpr, + projections, + qk, + slot, + start, + tokens, + tt, + valid, + vv, +): + # GDN precompute 路径：一次投影整个窗口，再按 token 顺序写入新 history。 + # 只有 corrected-value solve 依赖 token 顺序，避免重复计算历史投影。 + jj = tl.arange(0, BR) + fresh_tokens = start + jj - L + old = (jj < L) & (jj < n) + fresh = (jj >= L) & (jj < L + end - start) + # BR is rounded above the history length. Masked lanes must still address + # a real slot; otherwise they read uninitialized neighbor rows into the dot. + safe_jj = tl.where(jj < L, jj, 0) + safe_fresh = tl.where(fresh, fresh_tokens, start) + cached_keys = tl.load( + Keys + (key_slot * L) * K + safe_jj[:, None] * K + kk[None, :], + old[:, None] & (kk[None, :] < K), + 0, + ).to(tl.float32) + raw_keys = tl.load( + Kp + safe_fresh[:, None] * SK + h * K + kk[None, :], + fresh[:, None] & (kk[None, :] < K), + 0, + ).to(tl.float32) + fresh_keys = raw_keys * tl.rsqrt(tl.sum(raw_keys * raw_keys, axis=1) + 1.0e-6)[:, None] + fresh_keys = tl.where(fresh[:, None], fresh_keys, 0) + all_keys = cached_keys + fresh_keys + values = tl.load( + Vp + safe_fresh[:, None] * SV + hv * V + vv[None, :], + fresh[:, None] & (vv[None, :] < V), + 0, + ).to(tl.float32) + log_a, bias = tl.load(Alog + hv).to(tl.float32), tl.load(Bias + hv).to(tl.float32) + x = tl.load(A + safe_fresh * SA + hv, fresh, 0).to(tl.float32) + bias + fresh_gates = -tl.exp(log_a) * tl.where(x <= 20.0, tl.log(1.0 + tl.exp(x)), x) + gates = tl.load(Gates + slot * L + safe_jj, old, 0).to(tl.float32) + gates += tl.where(fresh, fresh_gates, 0.0) + prefix = tl.cumsum(gates, axis=0) + query_gates = tl.sum(tl.where(jj[None, :] <= L + (tt % BW)[:, None], gates[None, :], 0), 1) + coefficients = tl.dot( + qk.to(projection_dtype), tl.trans(all_keys).to(projection_dtype), input_precision=projection_precision + ) + causal = old[None, :] | (fresh[None, :] & (jj[None, :] <= L + (tt % BW)[:, None])) + coefficients *= tl.where(causal, tl.exp(tl.minimum(query_gates[:, None] - prefix[None, :], 0.0)), 0.0) + projections *= tl.exp(query_gates)[:, None] + deltas = tl.load( + Deltas + (slot * L) * V + safe_jj[:, None] * V + vv[None, :], + old[:, None] & (vv[None, :] < V), + 0, + ).to(tl.float32) + betas = tl.sigmoid(tl.load(B + safe_fresh * SB + hv, fresh, 0).to(tl.float32)) + for t in range(end - start): + checkpoint_sk = tl.sum(tl.where((tt == BW + t)[:, None], projections, 0), 0) + weights = tl.sum(tl.where((tt == BW + t)[:, None], coefficients, 0), 0) + sk = checkpoint_sk + _apply_deltas(weights, deltas, BR, BV) + value = tl.sum(tl.where((jj == L + t)[:, None], values, 0), 0) + beta = tl.sum(tl.where(jj == L + t, betas, 0)) + delta = beta * (value - sk) + deltas = tl.where((jj == L + t)[:, None], delta[None, :], deltas) + outputs = projections + tl.dot( + coefficients.to(projection_dtype), deltas.to(projection_dtype), input_precision=projection_precision + ) + out_tokens = tl.where(valid, tokens, start) + tl.store( + Out + (out_tokens[:, None] * HV + hv) * V + vv[None, :], + outputs, + (valid & (tt < BW))[:, None] & (vv[None, :] < V), + ) + records = n + jj - L + in_prefix = fresh & (records < L) + safe_records = tl.where(in_prefix, records, 0) + tl.store( + Deltas + (slot * L + safe_records[:, None]) * V + vv[None, :], + deltas, + in_prefix[:, None] & (vv[None, :] < V), + ) + tl.store( + RawValues + (slot * L + safe_records[:, None]) * V + vv[None, :], + values, + in_prefix[:, None] & (vv[None, :] < V), + ) + if iv == 0: + if hv % (HV // KH) == 0: + tl.store( + Keys + (key_slot * L + safe_records[:, None]) * K + kk[None, :], + fresh_keys, + in_prefix[:, None] & (kk[None, :] < K), + ) + tl.store( + RawKeys + (key_slot * L + safe_records[:, None]) * K + kk[None, :], + raw_keys, + in_prefix[:, None] & (kk[None, :] < K), + ) + tl.store(Gates + slot * L + safe_records, fresh_gates, in_prefix) + tl.store(Betas + slot * L + safe_records, betas, in_prefix) + return + + +@triton.jit +def _kda_precomputed_window( + Kp, + Vp, + A, + B, + Alog, + Bias, + Keys, + Deltas, + Gates, + RawKeys, + RawValues, + Betas, + Out, + SK: tl.constexpr, + SV: tl.constexpr, + SA: tl.constexpr, + SB: tl.constexpr, + HV: tl.constexpr, + KH: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + L: tl.constexpr, + BV: tl.constexpr, + BW: tl.constexpr, + BR: tl.constexpr, + BK: tl.constexpr, + LOWER: tl.constexpr, + end, + h, + hv, + iv, + key_slot, + kk, + n, + projections, + qk, + slot, + start, + tt, + vv, +): + # KDA 的衰减在每个 key 通道上，不能把标量 exp 提到点积外面。 + # 每个 token 用冻结 checkpoint 的投影，加上 history 的通道衰减点积重建输出。 + # exp(G_t - G_j) 对已可见的 token 不大于 1。新 delta 仍按 token 解。 + # 状态压缩继续用原始 k/v 在 fp32 里折回。 + jj = tl.arange(0, BR) + fresh_tokens = start + jj - L + old = (jj < L) & (jj < n) + fresh = (jj >= L) & (jj < L + end - start) + safe_jj = tl.where(jj < L, jj, 0) + safe_fresh = tl.where(fresh, fresh_tokens, start) + cached_keys = tl.load( + Keys + (key_slot * L) * K + safe_jj[:, None] * K + kk[None, :], + old[:, None] & (kk[None, :] < K), + 0, + ).to(tl.float32) + raw_keys = tl.load( + Kp + safe_fresh[:, None] * SK + h * K + kk[None, :], + fresh[:, None] & (kk[None, :] < K), + 0, + ).to(tl.float32) + fresh_keys = raw_keys * tl.rsqrt(tl.sum(raw_keys * raw_keys, axis=1) + 1.0e-6)[:, None] + fresh_keys = tl.where(fresh[:, None], fresh_keys, 0) + all_keys = tl.where(old[:, None], cached_keys, 0) + fresh_keys + values = tl.load( + Vp + safe_fresh[:, None] * SV + hv * V + vv[None, :], + fresh[:, None] & (vv[None, :] < V), + 0, + ).to(tl.float32) + hist_g = tl.load( + Gates + (slot * L + safe_jj[:, None]) * K + kk[None, :], + old[:, None] & (kk[None, :] < K), + 0, + ).to(tl.float32) + total_hist = tl.sum(hist_g, axis=0) + hist_prefix = tl.cumsum(hist_g, axis=0) + wi = tl.arange(0, BW) + gate_ok = start + wi < end + safe_wi = tl.where(gate_ok, wi, 0) + gate_inputs = tl.load( + A + (start + safe_wi[:, None]) * SA + hv * K + kk[None, :], + gate_ok[:, None] & (kk[None, :] < K), + 0, + ).to(tl.float32) + gate_inputs += tl.load(Bias + hv * K + kk, kk < K, 0).to(tl.float32)[None, :] + fresh_g_rows = LOWER * tl.sigmoid(tl.exp(tl.load(Alog + hv).to(tl.float32)) * gate_inputs) + fresh_g_rows = tl.where(gate_ok[:, None], fresh_g_rows, 0) + fresh_prefix = tl.cumsum(fresh_g_rows, axis=0) + fresh_index = tl.broadcast_to(tl.where(fresh, jj - L, 0)[:, None], (BR, BK)) + G_pos = tl.where(old[:, None], hist_prefix, 0.0) + G_pos = tl.where(fresh[:, None], total_hist[None, :] + tl.gather(fresh_prefix, fresh_index, 0), G_pos) + step_g = tl.where(fresh[:, None], tl.gather(fresh_g_rows, fresh_index, 0), 0.0) + G_query = total_hist[None, :] + tl.gather(fresh_prefix, tl.broadcast_to((tt % BW)[:, None], qk.shape), 0) + deltas = tl.load( + Deltas + (slot * L) * V + safe_jj[:, None] * V + vv[None, :], + old[:, None] & (vv[None, :] < V), + 0, + ).to(tl.float32) + betas = tl.sigmoid(tl.load(B + safe_fresh * SB + hv, fresh, 0).to(tl.float32)) + for t in range(start, end): + local = t - start + q = tl.sum(tl.where((tt == local)[:, None], qk, 0), 0) + k = tl.sum(tl.where((tt == BW + local)[:, None], qk, 0), 0) + g_now = tl.sum(tl.where((tt == local)[:, None], G_query, 0), 0) + causal = old | (fresh & (jj <= L + local)) + decay = tl.exp(tl.where(causal[:, None], g_now[None, :] - G_pos, 0.0)) + scaled = tl.where(causal[:, None], all_keys * decay, 0) + hk = _history_coeff(scaled, k, BK) + proj_k = tl.sum(tl.where((tt == BW + local)[:, None], projections, 0), 0) + sk = proj_k + _apply_deltas(hk, deltas, BR, BV) + value = tl.sum(tl.where((jj == L + local)[:, None], values, 0), 0) + beta = tl.sum(tl.where(jj == L + local, betas, 0)) + delta = beta * (value - sk) + deltas = tl.where((jj == L + local)[:, None], delta[None, :], deltas) + hq = _history_coeff(scaled, q, BK) + proj_q = tl.sum(tl.where((tt == local)[:, None], projections, 0), 0) + sq = proj_q + _apply_deltas(hq, deltas, BR, BV) + tl.store(Out + (t * HV + hv) * V + vv, sq, vv < V) + records = n + jj - L + in_prefix = fresh & (records < L) + safe_records = tl.where(in_prefix, records, 0) + tl.store( + Deltas + (slot * L + safe_records[:, None]) * V + vv[None, :], + deltas, + in_prefix[:, None] & (vv[None, :] < V), + ) + tl.store( + RawValues + (slot * L + safe_records[:, None]) * V + vv[None, :], + values, + in_prefix[:, None] & (vv[None, :] < V), + ) + if iv == 0: + if hv % (HV // KH) == 0: + tl.store( + Keys + (key_slot * L + safe_records[:, None]) * K + kk[None, :], + fresh_keys, + in_prefix[:, None] & (kk[None, :] < K), + ) + tl.store( + RawKeys + (key_slot * L + safe_records[:, None]) * K + kk[None, :], + raw_keys, + in_prefix[:, None] & (kk[None, :] < K), + ) + tl.store( + Gates + (slot * L + safe_records[:, None]) * K + kk[None, :], + step_g, + in_prefix[:, None] & (kk[None, :] < K), + ) + tl.store(Betas + slot * L + safe_records, betas, in_prefix) + return + + +@triton.jit +def _replay_each_token( + Q, + Kp, + Vp, + A, + B, + Alog, + Bias, + Keys, + Deltas, + Gates, + RawKeys, + RawValues, + Betas, + Out, + SQ: tl.constexpr, + SK: tl.constexpr, + SV: tl.constexpr, + SA: tl.constexpr, + SB: tl.constexpr, + H: tl.constexpr, + HV: tl.constexpr, + KH: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + L: tl.constexpr, + PRECOMPUTE_PROJECTIONS: tl.constexpr, + BW: tl.constexpr, + KDA: tl.constexpr, + LOWER: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + end, + hv, + iv, + key_slot, + kk, + ll, + n, + projections, + slot, + start, + state, + tt, + vv, +): + # KDA、inline 以及不满足 precompute 条件的路径使用逐 token 递推。 + history = ll < n + keys = tl.load( + Keys + (key_slot * L) * K + ll[:, None] * K + kk[None, :], + history[:, None] & (kk[None, :] < K), + 0, + ).to(tl.float32) + deltas = tl.load( + Deltas + (slot * L) * V + ll[:, None] * V + vv[None, :], + history[:, None] & (vv[None, :] < V), + 0, + ).to(tl.float32) + if KDA: + gates = tl.load( + Gates + (slot * L + ll[:, None]) * K + kk[None, :], + history[:, None] & (kk[None, :] < K), + 0, + ).to(tl.float32) + else: + gates = tl.load(Gates + slot * L + ll, history, 0).to(tl.float32) + gate_prefix = tl.cumsum(gates, axis=0) + total_g = tl.sum(gates, axis=0) + if KDA: + keys *= tl.where(history[:, None], tl.exp(total_g[None, :] - gate_prefix), 0.0) + h = hv // (HV // H) + log_a = tl.load(Alog + hv).to(tl.float32) + bias = tl.load(Bias + hv * K + kk, kk < K, 0).to(tl.float32) if KDA else tl.load(Bias + hv).to(tl.float32) + + for t in range(start, end): + q = tl.load(Q + t * SQ + h * K + kk, kk < K, 0).to(tl.float32) + raw_k = tl.load(Kp + t * SK + h * K + kk, kk < K, 0).to(tl.float32) + v = tl.load(Vp + t * SV + hv * V + vv, vv < V, 0).to(tl.float32) + q = q / tl.sqrt(tl.sum(q * q) + 1.0e-6) * (K ** -0.5) + k = raw_k / tl.sqrt(tl.sum(raw_k * raw_k) + 1.0e-6) + if KDA: + x = tl.load(A + t * SA + hv * K + kk, kk < K, 0).to(tl.float32) + bias + g = LOWER * tl.sigmoid(tl.exp(log_a) * x) + else: + x = tl.load(A + t * SA + hv).to(tl.float32) + bias + g = -tl.exp(log_a) * tl.where(x <= 20.0, tl.log(1.0 + tl.exp(x)), x) + beta = tl.sigmoid(tl.load(B + t * SB + hv).to(tl.float32)) + total_g += g + if KDA: + # Advance the decayed keys once per token, rather than + # evaluating L independent vector exponentials each time. + keys *= tl.exp(g)[None, :] + hk = _history_coeff(keys, k, BK) + hq = _history_coeff(keys, q, BK) + if PRECOMPUTE_PROJECTIONS: + checkpoint_sk = tl.sum(tl.where((tt == BW + t - start)[:, None], projections, 0), 0) + checkpoint_sq = tl.sum(tl.where((tt == t - start)[:, None], projections, 0), 0) + else: + checkpoint_sk = tl.sum(state * (k * tl.exp(total_g))[:, None], axis=0) + checkpoint_sq = tl.sum(state * (q * tl.exp(total_g))[:, None], axis=0) + sk = checkpoint_sk + _apply_deltas(hk, deltas, L, BV) + sq = checkpoint_sq + _apply_deltas(hq, deltas, L, BV) + else: + weights = tl.where(history, tl.exp(total_g - gate_prefix), 0.0) + hk = _history_coeff(keys, k, BK) * weights + hq = _history_coeff(keys, q, BK) * weights + checkpoint_sk = tl.sum(state * k[:, None], axis=0) + checkpoint_sq = tl.sum(state * q[:, None], axis=0) + sk = checkpoint_sk * tl.exp(total_g) + _apply_deltas(hk, deltas, L, BV) + sq = checkpoint_sq * tl.exp(total_g) + _apply_deltas(hq, deltas, L, BV) + d = beta * (v - sk) + tl.store(Out + (t * HV + hv) * V + vv, sq + d * tl.sum(k * q), vv < V) + record = n + in_prefix = record < L + # 被 mask 掉的 lane 仍会计算地址。cursor 已经顶到容量时必须改写 0 号槽，避免非法访问。 + safe_record = tl.where(in_prefix, record, 0) + tl.store(Deltas + (slot * L + safe_record) * V + vv, d, in_prefix & (vv < V)) + tl.store(RawValues + (slot * L + safe_record) * V + vv, v, in_prefix & (vv < V)) + if iv == 0: + # Only the group owner writes shared K. Every CTA computes this + # round's K locally; readers only consume earlier accepted records. + if hv % (HV // KH) == 0: + tl.store(Keys + (key_slot * L + safe_record) * K + kk, k, in_prefix & (kk < K)) + tl.store(RawKeys + (key_slot * L + safe_record) * K + kk, raw_k, in_prefix & (kk < K)) + if KDA: + tl.store(Gates + (slot * L + safe_record) * K + kk, g, in_prefix & (kk < K)) + else: + tl.store(Gates + slot * L + safe_record, g, in_prefix) + tl.store(Betas + slot * L + safe_record, beta, in_prefix) + keys = tl.where((ll == n)[:, None], k[None, :], keys) + deltas = tl.where((ll == n)[:, None], d[None, :], deltas) + if not KDA: + gate_prefix = tl.where(ll >= n, total_g, gate_prefix) + history |= ll == n + n += 1 diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_delta.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_delta.py new file mode 100644 index 0000000000..b3a6b2aec0 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_delta.py @@ -0,0 +1,21 @@ +# SPDX-License-Identifier: Apache-2.0 +"""One gated-delta step shared by prefix folding.""" + +import triton +import triton.language as tl + + +@triton.jit +def _delta_rule_step(state, raw_k, raw_v, g, beta, KDA: tl.constexpr): + # Folding applies this same update once per accepted history token. + raw_k /= tl.sqrt(tl.sum(raw_k * raw_k) + 1.0e-6) + decay = tl.exp(g) + if KDA: + state *= decay[:, None] + d = beta * (raw_v - tl.sum(state * raw_k[:, None], 0)) + state += raw_k[:, None] * d[None, :] + else: + sk = tl.sum(state * raw_k[:, None], 0) + d = beta * (raw_v - decay * sk) + state = state * decay + raw_k[:, None] * d[None, :] + return state diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py new file mode 100644 index 0000000000..7ac6a25c8b --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/linear_att/replayssm_fold.py @@ -0,0 +1,104 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Fold an accepted prefix into the canonical SSM.""" + +import triton +import triton.language as tl + +from .replayssm_delta import _delta_rule_step + + +@triton.jit(do_not_specialize=["N"]) +def _reset_folded_cursors( + Cursors, Reqs, N, HOLD: tl.constexpr, L: tl.constexpr, NEXT: tl.constexpr, BLOCK: tl.constexpr +): + row = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + req = tl.load(Reqs + row, row < N, HOLD) + valid = (row < N) & (req != HOLD) + n = tl.load(Cursors + req, valid, 0) + tl.store(Cursors + req, 0, valid & (n + NEXT > L)) + + +@triton.jit +def _fold_prefix( + State, + Output, + RawKeys, + RawValues, + Gates, + Betas, + Cursors, + Reqs, + STATE_STRIDES: tl.constexpr, + SLOTS: tl.constexpr, + HV: tl.constexpr, + KH: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + L: tl.constexpr, + HOLD: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + SNAPSHOT: tl.constexpr, + KDA: tl.constexpr, + FLUSH_EXPORT: tl.constexpr = False, + OUTPUT_STRIDES: tl.constexpr = None, + WHEN_FULL: tl.constexpr = False, + NEXT: tl.constexpr = 1, +): + # 将 cursor 指向的 accepted history 应用到 canonical state。SNAPSHOT + # 只写独立输出。WHEN_FULL 只压缩「长度 + 下一次最大宽度」放不下的请求， + # cursor 在后续 kernel 中清零；这里的所有 layer/head tile 都需要读到原始长度。 + iv, row, lh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + req = tl.load(Reqs + row).to(tl.int64) + if req == HOLD: + return + n = tl.load(Cursors + req) + if WHEN_FULL: + if n == 0 or n + NEXT <= L: + return + elif n == 0 and not SNAPSHOT and not FLUSH_EXPORT: + return + linear_layer_idx, head = lh // HV, lh % HV + slot = (linear_layer_idx * SLOTS + req) * HV + head + key_slot = (linear_layer_idx * SLOTS + req) * KH + head // (HV // KH) + kk = tl.arange(0, BK) + vv = iv * BV + tl.arange(0, BV) + state_offset = ( + linear_layer_idx.to(tl.int64) * STATE_STRIDES[0] + + req * STATE_STRIDES[1] + + head.to(tl.int64) * STATE_STRIDES[2] + + kk[:, None].to(tl.int64) * STATE_STRIDES[3] + + vv[None, :].to(tl.int64) * STATE_STRIDES[4] + ) + sp = State + state_offset + state = tl.load(sp, (kk[:, None] < K) & (vv[None, :] < V), 0).to(tl.float32) + # cursor 可能暂时超过容量。循环只走容量内的槽，越界步不读 history，也不更新 state。 + for j in range(L): + active = j < n + raw_k = tl.load(RawKeys + (key_slot * L + j) * K + kk, active & (kk < K), 0).to(tl.float32) + raw_v = tl.load(RawValues + (slot * L + j) * V + vv, active & (vv < V), 0).to(tl.float32) + if KDA: + g = tl.load(Gates + (slot * L + j) * K + kk, active & (kk < K), 0).to(tl.float32) + else: + g = tl.load(Gates + slot * L + j, active, 0).to(tl.float32) + beta = tl.load(Betas + slot * L + j, active, 0).to(tl.float32) + updated = _delta_rule_step(state, raw_k, raw_v, g, beta, KDA) + state = tl.where(active, updated, state) + mask = (kk[:, None] < K) & (vv[None, :] < V) + if OUTPUT_STRIDES is not None: + target = tl.load(Output + row).to(tl.pointer_type(State.dtype.element_ty)) + out = ( + target + + linear_layer_idx.to(tl.int64) * OUTPUT_STRIDES[0] + + head * OUTPUT_STRIDES[1] + + kk[:, None] * OUTPUT_STRIDES[2] + + vv[None, :] * OUTPUT_STRIDES[3] + ) + else: + if SNAPSHOT: + out = Output + (linear_layer_idx * HV + head) * K * V + kk[:, None] * V + vv[None, :] + else: + out = Output + state_offset + tl.store(out, state, mask) + if FLUSH_EXPORT: + tl.store(sp, state, mask) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_acceptance.py b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_acceptance.py new file mode 100644 index 0000000000..5244dada83 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_acceptance.py @@ -0,0 +1,55 @@ +from __future__ import annotations + +import torch +import triton +import triton.language as tl + + +def commit_ssm_prefix( + indices: torch.Tensor, + reqs: torch.Tensor, + lengths: torch.Tensor, + replay_cursors: torch.Tensor | None = None, +) -> None: + """把每个请求的 accepted length 写入 MTP 索引并推进 Replay cursor。""" + if not reqs.is_cuda: + indices[reqs] = (lengths - 1).to(indices.dtype) + if replay_cursors is not None: + valid = reqs != replay_cursors.numel() - 1 + replay_cursors[reqs[valid]] += lengths[valid].to(replay_cursors.dtype) + return + if reqs.numel(): + _commit_prefix[(triton.cdiv(reqs.numel(), 256),)]( + indices, + reqs, + lengths, + replay_cursors, + reqs.numel(), + indices.numel() - 1, + reqs.stride(0), + lengths.stride(0), + 256, + num_warps=4, + ) + + +@triton.jit(do_not_specialize=["N"]) +def _commit_prefix( + Indices, + Reqs, + Lengths, + Cursors, + N, + HOLD: tl.constexpr, + REQ_STRIDE: tl.constexpr, + LENGTH_STRIDE: tl.constexpr, + BLOCK: tl.constexpr, +): + row = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + req = tl.load(Reqs + row * REQ_STRIDE, row < N, 0) + length = tl.load(Lengths + row * LENGTH_STRIDE, row < N, 0) + tl.store(Indices + req, length - 1, row < N) + if Cursors is not None: + valid = (row < N) & (req != HOLD) + cursor = tl.load(Cursors + req, valid, 0) + tl.store(Cursors + req, cursor + length, valid) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py new file mode 100644 index 0000000000..7ca84503c1 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/linear_att/ssm_autotune.py @@ -0,0 +1,180 @@ +"""使用共享 autotuner 调优单层 SSM update cycle。""" + +from __future__ import annotations + +import ast + +import torch +import triton + +from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType, autotune +from lightllm.utils.log_utils import init_logger + +logger = init_logger(__name__) + +_DEFAULT_REPLAY_CONFIG = {"BV": 32, "num_warps": 1, "num_stages": 3} + + +def resolve_run_config(ssm_backend, mode, q, k, v, a, b, a_log, bias, cu_seqlens) -> dict[str, int]: + """返回当前 batch 的固定或已调优 SSM kernel 配置。""" + if ssm_backend._config_is_fixed and ssm_backend.run_config is not None: + return ssm_backend.run_config + default = _replay_config(ssm_backend.run_config) + key = (tuple(static_key(ssm_backend, mode, q, cu_seqlens).items()), run_key(q, cu_seqlens)) + if key in ssm_backend.batch_configs: + return ssm_backend.batch_configs[key] + if not Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + return default + config = _replay_config(select_config(ssm_backend, mode, q, k, v, a, b, a_log, bias, cu_seqlens)) + # 每个 captured graph 保留自己的布局，即使其他 bucket 后续完成 warmup。 + ssm_backend.batch_configs[key] = config + logger.info(f"SSM {mode} update-cycle config: {config}, warmup batch: {key[1]}") + return config + + +def get_configs() -> list[dict[str, int]]: + """返回 ReplaySSM update cycle 的候选 Triton 配置。""" + return [ + {"BV": bv, "num_warps": warps, "num_stages": stages} + for bv, warps in [(8, 1), (16, 1), (16, 2), (32, 1), (32, 2), (32, 4), (64, 2), (64, 4), (128, 4)] + for stages in (1, 2, 3) + ] + + +def static_key(ssm_backend, mode, q, cu_seqlens) -> dict[str, object]: + """构造区分状态布局、数据类型和投影模式的静态调优键。""" + _, _, hv, kd, vd = ssm_backend.ssm_state.shape + # 保持文件名较短，同时区分 packed projection、TP 布局、状态精度、递推规则和 history 容量。 + key = { + "mode": mode, + "v": 7, + "shape": str((q.shape[-2], hv, ssm_backend.num_key_heads, kd, vd, ssm_backend.verify_width)), + "dtype": str((q.dtype, ssm_backend.ssm_state.dtype)), + "history": ssm_backend.capacity, + "gate": int(ssm_backend.kda), + "varlen": cu_seqlens is not None, + } + if ssm_backend.kda: + key["lower"] = ssm_backend.gate_lower_bound + key["history_dtype"] = str(ssm_backend.keys.dtype) + key["delta_dtype"] = str(ssm_backend.deltas.dtype) + key["projection"] = ( + ("kda_mma" if ssm_backend.kda else "gdn_mma") if ssm_backend.compute_mode == "precompute" else "inline" + ) + return key + + +def run_key(q, cu_seqlens) -> tuple[int, int]: + """返回调优使用的 token 数和请求数键。""" + tokens = q.shape[1 if cu_seqlens is not None else 0] + return tokens, cu_seqlens.numel() - 1 if cu_seqlens is not None else tokens + + +def rebuild_inputs( + ssm_backend, mode, q, k, v, a, b, a_log, bias, cu_seqlens=None, workload=None, run_config=None +) -> tuple[tuple[object, ...], dict[str, object]]: + """为一层 scratch 和同形状的紧凑激活分配调优输入。不复制服务状态池。""" + tokens, sequences = run_key(q, cu_seqlens) + batch = min(sequences, triton.cdiv(tokens, ssm_backend.verify_width)) if cu_seqlens is not None else tokens + _, _, hv, kd, vd = ssm_backend.ssm_state.shape + from lightllm.common.attention_state.bounded_att.conv_ssm_types import ConvSSMStateConfig + + # 计时只需要一层、本 batch 的槽位。正式执行仍用调用方原来的 state。 + scratch = type(ssm_backend)( + ConvSSMStateConfig( + logical_layer_ids=(0,), + global_k_heads=ssm_backend.num_key_heads, + global_v_heads=hv, + key_dim=kd, + value_dim=vd, + conv_width=1, + conv_dtype=q.dtype, + ssm_dtype=ssm_backend.ssm_state.dtype, + ), + batch + 1, + mtp_step=ssm_backend.verify_width - 1, + device=q.device, + ssm_replay_config={"cache_len": ssm_backend.capacity, "compute_mode": ssm_backend.compute_mode}, + gate_lower_bound=ssm_backend.gate_lower_bound, + ) + scratch.ssm_state.normal_(0, 0.02) + # benchmark 回调不能再次触发配置选择。 + scratch._config_is_fixed = True + inputs = [ + torch.empty(x.shape, device=x.device, dtype=x.dtype).normal_(0, 0.2) for x in (q, k, v, a, b, a_log, bias) + ] + reqs = torch.cat( + ( + torch.arange(batch, dtype=torch.int32, device=q.device), + torch.full((sequences - batch,), batch, dtype=torch.int32, device=q.device), + ) + ) + cu = (torch.arange(sequences + 1, dtype=torch.int32, device=q.device) * ssm_backend.verify_width).clamp_max(tokens) + lengths = cu[1 : batch + 1] - cu[:batch] + accepted = torch.cat((lengths - 1, lengths.new_zeros(1))) + return (), dict( + ssm_backend=scratch, + mode=mode, + q=inputs[0], + k=inputs[1], + v=inputs[2], + a=inputs[3], + b=inputs[4], + a_log=inputs[5], + bias=inputs[6], + cu_seqlens=cu if cu_seqlens is not None else None, + workload=(reqs, accepted), + ) + + +def replay_layer(history, *args, **kwargs): + """按 history 的具体类型调用对应的 replay decode。""" + + 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 + + if isinstance(history, ReplayKDAState): + from lightllm.common.basemodel.attention.linear.replay_kda_decode import replay + elif isinstance(history, ReplayGDNState): + from lightllm.common.basemodel.attention.linear.replay_gdn_decode import replay + else: + raise TypeError(f"replay_layer expects ReplayGDNState or ReplayKDAState, got {type(history).__name__}") + return replay(history, *args, **kwargs) + + +@autotune( + kernel_name="ssm_update_cycle:v2", + configs_gen_func=get_configs, + static_key_func=static_key, + run_key_func=run_key, + run_key_distance_func=lambda a, b: sum(abs(x - y) for x, y in zip(ast.literal_eval(a), ast.literal_eval(b))), + kernel_type=AutotuneKernelType.DECODE_ATTENTION, + rebuild_input_func=rebuild_inputs, + warmup_all_exist_config=False, +) +def select_config( + ssm_backend, mode, q, k, v, a, b, a_log, bias, cu_seqlens=None, workload=None, run_config=None +) -> dict[str, int] | None: + """搜索 update cycle 布局；计时只使用重建的临时输入。""" + if workload is None: + return run_config + ssm_backend.run_config = run_config + reqs, accepted = workload + ssm_backend.cursors.zero_() + # 把前缀写满但不越出 capacity。replay 自己读 cursor，正式 decode 也不经过 prepare_decode。 + tokens = q.shape[1 if cu_seqlens is not None else 0] + rounds = max(1, ssm_backend.capacity // min(ssm_backend.verify_width, tokens)) + for _ in range(rounds): + replay_layer(ssm_backend, 0, q, k, v, a, b, a_log, bias, reqs, cu_seqlens) + ssm_backend.accept_updates(reqs, accepted) + return run_config + + +def _replay_config(config: dict | None) -> dict[str, int]: + if not config: + return dict(_DEFAULT_REPLAY_CONFIG) + return { + "BV": int(config["BV"]), + "num_warps": int(config["num_warps"]), + "num_stages": int(config.get("num_stages", 3)), + } diff --git a/lightllm/common/basemodel/triton_kernel/mtp_utils.py b/lightllm/common/basemodel/triton_kernel/mtp_utils.py index 7525a754df..85ac266691 100644 --- a/lightllm/common/basemodel/triton_kernel/mtp_utils.py +++ b/lightllm/common/basemodel/triton_kernel/mtp_utils.py @@ -309,12 +309,11 @@ def mtp_scatter_next_token_ids( ) -@triton.jit +@triton.jit(do_not_specialize=["batch_size"]) def _fwd_kernel_gen_b_req_mtp_start_loc( b_mtp_index, b_req_mtp_start_loc, - num_reqs: tl.constexpr, - batch_size: tl.constexpr, + batch_size, BLOCK_SIZE: tl.constexpr, ): offset = tl.arange(0, BLOCK_SIZE) @@ -333,7 +332,6 @@ def gen_b_req_mtp_start_loc(b_mtp_index: torch.Tensor, num_reqs: int): _fwd_kernel_gen_b_req_mtp_start_loc[grid]( b_mtp_index=b_mtp_index, b_req_mtp_start_loc=b_req_mtp_start_loc, - num_reqs=num_reqs, batch_size=batch_size, BLOCK_SIZE=BLOCK_SIZE, num_warps=8, diff --git a/lightllm/common/basemodel/triton_kernel/scalar_linear.py b/lightllm/common/basemodel/triton_kernel/scalar_linear.py new file mode 100644 index 0000000000..e78678c40a --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/scalar_linear.py @@ -0,0 +1,42 @@ +import torch +import triton +import triton.language as tl + + +@triton.jit +def _scalar_linear(X, W, Y, K: tl.constexpr, XS0, XS1, WS, YS, BLOCK: tl.constexpr): + row = tl.program_id(0) + column = tl.arange(0, BLOCK) + value = tl.load(X + row * XS0 + column * XS1, column < K, other=0).to(tl.float64) + weight = tl.load(W + column * WS, column < K, other=0).to(tl.float64) + tl.store(Y + row * YS, tl.sum(value * weight, 0)) + + +def scalar_linear(input: torch.Tensor, weight: torch.Tensor, out=None, alloc_func=torch.empty) -> torch.Tensor: + if input.ndim != 2 or weight.shape != (1, input.shape[1]): + raise ValueError("Scalar projection requires (tokens, hidden) input and (1, hidden) weight") + if not input.is_cuda or weight.device != input.device or weight.dtype != input.dtype: + raise ValueError("Scalar projection requires CUDA input and matching weight dtype/device") + if not input.is_floating_point(): + raise ValueError("Scalar projection requires floating-point activations") + shape = (input.shape[0], 1) + if out is None: + out = alloc_func(shape, dtype=input.dtype, device=input.device) + elif out.shape != shape or out.dtype != input.dtype or out.device != input.device: + raise ValueError("Scalar projection output must match the unpadded shape, dtype and device") + if input.shape[0]: + block = triton.next_power_of_2(max(1, input.shape[1])) + _scalar_linear[(input.shape[0],)]( + input, + weight, + out, + input.shape[1], + input.stride(0), + input.stride(1), + weight.stride(1), + out.stride(0), + block, + num_warps=4 if block <= 8192 else 8, + enable_fp_fusion=False, + ) + return out diff --git a/lightllm/common/gated_residual/__init__.py b/lightllm/common/gated_residual/__init__.py new file mode 100644 index 0000000000..7f0da1f957 --- /dev/null +++ b/lightllm/common/gated_residual/__init__.py @@ -0,0 +1,85 @@ +from dataclasses import dataclass +from typing import Mapping, Optional + +import torch + +from .kernels import grouped_norm, inject_streams, injection_gates +from .projection import project_and_mix + + +@dataclass(frozen=True) +class GatedResidual: + norm_weight: torch.Tensor + down_weight: torch.Tensor + up_weight: torch.Tensor + inject_weight: Optional[torch.Tensor] + branches: int + eps: float = 1e-6 + + def __post_init__(self): + width = self.norm_weight.numel() + if self.branches < 2 or width == 0 or width % self.branches: + raise ValueError("Residual width must be a positive multiple of branches >= 2") + if self.norm_weight.shape != (width,) or self.down_weight.ndim != 2: + raise ValueError("Expected a norm vector and a matrix for the down projection") + rank = self.down_weight.shape[0] + if rank == 0 or self.down_weight.shape[1] != width or self.up_weight.shape != (width, rank): + raise ValueError("Residual projection dimensions do not match") + if self.inject_weight is not None and self.inject_weight.shape != (self.branches, width): + raise ValueError("Injection projection must have shape (branches, residual width)") + weights = [self.norm_weight, self.down_weight, self.up_weight] + if self.inject_weight is not None: + weights.append(self.inject_weight) + if self.eps <= 0: + raise ValueError("Normalization epsilon must be positive") + for weight in weights: + if not weight.is_cuda or not weight.is_contiguous(): + raise ValueError("Gated residual weights must be contiguous CUDA tensors") + if weight.device != self.norm_weight.device or weight.dtype != self.norm_weight.dtype: + raise ValueError("Gated residual weights must share a device and dtype") + if self.norm_weight.dtype not in (torch.float16, torch.bfloat16, torch.float32): + raise ValueError("Gated residuals require floating point weights") + + @classmethod + def from_checkpoint( + cls, + weights: Mapping[str, torch.Tensor], + prefix: str, + branches: int, + eps: float = 1e-6, + use_combine: bool = True, + ) -> "GatedResidual": + return cls( + weights[f"{prefix}.hc_norm.weight"], + weights[f"{prefix}.input_mix_weight_down.weight"], + weights[f"{prefix}.input_mix_weight_up.weight"], + weights[f"{prefix}.block_inject_weight.weight"] if use_combine else None, + branches, + eps, + ) + + def _check_input(self, x: torch.Tensor, width: int): + if x.ndim < 2 or x.shape[-1] != width or not x.is_contiguous(): + raise ValueError(f"Expected contiguous input with last dimension {width}") + if x.device != self.norm_weight.device or x.dtype != self.norm_weight.dtype: + raise ValueError("Input must share the weights' device and dtype") + + @torch.no_grad() + def read(self, residual: torch.Tensor): + self._check_input(residual, self.norm_weight.numel()) + normalized = grouped_norm(residual, self.norm_weight, self.branches, self.eps) + mixed = project_and_mix(normalized, self.down_weight, self.up_weight, self.branches) + if self.inject_weight is None: + return mixed + gates = injection_gates(normalized, self.inject_weight, self.branches) + return mixed, gates + + @torch.no_grad() + def write(self, residual: torch.Tensor, update: torch.Tensor, gates: torch.Tensor) -> torch.Tensor: + width = self.norm_weight.numel() + self._check_input(residual, width) + self._check_input(update, width // self.branches) + self._check_input(gates, self.branches) + if update.shape[:-1] != residual.shape[:-1] or gates.shape[:-1] != residual.shape[:-1]: + raise ValueError("Residual, update, and gates must have matching token dimensions") + return inject_streams(residual, update, gates) diff --git a/lightllm/common/gated_residual/kernels.py b/lightllm/common/gated_residual/kernels.py new file mode 100644 index 0000000000..96f7271db0 --- /dev/null +++ b/lightllm/common/gated_residual/kernels.py @@ -0,0 +1,103 @@ +import torch +import triton +import triton.language as tl + + +@triton.jit +def _norm(X, W, Y, D: tl.constexpr, C: tl.constexpr, EPS: tl.constexpr, B: tl.constexpr): + row = tl.program_id(0) + branch = tl.program_id(1) + col = tl.arange(0, B) + offset = row * C * D + branch * D + col + x = tl.load(X + offset, col < D, other=0).to(tl.float32) + weight = tl.load(W + branch * D + col, col < D, other=0).to(tl.float32) + inv = tl.rsqrt(tl.sum(x * x, 0) / D + EPS) + tl.store(Y + offset, (x * inv) * (1.0 + weight), col < D) + + +@triton.jit +def _mix(X, G, Y, D: tl.constexpr, C: tl.constexpr, BC: tl.constexpr, B: tl.constexpr): + row = tl.program_id(0) + col = tl.program_id(1) * B + tl.arange(0, B) + branch = tl.arange(0, BC) + offset = row * C * D + branch[:, None] * D + col[None, :] + mask = (branch[:, None] < C) & (col[None, :] < D) + x = tl.load(X + offset, mask, other=0).to(tl.float32) + logits = tl.load(G + offset, mask, other=0).to(tl.float32) + + gate = tl.sigmoid(logits).to(G.dtype.element_ty).to(tl.float32) + product = (x * gate).to(X.dtype.element_ty).to(tl.float32) + product = tl.where(branch[:, None] < C, product, 0.0) + result = tl.sum(product, 0) / C + tl.store(Y + row * D + col, result, col < D) + + +@triton.jit +def _gates(X, W, Out, D: tl.constexpr, C: tl.constexpr, B: tl.constexpr): + row = tl.program_id(0) + branch = tl.program_id(1) + col = tl.arange(0, B) + x = tl.load(X + row * D + col, col < D, other=0).to(tl.float32) + weight = tl.load(W + branch * D + col, col < D, other=0).to(tl.float32) + + logits = tl.sum(x * weight, 0).to(X.dtype.element_ty).to(tl.float32) + scaled = (logits / C).to(X.dtype.element_ty).to(tl.float32) + gate = tl.sigmoid(scaled).to(X.dtype.element_ty).to(tl.float32) + tl.store(Out + row * C + branch, 2 * gate) + + +@triton.jit +def _inject(R, X, G, Y, D: tl.constexpr, C: tl.constexpr, B: tl.constexpr): + row = tl.program_id(0) + branch = tl.program_id(1) + col = tl.program_id(2) * B + tl.arange(0, B) + offset = row * C * D + branch * D + col + x = tl.load(X + row * D + col, col < D, other=0).to(tl.float32) + gate = tl.load(G + row * C + branch).to(tl.float32) + residual = tl.load(R + offset, col < D, other=0).to(tl.float32) + + injection = (x * gate).to(R.dtype.element_ty).to(tl.float32) + tl.store(Y + offset, residual + injection, col < D) + + +def grouped_norm(x: torch.Tensor, weight: torch.Tensor, branches: int, eps: float) -> torch.Tensor: + output = torch.empty_like(x) + rows = x.numel() // x.shape[-1] + dim = x.shape[-1] // branches + if rows: + _norm[(rows, branches)]( + x, weight, output, dim, branches, eps, triton.next_power_of_2(dim), enable_fp_fusion=False + ) + return output + + +def mix_streams(x: torch.Tensor, logits: torch.Tensor, branches: int) -> torch.Tensor: + dim = x.shape[-1] // branches + output = torch.empty((*x.shape[:-1], dim), dtype=x.dtype, device=x.device) + rows = x.numel() // x.shape[-1] + if rows: + _mix[(rows, triton.cdiv(dim, 256))](x, logits, output, dim, branches, triton.next_power_of_2(branches), 256) + return output + + +def inject_streams(residual: torch.Tensor, update: torch.Tensor, gates: torch.Tensor) -> torch.Tensor: + output = torch.empty_like(residual) + branches = gates.shape[-1] + dim = update.shape[-1] + rows = update.numel() // dim + if rows: + _inject[(rows, branches, triton.cdiv(dim, 256))]( + residual, update, gates, output, dim, branches, 256, enable_fp_fusion=False + ) + return output + + +def injection_gates(x: torch.Tensor, weight: torch.Tensor, branches: int) -> torch.Tensor: + output = torch.empty((*x.shape[:-1], branches), dtype=x.dtype, device=x.device) + rows = x.numel() // x.shape[-1] + if rows: + block = triton.next_power_of_2(x.shape[-1]) + _gates[(rows, branches)]( + x, weight, output, x.shape[-1], branches, block, num_warps=8 if block >= 8192 else 4, enable_fp_fusion=False + ) + return output diff --git a/lightllm/common/gated_residual/projection.py b/lightllm/common/gated_residual/projection.py new file mode 100644 index 0000000000..8d498650b1 --- /dev/null +++ b/lightllm/common/gated_residual/projection.py @@ -0,0 +1,156 @@ +import torch +import torch.nn.functional as F +import triton +import triton.language as tl + +from .kernels import mix_streams + + +@triton.jit +def _activate(value, DTYPE: tl.constexpr, C: tl.constexpr): + value = value.to(DTYPE).to(tl.float32) + value = (value / C).to(DTYPE).to(tl.float32) + return value * tl.sigmoid(value) + + +@triton.jit +def _down( + X, + W, + P, + Y, + M: tl.constexpr, + N: tl.constexpr, + K: tl.constexpr, + C: tl.constexpr, + BM: tl.constexpr, + BN: tl.constexpr, + BK: tl.constexpr, + S: tl.constexpr, +): + r = tl.program_id(0) * BM + tl.arange(0, BM) + c = tl.program_id(1) * BN + tl.arange(0, BN) + split = tl.program_id(2) + part = tl.cdiv(K, BK * S) * BK + k = split * part + tl.arange(0, BK) + acc = tl.full((BM, BN), 0, tl.float32) + for _ in range(tl.cdiv(K, BK * S)): + x = tl.load(X + r[:, None] * K + k[None, :], (r[:, None] < M) & (k[None, :] < K), other=0) + w = tl.load(W + c[None, :] * K + k[:, None], (c[None, :] < N) & (k[:, None] < K), other=0) + acc = tl.dot(x, w, acc, input_precision="ieee") + k += BK + if S == 1: + tl.store( + Y + r[:, None] * N + c[None, :], _activate(acc, Y.dtype.element_ty, C), (r[:, None] < M) & (c[None, :] < N) + ) + else: + tl.store(P + (split * M + r[:, None]) * N + c[None, :], acc, (r[:, None] < M) & (c[None, :] < N)) + + +@triton.jit +def _merge(P, Y, M: tl.constexpr, N: tl.constexpr, C: tl.constexpr, S: tl.constexpr, B: tl.constexpr): + row = tl.program_id(0) + col = tl.program_id(1) * B + tl.arange(0, B) + split = tl.arange(0, S) + value = tl.load(P + (split[:, None] * M + row) * N + col[None, :], col[None, :] < N, other=0) + tl.store(Y + row * N + col, _activate(tl.sum(value, 0), Y.dtype.element_ty, C), col < N) + + +@triton.jit +def _up_mix( + X, + W, + Normalized, + Y, + M: tl.constexpr, + D: tl.constexpr, + K: tl.constexpr, + C: tl.constexpr, + BC: tl.constexpr, + BM: tl.constexpr, + BD: tl.constexpr, + BK: tl.constexpr, +): + row = tl.program_id(0) * BM + tl.arange(0, BM) + n = tl.arange(0, BC * BD) + branch = n // BD + col = tl.program_id(1) * BD + n % BD + weight_row = branch * D + col + k = tl.arange(0, BK) + acc = tl.full((BM, BC * BD), 0, tl.float32) + for _ in range(tl.cdiv(K, BK)): + x = tl.load(X + row[:, None] * K + k[None, :], (row[:, None] < M) & (k[None, :] < K), other=0) + w = tl.load( + W + weight_row[None, :] * K + k[:, None], + (branch[None, :] < C) & (col[None, :] < D) & (k[:, None] < K), + other=0, + ) + acc = tl.dot(x, w, acc, input_precision="ieee") + k += BK + logits = acc.to(X.dtype.element_ty).to(tl.float32) + gate = tl.sigmoid(logits).to(X.dtype.element_ty).to(tl.float32) + normalized = tl.load( + Normalized + row[:, None] * (C * D) + weight_row[None, :], + (row[:, None] < M) & (branch[None, :] < C) & (col[None, :] < D), + other=0, + ).to(tl.float32) + products = (normalized * gate).to(X.dtype.element_ty).to(tl.float32) + mixed = tl.sum(products.reshape(BM, BC, BD), 1) / C + out_col = tl.program_id(1) * BD + tl.arange(0, BD) + tl.store(Y + row[:, None] * D + out_col[None, :], mixed, (row[:, None] < M) & (out_col[None, :] < D)) + + +def _lowrank(x, w, branches): + m, k = x.shape + n = w.shape[0] + out = torch.empty((m, n), device=x.device, dtype=x.dtype) + if not m: + return out + + splits = 16 if k >= 8192 and n <= 512 else 1 + p = torch.empty((splits, m, n), device=x.device, dtype=torch.float32) if splits > 1 else out + bm = 16 if m <= 16 else 32 if m <= 64 else 64 if m < 128 else 128 + _down[(triton.cdiv(m, bm), triton.cdiv(n, 64), splits)]( + x, w, p, out, m, n, k, branches, bm, 64, 64, splits, num_warps=8 if bm == 128 else 4, enable_fp_fusion=False + ) + if splits > 1: + _merge[(m, triton.cdiv(n, 128))](p, out, m, n, branches, splits, 128, num_warps=4, enable_fp_fusion=False) + return out + + +def _mix(normalized, low, w, branches): + m, k = low.shape + dim = normalized.shape[1] // branches + out = torch.empty((m, dim), device=low.device, dtype=low.dtype) + if not m: + return out + bm = 16 if m <= 16 else 32 if m <= 64 else 64 if m < 128 else 128 + _up_mix[(triton.cdiv(m, bm), triton.cdiv(dim, 32))]( + low, + w, + normalized, + out, + m, + dim, + k, + branches, + triton.next_power_of_2(branches), + bm, + 32, + 64, + num_warps=8 if bm == 128 else 4, + enable_fp_fusion=False, + ) + return out + + +def project_and_mix( + normalized: torch.Tensor, down_weight: torch.Tensor, up_weight: torch.Tensor, branches: int +) -> torch.Tensor: + if normalized.dtype == torch.float32: + low = F.silu(F.linear(normalized, down_weight) / branches) + return mix_streams(normalized, F.linear(low, up_weight), branches) + rows = normalized.reshape(-1, normalized.shape[-1]) + low = _lowrank(rows, down_weight, branches) + mixed = _mix(rows, low, up_weight, branches) + return mixed.reshape((*normalized.shape[:-1], normalized.shape[-1] // branches)) diff --git a/lightllm/common/ngram_embedding/__init__.py b/lightllm/common/ngram_embedding/__init__.py new file mode 100644 index 0000000000..9b10a8e5a1 --- /dev/null +++ b/lightllm/common/ngram_embedding/__init__.py @@ -0,0 +1,132 @@ +from dataclasses import dataclass +from typing import Mapping, Optional + +import torch +import triton + +from .kernels import _hash, _history + + +@dataclass(frozen=True) +class NGramHasher: + multipliers: torch.Tensor + head_vocab_sizes: torch.Tensor + head_offsets: torch.Tensor + eos_token_id: int + + def __post_init__(self): + if self.multipliers.ndim != 1 or self.multipliers.numel() < 2: + raise ValueError("N-gram hashing requires at least two multipliers") + heads = self.head_vocab_sizes.numel() + if heads == 0 or heads % self.context_len or self.head_vocab_sizes.shape != (heads,): + raise ValueError("Head count must be a positive multiple of ngram_size - 1") + if self.head_offsets.shape != (heads,): + raise ValueError("Head offsets must match head vocabulary sizes") + for tensor in (self.multipliers, self.head_vocab_sizes, self.head_offsets): + if not tensor.is_cuda or not tensor.is_contiguous() or tensor.dtype != torch.int64: + raise ValueError("Hash buffers must be contiguous CUDA int64 tensors") + if tensor.device != self.multipliers.device: + raise ValueError("Hash buffers must share a device") + + sizes = self.head_vocab_sizes.tolist() + offsets = self.head_offsets.tolist() + previous_end = 0 + for size, offset in zip(sizes, offsets): + if size <= 0 or offset < previous_end or offset + size > torch.iinfo(torch.int64).max: + raise ValueError("Head vocabularies must have positive sizes and non-overlapping ordered offsets") + previous_end = offset + size + if not 0 <= self.eos_token_id <= torch.iinfo(torch.int64).max: + raise ValueError("EOS token ID must be a nonnegative int64 value") + + @property + def context_len(self) -> int: + return self.multipliers.numel() - 1 + + @property + def heads(self) -> int: + return self.head_vocab_sizes.numel() + + @classmethod + def from_checkpoint(cls, weights: Mapping[str, torch.Tensor], prefix: str, eos_token_id: int) -> "NGramHasher": + return cls( + weights[f"{prefix}.layer_multipliers"], + weights[f"{prefix}.ngram_heads_vocab_sizes"], + weights[f"{prefix}.ngram_heads_offsets"], + eos_token_id, + ) + + def empty_history(self, requests: int) -> torch.Tensor: + return torch.full( + (requests, self.context_len), self.eos_token_id, dtype=torch.int64, device=self.multipliers.device + ) + + def _validate(self, tokens: torch.Tensor, starts: torch.Tensor, history: torch.Tensor): + if tokens.ndim != 1 or tokens.dtype not in (torch.int32, torch.int64): + raise ValueError("Tokens must be a packed int32 or int64 vector") + if starts.ndim != 1 or starts.numel() == 0 or starts.dtype != torch.int32: + raise ValueError("Starts must be an int32 cumulative-length vector of length requests + 1") + requests = starts.numel() - 1 + if history.shape != (requests, self.context_len) or history.dtype != torch.int64: + raise ValueError("History must have shape (requests, ngram_size - 1) and dtype int64") + if requests == 0 and tokens.numel(): + raise ValueError("Nonempty tokens require at least one request") + for tensor in (tokens, starts, history): + if tensor.device != self.multipliers.device or not tensor.is_contiguous(): + raise ValueError("Inputs must be contiguous and on the hash buffers' device") + + @torch.no_grad() + def hash(self, tokens: torch.Tensor, starts: torch.Tensor, history: torch.Tensor): + self._validate(tokens, starts, history) + requests = starts.numel() - 1 + indices = torch.empty((tokens.numel(), self.heads), dtype=torch.int64, device=tokens.device) + if tokens.numel(): + _hash[(triton.cdiv(tokens.numel(), 64),)]( + tokens, + starts, + history, + self.multipliers, + self.head_vocab_sizes, + self.head_offsets, + indices, + tokens.numel(), + requests, + requests.bit_length(), + self.context_len, + self.heads, + self.heads // self.context_len, + self.eos_token_id, + 64, + triton.next_power_of_2(self.heads), + ) + return indices, self.history_after(tokens, starts, history) + + @torch.no_grad() + def history_after( + self, + tokens: torch.Tensor, + starts: torch.Tensor, + history: torch.Tensor, + accepted_lengths: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + self._validate(tokens, starts, history) + requests = starts.numel() - 1 + if accepted_lengths is not None: + if accepted_lengths.shape != (requests,) or accepted_lengths.dtype != torch.int32: + raise ValueError("Accepted lengths must be an int32 vector with one count per request") + if accepted_lengths.device != tokens.device or not accepted_lengths.is_contiguous(): + raise ValueError("Accepted lengths must be contiguous and on the input device") + output = torch.empty_like(history) + if requests: + _history[(requests,)]( + tokens, + starts, + history, + accepted_lengths, + output, + self.context_len, + accepted_lengths is not None, + self.eos_token_id, + triton.next_power_of_2(self.context_len), + num_warps=1, + ) + return output diff --git a/lightllm/common/ngram_embedding/kernels.py b/lightllm/common/ngram_embedding/kernels.py new file mode 100644 index 0000000000..f20f28c5b5 --- /dev/null +++ b/lightllm/common/ngram_embedding/kernels.py @@ -0,0 +1,94 @@ +import triton +import triton.language as tl + + +@triton.jit +def request_for_token(token, Starts, valid, REQUESTS: tl.constexpr, SEARCH_STEPS: tl.constexpr): + lo = tl.full(token.shape, 0, tl.int32) + hi = tl.full(token.shape, REQUESTS, tl.int32) + for _ in tl.static_range(SEARCH_STEPS): + mid = (lo + hi) // 2 + start = tl.load(Starts + mid, (mid < REQUESTS) & valid, other=2147483647) + right = start <= token + active = lo < hi + lo = tl.where(active & right, mid + 1, lo) + hi = tl.where(active & ~right, mid, hi) + request = tl.maximum(lo - 1, 0) + start = tl.load(Starts + request, valid, other=0) + return request, start + + +@triton.jit +def _hash( + Tokens, + Starts, + History, + Multipliers, + Sizes, + Offsets, + Output, + TOKENS: tl.constexpr, + REQUESTS: tl.constexpr, + SEARCH_STEPS: tl.constexpr, + CONTEXT: tl.constexpr, + HEADS: tl.constexpr, + HEADS_PER_NGRAM: tl.constexpr, + EOS: tl.constexpr, + BLOCK_T: tl.constexpr, + BLOCK_H: tl.constexpr, +): + token = tl.program_id(0) * BLOCK_T + tl.arange(0, BLOCK_T) + request, start = request_for_token(token, Starts, token < TOKENS, REQUESTS, SEARCH_STEPS) + position = token - start + head = tl.arange(0, BLOCK_H) + ngram_order = head // HEADS_PER_NGRAM + 2 + current = tl.load(Tokens + token, token < TOKENS, other=EOS).to(tl.int64) + multiplier = tl.load(Multipliers) + mixed = tl.broadcast_to((current * multiplier)[:, None], (BLOCK_T, BLOCK_H)) + alive = tl.full((BLOCK_T,), True, tl.int1) + for shift in tl.static_range(1, CONTEXT + 1): + in_chunk = position >= shift + source = tl.load(Tokens + token - shift, (token < TOKENS) & in_chunk, other=EOS).to(tl.int64) + from_history = tl.load( + History + request * CONTEXT + CONTEXT + position - shift, + (token < TOKENS) & ~in_chunk, + other=EOS, + ) + source = tl.where(in_chunk, source, from_history) + shifted = tl.where(alive, source, EOS) + multiplier = tl.load(Multipliers + shift) + mixed = mixed ^ tl.where(ngram_order[None, :] > shift, (shifted * multiplier)[:, None], 0) + + alive = alive & (source != EOS) + sizes = tl.load(Sizes + head, head < HEADS, other=1) + offsets = tl.load(Offsets + head, head < HEADS, other=0) + remainder = mixed % sizes[None, :] + remainder = tl.where(remainder < 0, remainder + sizes[None, :], remainder) + tl.store( + Output + token[:, None] * HEADS + head[None, :], + remainder + offsets[None, :], + (token[:, None] < TOKENS) & (head[None, :] < HEADS), + ) + + +@triton.jit +def _history( + Tokens, + Starts, + Previous, + Accepted, + Output, + CONTEXT: tl.constexpr, + USE_ACCEPTED: tl.constexpr, + EOS: tl.constexpr, + BLOCK: tl.constexpr, +): + request = tl.program_id(0) + col = tl.arange(0, BLOCK) + start = tl.load(Starts + request) + end = tl.load(Starts + request + 1) + count = tl.load(Accepted + request) if USE_ACCEPTED else end - start + source = count - CONTEXT + col + old = tl.load(Previous + request * CONTEXT + count + col, (col < CONTEXT) & (source < 0), other=EOS) + new = tl.load(Tokens + start + source, (col < CONTEXT) & (source >= 0), other=EOS).to(tl.int64) + tl.store(Output + request * CONTEXT + col, tl.where(source >= 0, new, old), col < CONTEXT) diff --git a/lightllm/common/ngram_embedding/ple.py b/lightllm/common/ngram_embedding/ple.py new file mode 100644 index 0000000000..042002cb8c --- /dev/null +++ b/lightllm/common/ngram_embedding/ple.py @@ -0,0 +1,198 @@ +from dataclasses import dataclass +from typing import Mapping, Optional + +import torch +import torch.nn.functional as F +import triton + +from lightllm.common.gated_residual.kernels import grouped_norm +from .ple_kernels import _gate, _conv, _conv_history + + +@dataclass(frozen=True) +class PLELayer: + key_weight: torch.Tensor + value_weight: torch.Tensor + norm_key: torch.Tensor + norm_query: torch.Tensor + norm_conv: torch.Tensor + conv_weight: torch.Tensor + branches: int + dilation: int + eps: float = 1e-6 + + def __post_init__(self): + width = self.norm_key.numel() + if self.branches < 1 or not width or width % self.branches or self.dilation < 1 or self.eps <= 0: + raise ValueError("PLE requires positive dimensions, dilation and epsilon") + if self.key_weight.ndim != 2 or self.key_weight.shape[1] == 0: + raise ValueError("PLE key projection must be a nonempty matrix") + if self.key_weight.shape[0] != width or self.value_weight.shape != (self.hidden_size, self.key_weight.shape[1]): + raise ValueError("PLE projection dimensions do not match residual/embedding widths") + if any(weight.shape != (width,) for weight in (self.norm_key, self.norm_query, self.norm_conv)): + raise ValueError("PLE normalization vectors must match the residual width") + if self.conv_weight.ndim != 2 or self.conv_weight.shape[0] != width or not self.conv_weight.shape[1]: + raise ValueError("PLE convolution weight must have shape (residual width, kernel size)") + for weight in ( + self.key_weight, + self.value_weight, + self.norm_key, + self.norm_query, + self.norm_conv, + self.conv_weight, + ): + if not weight.is_cuda or not weight.is_contiguous(): + raise ValueError("PLE weights must be contiguous CUDA tensors") + if weight.device != self.key_weight.device or weight.dtype != self.key_weight.dtype: + raise ValueError("PLE weights must share a device and dtype") + if self.key_weight.dtype not in (torch.float16, torch.bfloat16, torch.float32): + raise ValueError("PLE weights require a floating point dtype") + + @property + def hidden_size(self) -> int: + return self.norm_key.numel() // self.branches + + @property + def context_len(self) -> int: + return (self.conv_weight.shape[1] - 1) * self.dilation + + @classmethod + def from_checkpoint( + cls, weights: Mapping[str, torch.Tensor], prefix: str, branches: int, dilation: int, eps: float = 1e-6 + ) -> "PLELayer": + conv = weights[f"{prefix}.conv1d.weight"] + if conv.ndim != 3 or conv.shape[1] != 1: + raise ValueError("Checkpoint PLE convolution must be depthwise with shape (width, 1, kernel)") + return cls( + weights[f"{prefix}.key_proj.weight"], + weights[f"{prefix}.value_proj.weight"], + weights[f"{prefix}.norm_key.weight"], + weights[f"{prefix}.norm_query.weight"], + weights[f"{prefix}.norm_conv.weight"], + conv.squeeze(1), + branches, + dilation, + eps, + ) + + def empty_history(self, requests: int) -> torch.Tensor: + return torch.zeros( + (requests, self.context_len, self.norm_key.numel()), + dtype=self.key_weight.dtype, + device=self.key_weight.device, + ) + + def _check_tensor(self, tensor: torch.Tensor): + if ( + tensor.device != self.key_weight.device + or tensor.dtype != self.key_weight.dtype + or not tensor.is_contiguous() + ): + raise ValueError("PLE inputs must be contiguous and share the weights' device and dtype") + + def _validate_state(self, values: torch.Tensor, starts: torch.Tensor, history: torch.Tensor): + width = self.norm_key.numel() + self._check_tensor(values) + self._check_tensor(history) + if values.ndim != 2 or values.shape[1] != width: + raise ValueError("PLE values must have shape (packed tokens, residual width)") + if starts.ndim != 1 or starts.numel() == 0 or starts.dtype != torch.int32: + raise ValueError("PLE starts must be an int32 cumulative-length vector") + if starts.device != values.device or not starts.is_contiguous(): + raise ValueError("PLE starts must be contiguous and on the input device") + requests = starts.numel() - 1 + if history.shape != (requests, self.context_len, width): + raise ValueError("PLE history must have shape (requests, context length, residual width)") + if requests == 0 and values.shape[0]: + raise ValueError("Nonempty PLE input requires at least one request") + + @torch.no_grad() + def forward( + self, + residual: torch.Tensor, + embeddings: torch.Tensor, + starts: torch.Tensor, + history: torch.Tensor, + mask: Optional[torch.Tensor] = None, + ): + self._validate_state(residual, starts, history) + self._check_tensor(embeddings) + tokens = residual.shape[0] + if embeddings.shape != (tokens, self.key_weight.shape[1]): + raise ValueError("PLE embeddings must have shape (packed tokens, embedding width)") + if mask is not None: + if ( + mask.shape != (tokens,) + or mask.dtype != torch.bool + or mask.device != residual.device + or not mask.is_contiguous() + ): + raise ValueError("PLE mask must be a contiguous boolean vector on the input device") + key = grouped_norm(F.linear(embeddings, self.key_weight), self.norm_key, self.branches, self.eps) + query = grouped_norm(residual, self.norm_query, self.branches, self.eps) + value = F.linear(embeddings, self.value_weight) + gated = torch.empty_like(residual) + if tokens: + _gate[(tokens, self.branches)]( + key, + query, + value, + mask, + gated, + self.hidden_size, + self.branches, + mask is not None, + triton.next_power_of_2(self.hidden_size), + enable_fp_fusion=False, + ) + + conv_inputs = grouped_norm(gated, self.norm_conv, self.branches, self.eps) + output = torch.empty_like(residual) + if tokens: + requests = starts.numel() - 1 + _conv[(tokens, triton.cdiv(residual.shape[1], 256))]( + conv_inputs, + gated, + self.conv_weight, + starts, + history, + output, + residual.shape[1], + self.conv_weight.shape[1], + self.dilation, + requests, + requests.bit_length(), + 256, + enable_fp_fusion=False, + ) + return output, conv_inputs + + @torch.no_grad() + def history_after( + self, + conv_inputs: torch.Tensor, + starts: torch.Tensor, + history: torch.Tensor, + accepted_lengths: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + self._validate_state(conv_inputs, starts, history) + requests = starts.numel() - 1 + if accepted_lengths is not None: + if accepted_lengths.shape != (requests,) or accepted_lengths.dtype != torch.int32: + raise ValueError("PLE accepted lengths must be an int32 vector with one count per request") + if accepted_lengths.device != conv_inputs.device or not accepted_lengths.is_contiguous(): + raise ValueError("PLE accepted lengths must be contiguous and on the input device") + result = torch.empty_like(history) + if requests and self.context_len: + _conv_history[(requests, triton.cdiv(self.context_len * self.norm_key.numel(), 256))]( + conv_inputs, + starts, + history, + accepted_lengths, + result, + self.norm_key.numel(), + self.context_len, + accepted_lengths is not None, + 256, + ) + return result diff --git a/lightllm/common/ngram_embedding/ple_kernels.py b/lightllm/common/ngram_embedding/ple_kernels.py new file mode 100644 index 0000000000..79f7e13bac --- /dev/null +++ b/lightllm/common/ngram_embedding/ple_kernels.py @@ -0,0 +1,100 @@ +import triton +import triton.language as tl + +from .kernels import request_for_token + + +@triton.jit +def _gate(Key, Query, Value, Mask, Output, D: tl.constexpr, C: tl.constexpr, HAS_MASK: tl.constexpr, B: tl.constexpr): + token = tl.program_id(0).to(tl.int64) + branch = tl.program_id(1) + col = tl.arange(0, B) + offset = token * C * D + branch * D + col + key = tl.load(Key + offset, col < D, other=0).to(tl.float32) + query = tl.load(Query + offset, col < D, other=0).to(tl.float32) + + product = (key * query).to(Key.dtype.element_ty).to(tl.float32) + dot = tl.sum(product, 0).to(Key.dtype.element_ty).to(tl.float32) + dot = (dot / tl.sqrt(D * 1.0)).to(Key.dtype.element_ty).to(tl.float32) + magnitude = tl.maximum(tl.abs(dot), 1e-6).to(Key.dtype.element_ty).to(tl.float32) + root = tl.sqrt(magnitude).to(Key.dtype.element_ty).to(tl.float32) + signed = tl.where(dot > 0, root, tl.where(dot < 0, -root, 0.0)) + gate = tl.sigmoid(signed).to(Key.dtype.element_ty).to(tl.float32) + value = tl.load(Value + token * D + col, col < D, other=0).to(tl.float32) + result = value * gate + if HAS_MASK: + keep = tl.load(Mask + token) + result = tl.where(keep, result, 0.0) + tl.store(Output + offset, result, col < D) + + +@triton.jit +def _conv( + Input, + Gated, + Weight, + Starts, + History, + Output, + WIDTH: tl.constexpr, + KERNEL: tl.constexpr, + DILATION: tl.constexpr, + REQUESTS: tl.constexpr, + SEARCH_STEPS: tl.constexpr, + BLOCK: tl.constexpr, +): + token = tl.program_id(0).to(tl.int64) + col = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + request, start = request_for_token(token, Starts, True, REQUESTS, SEARCH_STEPS) + context: tl.constexpr = (KERNEL - 1) * DILATION + position = token - start + acc = tl.full((BLOCK,), 0, tl.float32) + for tap in tl.static_range(KERNEL): + source = position - context + tap * DILATION + current = tl.load(Input + (start + source) * WIDTH + col, (col < WIDTH) & (source >= 0), other=0) + previous = tl.load( + History + (request.to(tl.int64) * context + context + source) * WIDTH + col, + (col < WIDTH) & (source < 0), + other=0, + ) + value = tl.where(source >= 0, current, previous).to(tl.float32) + weight = tl.load(Weight + col * KERNEL + tap, col < WIDTH, other=0).to(tl.float32) + acc = acc + value * weight + + conv = acc.to(Input.dtype.element_ty).to(tl.float32) + activated = (conv * tl.sigmoid(conv)).to(Input.dtype.element_ty).to(tl.float32) + gated = tl.load(Gated + token * WIDTH + col, col < WIDTH, other=0).to(tl.float32) + tl.store(Output + token * WIDTH + col, gated + activated, col < WIDTH) + + +@triton.jit +def _conv_history( + Input, + Starts, + Previous, + Accepted, + Output, + WIDTH: tl.constexpr, + CONTEXT: tl.constexpr, + USE_ACCEPTED: tl.constexpr, + BLOCK: tl.constexpr, +): + request = tl.program_id(0).to(tl.int64) + flat = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + time = flat // WIDTH + col = flat % WIDTH + start = tl.load(Starts + request) + end = tl.load(Starts + request + 1) + count = tl.load(Accepted + request) if USE_ACCEPTED else end - start + source = count - CONTEXT + time + old = tl.load( + Previous + (request * CONTEXT + count + time) * WIDTH + col, + (time < CONTEXT) & (source < 0), + other=0, + ) + new = tl.load( + Input + (start + source.to(tl.int64)) * WIDTH + col, + (time < CONTEXT) & (source >= 0), + other=0, + ) + tl.store(Output + request * CONTEXT * WIDTH + flat, tl.where(source >= 0, new, old), time < CONTEXT) diff --git a/lightllm/common/qsa/__init__.py b/lightllm/common/qsa/__init__.py new file mode 100644 index 0000000000..2ae28399f5 --- /dev/null +++ b/lightllm/common/qsa/__init__.py @@ -0,0 +1 @@ +pass diff --git a/lightllm/common/qsa/attention.py b/lightllm/common/qsa/attention.py new file mode 100644 index 0000000000..91815a0d63 --- /dev/null +++ b/lightllm/common/qsa/attention.py @@ -0,0 +1,120 @@ +from typing import Optional + +import torch +import triton + +from .attention_kernels import _attention, _reduce + + +@torch.no_grad() +def indexed_attention( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + indices: torch.Tensor, + scale: Optional[float] = None, + splits: Optional[int] = None, + selection_lengths: Optional[torch.Tensor] = None, + short_rows_direct: bool = False, + skip_empty_splits: bool = False, + block_size: int = 64, + num_stages: Optional[int] = None, +) -> torch.Tensor: + if query.ndim != 3 or key.ndim != 3 or value.ndim != 3 or indices.ndim != 2: + raise ValueError("Indexed attention requires rank-three Q/K/V and rank-two indices") + rows, q_heads, dim = query.shape + if not q_heads or not dim or not key.shape[1] or q_heads % key.shape[1]: + raise ValueError("Indexed attention query heads must be divisible by KV heads") + if key.shape[0:2] != value.shape[0:2] or key.shape[2] != dim or not value.shape[2]: + raise ValueError("Indexed attention key/value dimensions do not match") + for tensor in (query, key, value): + if not tensor.is_cuda or tensor.dtype != query.dtype or tensor.device != query.device: + raise ValueError("Indexed attention Q/K/V must be CUDA tensors with matching dtype/device") + if not query.is_contiguous(): + raise ValueError("Indexed attention query must be contiguous") + if query.dtype not in (torch.float16, torch.bfloat16, torch.float32): + raise ValueError("Indexed attention requires floating point inputs") + if ( + indices.shape[0] != rows + or indices.dtype != torch.int32 + or indices.device != query.device + or not indices.is_contiguous() + ): + raise ValueError("Indexed attention indices must be a contiguous int32 matrix on the query device") + selected = indices.shape[1] + if selection_lengths is not None and ( + selection_lengths.shape != (rows,) + or selection_lengths.dtype != torch.int32 + or selection_lengths.device != query.device + or not selection_lengths.is_contiguous() + ): + raise ValueError("Selection lengths must be a contiguous int32 vector on the query device") + kv_heads, value_dim = value.shape[1:] + output = torch.empty((rows, q_heads, value_dim), dtype=query.dtype, device=query.device) + if not rows: + return output + if not selected: + return output.zero_() + if block_size < 16 or block_size & (block_size - 1): + raise ValueError("Indexed attention block size must be a power of two >= 16") + if num_stages is not None and num_stages < 1: + raise ValueError("Indexed attention num_stages must be positive") + if splits is None: + splits = min(32, max(1, triton.cdiv(block_size, rows * kv_heads)), triton.cdiv(selected, block_size)) + if splits < 1: + raise ValueError("Indexed attention split count must be positive") + partial = ( + torch.empty((rows, q_heads, splits, value_dim), dtype=torch.float32, device=query.device) + if splits > 1 + else None + ) + lse = torch.empty((rows, q_heads, splits), dtype=torch.float32, device=query.device) if splits > 1 else None + launch_kwargs = {} + if num_stages is not None: + launch_kwargs["num_stages"] = num_stages + _attention[(rows, kv_heads, splits)]( + query, + key, + value, + indices, + selection_lengths, + partial, + lse, + output, + q_heads, + kv_heads, + *key.stride(), + *value.stride(), + q_heads // kv_heads, + dim, + value_dim, + selected, + splits, + dim ** -0.5 if scale is None else scale, + max(16, triton.next_power_of_2(q_heads // kv_heads)), + max(16, triton.next_power_of_2(dim)), + max(16, triton.next_power_of_2(value_dim)), + block_size, + short_rows_direct, + skip_empty_splits, + "ieee" if query.dtype == torch.float32 else "tf32", + **launch_kwargs, + ) + if splits > 1: + _reduce[(rows, q_heads)]( + selection_lengths, + partial, + lse, + output, + q_heads, + value_dim, + splits, + triton.next_power_of_2(splits), + triton.next_power_of_2(value_dim), + selected, + block_size, + short_rows_direct, + skip_empty_splits, + **launch_kwargs, + ) + return output diff --git a/lightllm/common/qsa/attention_kernels.py b/lightllm/common/qsa/attention_kernels.py new file mode 100644 index 0000000000..7ce95104b0 --- /dev/null +++ b/lightllm/common/qsa/attention_kernels.py @@ -0,0 +1,137 @@ +import triton +import triton.language as tl + + +@triton.jit +def _attention( + Q, + K, + V, + Indices, + Lengths, + Partial, + LSE, + Output, + Q_HEADS: tl.constexpr, + KV_HEADS: tl.constexpr, + KS: tl.constexpr, + KH: tl.constexpr, + KD: tl.constexpr, + VS: tl.constexpr, + VH: tl.constexpr, + VD: tl.constexpr, + GROUP: tl.constexpr, + D: tl.constexpr, + DV: tl.constexpr, + SELECTED: tl.constexpr, + SPLITS: tl.constexpr, + SCALE: tl.constexpr, + BG: tl.constexpr, + BD: tl.constexpr, + BV: tl.constexpr, + BN: tl.constexpr, + SHORT_DIRECT: tl.constexpr, + SKIP_EMPTY: tl.constexpr, + INPUT_PRECISION: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + kv_head = tl.program_id(1) + split = tl.program_id(2) + limit = SELECTED + if Lengths is not None: + limit = tl.minimum(SELECTED, tl.maximum(0, tl.load(Lengths + row))) + if (not SHORT_DIRECT or limit > BN or split == 0) and ( + not SKIP_EMPTY or SPLITS == 1 or split * BN < limit or (SHORT_DIRECT and limit <= BN and split == 0) + ): + group = tl.arange(0, BG) + head = kv_head * GROUP + group + dim = tl.arange(0, BD) + value_dim = tl.arange(0, BV) + q = tl.load( + Q + (row * Q_HEADS + head[:, None]) * D + dim[None, :], + (group[:, None] < GROUP) & (dim[None, :] < D), + other=0, + ) + maximum = tl.full((BG,), -float("inf"), tl.float32) + denominator = tl.full((BG,), 0, tl.float32) + acc = tl.full((BG, BV), 0, tl.float32) + for start in tl.range(split * BN, limit, SPLITS * BN): + selected = start + tl.arange(0, BN) + slot = tl.load(Indices + row * SELECTED + selected, selected < limit, other=-1).to(tl.int64) + valid = (selected < limit) & (slot >= 0) + key = tl.load( + K + slot[None, :] * KS + kv_head * KH + dim[:, None] * KD, valid[None, :] & (dim[:, None] < D), other=0 + ) + score = tl.dot(q, key, input_precision=INPUT_PRECISION) * SCALE + score = tl.where(valid[None, :], score, -float("inf")) + new_max = tl.maximum(maximum, tl.max(score, 1)) + + safe_max = tl.where(new_max == -float("inf"), 0.0, new_max) + alpha = tl.exp(maximum - safe_max) + prob = tl.exp(score - safe_max[:, None]) + value = tl.load( + V + slot[:, None] * VS + kv_head * VH + value_dim[None, :] * VD, + valid[:, None] & (value_dim[None, :] < DV), + other=0, + ) + acc = acc * alpha[:, None] + acc = tl.dot(prob.to(V.dtype.element_ty), value, acc, input_precision=INPUT_PRECISION) + denominator = denominator * alpha + tl.sum(prob, 1) + maximum = new_max + safe_denominator = tl.where(denominator > 0, denominator, 1.0) + result = acc / safe_denominator[:, None] + if SPLITS == 1 or (SHORT_DIRECT and limit <= BN): + tl.store( + Output + (row * Q_HEADS + head[:, None]) * DV + value_dim[None, :], + result, + (group[:, None] < GROUP) & (value_dim[None, :] < DV), + ) + else: + tl.store( + Partial + ((row * Q_HEADS + head[:, None]) * SPLITS + split) * DV + value_dim[None, :], + result, + (group[:, None] < GROUP) & (value_dim[None, :] < DV), + ) + lse = tl.where(denominator > 0, maximum + tl.log(safe_denominator), -float("inf")) + tl.store(LSE + (row * Q_HEADS + head) * SPLITS + split, lse, group < GROUP) + + +@triton.jit +def _reduce( + Lengths, + Partial, + LSE, + Output, + HEADS: tl.constexpr, + DV: tl.constexpr, + SPLITS: tl.constexpr, + BS: tl.constexpr, + BV: tl.constexpr, + SELECTED: tl.constexpr, + BN: tl.constexpr, + SHORT_DIRECT: tl.constexpr, + SKIP_EMPTY: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + limit = SELECTED + if Lengths is not None: + limit = tl.minimum(SELECTED, tl.maximum(0, tl.load(Lengths + row))) + if not SHORT_DIRECT or limit > BN: + head = tl.program_id(1) + split = tl.arange(0, BS) + dim = tl.arange(0, BV) + valid_split = split < SPLITS + if SKIP_EMPTY: + valid_split = valid_split & (split * BN < limit) + lse = tl.load(LSE + (row * HEADS + head) * SPLITS + split, valid_split, other=-float("inf")) + maximum = tl.max(lse, 0) + maximum = tl.where(maximum == -float("inf"), 0.0, maximum) + weights = tl.exp(lse - maximum) + partial = tl.load( + Partial + ((row * HEADS + head) * SPLITS + split[:, None]) * DV + dim[None, :], + valid_split[:, None] & (dim[None, :] < DV), + other=0, + ) + total = tl.sum(weights, 0) + result = tl.sum(partial * weights[:, None], 0) / tl.where(total > 0, total, 1.0) + tl.store(Output + (row * HEADS + head) * DV + dim, result, dim < DV) diff --git a/lightllm/common/qsa/indexer.py b/lightllm/common/qsa/indexer.py new file mode 100644 index 0000000000..eda8a9158b --- /dev/null +++ b/lightllm/common/qsa/indexer.py @@ -0,0 +1,279 @@ +from dataclasses import dataclass +from typing import Mapping, Optional + +import torch +import torch.nn.functional as F +import triton + +from lightllm.common.gated_residual.kernels import grouped_norm +from .indexer_kernels import _rotate, _pool, _scores, _expand +from .topk import cuda_topk_indices, cuda_topk_supported, flashinfer_topk_indices, flashinfer_topk_supported + + +@dataclass(frozen=True) +class QSAIndexer: + projection: torch.Tensor + query_norm: torch.Tensor + key_norm: torch.Tensor + query_heads: int = 4 + token_budget: int = 2048 + compress_ratio: int = 4 + eps: float = 1e-6 + use_cuda_topk: bool = False + + def __post_init__(self): + dim = self.key_norm.numel() + if not dim or self.query_heads < 1 or self.compress_ratio < 1 or self.token_budget < 1 or self.eps <= 0: + raise ValueError("QSA dimensions, budget, compression ratio and epsilon must be positive") + if self.token_budget % self.compress_ratio: + raise ValueError("QSA token budget must be divisible by compression ratio") + if self.key_norm.shape != (dim,) or self.query_norm.shape != (dim,): + raise ValueError("QSA normalization weights must be matching vectors") + if ( + self.projection.ndim != 2 + or self.projection.shape[0] != (self.query_heads + 1) * dim + or not self.projection.shape[1] + ): + raise ValueError("QSA projection must contain query heads and one raw key head") + for weight in (self.projection, self.query_norm, self.key_norm): + self._activation(weight) + + @property + def head_dim(self) -> int: + return self.key_norm.numel() + + @classmethod + def from_checkpoint(cls, weights: Mapping[str, torch.Tensor], prefix: str, **kwargs) -> "QSAIndexer": + return cls( + weights[f"{prefix}.index_qk_proj.weight"], + weights[f"{prefix}.q_layernorm.weight"], + weights[f"{prefix}.k_layernorm.weight"], + **kwargs, + ) + + def _activation(self, x: torch.Tensor): + if ( + not x.is_cuda + or not x.is_contiguous() + or x.device != self.projection.device + or x.dtype != self.projection.dtype + ): + raise ValueError("QSA activations/weights must be contiguous CUDA tensors on the same device and dtype") + if x.dtype not in (torch.float16, torch.bfloat16, torch.float32): + raise ValueError("QSA requires floating point activations") + + def _integer(self, x: torch.Tensor): + if x.dtype != torch.int32 or x.device != self.projection.device or not x.is_contiguous(): + raise ValueError("QSA metadata must be contiguous int32 tensors on the input device") + + def _rotary(self, cos: torch.Tensor, sin: torch.Tensor, rows: int) -> int: + self._activation(cos) + self._activation(sin) + if cos.ndim != 2 or sin.shape != cos.shape or cos.shape[0] != rows: + raise ValueError("QSA rotary tensors must have matching shapes (tokens, rotary dimensions)") + if cos.shape[1] > self.head_dim or cos.shape[1] % 2: + raise ValueError("QSA rotary dimensions must be even and fit the index head") + return cos.shape[1] + + @torch.no_grad() + def project(self, hidden: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor): + self._activation(hidden) + if hidden.ndim != 2 or hidden.shape[1] != self.projection.shape[1]: + raise ValueError("QSA hidden input must have shape (tokens, hidden size)") + rotary_dim = self._rotary(cos, sin, hidden.shape[0]) + projected = F.linear(hidden, self.projection) + query = projected[:, : self.query_heads * self.head_dim].contiguous().view(-1, self.head_dim) + query = grouped_norm(query, self.query_norm, 1, self.eps).view(hidden.shape[0], self.query_heads, self.head_dim) + output = torch.empty_like(query) + if hidden.shape[0]: + _rotate[(hidden.shape[0], self.query_heads)]( + query, + cos, + sin, + output, + self.query_heads, + self.head_dim, + rotary_dim, + triton.next_power_of_2(self.head_dim), + enable_fp_fusion=False, + ) + # Keep the raw-key projection as a view. The state writer accepts a + # strided source, so making this slice contiguous only adds a full + # token-by-head-dimension copy on every QSA layer. + return output, projected[:, self.query_heads * self.head_dim :] + + def empty_cache(self, requests: int, max_tokens: int) -> torch.Tensor: + return torch.empty( + (requests, max_tokens // self.compress_ratio, self.head_dim), + dtype=self.projection.dtype, + device=self.projection.device, + ) + + def empty_physical_cache(self, slots: int) -> torch.Tensor: + if slots < 0: + raise ValueError("QSA physical cache size cannot be negative") + return torch.empty((slots, self.head_dim), dtype=self.projection.dtype, device=self.projection.device) + + @torch.no_grad() + def pool( + self, + raw_keys: torch.Tensor, + indices: torch.Tensor, + lengths: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + cache: torch.Tensor, + first_blocks: torch.Tensor, + max_new_blocks: int, + request_rows: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + self._activation(raw_keys) + self._activation(cache) + for tensor in (indices, lengths, first_blocks): + self._integer(tensor) + if raw_keys.ndim != 2 or raw_keys.shape[1] != self.head_dim or indices.ndim != 2: + raise ValueError("QSA raw keys and slot indices must be matrices") + requests, capacity = indices.shape + active = requests + if request_rows is not None: + self._integer(request_rows) + if request_rows.ndim != 1: + raise ValueError("QSA active request rows must be a vector") + active = request_rows.numel() + if lengths.shape != (active,) or first_blocks.shape != (active,): + raise ValueError("QSA lengths and first blocks must have one entry per request") + physical = cache.ndim == 2 + if physical: + if cache.shape != raw_keys.shape: + raise ValueError("QSA physical pooled cache must match the raw key namespace") + elif ( + cache.ndim != 3 + or cache.shape[0] != requests + or cache.shape[1] < capacity // self.compress_ratio + or cache.shape[2] != self.head_dim + ): + raise ValueError("QSA cache must hold all complete blocks in each request's slot table") + if max_new_blocks < 0: + raise ValueError("QSA maximum new block count cannot be negative") + rotary_dim = self._rotary(cos, sin, raw_keys.shape[0]) + if active and max_new_blocks: + _pool[(active, max_new_blocks)]( + raw_keys, + indices, + lengths, + first_blocks, + self.key_norm, + cos, + sin, + cache, + request_rows, + request_rows is not None, + physical, + capacity, + cache.shape[1], + self.head_dim, + rotary_dim, + self.compress_ratio, + self.eps, + triton.next_power_of_2(self.head_dim), + triton.next_power_of_2(self.compress_ratio), + enable_fp_fusion=False, + ) + return cache + + @torch.no_grad() + def select( + self, + query: torch.Tensor, + cache: torch.Tensor, + requests: torch.Tensor, + visible_lengths: torch.Tensor, + indices: torch.Tensor, + query_chunk_size: int = 64, + ) -> torch.Tensor: + self._activation(query) + self._activation(cache) + for tensor in (requests, visible_lengths, indices): + self._integer(tensor) + if query.ndim != 3 or query.shape[1:] != (self.query_heads, self.head_dim): + raise ValueError("QSA query shape must be (tokens, index query heads, index head dimension)") + tokens = query.shape[0] + if requests.shape != (tokens,) or visible_lengths.shape != (tokens,) or indices.ndim != 2: + raise ValueError("QSA requires one request ID and visible length per query") + physical = cache.ndim == 2 + if physical: + if cache.shape[1] != self.head_dim: + raise ValueError("QSA physical pooled cache has the wrong head dimension") + elif ( + cache.ndim != 3 + or cache.shape[0] != indices.shape[0] + or cache.shape[2] != self.head_dim + or cache.shape[1] < indices.shape[1] // self.compress_ratio + ): + raise ValueError("QSA pooled cache and physical slot table dimensions do not match") + if query_chunk_size < 1: + raise ValueError("QSA query chunk size must be positive") + width = self.token_budget + self.compress_ratio - 1 + output = torch.empty((tokens, width), dtype=torch.int32, device=query.device) + blocks = min(indices.shape[1], cache.shape[0]) // self.compress_ratio if physical else cache.shape[1] + block_budget = self.token_budget // self.compress_ratio + topk = min(blocks, block_budget) + cuda_topk = self.use_cuda_topk and cuda_topk_supported(query.device, blocks, topk) + flashinfer_topk = self.use_cuda_topk and not cuda_topk and flashinfer_topk_supported(query.device, blocks, topk) + max_score_elements = 16 * 1024 * 1024 + query_chunk_size = min(query_chunk_size, max(1, max_score_elements // max(blocks, 1))) + score_buffer = torch.empty((min(tokens, query_chunk_size), blocks), dtype=torch.float32, device=query.device) + group_size = 4 if physical and self.query_heads == 4 else 1 + score_block = 128 if group_size == 4 else 64 + for first in range(0, tokens, query_chunk_size): + count = min(query_chunk_size, tokens - first) + scores = score_buffer[:count] + if blocks: + _scores[(triton.cdiv(count, group_size), triton.cdiv(blocks, score_block))]( + query, + cache, + requests, + visible_lengths, + indices, + scores, + first, + self.query_heads, + self.head_dim, + blocks, + indices.shape[1], + physical, + self.compress_ratio, + max(16, triton.next_power_of_2(self.query_heads)), + max(16, triton.next_power_of_2(self.head_dim)), + score_block, + block_budget if cuda_topk or flashinfer_topk else 0, + "ieee" if query.dtype == torch.float32 else "tf32", + GROUP_SIZE=group_size, + TOKENS=count, + num_warps=4, + ) + if cuda_topk: + selected = cuda_topk_indices( + scores, topk, visible_lengths=visible_lengths, first=first, compress_ratio=self.compress_ratio + ) + elif flashinfer_topk: + selected = flashinfer_topk_indices( + scores, topk, visible_lengths=visible_lengths, first=first, compress_ratio=self.compress_ratio + ) + else: + selected = torch.topk(scores, topk, dim=-1, sorted=False).indices + _expand[(count,)]( + selected, + requests, + visible_lengths, + indices, + output, + first, + topk, + block_budget, + indices.shape[1], + self.compress_ratio, + width, + triton.next_power_of_2(width), + ) + return output diff --git a/lightllm/common/qsa/indexer_kernels.py b/lightllm/common/qsa/indexer_kernels.py new file mode 100644 index 0000000000..60a98de0dc --- /dev/null +++ b/lightllm/common/qsa/indexer_kernels.py @@ -0,0 +1,216 @@ +import triton +import triton.language as tl + + +@triton.jit +def _rotate(X, Cos, Sin, Output, HEADS: tl.constexpr, D: tl.constexpr, R: tl.constexpr, BLOCK: tl.constexpr): + row = tl.program_id(0).to(tl.int64) + head = tl.program_id(1) + col = tl.arange(0, BLOCK) + x = tl.load(X + (row * HEADS + head) * D + col, col < D, other=0) + if R > 0: + partner = tl.where(col < R // 2, col + R // 2, col - R // 2) + other = tl.load(X + (row * HEADS + head) * D + partner, col < R, other=0).to(tl.float32) + other = tl.where(col < R // 2, -other, other) + cos = tl.load(Cos + row * R + col, col < R, other=0).to(tl.float32) + sin = tl.load(Sin + row * R + col, col < R, other=0).to(tl.float32) + left = (x.to(tl.float32) * cos).to(X.dtype.element_ty).to(tl.float32) + right = (other * sin).to(X.dtype.element_ty).to(tl.float32) + x = tl.where(col < R, left + right, x.to(tl.float32)) + tl.store(Output + (row * HEADS + head) * D + col, x, col < D) + + +@triton.jit +def _pool( + Raw, + Indices, + Lengths, + First, + Norm, + Cos, + Sin, + Cache, + RequestRows, + MAPPED_ROWS: tl.constexpr, + PHYSICAL_CACHE: tl.constexpr, + CAPACITY: tl.constexpr, + BLOCKS: tl.constexpr, + D: tl.constexpr, + R: tl.constexpr, + RATIO: tl.constexpr, + EPS: tl.constexpr, + BD: tl.constexpr, + BR: tl.constexpr, +): + batch_row = tl.program_id(0) + request = tl.load(RequestRows + batch_row).to(tl.int64) if MAPPED_ROWS else batch_row.to(tl.int64) + block = tl.load(First + batch_row) + tl.program_id(1) + count = tl.load(Lengths + batch_row) + if (request >= 0) & (block < count // RATIO): + token = tl.arange(0, BR) + col = tl.arange(0, BD) + slots = tl.load(Indices + request * CAPACITY + block * RATIO + token, token < RATIO, other=0).to(tl.int64) + values = tl.load( + Raw + slots[:, None] * D + col[None, :], (token[:, None] < RATIO) & (col[None, :] < D), other=0 + ).to(tl.float32) + pooled = (tl.sum(values, 0) / RATIO).to(Raw.dtype.element_ty).to(tl.float32) + inv = tl.rsqrt(tl.sum(pooled * pooled, 0) / D + EPS) + weight = tl.load(Norm + col, col < D, other=0).to(tl.float32) + normalized = ((pooled * inv) * (1 + weight)).to(Raw.dtype.element_ty).to(tl.float32) + if R > 0: + slot = tl.load(Indices + request * CAPACITY + block * RATIO).to(tl.int64) + partner = tl.where(col < R // 2, col + R // 2, col - R // 2) + + partner = tl.minimum(tl.maximum(partner, 0), BD - 1) + other = tl.gather(normalized, partner, 0) + other = tl.where(col < R // 2, -other, other) + cos = tl.load(Cos + slot * R + col, col < R, other=0).to(tl.float32) + sin = tl.load(Sin + slot * R + col, col < R, other=0).to(tl.float32) + left = (normalized * cos).to(Raw.dtype.element_ty).to(tl.float32) + right = (other * sin).to(Raw.dtype.element_ty).to(tl.float32) + normalized = tl.where(col < R, left + right, normalized) + if PHYSICAL_CACHE: + destination = tl.load(Indices + request * CAPACITY + block * RATIO + RATIO - 1).to(tl.int64) + else: + destination = request * BLOCKS + block + tl.store(Cache + destination * D + col, normalized, col < D) + + +@triton.jit +def _scores( + Query, + Cache, + Requests, + Lengths, + Indices, + Output, + QUERY_START: tl.constexpr, + HEADS: tl.constexpr, + D: tl.constexpr, + BLOCKS: tl.constexpr, + CAPACITY: tl.constexpr, + PHYSICAL_CACHE: tl.constexpr, + RATIO: tl.constexpr, + BH: tl.constexpr, + BD: tl.constexpr, + BN: tl.constexpr, + MIN_BLOCKS: tl.constexpr = 0, + INPUT_PRECISION: tl.constexpr = "tf32", + GROUP_SIZE: tl.constexpr = 1, + TOKENS: tl.constexpr = 0, +): + tl.static_assert(GROUP_SIZE == 1 or (GROUP_SIZE == 4 and HEADS == 4 and PHYSICAL_CACHE)) + group = tl.program_id(0) + if GROUP_SIZE == 4: + rows = group * 4 + tl.arange(0, 4) + valid = rows < TOKENS + requests = tl.load(Requests + QUERY_START + rows, valid, 0).to(tl.int64) + lengths = tl.load(Lengths + QUERY_START + rows, valid, 0) // RATIO + shared = tl.min(tl.where(valid, requests, 0x7FFFFFFF)) == tl.max(tl.where(valid, requests, 0)) + else: + shared = tl.full((), False, tl.int1) + if GROUP_SIZE == 4 and shared: + # Four queries fill the MMA rows and share the same pooled keys. + request = tl.sum(tl.where(tl.arange(0, 4) == 0, requests, 0)) + maximum = tl.max(lengths) + block = tl.program_id(1) * BN + tl.arange(0, BN) + scores = tl.full((4, BN), -float("inf"), tl.float32) + if (tl.program_id(1) * BN < maximum) & (maximum >= MIN_BLOCKS): + head = tl.arange(0, 16) + col = tl.arange(0, BD) + query_rows = QUERY_START + group * 4 + head // HEADS + q = tl.load( + Query + (query_rows[:, None].to(tl.int64) * HEADS + (head % HEADS)[:, None]) * D + col[None, :], + (query_rows[:, None] < QUERY_START + TOKENS) & (col[None, :] < D), + 0, + ) + key_rows = tl.load( + Indices + request * CAPACITY + block * RATIO + RATIO - 1, + (block < maximum) & (block < BLOCKS), + 0, + ).to(tl.int64) + k = tl.load( + Cache + key_rows[None, :] * D + col[:, None], (block[None, :] < maximum) & (col[:, None] < D), 0 + ) + logits = tl.dot(q, k, input_precision=INPUT_PRECISION) + scores = tl.sum(tl.maximum(logits, 0).reshape(4, HEADS, BN), 1) * (D ** -0.5) + scores = tl.where(block[None, :] < lengths[:, None], scores, -float("inf")) + tl.store( + Output + rows[:, None].to(tl.int64) * BLOCKS + block[None, :], + scores, + valid[:, None] & (block[None, :] < BLOCKS) & ((MIN_BLOCKS == 0) | (lengths[:, None] >= MIN_BLOCKS)), + ) + else: + for offset in range(GROUP_SIZE): + row = group * GROUP_SIZE + offset + if GROUP_SIZE == 1 or row < TOKENS: + query_row = row + QUERY_START + request = tl.load(Requests + query_row).to(tl.int64) + visible = tl.load(Lengths + query_row) // RATIO + block = tl.program_id(1) * BN + tl.arange(0, BN) + + scores = tl.full((BN,), -float("inf"), tl.float32) + if (tl.program_id(1) * BN < visible) & (visible >= MIN_BLOCKS): + head = tl.arange(0, BH) + col = tl.arange(0, BD) + q = tl.load( + Query + (query_row.to(tl.int64) * HEADS + head[:, None]) * D + col[None, :], + (head[:, None] < HEADS) & (col[None, :] < D), + other=0, + ) + if PHYSICAL_CACHE: + key_rows = tl.load( + Indices + request * CAPACITY + block * RATIO + RATIO - 1, + (block < visible) & (block < BLOCKS), + other=0, + ).to(tl.int64) + else: + key_rows = request * BLOCKS + block + k = tl.load( + Cache + key_rows[None, :] * D + col[:, None], + (block[None, :] < visible) & (col[:, None] < D), + other=0, + ) + logits = tl.dot(q, k, input_precision=INPUT_PRECISION) + logits = tl.where(head[:, None] < HEADS, tl.maximum(logits, 0), 0) + scores = tl.sum(logits, 0) * (D ** -0.5) + scores = tl.where(block < visible, scores, -float("inf")) + tl.store( + Output + row.to(tl.int64) * BLOCKS + block, + scores, + (block < BLOCKS) & ((MIN_BLOCKS == 0) | (visible >= MIN_BLOCKS)), + ) + + +@triton.jit +def _expand( + Selected, + Requests, + Lengths, + Indices, + Output, + QUERY_START: tl.constexpr, + TOPK: tl.constexpr, + BLOCK_BUDGET: tl.constexpr, + CAPACITY: tl.constexpr, + RATIO: tl.constexpr, + WIDTH: tl.constexpr, + B: tl.constexpr, +): + row = tl.program_id(0) + query_row = row + QUERY_START + request = tl.load(Requests + query_row).to(tl.int64) + visible = tl.load(Lengths + query_row) + complete = visible // RATIO + chosen = tl.minimum(complete, BLOCK_BUDGET) + col = tl.arange(0, B) + rank = col // RATIO + block = tl.full((B,), 0, tl.int64) + if TOPK > 0: + block = tl.load(Selected + row.to(tl.int64) * TOPK + rank, (rank < chosen) & (col < WIDTH), other=0) + full = col < chosen * RATIO + tail_col = col - chosen * RATIO + tail = (tail_col >= 0) & (tail_col < visible % RATIO) + position = tl.where(full, block * RATIO + col % RATIO, complete * RATIO + tail_col) + slot = tl.load(Indices + request * CAPACITY + position, (col < WIDTH) & (full | tail), other=-1) + tl.store(Output + query_row.to(tl.int64) * WIDTH + col, slot, col < WIDTH) diff --git a/lightllm/common/qsa/topk.py b/lightllm/common/qsa/topk.py new file mode 100644 index 0000000000..c6c983d860 --- /dev/null +++ b/lightllm/common/qsa/topk.py @@ -0,0 +1,127 @@ +from functools import lru_cache +from inspect import signature + +import torch +import triton +import triton.language as tl +from lightllm.utils.log_utils import init_logger + +logger = init_logger(__name__) + + +_CUDA_TOPK_RUNTIME = str(torch.__version__) == "2.9.1+cu128" and triton.__version__ == "3.5.1" + + +@lru_cache(maxsize=16) +def _supports_device(device): + supported = torch.cuda.get_device_capability(device) == (9, 0) and "H100" in torch.cuda.get_device_name(device) + if supported: + logger.info("QSA CUDA TopK selection enabled") + return supported + + +@triton.jit +def _key(value): + bits = value.to(tl.uint32, bitcast=True) + mask = tl.where((bits & 0x80000000) != 0, 0xFFFFFFFF, 0x80000000).to(tl.uint32) + return tl.where(value == value, bits ^ mask, 0xFFFFFFFF).to(tl.uint32) + + +@triton.jit +def _topk( + Scores, + Selected, + Lengths, + FIRST: tl.constexpr, + RATIO: tl.constexpr, + N: tl.constexpr, + K: tl.constexpr, + B: tl.constexpr, +): + tl.static_assert(N == 1027 and K == 512 and B == 2048) + row = tl.program_id(0).to(tl.int64) + if Lengths is None: + short = tl.full((), False, tl.int1) + else: + short = tl.load(Lengths + row + FIRST) // RATIO < K + if short: + short_col = tl.arange(0, K) + tl.store(Selected + row * K + short_col, short_col.to(tl.int64)) + else: + prefix_col = tl.arange(0, 1024) + prefix = tl.topk(_key(tl.load(Scores + row * N + prefix_col)), K) + merge_col = tl.arange(0, 4) + tail = tl.load(Scores + row * N + 1024 + merge_col, merge_col < 3, other=0.0) + tail = tl.sort(tl.where(merge_col < 3, _key(tail), 0).to(tl.uint32), descending=True) + + prefix_cutoffs = tl.gather(prefix, K - 1 - merge_col, 0) + tail_cutoffs = tl.gather(tail, tl.maximum(merge_col - 1, 0), 0) + tail_cutoffs = tl.where(merge_col == 0, 0xFFFFFFFF, tail_cutoffs).to(tl.uint32) + threshold = tl.max(tl.minimum(prefix_cutoffs, tail_cutoffs)) + col = tl.arange(0, B) + value = tl.load(Scores + row * N + col, col < N, other=-float("inf")) + key = _key(value) + greater = (col < N) & (key > threshold) + equal = (col < N) & (key == threshold) + greater_rank = tl.cumsum(greater.to(tl.int32)) - 1 + equal_rank = tl.cumsum(equal.to(tl.int32)) - 1 + rank = tl.where(greater, greater_rank, tl.sum(greater.to(tl.int32)) + equal_rank) + selected = greater | (equal & (rank < K)) + tl.store(Selected + row * K + rank, col.to(tl.int64), selected) + + +def cuda_topk_supported(device, blocks: int, k: int) -> bool: + return _CUDA_TOPK_RUNTIME and device.type == "cuda" and blocks == 1027 and k == 512 and _supports_device(device) + + +@lru_cache(maxsize=1) +def _flashinfer_topk(): + try: + from flashinfer import top_k_ragged_transform + except ImportError: + return None + if "dsa_graph_safe" not in signature(top_k_ragged_transform).parameters: + return None + return top_k_ragged_transform + + +@lru_cache(maxsize=16) +def flashinfer_topk_supported(device, blocks: int, k: int) -> bool: + return ( + device.type == "cuda" + and k == 512 + and blocks >= k + and torch.cuda.get_device_capability(device) == (9, 0) + and _flashinfer_topk() is not None + ) + + +def flashinfer_topk_indices(scores, k, *, visible_lengths, first=0, compress_ratio=4): + lengths = visible_lengths[first : first + scores.shape[0]] // compress_ratio + return _flashinfer_topk()(scores, torch.zeros_like(lengths), lengths, k, dsa_graph_safe=True) + + +def cuda_topk_indices(scores: torch.Tensor, k: int, *, visible_lengths=None, first=0, compress_ratio=4) -> torch.Tensor: + if not ( + scores.is_cuda + and scores.dtype == torch.float32 + and scores.ndim == 2 + and scores.is_contiguous() + and cuda_topk_supported(scores.device, scores.shape[1], k) + ): + return torch.topk(scores, k, dim=-1, sorted=False).indices + if visible_lengths is not None: + if ( + visible_lengths.dtype != torch.int32 + or visible_lengths.device != scores.device + or visible_lengths.ndim != 1 + or not visible_lengths.is_contiguous() + or first < 0 + or first + scores.shape[0] > visible_lengths.numel() + or compress_ratio < 1 + ): + raise ValueError("QSA TopK requires matching contiguous int32 visible lengths and valid bounds") + output = torch.empty((scores.shape[0], k), dtype=torch.int64, device=scores.device) + if scores.shape[0]: + _topk[(scores.shape[0],)](scores, output, visible_lengths, first, compress_ratio, 1027, 512, 2048, num_warps=16) + return output diff --git a/lightllm/common/quantization/aligned_projection.py b/lightllm/common/quantization/aligned_projection.py new file mode 100644 index 0000000000..06e1f9a287 --- /dev/null +++ b/lightllm/common/quantization/aligned_projection.py @@ -0,0 +1,51 @@ +import torch + +from lightllm.common.quantization.no_quant import NoQuantization +from lightllm.common.quantization.quantize_method import WeightPack + + +class OutputAlignedProjection(NoQuantization): + def __init__(self, alignment=8): + super().__init__() + if type(alignment) is not int or alignment < 1 or alignment & (alignment - 1): + raise ValueError("Projection alignment must be a positive power of two") + self.alignment = alignment + self._padded_weight = None + + def _create_weight(self, out_dims, in_dim, dtype, device_id, num_experts=1): + width = sum(out_dims) if isinstance(out_dims, list) else out_dims + physical_width = (width + self.alignment - 1) // self.alignment * self.alignment + self._padded_weight = None + if dtype not in (torch.float16, torch.bfloat16) or physical_width == width or num_experts != 1: + return super()._create_weight(out_dims, in_dim, dtype, device_id, num_experts) + self._padded_weight = torch.zeros((physical_width, in_dim), dtype=dtype, device=device_id) + pack = WeightPack(weight=self._padded_weight[:width]) + parts = self._split_weight_pack(pack, out_dims if isinstance(out_dims, list) else [out_dims], -2) + return pack, parts + + def apply(self, input_tensor, weight_pack, out=None, workspace=None, use_custom_tensor_mananger=True, bias=None): + if self._padded_weight is None or bias is not None or input_tensor.shape[0] == 0: + return super().apply( + input_tensor, + weight_pack, + out, + workspace, + use_custom_tensor_mananger=use_custom_tensor_mananger, + bias=bias, + ) + from lightllm.common.basemodel.layer_infer.cache_tensor_manager import g_cache_manager + + def allocate(shape): + if use_custom_tensor_mananger: + return g_cache_manager.alloc_tensor(shape, input_tensor.dtype, device=input_tensor.device) + return torch.empty(shape, dtype=input_tensor.dtype, device=input_tensor.device) + + rows, width = input_tensor.shape[0], weight_pack.weight.shape[0] + physical_output = allocate((rows, self._padded_weight.shape[0])) + torch.mm(input_tensor, self._padded_weight.T, out=physical_output) + if out is None: + out = allocate((rows, width)) + elif out.shape != (rows, width) or out.dtype != input_tensor.dtype or out.device != input_tensor.device: + raise ValueError("Projection output must match logical shape, dtype and device") + out.copy_(physical_output[:, :width]) + return out diff --git a/lightllm/common/quantization/scalar_projection.py b/lightllm/common/quantization/scalar_projection.py new file mode 100644 index 0000000000..cef317e4a2 --- /dev/null +++ b/lightllm/common/quantization/scalar_projection.py @@ -0,0 +1,21 @@ +import torch + +from lightllm.common.basemodel.triton_kernel.scalar_linear import scalar_linear +from lightllm.common.quantization.no_quant import NoQuantization + + +class StableScalarProjection(NoQuantization): + def apply(self, input_tensor, weight_pack, out=None, workspace=None, use_custom_tensor_mananger=True, bias=None): + if not input_tensor.is_cuda or weight_pack.weight.shape[0] != 1 or bias is not None: + return super().apply( + input_tensor, + weight_pack, + out=out, + workspace=workspace, + use_custom_tensor_mananger=use_custom_tensor_mananger, + bias=bias, + ) + from lightllm.common.basemodel.layer_infer.cache_tensor_manager import g_cache_manager + + allocate = g_cache_manager.empty if use_custom_tensor_mananger else torch.empty + return scalar_linear(input_tensor, weight_pack.weight, out=out, alloc_func=allocate) diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_H100_80GB_HBM3/grouped_matmul:v1/{K=2560,N=160,expert_num=513,mul_routed_weight=false,out_dtype=torch.bfloat16,topk_num=11,use_fp8_w8a8=false}_NVIDIA_H100_80GB_HBM3.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_H100_80GB_HBM3/grouped_matmul:v1/{K=2560,N=160,expert_num=513,mul_routed_weight=false,out_dtype=torch.bfloat16,topk_num=11,use_fp8_w8a8=false}_NVIDIA_H100_80GB_HBM3.json new file mode 100644 index 0000000000..a4e6e1885e --- /dev/null +++ b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_H100_80GB_HBM3/grouped_matmul:v1/{K=2560,N=160,expert_num=513,mul_routed_weight=false,out_dtype=torch.bfloat16,topk_num=11,use_fp8_w8a8=false}_NVIDIA_H100_80GB_HBM3.json @@ -0,0 +1,92 @@ +{ + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 3 + }, + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 5 + }, + "3": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 5 + }, + "513": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 1 + }, + "514": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 1 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 5 + } +} \ No newline at end of file diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_H100_80GB_HBM3/grouped_matmul:v1/{K=80,N=2560,expert_num=513,mul_routed_weight=true,out_dtype=torch.bfloat16,topk_num=1,use_fp8_w8a8=false}_NVIDIA_H100_80GB_HBM3.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_H100_80GB_HBM3/grouped_matmul:v1/{K=80,N=2560,expert_num=513,mul_routed_weight=true,out_dtype=torch.bfloat16,topk_num=1,use_fp8_w8a8=false}_NVIDIA_H100_80GB_HBM3.json new file mode 100644 index 0000000000..444e3b96fe --- /dev/null +++ b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_H100_80GB_HBM3/grouped_matmul:v1/{K=80,N=2560,expert_num=513,mul_routed_weight=true,out_dtype=torch.bfloat16,topk_num=1,use_fp8_w8a8=false}_NVIDIA_H100_80GB_HBM3.json @@ -0,0 +1,92 @@ +{ + "88": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 4 + }, + "264": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 4 + }, + "5632": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 4 + }, + "11": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 2 + }, + "22": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 2 + }, + "33": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 4 + }, + "44": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 2 + }, + "513": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 2 + }, + "514": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 1 + }, + "5633": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 1 + } +} \ No newline at end of file diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_H200/grouped_matmul:v1/{K=160,N=2560,expert_num=513,mul_routed_weight=true,out_dtype=torch.bfloat16,topk_num=1,use_fp8_w8a8=false}_NVIDIA_H200.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_H200/grouped_matmul:v1/{K=160,N=2560,expert_num=513,mul_routed_weight=true,out_dtype=torch.bfloat16,topk_num=1,use_fp8_w8a8=false}_NVIDIA_H200.json new file mode 100644 index 0000000000..967fd15e61 --- /dev/null +++ b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_H200/grouped_matmul:v1/{K=160,N=2560,expert_num=513,mul_routed_weight=true,out_dtype=torch.bfloat16,topk_num=1,use_fp8_w8a8=false}_NVIDIA_H200.json @@ -0,0 +1,47 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 2 + }, + "1408": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 2 + }, + "1409": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 2 + }, + "22528": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 2 + }, + "22529": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 1 + } +} diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_H200/grouped_matmul:v1/{K=2560,N=320,expert_num=513,mul_routed_weight=false,out_dtype=torch.bfloat16,topk_num=11,use_fp8_w8a8=false}_NVIDIA_H200.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_H200/grouped_matmul:v1/{K=2560,N=320,expert_num=513,mul_routed_weight=false,out_dtype=torch.bfloat16,topk_num=11,use_fp8_w8a8=false}_NVIDIA_H200.json new file mode 100644 index 0000000000..ace348c88f --- /dev/null +++ b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_H200/grouped_matmul:v1/{K=2560,N=320,expert_num=513,mul_routed_weight=false,out_dtype=torch.bfloat16,topk_num=11,use_fp8_w8a8=false}_NVIDIA_H200.json @@ -0,0 +1,65 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 3 + }, + "33": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 1 + }, + "127": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 1 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 2 + }, + "513": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 2 + }, + "514": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 1 + } +} diff --git a/lightllm/distributed/symm_mem_all_reduce.py b/lightllm/distributed/symm_mem_all_reduce.py index 2256a5093f..7a93318313 100644 --- a/lightllm/distributed/symm_mem_all_reduce.py +++ b/lightllm/distributed/symm_mem_all_reduce.py @@ -90,3 +90,22 @@ def all_reduce(self, inp: torch.Tensor) -> None: else: torch.ops.symm_mem.two_shot_all_reduce_(self.buffer[:n], "sum", self.group.group_name) inp.view(-1).copy_(self.buffer[:n]) + + def ordered_all_gather(self, inp: torch.Tensor): + if not self.should_use(inp) or inp.device != self.device: + return None + count = inp.numel() * self.world_size + if not count or count * inp.element_size() > self.max_size: + return None + if not hasattr(torch.ops.symm_mem, "multimem_all_gather_out"): + return None + if not hasattr(self, "ordered_gather_buffer"): + if torch.cuda.is_current_stream_capturing(): + return None + self.ordered_gather_buffer = torch_symm_mem.empty( + self.max_size // self.dtype.itemsize, device=self.device, dtype=self.dtype + ) + torch_symm_mem.rendezvous(self.ordered_gather_buffer, self.group.group_name) + gathered = self.ordered_gather_buffer[:count] + torch.ops.symm_mem.multimem_all_gather_out(inp.view(-1), self.group.group_name, gathered) + return gathered.view(self.world_size, *inp.shape) diff --git a/lightllm/models/builtin.py b/lightllm/models/builtin.py index d8e091d9aa..cb3bff203c 100644 --- a/lightllm/models/builtin.py +++ b/lightllm/models/builtin.py @@ -209,3 +209,11 @@ def _tarsier_text_model_is(model_type): for config in configs } MODEL_CLASS_PATHS.update({path.rsplit(":", 1)[1]: path for path in DraftModelRegistry._registry.values()}) + +ModelRegistry.register("qwen4_exp", "lightllm.models.qwen4_exp.model:Qwen4ExpTpPartModel", is_multimodal=True) +ModelRegistry.register("qwen4_exp_text", "lightllm.models.qwen4_exp.model:Qwen4ExpTpPartModel") +DraftModelRegistry.register( + ("qwen4_exp", "qwen4_exp_text"), "eagle_with_att", "lightllm.models.qwen4_exp.mtp.model:Qwen4ExpMTPModel" +) +MODEL_CLASS_PATHS["Qwen4ExpTpPartModel"] = "lightllm.models.qwen4_exp.model:Qwen4ExpTpPartModel" +MODEL_CLASS_PATHS["Qwen4ExpMTPModel"] = "lightllm.models.qwen4_exp.mtp.model:Qwen4ExpMTPModel" diff --git a/lightllm/models/glm5_next/config.py b/lightllm/models/glm5_next/config.py index 0e52e736a9..31d79697f6 100644 --- a/lightllm/models/glm5_next/config.py +++ b/lightllm/models/glm5_next/config.py @@ -14,15 +14,17 @@ @dataclasses.dataclass class Glm5NextResourceConfig(Qwen3NextResourceConfig): - """Replicated MLA/index KV plus TP-sharded KDA checkpoints.""" + """描述复制式 MLA/index KV 与 TP 分片 KDA checkpoint 的资源配置。""" # Indexer 池化步长，单位为原始 token；同时决定 K-pool tail 的基准 ring 长度和恢复边界。 + gate_lower_bound: float = -5.0 index_kpool: int = 4 # 单个 indexer key 的特征维度；原始 key/gate tail 的 channel 宽度为该值的两倍。 index_head_dim: int = 128 @classmethod def from_model_config(cls, config: dict[str, object], args: StartArgs | EasyDict) -> "Glm5NextResourceConfig": + """从模型配置和启动参数构造 GLM5.3 Flash 资源配置。""" config = config.get("text_config", config) linear = config["linear_attn_config"] tp = args.tp // args.dp @@ -45,6 +47,7 @@ def from_model_config(cls, config: dict[str, object], args: StartArgs | EasyDict ssm_state_dtype=get_torch_dtype(args.linear_att_ssm_data_type), full_attention_interval=4, all_layer_num=layers, + gate_lower_bound=linear.get("gate_lower_bound", -5.0), index_kpool=config["index_kpool"], index_head_dim=config["index_head_dim"], ) diff --git a/lightllm/models/glm5_next/model.py b/lightllm/models/glm5_next/model.py index e15efe0143..34d28152e8 100644 --- a/lightllm/models/glm5_next/model.py +++ b/lightllm/models/glm5_next/model.py @@ -79,6 +79,9 @@ def create_bounded_att_states( request_capacity=request_capacity, mtp_step=args.mtp_step, device=device, + ssm_state_mode=args.ssm_state_mode, + ssm_replay_config=args.ssm_replay_config, + gate_lower_bound=cfg.gate_lower_bound, ) ) full_layers = tuple(model_logical_layer_ids[index] for index in cfg.get_full_att_layer_indices()) @@ -114,37 +117,9 @@ def normalize_model_config(cls, model_config: dict[str, object]) -> dict[str, ob config["autotune_layer_num"] = 4 return config - def _init_config(self): - with open(os.path.join(self.weight_dir_, "config.json")) as f: - self.config = self.normalize_model_config(json.load(f)) - - def _verify_params(self): - super()._verify_params() - assert self.config["qk_rope_head_dim"] == 0, "GLM-5.3 Flash uses NoPE attention" - assert not self.args.enable_tpsp_mix_mode, "GLM-5.3 Flash does not support TP/SP mixed mode" - if self.args.llm_kv_type not in ("None", "fp8kv_dsa"): - raise ValueError("GLM-5.3 Flash supports only unquantized KV or dynamic fp8kv_dsa") - if self.args.llm_kv_type == "fp8kv_dsa": - # group64/E8M0 编码与稀疏 attention 算子要求 BF16 输入和固定的 512 维 latent。 - if self.data_type != torch.bfloat16 or self.config["kv_lora_rank"] != 512: - raise ValueError("GLM-5.3 Flash FP8 KV requires BF16 computation and kv_lora_rank=512") - def autotune_layers(self): return 4 - def _init_some_value(self): - self.layers_num = self.config["n_layer"] - self.vocab_size = self.config["vocab_size"] - # MLA stores one replicated latent KV vector per token on every TP rank. - self.tp_k_head_num_ = 1 - self.tp_v_head_num_ = 0 - self.qk_nope_head_dim = self.config["qk_nope_head_dim"] - self.qk_rope_head_dim = self.config["qk_rope_head_dim"] - self.q_lora_rank = self.config["q_lora_rank"] - self.kv_lora_rank = self.config["kv_lora_rank"] - self.v_head_dim = self.config.get("v_head_dim", self.qk_nope_head_dim) - self.head_dim_ = self.kv_lora_rank + self.qk_rope_head_dim - @classmethod def create_full_att_states( cls, @@ -155,7 +130,10 @@ def create_full_att_states( tp_world_size: int = 1, ) -> list["FullAttentionState"]: """GLM full-attention 层采用 NoPE packed MLA，并保存压缩 indexer 字段。""" - from lightllm.common.attention_state.full_att import PackedMLALogicalConfig, create_packed_mla_state + from lightllm.common.attention_state.full_att import ( + PackedMLALogicalConfig, + create_packed_mla_state, + ) from lightllm.utils.torch_dtype_utils import get_torch_dtype config = cls.normalize_model_config(model_config) @@ -178,18 +156,38 @@ def create_full_att_states( ) return [state] - def prepare_attention_states(self) -> None: - super().prepare_attention_states() - if self.args.llm_kv_type == "fp8kv_dsa" and self.att_state_manager.get_full_attention_states( - self.att_state_role - ): - from lightllm.common.basemodel.attention.nsa.fp8_glm5_next import Fp8Glm5NextSparseAttBackend + def _init_config(self): + with open(os.path.join(self.weight_dir_, "config.json")) as f: + self.config = self.normalize_model_config(json.load(f)) - # Reserve fixed scratch space before the shared KV capacity is profiled. - Fp8Glm5NextSparseAttBackend.allocate_prefill_workspaces(self) + def _verify_params(self): + super()._verify_params() + assert self.config["qk_rope_head_dim"] == 0, "GLM-5.3 Flash uses NoPE attention" + assert not self.args.enable_tpsp_mix_mode, "GLM-5.3 Flash does not support TP/SP mixed mode" + if self.args.llm_kv_type not in ("None", "fp8kv_dsa"): + raise ValueError("GLM-5.3 Flash supports only unquantized KV or dynamic fp8kv_dsa") + if self.args.llm_kv_type == "fp8kv_dsa": + # group64/E8M0 编码与稀疏 attention 算子要求 BF16 输入和固定的 512 维 latent。 + if self.data_type != torch.bfloat16 or self.config["kv_lora_rank"] != 512: + raise ValueError("GLM-5.3 Flash FP8 KV requires BF16 computation and kv_lora_rank=512") + + def _init_some_value(self): + self.layers_num = self.config["n_layer"] + self.vocab_size = self.config["vocab_size"] + # MLA stores one replicated latent KV vector per token on every TP rank. + self.tp_k_head_num_ = 1 + self.tp_v_head_num_ = 0 + self.qk_nope_head_dim = self.config["qk_nope_head_dim"] + self.qk_rope_head_dim = self.config["qk_rope_head_dim"] + self.q_lora_rank = self.config["q_lora_rank"] + self.kv_lora_rank = self.config["kv_lora_rank"] + self.v_head_dim = self.config.get("v_head_dim", self.qk_nope_head_dim) + self.head_dim_ = self.kv_lora_rank + self.qk_rope_head_dim def _init_att_backend(self): - from lightllm.common.basemodel.attention.nsa.fp8_glm5_next import Fp8Glm5NextSparseAttBackend + from lightllm.common.basemodel.attention.nsa.fp8_glm5_next import ( + Fp8Glm5NextSparseAttBackend, + ) backend_class = ( Fp8Glm5NextSparseAttBackend if self.args.llm_kv_type == "fp8kv_dsa" else Glm5NextSparseAttBackend @@ -198,7 +196,15 @@ def _init_att_backend(self): self.decode_att_backend = self.prefill_att_backend def _init_att_backend1(self): - self.prefill_att_backend1 = KDALinearAttBackend(model=self) + if getattr(self.args, "ssm_state_mode", "native") == "replay": + from lightllm.common.basemodel.attention.linear.replay_kda import ( + ReplayKDALinearAttBackend, + ) + + backend_cls = ReplayKDALinearAttBackend + else: + backend_cls = KDALinearAttBackend + self.prefill_att_backend1 = backend_cls(model=self) self.decode_att_backend1 = self.prefill_att_backend1 def _init_custom(self): diff --git a/lightllm/models/qwen2_vl/triton_kernel/rotary_pos_emb.py b/lightllm/models/qwen2_vl/triton_kernel/rotary_pos_emb.py index 8063856599..fa61e05c69 100644 --- a/lightllm/models/qwen2_vl/triton_kernel/rotary_pos_emb.py +++ b/lightllm/models/qwen2_vl/triton_kernel/rotary_pos_emb.py @@ -1,4 +1,3 @@ -import math import torch import triton import triton.language as tl @@ -13,6 +12,9 @@ def rotary_kernel( stride_l, stride_h, stride_d, + out_stride_l, + out_stride_h, + out_stride_d, stride_cos_l, stride_cos_d, stride_sin_l, @@ -53,7 +55,7 @@ def rotary_kernel( y = x * cos + rotated * sin - out_ptr_ = out_ptr + base + d * stride_d + out_ptr_ = out_ptr + pid_l * out_stride_l + pid_h * out_stride_h + d * out_stride_d tl.store(out_ptr_, y, mask=mask) @@ -89,6 +91,9 @@ def apply_rotary_pos_emb_triton( stride_l=x.stride(0), stride_h=x.stride(1), stride_d=x.stride(2), + out_stride_l=y.stride(0), + out_stride_h=y.stride(1), + out_stride_d=y.stride(2), stride_cos_l=cos.stride(0), stride_cos_d=cos.stride(1), stride_sin_l=sin.stride(0), diff --git a/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py b/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py index 1bc4ef914b..df30288562 100644 --- a/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py @@ -266,9 +266,7 @@ def context_attention_forward( gdn_out = self._linear_post(core_attn_out, z, layer_weight) - if self.tp_world_size_ > 1: - all_reduce(gdn_out, op=dist.ReduceOp.SUM, group=infer_state.dist_group, async_op=False) - return gdn_out + return self._reduce_linear_output(gdn_out, infer_state) def token_attention_forward( self, @@ -296,9 +294,12 @@ def token_attention_forward( ) gdn_out = self._linear_post(core_attn_out, z, layer_weight) + return self._reduce_linear_output(gdn_out, infer_state) + + def _reduce_linear_output(self, output, infer_state): if self.tp_world_size_ > 1: - all_reduce(gdn_out, op=dist.ReduceOp.SUM, group=infer_state.dist_group, async_op=False) - return gdn_out + all_reduce(output, op=dist.ReduceOp.SUM, group=infer_state.dist_group, async_op=False) + return output def _linear_in_proj( self, @@ -316,7 +317,6 @@ def _linear_post( ) -> torch.Tensor: num_tokens = z.shape[0] core_attn_out = core_attn_out.view(-1, core_attn_out.shape[-1]) - z = z.contiguous().view(-1, z.shape[-1]) norm_out = layer_weight.linear_norm(core_attn_out, z, self.eps_) core_attn_out = norm_out.view(num_tokens, -1) output = layer_weight.linear_out_proj.mm(core_attn_out) diff --git a/lightllm/models/qwen3next/model.py b/lightllm/models/qwen3next/model.py index f1a3a65748..9da87c6acb 100644 --- a/lightllm/models/qwen3next/model.py +++ b/lightllm/models/qwen3next/model.py @@ -81,26 +81,14 @@ def create_bounded_att_states( request_capacity=request_capacity, mtp_step=args.mtp_step, device=device, + ssm_state_mode=args.ssm_state_mode, + ssm_replay_config=args.ssm_replay_config, ) return [state] - def _init_triton(self): - def _triton_allocator(size: int, alignment: int, stream: Optional[int]) -> torch.Tensor: - return torch.empty(size, device="cuda", dtype=torch.int8) - - # Set Triton allocator for TMA descriptors - # This is required for kernels in common/.../linear_att/fla/ops/solve_tril.py - triton.set_allocator(_triton_allocator) - logger.info("Triton allocator set for Qwen3Next model") - return - def autotune_layers(self): return self.config["full_attention_interval"] - def _init_config(self): - super()._init_config() - self.num_kv_heads = max(self.config["num_key_value_heads"] // self.tp_world_size_, 1) - @classmethod def create_full_att_states( cls, @@ -135,6 +123,20 @@ def create_full_att_states( ) return [state] + def _init_triton(self): + def _triton_allocator(size: int, alignment: int, stream: Optional[int]) -> torch.Tensor: + return torch.empty(size, device="cuda", dtype=torch.int8) + + # Set Triton allocator for TMA descriptors + # This is required for kernels in common/.../linear_att/fla/ops/solve_tril.py + triton.set_allocator(_triton_allocator) + logger.info("Triton allocator set for Qwen3Next model") + return + + def _init_config(self): + super()._init_config() + self.num_kv_heads = max(self.config["num_key_value_heads"] // self.tp_world_size_, 1) + def _init_att_backend1(self): role = self.att_state_role if not self.att_state_manager.get_bounded_attention_states(role): diff --git a/lightllm/models/qwen4_exp/__init__.py b/lightllm/models/qwen4_exp/__init__.py new file mode 100644 index 0000000000..2ae28399f5 --- /dev/null +++ b/lightllm/models/qwen4_exp/__init__.py @@ -0,0 +1 @@ +pass diff --git a/lightllm/models/qwen4_exp/attention_state.py b/lightllm/models/qwen4_exp/attention_state.py new file mode 100644 index 0000000000..d07fa8de89 --- /dev/null +++ b/lightllm/models/qwen4_exp/attention_state.py @@ -0,0 +1,100 @@ +from dataclasses import replace + +import torch + +from lightllm.common.attention_state.full_att.gqa.gqa_state import GQAFullAttentionState, GQAAttentionConfig +from .state import QSAExecutionState + + +class Qwen4ExpAttentionState(GQAFullAttentionState): + def __init__(self, layout_config, indexer_dim, rotary_dim, page_size=1, device="meta"): + super().__init__(layout_config, page_size, device) + common = dict( + global_head_num=1, + dtype=layout_config.dtype, + tp_rank=layout_config.tp_rank, + tp_world_size=layout_config.tp_world_size, + replicated=True, + ) + self.indexer = GQAFullAttentionState( + GQAAttentionConfig(logical_layer_groups=layout_config.logical_layer_groups, head_dim=indexer_dim, **common), + page_size, + device, + ) + layers = tuple(layer for group in layout_config.logical_layer_groups for layer in group) + self.positions = GQAFullAttentionState( + GQAAttentionConfig(logical_layer_groups=(layers,), head_dim=rotary_dim, **common), page_size, device + ) + + def get_state_type(self): + return "qwen4_" + super().get_state_type() + + def allocate_storage(self, token_capacity): + super().allocate_storage(token_capacity) + for state in (self.indexer, self.positions): + state.device = self.device + state.allocate_storage(token_capacity) + + def get_token_bytes(self): + return super().get_token_bytes() + self.indexer.get_token_bytes() + self.positions.get_token_bytes() + + def qsa_view(self, layer): + raw, pooled = self.indexer.get_layer_view(layer) + cos, sin = self.positions.get_layer_view(layer) + return QSAExecutionState(raw[:, 0], pooled[:, 0], cos[:, 0], sin[:, 0]) + + def store_positions(self, indices, cos, sin): + self.positions.k_storage[0, :, 0].index_copy_(0, indices, cos) + self.positions.v_storage[0, :, 0].index_copy_(0, indices, sin) + + def _get_transfer_regions(self, page_tokens): + regions = list(super()._get_transfer_regions(page_tokens)) + for prefix, state in (("indexer", self.indexer), ("positions", self.positions)): + offset = max(region.byte_end for region in regions) + regions.extend( + replace( + region, + field=prefix + "." + region.field, + byte_start=offset + region.byte_start, + byte_end=offset + region.byte_end, + ) + for region in state._get_transfer_regions(page_tokens) + ) + return tuple(regions) + + def _layouts(self, layout): + for prefix, state in (("", super()), ("indexer.", self.indexer), ("positions.", self.positions)): + regions = tuple( + replace(region, field=region.field[len(prefix) :]) + for region in layout.regions + if (region.field.startswith(prefix) if prefix else "." not in region.field) + ) + yield state, replace(layout, regions=regions) + + def write_to_cpu(self, token_indices, cpu_page_buffers, cpu_page_indices, layout, cpu_page_ready_flags=None): + for state, part in self._layouts(layout): + state.write_to_cpu(token_indices, cpu_page_buffers, cpu_page_indices, part, cpu_page_ready_flags) + + def read_from_cpu(self, cpu_page_buffers, cpu_page_indices, token_indices, layout): + for state, part in self._layouts(layout): + state.read_from_cpu(cpu_page_buffers, cpu_page_indices, token_indices, part) + + def write_to_pd_page(self, token_indices, pd_page_buffers, pd_page_index, layout): + for state, part in self._layouts(layout): + state.write_to_pd_page(token_indices, pd_page_buffers, pd_page_index, part) + + def read_from_pd_page(self, pd_page_buffers, token_indices, pd_page_index, layout): + for state, part in self._layouts(layout): + state.read_from_pd_page(pd_page_buffers, token_indices, pd_page_index, part) + + def copy_full_att_tokens(self, src_indices, dst_indices): + super().copy_full_att_tokens(src_indices, dst_indices) + for state in (self.indexer, self.positions): + state.copy_full_att_tokens(src_indices, dst_indices) + + def copy_tokens_from_dp(self, src_token_indices, dst_token_indices, remote_dp_full_att_state): + super().copy_tokens_from_dp(src_token_indices, dst_token_indices, remote_dp_full_att_state) + for name in ("indexer", "positions"): + getattr(self, name).copy_tokens_from_dp( + src_token_indices, dst_token_indices, getattr(remote_dp_full_att_state, name) + ) diff --git a/lightllm/models/qwen4_exp/batch.py b/lightllm/models/qwen4_exp/batch.py new file mode 100644 index 0000000000..299e98e049 --- /dev/null +++ b/lightllm/models/qwen4_exp/batch.py @@ -0,0 +1,119 @@ +from dataclasses import dataclass + +import torch +import triton +import triton.language as tl + +from lightllm.common.basemodel.triton_kernel.gen_nsa_ks_ke import gen_same_req_mark + + +@triton.jit +def _batch_metadata( + Requests, + Lengths, + Starts, + Active, + KeyLengths, + FirstBlocks, + QueryRows, + Visible, + TokenMask, + HOLD: tl.constexpr, + RATIO: tl.constexpr, + BLOCK: tl.constexpr, +): + row = tl.program_id(0) + request = tl.load(Requests + row) + start = tl.load(Starts + row) + end = tl.load(Starts + row + 1) + count = end - start + length = tl.load(Lengths + row) + active = (request != HOLD) & (count > 0) + tl.store(Active + row, tl.where(active, request, -1)) + tl.store(KeyLengths + row, tl.where(active, length, 0)) + tl.store(FirstBlocks + row, tl.where(active, (length - count) // RATIO, 0)) + for first in range(0, tl.cdiv(count, BLOCK)): + offset = first * BLOCK + tl.arange(0, BLOCK) + mask = offset < count + tl.store(QueryRows + start + offset, request, mask) + tl.store(Visible + start + offset, tl.where(active, length - count + offset + 1, 0), mask) + tl.store(TokenMask + start + offset, active, mask) + + +@dataclass(frozen=True) +class PackedBatchMetadata: + starts: torch.Tensor + active_rows: torch.Tensor + key_lengths: torch.Tensor + first_blocks: torch.Tensor + query_rows: torch.Tensor + visible_lengths: torch.Tensor + token_mask: torch.Tensor + + @classmethod + @torch.no_grad() + def prepare( + cls, + requests, + lengths, + query_lengths, + token_count: int, + hold_request: int, + compress_ratio: int, + segment_capacity=None, + ): + if token_count < 0 or compress_ratio < 1: + raise ValueError("Token count must be nonnegative and compression ratio positive") + for tensor in (requests, lengths, query_lengths): + if ( + not tensor.is_cuda + or tensor.device != requests.device + or tensor.dtype != torch.int32 + or tensor.ndim != 1 + or not tensor.is_contiguous() + or tensor.shape != requests.shape + ): + raise ValueError("Batch metadata requires matching contiguous int32 CUDA vectors") + batch = requests.numel() + capacity = batch if segment_capacity is None else segment_capacity + if capacity < batch: + raise ValueError("Metadata segment capacity cannot be smaller than the input batch") + if not batch and token_count: + raise ValueError("An empty batch cannot contain input tokens") + + starts = torch.full((capacity + 1,), token_count, dtype=torch.int32, device=requests.device) + starts[:1].zero_() + if batch: + marks = gen_same_req_mark(requests, hold_req_idx=hold_request) + counts = torch.where(marks > 0, query_lengths + marks - 1, 0) + torch.cumsum(counts, dim=0, dtype=torch.int32, out=starts[1 : batch + 1]) + active = torch.full((capacity,), -1, dtype=torch.int32, device=requests.device) + key_lengths = torch.zeros_like(active) + first_blocks = torch.zeros_like(active) + query_rows = torch.empty(token_count, dtype=torch.int32, device=requests.device) + visible = torch.empty_like(query_rows) + token_mask = torch.empty(token_count, dtype=torch.bool, device=requests.device) + if batch: + _batch_metadata[(batch,)]( + requests, + lengths, + starts, + active, + key_lengths, + first_blocks, + query_rows, + visible, + token_mask, + hold_request, + compress_ratio, + 256, + ) + return cls(starts, active, key_lengths, first_blocks, query_rows, visible, token_mask) + + def copy_from(self, other: "PackedBatchMetadata"): + for name in self.__dataclass_fields__: + target, source = getattr(self, name), getattr(other, name) + if target.shape != source.shape or target.device != source.device or target.dtype != source.dtype: + raise ValueError(f"Incompatible graph batch metadata: {name}") + for name in self.__dataclass_fields__: + getattr(self, name).copy_(getattr(other, name), non_blocking=True) diff --git a/lightllm/models/qwen4_exp/checkpoint.py b/lightllm/models/qwen4_exp/checkpoint.py new file mode 100644 index 0000000000..e5a42a716b --- /dev/null +++ b/lightllm/models/qwen4_exp/checkpoint.py @@ -0,0 +1,38 @@ +import math +from functools import lru_cache + + +def normalize_text_weight_names(weights: dict): + prefix = "model.language_model." + for name in tuple(weights): + if name.startswith(prefix): + target = "model." + name[len(prefix) :] + if target in weights and weights[target] is not weights[name]: + raise ValueError(f"Conflicting text checkpoint names: {name} and {target}") + weights[target] = weights.pop(name) + + +def ngram_eos_token_id(config: dict) -> int: + eos = config["eos_token_id"] + if isinstance(eos, list): + if not eos: + raise ValueError("PLE requires an EOS token ID") + eos = eos[0] + return eos + + +@lru_cache(maxsize=None) +def ngram_table_shape(base: int, heads: int, ple_index: int, embedding_dim: int, padding: int): + if base < 1 or heads < 1 or ple_index < 0 or embedding_dim < 1 or embedding_dim % heads or padding < 1: + raise ValueError("Invalid n-gram table dimensions") + sizes = [] + candidate = base - 1 + for index in range((ple_index + 1) * heads): + while True: + candidate += 1 + if candidate >= 2 and all(candidate % divisor for divisor in range(2, math.isqrt(candidate) + 1)): + break + if index >= ple_index * heads: + sizes.append(candidate) + total = sum(sizes) + return ((total + padding - 1) // padding) * padding, embedding_dim // heads diff --git a/lightllm/models/qwen4_exp/configuration.py b/lightllm/models/qwen4_exp/configuration.py new file mode 100644 index 0000000000..a47e250b84 --- /dev/null +++ b/lightllm/models/qwen4_exp/configuration.py @@ -0,0 +1,24 @@ +from lightllm.common.build_utils import repair_config + + +def normalize_model_config(checkpoint, finetune_config=None): + config = dict(checkpoint.get("text_config", checkpoint)) + for aliases in ( + ["num_attention_heads", "n_head"], + ["hidden_size", "n_embd", "n_embed"], + ["num_hidden_layers", "n_layer"], + ): + repair_config(config, same_names=aliases) + rope = config.get("rope_parameters") + if isinstance(rope, dict): + for name in ("rope_theta", "partial_rotary_factor"): + if name in rope: + config.setdefault(name, rope[name]) + config.setdefault("rope_scaling", dict(rope)) + for name in ("image_token_id", "video_token_id", "vision_start_token_id", "vision_end_token_id"): + if name in checkpoint: + config[name] = checkpoint[name] + config.setdefault("norm_topk_prob", True) + if finetune_config is not None: + config["vocab_size"] = finetune_config.vocab_size + return config, checkpoint.get("vision_config") diff --git a/lightllm/models/qwen4_exp/image_processing.py b/lightllm/models/qwen4_exp/image_processing.py new file mode 100644 index 0000000000..902e0d9d36 --- /dev/null +++ b/lightllm/models/qwen4_exp/image_processing.py @@ -0,0 +1,9 @@ +from lightllm.models.qwen2_vl.vision_process import Qwen2VLImageProcessor + + +class Qwen4ExpImageProcessor(Qwen2VLImageProcessor): + @classmethod + def from_pretrained(cls, model_path, **kwargs): + config, overrides = cls.get_image_processor_dict(model_path, **kwargs) + config.update(overrides) + return cls(**config) diff --git a/lightllm/models/qwen4_exp/infer_struct.py b/lightllm/models/qwen4_exp/infer_struct.py new file mode 100644 index 0000000000..06062d7959 --- /dev/null +++ b/lightllm/models/qwen4_exp/infer_struct.py @@ -0,0 +1,76 @@ +from lightllm.models.qwen3_5.infer_struct import Qwen35InferStateInfo +from .batch import PackedBatchMetadata +from .position import combine_mrope +from .state import QSABatch + + +class Qwen4ExpInferStateInfo(Qwen35InferStateInfo): + def init_some_extra_state(self, model): + super().init_some_extra_state(model) + config = model.config + ratio = config["indexer_compress_ratio"] + self.qwen4_exp_metadata = PackedBatchMetadata.prepare( + self.b_req_idx, + self.b_seq_len, + self.b_q_seq_len, + self.input_ids.numel(), + self.req_manager.HOLD_REQUEST_ID, + ratio, + segment_capacity=model.max_req_num + 1 if self.is_prefill else None, + ) + cos, sin = combine_mrope(self.position_cos, self.position_sin, config["rope_scaling"]["mrope_section"]) + self.qwen4_exp_batch = QSABatch.from_metadata( + self.qwen4_exp_metadata, self.req_manager.req_to_token_indexs, self.mem_index, cos, sin, ratio + ) + from .attention_state import Qwen4ExpAttentionState + from .lexical_state import Qwen4ExpLexicalState + + states = self.att_state_manager.get_full_attention_states(self.att_state_role) + self.qwen4_exp_qsa = {} + for state in states: + if isinstance(state, Qwen4ExpAttentionState): + state.store_positions(self.qwen4_exp_batch.mem_indices, cos, sin) + for group in state.config.logical_layer_groups: + for layer in group: + self.qwen4_exp_qsa[layer] = state.qsa_view(layer) + self.qwen4_exp_lexical_state = next( + ( + state + for state in self.att_state_manager.get_bounded_attention_states(self.att_state_role) + if isinstance(state, Qwen4ExpLexicalState) + ), + None, + ) + self.qwen4_exp_ple = ( + {} + if self.qwen4_exp_lexical_state is None + else self.qwen4_exp_lexical_state.prepare_lexical( + self.qwen4_exp_metadata, reset_empty_prefix=self.is_prefill + ) + ) + self.qwen4_exp_auto_commit = self.is_prefill or model.args.mtp_step == 0 + + def commit_lexical(self, accepted_lengths): + self.qwen4_exp_lexical_state.commit_lexical(self.qwen4_exp_metadata, self.qwen4_exp_ple, accepted_lengths) + + def _copy_qwen_state(self, new_state): + if self.qwen4_exp_auto_commit != new_state.qwen4_exp_auto_commit: + raise ValueError("Qwen graph replay cannot change lexical commit mode") + self.qwen4_exp_metadata.copy_from(new_state.qwen4_exp_metadata) + old, new = self.qwen4_exp_batch, new_state.qwen4_exp_batch + if old.max_new_blocks != new.max_new_blocks: + raise ValueError("Qwen graph replay requires the same token capacity") + for name in ("mem_indices", "cos", "sin"): + getattr(old, name).copy_(getattr(new, name), non_blocking=True) + for layer, batch in self.qwen4_exp_ple.items(): + source = new_state.qwen4_exp_ple[layer] + batch.ngram_history.copy_(source.ngram_history, non_blocking=True) + batch.conv_history.copy_(source.conv_history, non_blocking=True) + + def copy_for_cuda_graph(self, new_infer_state): + super().copy_for_cuda_graph(new_infer_state) + self._copy_qwen_state(new_infer_state) + + def copy_for_prefill_cuda_graph(self, new_infer_state): + super().copy_for_prefill_cuda_graph(new_infer_state) + self._copy_qwen_state(new_infer_state) diff --git a/lightllm/models/qwen4_exp/layer_infer/__init__.py b/lightllm/models/qwen4_exp/layer_infer/__init__.py new file mode 100644 index 0000000000..2ae28399f5 --- /dev/null +++ b/lightllm/models/qwen4_exp/layer_infer/__init__.py @@ -0,0 +1 @@ +pass diff --git a/lightllm/models/qwen4_exp/layer_infer/post_layer_infer.py b/lightllm/models/qwen4_exp/layer_infer/post_layer_infer.py new file mode 100644 index 0000000000..a685363b1d --- /dev/null +++ b/lightllm/models/qwen4_exp/layer_infer/post_layer_infer.py @@ -0,0 +1,6 @@ +from lightllm.models.llama.layer_infer.post_layer_infer import LlamaPostLayerInfer + + +class Qwen4ExpPostLayerInfer(LlamaPostLayerInfer): + def _norm(self, input, infer_state, layer_weight): + return layer_weight.hyper_connection_mixer.component.read(input) diff --git a/lightllm/models/qwen4_exp/layer_infer/pre_layer_infer.py b/lightllm/models/qwen4_exp/layer_infer/pre_layer_infer.py new file mode 100644 index 0000000000..129ab20853 --- /dev/null +++ b/lightllm/models/qwen4_exp/layer_infer/pre_layer_infer.py @@ -0,0 +1,23 @@ +from lightllm.models.qwen3_vl.layer_infer.pre_layer_infer import Qwen3VLMultimodalPreLayerInfer +from ..multimodal import canonical_image_ids + + +class Qwen4ExpPreLayerInfer(Qwen3VLMultimodalPreLayerInfer): + def __init__(self, network_config): + super().__init__(network_config) + self.branches = network_config["hc_count"] + self.image_token_id = network_config.get("image_token_id") + + def context_forward(self, input_ids, infer_state, layer_weight): + if not any(params.get("images") or params.get("audios") for params in (infer_state.multimodal_params or [])): + return self.token_forward(input_ids, infer_state, layer_weight) + embeddings = super().context_forward(input_ids, infer_state, layer_weight) + infer_state.qwen4_exp_lexical_input_ids = canonical_image_ids( + input_ids, infer_state.img_start_token_ids, infer_state.img_token_lens, self.image_token_id + ) + return embeddings.repeat(1, self.branches) + + def token_forward(self, input_ids, infer_state, layer_weight): + infer_state.qwen4_exp_lexical_input_ids = input_ids + embeddings = super().token_forward(input_ids, infer_state, layer_weight) + return embeddings.repeat(1, self.branches) diff --git a/lightllm/models/qwen4_exp/layer_infer/transformer_layer_infer.py b/lightllm/models/qwen4_exp/layer_infer/transformer_layer_infer.py new file mode 100644 index 0000000000..7ccb2e98d3 --- /dev/null +++ b/lightllm/models/qwen4_exp/layer_infer/transformer_layer_infer.py @@ -0,0 +1,93 @@ +from lightllm.common.basemodel.tp_reduce import ordered_tp_sum +from lightllm.utils.envs_utils import get_env_start_args +from lightllm.models.qwen3_5.layer_infer.transformer_layer_infer import Qwen35TransformerLayerInfer + + +class Qwen4ExpTransformerLayerInfer(Qwen35TransformerLayerInfer): + def __init__(self, layer_num, network_config): + super().__init__(layer_num, network_config) + layer_types = network_config.get("layer_types") + if layer_types is not None: + self.is_linear_attention_layer = layer_types[layer_num] == "linear_attention" + + def _reduce_linear_output(self, output, infer_state): + return self._tpsp_reduce(output, infer_state) + + def _tpsp_reduce(self, input, infer_state): + if self.tp_world_size_ == 1: + return input + scatter_rank = self.tp_rank_ if get_env_start_args().enable_tpsp_mix_mode else None + return ordered_tp_sum(input, infer_state.dist_group, self.tp_world_size_, scatter_rank=scatter_rank) + + def _forward(self, residual, infer_state, layer_weight, is_prefill): + if hasattr(layer_weight, "ple"): + ple_batch = infer_state.qwen4_exp_ple[self.layer_num_] + lexical_tokens = getattr(infer_state, "qwen4_exp_lexical_input_ids", infer_state.input_ids) + lexical = ple_batch.apply( + self._tpsp_allgather(residual, infer_state), + lexical_tokens, + layer_weight, + self.tp_world_size_, + infer_state.dist_group, + ) + auto_commit = getattr(infer_state, "qwen4_exp_auto_commit", None) + if auto_commit: + infer_state.qwen4_exp_lexical_state.lexical_caches[self.layer_num_].commit( + infer_state.qwen4_exp_metadata, ple_batch + ) + elif auto_commit is False: + infer_state.qwen4_exp_lexical_state.stage_speculative_lexical( + self.layer_num_, infer_state.qwen4_exp_metadata, ple_batch + ) + residual = residual + self._tpsp_sp_split(lexical, infer_state) + attention_residual = layer_weight.attn_hyper_connection.component + mixed, gates = attention_residual.read(residual) + if is_prefill: + update = self.context_attention_forward(mixed, infer_state, layer_weight) + else: + update = self.token_attention_forward(mixed, infer_state, layer_weight) + residual = attention_residual.write(residual, update, gates) + mlp_residual = layer_weight.mlp_hyper_connection.component + mixed, gates = mlp_residual.read(residual) + update = self._ffn(mixed, infer_state, layer_weight) + return mlp_residual.write(residual, update, gates) + + def context_forward(self, input_embdings, infer_state, layer_weight): + return self._forward(input_embdings, infer_state, layer_weight, True) + + def token_forward(self, input_embdings, infer_state, layer_weight): + return self._forward(input_embdings, infer_state, layer_weight, False) + + def _qsa_attention_forward(self, input_embdings, infer_state, layer_weight): + query, kv = self._get_qkv(input_embdings, infer_state, layer_weight) + self._post_cache_kv(kv, infer_state, layer_weight) + from ..attention_state import Qwen4ExpAttentionState + + state = infer_state.att_state_manager.get_state( + infer_state.att_state_role, "qwen4_gqa", self.layer_num_, Qwen4ExpAttentionState + ) + key, value = state.get_layer_view(self.layer_num_) + hidden = self._tpsp_allgather(input_embdings, infer_state) + output = infer_state.qwen4_exp_qsa[self.layer_num_].attend( + layer_weight.indexer.component, + hidden, + query.view(-1, self.tp_q_head_num_, self.head_dim_), + key, + value, + infer_state.qwen4_exp_batch, + decode=not infer_state.is_prefill, + attention_splits=None if infer_state.is_prefill else 1, + ) + return self._get_o(output, infer_state, layer_weight) + + def context_attention_forward(self, input_embdings, infer_state, layer_weight): + if self.is_linear_attention_layer: + input_embdings = self._tpsp_allgather(input_embdings, infer_state) + return super().context_attention_forward(input_embdings, infer_state, layer_weight) + return self._qsa_attention_forward(input_embdings, infer_state, layer_weight) + + def token_attention_forward(self, input_embdings, infer_state, layer_weight): + if self.is_linear_attention_layer: + input_embdings = self._tpsp_allgather(input_embdings, infer_state) + return super().token_attention_forward(input_embdings, infer_state, layer_weight) + return self._qsa_attention_forward(input_embdings, infer_state, layer_weight) diff --git a/lightllm/models/qwen4_exp/layer_weights/__init__.py b/lightllm/models/qwen4_exp/layer_weights/__init__.py new file mode 100644 index 0000000000..2ae28399f5 --- /dev/null +++ b/lightllm/models/qwen4_exp/layer_weights/__init__.py @@ -0,0 +1 @@ +pass diff --git a/lightllm/models/qwen4_exp/layer_weights/components.py b/lightllm/models/qwen4_exp/layer_weights/components.py new file mode 100644 index 0000000000..8c0e9d5717 --- /dev/null +++ b/lightllm/models/qwen4_exp/layer_weights/components.py @@ -0,0 +1,129 @@ +from functools import partial +from typing import Callable, Mapping + +import torch + +from lightllm.common.basemodel.layer_weights.meta_weights.base_weight import BaseWeight +from lightllm.common.basemodel.layer_weights.meta_weights.parameter_weight import ParameterWeight +from lightllm.common.gated_residual import GatedResidual +from lightllm.common.ngram_embedding import NGramHasher +from lightllm.common.ngram_embedding.ple import PLELayer +from lightllm.common.qsa.indexer import QSAIndexer +from lightllm.models.qwen4_exp.checkpoint import ngram_eos_token_id + + +class ComponentWeight(BaseWeight): + def __init__(self, specs: Mapping[str, tuple], builder: Callable): + super().__init__() + self._specs = specs + self._builder = builder + self._component = None + self._create_weight() + + def _create_weight(self): + self.parameters = {name: ParameterWeight(name, dtype, shape) for name, (shape, dtype) in self._specs.items()} + + def load_hf_weights(self, weights): + for name, parameter in self.parameters.items(): + if name not in weights: + continue + source = weights[name] + shape, dtype = self._specs[name] + if source.shape != shape: + raise ValueError(f"Component checkpoint shape mismatch for {name}: {tuple(source.shape)} != {shape}") + if dtype == torch.int64 and source.dtype != torch.int64: + raise ValueError(f"Integer component buffer {name} must use int64") + if dtype != torch.int64 and source.dtype not in (torch.float16, torch.bfloat16, torch.float32): + raise ValueError(f"Component tensor {name} requires unquantized floating point weights") + parameter.load_hf_weights(weights) + + def verify_load(self): + ready = all(parameter.verify_load() for parameter in self.parameters.values()) + if ready and self._component is None: + self._component = self._builder({name: parameter.weight for name, parameter in self.parameters.items()}) + return ready + + @property + def component(self): + if self._component is None and not self.verify_load(): + raise RuntimeError("Component weights are not fully loaded") + return self._component + + +def gated_residual_weight(prefix: str, config: dict, dtype: torch.dtype, use_combine: bool = True): + width = config["hc_count"] * config["hidden_size"] + rank = config["hc_lowrank"] + specs = { + f"{prefix}.hc_norm.weight": ((width,), dtype), + f"{prefix}.input_mix_weight_down.weight": ((rank, width), dtype), + f"{prefix}.input_mix_weight_up.weight": ((width, rank), dtype), + } + if use_combine: + specs[f"{prefix}.block_inject_weight.weight"] = ((config["hc_count"], width), dtype) + return ComponentWeight( + specs, + partial( + GatedResidual.from_checkpoint, + prefix=prefix, + branches=config["hc_count"], + eps=config["rms_norm_eps"], + use_combine=use_combine, + ), + ) + + +def qsa_indexer_weight(prefix: str, config: dict, dtype: torch.dtype): + dim, heads = config["indexer_head_dim"], config["indexer_n_heads"] + if config["indexer_kv_heads"] != 1: + raise ValueError("Qwen4Exp QSA requires one index key head") + specs = { + f"{prefix}.index_qk_proj.weight": (((heads + 1) * dim, config["hidden_size"]), dtype), + f"{prefix}.q_layernorm.weight": ((dim,), dtype), + f"{prefix}.k_layernorm.weight": ((dim,), dtype), + } + return ComponentWeight( + specs, + partial( + QSAIndexer.from_checkpoint, + prefix=prefix, + query_heads=heads, + token_budget=config["indexer_budget"], + use_cuda_topk=True, + compress_ratio=config["indexer_compress_ratio"], + eps=config["rms_norm_eps"], + ), + ) + + +def ple_weight(prefix: str, config: dict, dtype: torch.dtype): + width = config["hc_count"] * config["hidden_size"] + embed = config.get("ple_embed_dim", config["hidden_size"]) + specs = { + f"{prefix}.key_proj.weight": ((width, embed), dtype), + f"{prefix}.value_proj.weight": ((config["hidden_size"], embed), dtype), + f"{prefix}.conv1d.weight": ((width, 1, config["ple_conv_kernel_size"]), dtype), + } + for norm in ("norm_key", "norm_query", "norm_conv"): + specs[f"{prefix}.{norm}.weight"] = ((width,), dtype) + return ComponentWeight( + specs, + partial( + PLELayer.from_checkpoint, + prefix=prefix, + branches=config["hc_count"], + dilation=config["ngram_size"], + eps=config["rms_norm_eps"], + ), + ) + + +def ngram_hasher_weight(prefix: str, config: dict): + heads = (config["ngram_size"] - 1) * config["heads_per_ngram"] + specs = { + f"{prefix}.layer_multipliers": ((config["ngram_size"],), torch.int64), + f"{prefix}.ngram_heads_vocab_sizes": ((heads,), torch.int64), + f"{prefix}.ngram_heads_offsets": ((heads,), torch.int64), + } + return ComponentWeight( + specs, partial(NGramHasher.from_checkpoint, prefix=prefix, eos_token_id=ngram_eos_token_id(config)) + ) diff --git a/lightllm/models/qwen4_exp/layer_weights/embedding_weight.py b/lightllm/models/qwen4_exp/layer_weights/embedding_weight.py new file mode 100644 index 0000000000..ef2f5980dd --- /dev/null +++ b/lightllm/models/qwen4_exp/layer_weights/embedding_weight.py @@ -0,0 +1,14 @@ +from lightllm.common.basemodel.layer_weights.meta_weights.row_sharded_embedding_weight import RowShardedEmbeddingWeight +from lightllm.common.basemodel.layer_weights.meta_weights.pinned_embedding_weight import PinnedRowShardedEmbeddingWeight +from lightllm.utils.envs_utils import get_env_start_args + + +def ngram_embedding_weight(dim, rows, name, dtype, shard_count): + args = get_env_start_args() + if not args.get("offload_ngram_embedding", False): + return RowShardedEmbeddingWeight(dim, rows, name, dtype, shard_count) + if args.get("hardware_platform", "cuda") != "cuda": + raise ValueError("N-gram embedding offload requires CUDA Triton lookup") + from ..pinned_memory import pinned_zeros + + return PinnedRowShardedEmbeddingWeight(dim, rows, name, dtype, shard_count, cpu_allocator=pinned_zeros) diff --git a/lightllm/models/qwen4_exp/layer_weights/lm_head_weight.py b/lightllm/models/qwen4_exp/layer_weights/lm_head_weight.py new file mode 100644 index 0000000000..72dd4d1f50 --- /dev/null +++ b/lightllm/models/qwen4_exp/layer_weights/lm_head_weight.py @@ -0,0 +1,25 @@ +from typing import Optional + +import torch + +from lightllm.common.basemodel.layer_weights.meta_weights import LMHeadWeight + + +class Qwen4ExpLMHeadWeight(LMHeadWeight): + def _cuda_forward( + self, input: torch.Tensor, out: Optional[torch.Tensor] = None, alloc_func=torch.empty + ) -> torch.Tensor: + rows = input.shape[1] + if rows <= 1 or rows % 8 == 0 or input.dtype not in (torch.float16, torch.bfloat16): + return super()._cuda_forward(input, out=out, alloc_func=alloc_func) + + tokens = input.transpose(0, 1) + padded = torch.cat((tokens, tokens[:1].expand((-rows) % 8, -1)), dim=0) + scores = super()._cuda_forward(padded.transpose(0, 1), alloc_func=alloc_func) + shape = (self.weight.shape[0], rows) + if out is None: + out = alloc_func(shape, dtype=input.dtype, device=input.device) + elif out.shape != shape or out.dtype != input.dtype or out.device != input.device: + raise ValueError("LM-head output must match the unpadded shape, dtype and device") + out.copy_(scores[:, :rows]) + return out diff --git a/lightllm/models/qwen4_exp/layer_weights/norm_weight.py b/lightllm/models/qwen4_exp/layer_weights/norm_weight.py new file mode 100644 index 0000000000..1824e4948a --- /dev/null +++ b/lightllm/models/qwen4_exp/layer_weights/norm_weight.py @@ -0,0 +1,30 @@ +import torch + +from lightllm.common.basemodel.layer_weights.meta_weights import GatedRMSNormWeight +from lightllm.models.qwen4_exp.triton_kernel.gated_norm import gated_norm + + +class Qwen4ExpGatedNormWeight(GatedRMSNormWeight): + def __init__(self, dim, weight_name, data_type, activation="sigmoid"): + if activation not in ("sigmoid", "silu"): + raise ValueError("Qwen4Exp GDN output gate must be sigmoid or silu") + self.activation = activation + super().__init__(dim, weight_name, data_type) + + def _triton_forward(self, input, gate_value, eps, out=None, alloc_func=torch.empty): + if out is None: + out = alloc_func(input.shape, dtype=input.dtype, device=input.device) + return gated_norm(input, gate_value, self.weight, eps, self.activation, out) + + def _native_forward(self, input, gate_value, eps, out=None, alloc_func=torch.empty): + values = input.float() + normalized = values * torch.rsqrt(values.square().mean(-1, keepdim=True) + eps) + weighted = self.weight * normalized.to(input.dtype) + gate = gate_value.float().sigmoid() + if self.activation == "silu": + gate = gate * gate_value.float() + result = (weighted * gate).to(input.dtype) + if out is not None: + out.copy_(result) + return out + return result diff --git a/lightllm/models/qwen4_exp/layer_weights/pre_and_post_layer_weight.py b/lightllm/models/qwen4_exp/layer_weights/pre_and_post_layer_weight.py new file mode 100644 index 0000000000..dabedaed88 --- /dev/null +++ b/lightllm/models/qwen4_exp/layer_weights/pre_and_post_layer_weight.py @@ -0,0 +1,26 @@ +from lightllm.common.basemodel import PreAndPostLayerWeight +from lightllm.common.basemodel.layer_weights.meta_weights import EmbeddingWeight +from lightllm.models.qwen4_exp.checkpoint import normalize_text_weight_names +from .components import gated_residual_weight +from .lm_head_weight import Qwen4ExpLMHeadWeight + + +class Qwen4ExpPreAndPostLayerWeight(PreAndPostLayerWeight): + def __init__(self, data_type, network_config): + super().__init__(data_type, network_config) + hidden, vocab = network_config["hidden_size"], network_config["vocab_size"] + self.wte_weight_ = EmbeddingWeight(hidden, vocab, "model.embed_tokens.weight", data_type) + self.lm_head_weight_ = Qwen4ExpLMHeadWeight( + hidden, + vocab, + "lm_head.weight", + data_type, + self.wte_weight_ if network_config.get("tie_word_embeddings", False) else None, + ) + self.hyper_connection_mixer = gated_residual_weight( + "model.hyper_connection_mixer", network_config, data_type, use_combine=False + ) + + def load_hf_weights(self, weights): + normalize_text_weight_names(weights) + super().load_hf_weights(weights) diff --git a/lightllm/models/qwen4_exp/layer_weights/transformer_layer_weight.py b/lightllm/models/qwen4_exp/layer_weights/transformer_layer_weight.py new file mode 100644 index 0000000000..25e2d082b2 --- /dev/null +++ b/lightllm/models/qwen4_exp/layer_weights/transformer_layer_weight.py @@ -0,0 +1,85 @@ +import torch + +from lightllm.models.qwen3_5_moe.layer_weights.transformer_layer_weight import Qwen35MOETransformerLayerWeight +from lightllm.common.basemodel.layer_weights.meta_weights import QKGEMMANormWeight +from lightllm.common.quantization.scalar_projection import StableScalarProjection +from lightllm.common.quantization.aligned_projection import OutputAlignedProjection +from lightllm.common.quantization.no_quant import NoQuantization +from lightllm.models.qwen4_exp.checkpoint import normalize_text_weight_names, ngram_table_shape +from .components import gated_residual_weight, qsa_indexer_weight, ple_weight, ngram_hasher_weight +from .norm_weight import Qwen4ExpGatedNormWeight +from .embedding_weight import ngram_embedding_weight + + +class Qwen4ExpTransformerLayerWeight(Qwen35MOETransformerLayerWeight): + def get_quant_method(self, name): + method = super().get_quant_method(name) + if name == "in_proj_weight" and type(method) is NoQuantization: + aligned = OutputAlignedProjection() + aligned.hf_quantization_config = method.hf_quantization_config + return aligned + return method + + def _init_moe(self): + super()._init_moe() + + self.shared_expert_gate.quant_method = StableScalarProjection() + + def _init_weight(self): + config = self.network_config_ + layer_types = config.get("layer_types") + if layer_types is not None: + layer_type = layer_types[self.layer_num_] + if layer_type not in ("linear_attention", "full_attention", "indexed_attention"): + raise ValueError(f"Unsupported Qwen4Exp layer type {layer_type}") + self.is_linear_attention_layer = layer_type == "linear_attention" + super()._init_weight() + prefix = f"model.layers.{self.layer_num_}" + self.attn_hyper_connection = gated_residual_weight(f"{prefix}.attn_hyper_connection", config, self.data_type_) + self.mlp_hyper_connection = gated_residual_weight(f"{prefix}.mlp_hyper_connection", config, self.data_type_) + if not self.is_linear_attention_layer: + self.indexer = qsa_indexer_weight(f"{prefix}.self_attn.indexer", config, self.data_type_) + ple_layers = sorted(set(config.get("ple_layer_ids", []))) + if self.layer_num_ + 1 in ple_layers: + if not self.is_linear_attention_layer: + raise ValueError("Qwen4Exp PLE must be attached to a linear-attention layer") + self.ple = ple_weight(f"{prefix}.ple", config, self.data_type_) + self.ngram_hasher = ngram_hasher_weight(f"{prefix}.ple.ple_embedding", config) + heads = (config["ngram_size"] - 1) * config["heads_per_ngram"] + rows, dim = ngram_table_shape( + config["ngram_vocab_size_base"], + heads, + ple_layers.index(self.layer_num_ + 1), + config.get("ple_embed_dim", config["hidden_size"]), + config["make_ngram_vocab_size_divisible_by"], + ) + self.ngram_embedding = ngram_embedding_weight( + dim, + rows, + f"{prefix}.ple.ple_embedding.ngram_embedding.weight", + self.data_type_, + config["split_ngram_parts"], + ) + + def _init_gdn_weight(self): + super()._init_gdn_weight() + config = self.network_config_ + self.linear_norm = Qwen4ExpGatedNormWeight( + self.linear_v_head_dim, + f"model.layers.{self.layer_num_}.linear_attn.norm.weight", + self.data_type_, + config.get("output_gate_type") or config["hidden_act"], + ) + + def _init_norm(self): + if not self.is_linear_attention_layer: + self.qk_norm_weight_ = QKGEMMANormWeight( + self.head_dim, + self._q_norm_name, + self._k_norm_name, + torch.float32, + ) + + def load_hf_weights(self, weights): + normalize_text_weight_names(weights) + super().load_hf_weights(weights) diff --git a/lightllm/models/qwen4_exp/lexical_cache.py b/lightllm/models/qwen4_exp/lexical_cache.py new file mode 100644 index 0000000000..a5cdedab4b --- /dev/null +++ b/lightllm/models/qwen4_exp/lexical_cache.py @@ -0,0 +1,89 @@ +from dataclasses import dataclass + +import torch +import triton +import triton.language as tl + +from .state import PLEBatch +from .checkpoint import ngram_eos_token_id +from lightllm.utils.dist_utils import get_current_device_id + + +@triton.jit +def _scatter_history(Source, Destinations, Target, WIDTH: tl.constexpr, BLOCK: tl.constexpr): + row = tl.program_id(0) + destination = tl.load(Destinations + row).to(tl.int64) + if destination >= 0: + col = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + values = tl.load(Source + row.to(tl.int64) * WIDTH + col, col < WIDTH, other=0) + tl.store(Target + destination * WIDTH + col, values, col < WIDTH) + + +@dataclass(frozen=True) +class LexicalSnapshot: + ngram: torch.Tensor + convolution: torch.Tensor + + +class LexicalRequestCache: + def __init__(self, layer_weight, requests: int): + if requests < 1: + raise ValueError("Lexical cache requires at least one request slot") + self.layer_weight = layer_weight + config = layer_weight.network_config_ + self.eos_token_id = ngram_eos_token_id(config) + + device = torch.device("cuda", get_current_device_id()) + self.ngram = torch.full( + (requests, config["ngram_size"] - 1), self.eos_token_id, dtype=torch.int64, device=device + ) + self.convolution = torch.zeros( + ( + requests, + (config["ple_conv_kernel_size"] - 1) * config["ngram_size"], + config["hc_count"] * config["hidden_size"], + ), + dtype=layer_weight.data_type_, + device=device, + ) + + def reset(self, request: int): + self.ngram[request].fill_(self.eos_token_id) + self.convolution[request].zero_() + + @torch.no_grad() + def prepare(self, metadata, reset_empty_prefix=False) -> PLEBatch: + rows = metadata.active_rows + indices = rows.clamp_min(0).to(torch.int64) + ngram = self.ngram.index_select(0, indices) + convolution = self.convolution.index_select(0, indices) + empty = rows < 0 + if reset_empty_prefix: + counts = metadata.starts[1:] - metadata.starts[:-1] + empty = empty | (metadata.key_lengths == counts) + ngram.masked_fill_(empty[:, None], self.eos_token_id) + convolution.masked_fill_(empty[:, None, None], 0) + return PLEBatch(metadata.starts, ngram, convolution, metadata.token_mask) + + @torch.no_grad() + def commit(self, metadata, batch: PLEBatch, accepted_lengths=None): + ngram, convolution = batch.histories_after(self.layer_weight, accepted_lengths) + for values, storage in ((ngram, self.ngram), (convolution, self.convolution)): + width = storage[0].numel() + if metadata.active_rows.numel() and width: + _scatter_history[(metadata.active_rows.numel(), triton.cdiv(width, 256))]( + values, metadata.active_rows, storage, width, 256 + ) + + def snapshot(self, request: int) -> LexicalSnapshot: + return LexicalSnapshot(self.ngram[request].clone(), self.convolution[request].clone()) + + def restore(self, request: int, snapshot: LexicalSnapshot): + for target, source in ( + (self.ngram[request], snapshot.ngram), + (self.convolution[request], snapshot.convolution), + ): + if target.shape != source.shape or target.dtype != source.dtype: + raise ValueError("Lexical snapshot has incompatible shape or dtype") + self.ngram[request].copy_(snapshot.ngram, non_blocking=True) + self.convolution[request].copy_(snapshot.convolution, non_blocking=True) diff --git a/lightllm/models/qwen4_exp/lexical_state.py b/lightllm/models/qwen4_exp/lexical_state.py new file mode 100644 index 0000000000..cf13172468 --- /dev/null +++ b/lightllm/models/qwen4_exp/lexical_state.py @@ -0,0 +1,257 @@ +from dataclasses import dataclass +from functools import lru_cache +from math import prod + +import torch + +from lightllm.common.attention_state.bounded_att import BoundedAttentionState +from lightllm.common.attention_state.bounded_att.copy_kernel import copy_bounded_att_tensors +from lightllm.common.attention_state.bounded_att.transfer_utils import bounded_att_buffer_copies +from lightllm.common.attention_state.bounded_att.utils import create_page_pool_tensor, register_page_pool_tensor +from lightllm.common.attention_state.types import AxisMapping, TensorField, TransferRegion +from lightllm.common.attention_state.utils import deepcopy_result +from .lexical_cache import LexicalRequestCache +from .speculative_cache import LexicalSpeculativeCache + + +@dataclass(frozen=True) +class LexicalStateConfig: + logical_layer_ids: tuple[int, ...] + ngram_context: int + conv_context: int + width: int + eos_token_id: int + dtype: torch.dtype + tp_rank: int = 0 + + +class LexicalPages: + def __init__(self, capacity, config, device, pin_memory): + count = len(config.logical_layer_ids) + self.ngram = create_page_pool_tensor( + size=(capacity, count, config.ngram_context), + dtype=torch.int64, + device=torch.device(device), + pin_memory=pin_memory, + ) + self.convolution = create_page_pool_tensor( + size=(capacity, count, config.conv_context, config.width), + dtype=config.dtype, + device=torch.device(device), + pin_memory=pin_memory, + ) + + def register_pin_memory(self): + for tensor in (self.ngram, self.convolution): + register_page_pool_tensor(tensor=tensor) + + +class Qwen4ExpLexicalState(BoundedAttentionState): + def __init__(self, config, request_capacity, mtp_step=0, device="meta"): + self.config = config + self.state_name = None + self.request_capacity = request_capacity + self.mtp_step = mtp_step + count = len(config.logical_layer_ids) + self.ngram = torch.full( + (count, request_capacity, config.ngram_context), config.eos_token_id, dtype=torch.int64, device=device + ) + self.convolution = torch.zeros( + (count, request_capacity, config.conv_context, config.width), dtype=config.dtype, device=device + ) + self.big_pages_pool = self.small_pages_pool = None + self.lexical_caches = {} + self.speculative_lexical_caches = {} + + def __getstate__(self): + state = self.__dict__.copy() + state["lexical_caches"] = {} + state["speculative_lexical_caches"] = {} + return state + + def get_state_type(self): + return "qwen4_lexical" + + def contains_logical_layer(self, logical_layer_id): + return logical_layer_id in self.config.logical_layer_ids + + def get_layer_view(self, logical_layer_id): + layer = self.config.logical_layer_ids.index(logical_layer_id) + return self.ngram[layer], self.convolution[layer] + + def bind_weights(self, layer_weights): + for weight in layer_weights: + if not self.contains_logical_layer(weight.layer_num_): + continue + cache = LexicalRequestCache.__new__(LexicalRequestCache) + cache.layer_weight = weight + cache.eos_token_id = self.config.eos_token_id + cache.ngram, cache.convolution = self.get_layer_view(weight.layer_num_) + self.lexical_caches[weight.layer_num_] = cache + if self.mtp_step: + self.speculative_lexical_caches[weight.layer_num_] = LexicalSpeculativeCache(cache, self.mtp_step + 1) + + def prepare_lexical(self, metadata, reset_empty_prefix=False): + return {layer: cache.prepare(metadata, reset_empty_prefix) for layer, cache in self.lexical_caches.items()} + + def commit_lexical(self, metadata, batches, accepted_lengths=None): + for layer, cache in self.lexical_caches.items(): + cache.commit(metadata, batches[layer], accepted_lengths) + + def stage_speculative_lexical(self, layer, metadata, batch): + self.speculative_lexical_caches[layer].stage(metadata, batch) + + def initialize_page_pools(self, big_page_capacity, small_page_capacity=0, device="cpu", pin_memory=True): + self.big_pages_pool = ( + LexicalPages(big_page_capacity, self.config, device, pin_memory) if big_page_capacity else None + ) + self.small_pages_pool = ( + LexicalPages(small_page_capacity, self.config, device, pin_memory) if small_page_capacity else None + ) + + def reset_bounded_att_state(self, req_indices): + for request in req_indices: + self.ngram[:, request].fill_(self.config.eos_token_id) + self.convolution[:, request].zero_() + + def _pool(self, page_kind): + assert page_kind in ("big", "small") + pool = self.big_pages_pool if page_kind == "big" else self.small_pages_pool + assert pool is not None, f"{page_kind} lexical pages are disabled" + return pool + + def _copy_pages(self, req_indices, state_token_lens, page_kind, page_slot_ids, write): + assert len(req_indices) == len(state_token_lens) == len(page_slot_ids) + pool = self._pool(page_kind) + for name in ("ngram", "convolution"): + copies = [ + (getattr(self, name)[:, req], getattr(pool, name)[slot]) + for req, slot in zip(req_indices, page_slot_ids) + ] + if not write: + copies = [(dst, src) for src, dst in copies] + copy_bounded_att_tensors(copies=copies) + + def save_bounded_att_runtime_to_pages(self, req_indices, state_token_lens, page_kind, page_slot_ids): + self._copy_pages(req_indices, state_token_lens, page_kind, page_slot_ids, True) + + def restore_bounded_att_runtime_from_pages(self, req_indices, state_token_lens, page_kind, page_slot_ids): + self._copy_pages(req_indices, state_token_lens, page_kind, page_slot_ids, False) + + def commit_bounded_att_state(self, req_indices, accepted_steps): + starts = torch.arange(req_indices.numel(), dtype=torch.int32, device=req_indices.device) + for cache in self.speculative_lexical_caches.values(): + cache.commit(req_indices, starts, accepted_steps, self.request_capacity - 1) + + @deepcopy_result + @lru_cache(maxsize=None) + def get_cpu_transfer_regions(self): + return self._regions() + + @deepcopy_result + @lru_cache(maxsize=None) + def get_pd_transfer_regions(self): + return self._regions() + + def _regions(self): + regions = [] + offset = 0 + for name, tensor in (("ngram", self.ngram), ("convolution", self.convolution)): + shape = tensor.shape[:1] + tensor.shape[2:] + axes = ("layer", "context") if name == "ngram" else ("layer", "context", "dim") + end = offset + prod(shape) * tensor.element_size() + regions.append( + TransferRegion( + field=name, + byte_start=offset, + byte_end=end, + dtype=tensor.dtype, + axes=axes, + shape=tuple(shape), + axis_mappings=tuple( + AxisMapping(axis=axis, local_start=0, external_start=0, length=size) + for axis, size in zip(axes[1:], shape[1:]) + ), + write_enabled=self.config.tp_rank == 0, + ) + ) + offset = end + return tuple(regions) + + def _fields(self, ngram, convolution): + return { + "ngram": TensorField(ngram, ("layer", "context")), + "convolution": TensorField(convolution, ("layer", "context", "dim")), + } + + def _transfer(self, fields, buffers, pages, layout, write): + assert layout.state_name == self.state_name + for copies in bounded_att_buffer_copies(fields, buffers, layout, pages, write): + copy_bounded_att_tensors(copies=copies) + + def write_state_pages_to_cpu( + self, page_kind, page_slot_ids, cpu_page_buffers, cpu_page_indices, layout, cpu_page_ready_flags=None + ): + assert len(page_slot_ids) == len(cpu_page_indices) + if cpu_page_ready_flags is not None: + assert len(cpu_page_ready_flags) == len(cpu_page_indices) + pool = self._pool(page_kind) + selected = [i for i in range(len(page_slot_ids)) if cpu_page_ready_flags is None or not cpu_page_ready_flags[i]] + self._transfer( + [self._fields(pool.ngram[page_slot_ids[i]], pool.convolution[page_slot_ids[i]]) for i in selected], + cpu_page_buffers, + [cpu_page_indices[i] for i in selected], + layout, + True, + ) + + def read_state_pages_from_cpu(self, cpu_page_buffers, cpu_page_indices, big_page_slot_ids, layout): + assert len(cpu_page_indices) == len(big_page_slot_ids) + pool = self._pool("big") + self._transfer( + [self._fields(pool.ngram[slot], pool.convolution[slot]) for slot in big_page_slot_ids], + cpu_page_buffers, + cpu_page_indices, + layout, + False, + ) + + def copy_big_pages_from_dp(self, src_page_ids, dst_page_ids, remote_dp_bounded_state): + assert self.config == remote_dp_bounded_state.config + assert len(src_page_ids) == len(dst_page_ids) + for name in ("ngram", "convolution"): + copy_bounded_att_tensors( + copies=[ + (getattr(remote_dp_bounded_state._pool("big"), name)[src], getattr(self._pool("big"), name)[dst]) + for src, dst in zip(src_page_ids, dst_page_ids) + ] + ) + + def copy_runtime_state_from_dp(self, src_req_indices, dst_req_indices, remote_dp_bounded_state): + assert self.config == remote_dp_bounded_state.config + assert len(src_req_indices) == len(dst_req_indices) + for name in ("ngram", "convolution"): + copy_bounded_att_tensors( + copies=[ + (getattr(remote_dp_bounded_state, name)[:, src], getattr(self, name)[:, dst]) + for src, dst in zip(src_req_indices, dst_req_indices) + ] + ) + + def write_to_pd_page(self, req_idx, pd_page_buffers, pd_page_index, state_token_len, layout): + self._transfer( + [self._fields(self.ngram[:, req_idx], self.convolution[:, req_idx])], + pd_page_buffers, + [pd_page_index], + layout, + True, + ) + + def read_from_pd_page(self, pd_page_buffers, req_idx, pd_page_index, state_token_len, layout): + self._transfer( + [self._fields(self.ngram[:, req_idx], self.convolution[:, req_idx])], + pd_page_buffers, + [pd_page_index], + layout, + False, + ) diff --git a/lightllm/models/qwen4_exp/model.py b/lightllm/models/qwen4_exp/model.py new file mode 100644 index 0000000000..e5112afed1 --- /dev/null +++ b/lightllm/models/qwen4_exp/model.py @@ -0,0 +1,114 @@ +import json +from pathlib import Path + +from lightllm.models.qwen3_5_moe.model import Qwen3_5MOETpPartModel +from lightllm.models.registry import ModelRegistry +from lightllm.common.basemodel.attention.triton.fp import TritonAttBackend +from lightllm.utils.torch_dtype_utils import get_torch_dtype +from .configuration import normalize_model_config +from .infer_struct import Qwen4ExpInferStateInfo +from .layer_weights.pre_and_post_layer_weight import Qwen4ExpPreAndPostLayerWeight +from .layer_weights.transformer_layer_weight import Qwen4ExpTransformerLayerWeight +from .layer_infer.pre_layer_infer import Qwen4ExpPreLayerInfer +from .layer_infer.post_layer_infer import Qwen4ExpPostLayerInfer +from .layer_infer.transformer_layer_infer import Qwen4ExpTransformerLayerInfer + + +class Qwen4ExpTpPartModel(Qwen3_5MOETpPartModel): + linear_att_mtp_match_decode = True + pre_and_post_weight_class = Qwen4ExpPreAndPostLayerWeight + transformer_weight_class = Qwen4ExpTransformerLayerWeight + pre_layer_infer_class = Qwen4ExpPreLayerInfer + post_layer_infer_class = Qwen4ExpPostLayerInfer + transformer_layer_infer_class = Qwen4ExpTransformerLayerInfer + infer_state_class = Qwen4ExpInferStateInfo + + def _init_config(self): + with (Path(self.weight_dir_) / "config.json").open() as stream: + checkpoint = json.load(stream) + self.config, self.vision_config = normalize_model_config(checkpoint, self.finetune_config) + self.num_kv_heads = max(self.config["num_key_value_heads"] // self.tp_world_size_, 1) + + @classmethod + def normalize_model_config(cls, model_config): + config = normalize_model_config(model_config)[0] + interval = config["full_attention_interval"] + expected = [ + "full_attention" if (i + 1) % interval == 0 else "linear_attention" + for i in range(config["num_hidden_layers"]) + ] + actual = [ + "full_attention" if kind == "indexed_attention" else kind for kind in config.get("layer_types", expected) + ] + if actual != expected: + raise ValueError("Qwen4Exp requires the configured periodic attention layout") + config["layer_types"] = expected + return config + + @classmethod + def create_full_att_states(cls, model_config, args, model_logical_layer_ids, tp_rank=0, tp_world_size=1): + from lightllm.common.attention_state.full_att.gqa.gqa_state import GQAAttentionConfig + from .attention_state import Qwen4ExpAttentionState + + config = cls.normalize_model_config(model_config) + if args.llm_kv_type != "None": + raise ValueError("Qwen4Exp QSA currently requires unquantized KV storage") + interval = config["full_attention_interval"] + layers = tuple((layer,) for index, layer in enumerate(model_logical_layer_ids) if (index + 1) % interval == 0) + return [ + Qwen4ExpAttentionState( + GQAAttentionConfig( + logical_layer_groups=layers, + global_head_num=config["num_key_value_heads"], + head_dim=config["head_dim"], + dtype=get_torch_dtype(args.data_type), + tp_rank=tp_rank, + tp_world_size=tp_world_size, + ), + config["indexer_head_dim"], + int(config["head_dim"] * config["partial_rotary_factor"]), + args.page_size, + ) + ] + + @classmethod + def create_bounded_att_states( + cls, model_config, args, model_logical_layer_ids, request_capacity=1, tp_rank=0, tp_world_size=1, device="meta" + ): + from .lexical_state import LexicalStateConfig, Qwen4ExpLexicalState + from .checkpoint import ngram_eos_token_id + + config = cls.normalize_model_config(model_config) + states = super().create_bounded_att_states( + config, args, model_logical_layer_ids, request_capacity, tp_rank, tp_world_size, device + ) + layers = tuple(model_logical_layer_ids[layer - 1] for layer in config.get("ple_layer_ids", [])) + if layers: + states.append( + Qwen4ExpLexicalState( + LexicalStateConfig( + logical_layer_ids=layers, + ngram_context=config["ngram_size"] - 1, + conv_context=(config["ple_conv_kernel_size"] - 1) * config["ngram_size"], + width=config["hc_count"] * config["hidden_size"], + eos_token_id=ngram_eos_token_id(config), + dtype=get_torch_dtype(args.data_type), + tp_rank=tp_rank, + ), + request_capacity, + args.mtp_step, + device, + ) + ) + return states + + def prepare_attention_states(self): + super().prepare_attention_states() + from .lexical_state import Qwen4ExpLexicalState + + state = getattr(self.att_state_manager, f"{self.att_state_role}_qwen4_lexical_state", None) + if isinstance(state, Qwen4ExpLexicalState): + state.bind_weights(self.trans_layers_weight) + + def _init_att_backend(self): + self.prefill_att_backend = self.decode_att_backend = TritonAttBackend(model=self) diff --git a/lightllm/models/qwen4_exp/mtp/__init__.py b/lightllm/models/qwen4_exp/mtp/__init__.py new file mode 100644 index 0000000000..2ae28399f5 --- /dev/null +++ b/lightllm/models/qwen4_exp/mtp/__init__.py @@ -0,0 +1 @@ +pass diff --git a/lightllm/models/qwen4_exp/mtp/infer.py b/lightllm/models/qwen4_exp/mtp/infer.py new file mode 100644 index 0000000000..6e3ff96e3d --- /dev/null +++ b/lightllm/models/qwen4_exp/mtp/infer.py @@ -0,0 +1,29 @@ +import torch + +from lightllm.models.llama.layer_infer.pre_layer_infer import LlamaPreLayerInfer +from ..layer_infer.transformer_layer_infer import Qwen4ExpTransformerLayerInfer +from ..multimodal import canonical_image_ids + + +class Qwen4ExpMTPPreLayerInfer(LlamaPreLayerInfer): + def __init__(self, network_config): + super().__init__(network_config) + self.image_token_id = network_config.get("image_token_id") + + def context_forward(self, input_ids, infer_state, layer_weight): + images = [image for params in (infer_state.multimodal_params or []) for image in params.get("images", [])] + if images: + starts = torch.tensor([image["token_id"] for image in images], dtype=torch.int64, device=input_ids.device) + lengths = torch.tensor([image["token_num"] for image in images], dtype=torch.int64, device=input_ids.device) + input_ids = canonical_image_ids(input_ids, starts, lengths, self.image_token_id) + return self.token_forward(input_ids, infer_state, layer_weight) + + def token_forward(self, input_ids, infer_state, layer_weight): + embedding = super().token_forward(input_ids, infer_state, layer_weight) + return layer_weight.fusion.component(embedding, infer_state.mtp_draft_input_hiddens) + + +class Qwen4ExpMTPTransformerLayerInfer(Qwen4ExpTransformerLayerInfer): + def __init__(self, layer_num, network_config): + config = dict(network_config, layer_types=None, full_attention_interval=1) + super().__init__(layer_num, config) diff --git a/lightllm/models/qwen4_exp/mtp/model.py b/lightllm/models/qwen4_exp/mtp/model.py new file mode 100644 index 0000000000..7c28fc649d --- /dev/null +++ b/lightllm/models/qwen4_exp/mtp/model.py @@ -0,0 +1,51 @@ +import torch + +from lightllm.models.draft_registry import DraftModelRegistry +from lightllm.models.qwen3_5_mtp.model import Qwen3_5MTPModel +from ..model import Qwen4ExpTpPartModel +from ..infer_struct import Qwen4ExpInferStateInfo +from ..layer_infer.post_layer_infer import Qwen4ExpPostLayerInfer +from .weights import Qwen4ExpMTPPreAndPostLayerWeight, Qwen4ExpMTPTransformerLayerWeight +from .infer import Qwen4ExpMTPPreLayerInfer, Qwen4ExpMTPTransformerLayerInfer + + +class Qwen4ExpMTPModel(Qwen3_5MTPModel): + pre_and_post_weight_class = Qwen4ExpMTPPreAndPostLayerWeight + transformer_weight_class = Qwen4ExpMTPTransformerLayerWeight + pre_layer_infer_class = Qwen4ExpMTPPreLayerInfer + post_layer_infer_class = Qwen4ExpPostLayerInfer + transformer_layer_infer_class = Qwen4ExpMTPTransformerLayerInfer + infer_state_class = Qwen4ExpInferStateInfo + _init_att_backend = Qwen4ExpTpPartModel._init_att_backend + create_full_att_states = classmethod(Qwen4ExpTpPartModel.create_full_att_states.__func__) + + @classmethod + def normalize_model_config(cls, model_config): + config = Qwen4ExpTpPartModel.normalize_model_config(model_config) + config.update( + full_attention_interval=1, num_hidden_layers=1, n_layer=1, ple_layer_ids=[], layer_types=["full_attention"] + ) + return config + + def _init_config(self): + Qwen4ExpTpPartModel._init_config(self) + if self.config.get("mtp_num_hidden_layers", 1) != 1: + raise ValueError("Qwen4Exp recurrent MTP requires one trained draft layer") + self.config.update( + full_attention_interval=1, num_hidden_layers=1, n_layer=1, ple_layer_ids=[], layer_types=["full_attention"] + ) + + def _init_weights(self, start_layer_index=None): + if start_layer_index is not None or self.mtp_previous_draft_models: + raise ValueError("Qwen4Exp MTP uses one recurrent predictor") + self.pre_post_weight = self.pre_and_post_weight_class(self.data_type, self.config) + self.pre_post_weight.wte_weight_ = self.main_model.pre_post_weight.wte_weight_ + self.pre_post_weight.lm_head_weight_ = self.main_model.pre_post_weight.lm_head_weight_ + self.trans_layers_weight = [self.transformer_weight_class(0, self.data_type, self.config, self.quant_cfg)] + + def _gen_special_model_input(self, token_num): + return { + "mtp_draft_input_hiddens": torch.randn( + token_num, self.config["hidden_size"] * self.config["hc_count"], dtype=self.data_type, device="cuda" + ) + } diff --git a/lightllm/models/qwen4_exp/mtp/weights.py b/lightllm/models/qwen4_exp/mtp/weights.py new file mode 100644 index 0000000000..d0edb3e5ea --- /dev/null +++ b/lightllm/models/qwen4_exp/mtp/weights.py @@ -0,0 +1,55 @@ +from dataclasses import dataclass +from functools import partial + +import torch.nn.functional as F + +from lightllm.common.basemodel import PreAndPostLayerWeight +from lightllm.common.gated_residual.kernels import grouped_norm +from ..layer_weights.components import ComponentWeight, gated_residual_weight +from ..layer_weights.transformer_layer_weight import Qwen4ExpTransformerLayerWeight + + +@dataclass(frozen=True) +class MTPFusion: + parameters: dict + branches: int + eps: float + + def __call__(self, embedding, hidden): + if embedding.ndim != 2 or not embedding.is_contiguous() or not hidden.is_contiguous(): + raise ValueError("Qwen MTP fusion requires contiguous embedding and hidden matrices") + width = embedding.shape[-1] + if hidden.shape != (embedding.shape[0], self.branches * width): + raise ValueError("Qwen MTP requires the full pre-mixer HC hidden stream") + weights = self.parameters + embedding = grouped_norm(embedding, weights["mtp.pre_fc_norm_embedding.weight"], 1, self.eps) + hidden = grouped_norm(hidden, weights["mtp.pre_fc_norm_hidden.weight"], 1, self.eps) + embedding = F.linear(embedding, weights["mtp.fc_embedding.weight"]) + hidden = F.linear(hidden.view(-1, self.branches, width), weights["mtp.fc_hidden.weight"]) + return (hidden + embedding[:, None]).flatten(1) + + +class Qwen4ExpMTPPreAndPostLayerWeight(PreAndPostLayerWeight): + def __init__(self, data_type, network_config): + super().__init__(data_type, network_config) + hidden, branches = network_config["hidden_size"], network_config["hc_count"] + specs = { + "mtp.pre_fc_norm_embedding.weight": ((hidden,), data_type), + "mtp.pre_fc_norm_hidden.weight": ((hidden * branches,), data_type), + "mtp.fc_embedding.weight": ((hidden, hidden), data_type), + "mtp.fc_hidden.weight": ((hidden, hidden), data_type), + } + self.fusion = ComponentWeight(specs, partial(MTPFusion, branches=branches, eps=network_config["rms_norm_eps"])) + self.hyper_connection_mixer = gated_residual_weight( + "mtp.hyper_connection_mixer", network_config, data_type, use_combine=False + ) + self.wte_weight_ = self.lm_head_weight_ = None + + +class Qwen4ExpMTPTransformerLayerWeight(Qwen4ExpTransformerLayerWeight): + def load_hf_weights(self, weights): + prefix = "mtp.layers.0." + mapped = { + "model.layers.0." + name[len(prefix) :]: value for name, value in weights.items() if name.startswith(prefix) + } + super().load_hf_weights(mapped) diff --git a/lightllm/models/qwen4_exp/multimodal.py b/lightllm/models/qwen4_exp/multimodal.py new file mode 100644 index 0000000000..3978fd502a --- /dev/null +++ b/lightllm/models/qwen4_exp/multimodal.py @@ -0,0 +1,43 @@ +import torch +import triton +import triton.language as tl + + +@triton.jit +def _canonical_image_ids( + Tokens, + Starts, + Lengths, + Output, + COUNT: tl.constexpr, + IMAGES: tl.constexpr, + IMAGE_TOKEN: tl.constexpr, + BLOCK: tl.constexpr, +): + offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + values = tl.load(Tokens + offsets, offsets < COUNT, other=0).to(tl.int64) + canonical = values + for image in range(IMAGES): + start = tl.load(Starts + image).to(tl.int64) + length = tl.load(Lengths + image).to(tl.int64) + canonical = tl.where((values >= start) & (values < start + length), IMAGE_TOKEN, canonical) + tl.store(Output + offsets, canonical, offsets < COUNT) + + +def canonical_image_ids(tokens, starts, lengths, image_token_id): + if image_token_id is None: + raise ValueError("Qwen4 image hashing requires the checkpoint image token ID") + for value in (tokens, starts, lengths): + if value.ndim != 1 or not value.is_cuda or not value.is_contiguous() or value.device != tokens.device: + raise ValueError("Image token mapping requires contiguous CUDA vectors on one device") + if value.dtype not in (torch.int32, torch.int64): + raise ValueError("Image token mapping requires integer IDs and ranges") + if starts.shape != lengths.shape: + raise ValueError("Image cache starts and lengths must have matching shapes") + if not starts.numel() or not tokens.numel(): + return tokens + output = torch.empty_like(tokens) + _canonical_image_ids[(triton.cdiv(tokens.numel(), 256),)]( + tokens, starts, lengths, output, tokens.numel(), starts.numel(), image_token_id, 256 + ) + return output diff --git a/lightllm/models/qwen4_exp/pinned_memory.py b/lightllm/models/qwen4_exp/pinned_memory.py new file mode 100644 index 0000000000..cfa4f14269 --- /dev/null +++ b/lightllm/models/qwen4_exp/pinned_memory.py @@ -0,0 +1,72 @@ +import math +import mmap + +import numpy as np +import torch + +# Keep the number of CUDA host-registration ranges bounded when DP replicates +# the offloaded table on every rank. The registration unit still aligns each +# range to a complete recurrent snapshot, so a larger chunk cannot split a DMA +# snapshot across registrations. +_CHUNK_BYTES = 1024 * 1024 * 1024 +_LARGE_BYTES = 512 * 1024 * 1024 + + +class _RegisteredArray(np.ndarray): + def __array_finalize__(self, source): + self.allocation = getattr(source, "allocation", None) + + +class _Registration: + def __init__(self, mapping, device, runtime): + self.mapping = mapping + self.device = device + self.runtime = runtime + self.addresses = [] + + def __del__(self): + if not self.addresses: + return + + try: + with torch.cuda.device(self.device): + torch.cuda.synchronize() + except Exception: + pass + for address in reversed(self.addresses): + try: + self.runtime.cudaHostUnregister(address) + except Exception: + pass + + +def pinned_zeros(shape, *, dtype, device="cpu", pin_memory=True): + size = math.prod(shape) * dtype.itemsize + if device != "cpu" or not pin_memory: + raise ValueError("Qwen recurrent snapshots require pinned CPU storage") + if size < _LARGE_BYTES: + return torch.zeros(shape, dtype=dtype, device="cpu", pin_memory=True) + + cell_bytes = math.prod(shape[1:]) * dtype.itemsize + registration_unit = math.lcm(cell_bytes, mmap.PAGESIZE) + if registration_unit > _CHUNK_BYTES: + raise ValueError("A recurrent snapshot must fit within one bounded host registration") + chunk_bytes = (_CHUNK_BYTES // registration_unit) * registration_unit + runtime = torch.cuda.cudart() + mapping = mmap.mmap(-1, size) + allocation = _Registration(mapping, torch.cuda.current_device(), runtime) + array = np.frombuffer(mapping, dtype=np.uint8).view(_RegisteredArray) + array.allocation = allocation + value = torch.from_numpy(array).view(dtype).reshape(shape) + value.zero_() + for offset in range(0, size, chunk_bytes): + address = value.data_ptr() + offset + status = runtime.cudaHostRegister(address, min(chunk_bytes, size - offset), 0) + if int(status): + raise RuntimeError( + f"Cannot register Qwen pinned snapshot at byte {offset}: {runtime.cudaGetErrorString(status)}" + ) + allocation.addresses.append(address) + + value._qwen4_pinned_allocation = allocation + return value diff --git a/lightllm/models/qwen4_exp/position.py b/lightllm/models/qwen4_exp/position.py new file mode 100644 index 0000000000..e5fcaa054a --- /dev/null +++ b/lightllm/models/qwen4_exp/position.py @@ -0,0 +1,18 @@ +import torch + + +def combine_mrope(cos: torch.Tensor, sin: torch.Tensor, sections): + if cos.shape != sin.shape or cos.device != sin.device or cos.dtype != sin.dtype: + raise ValueError("Rotary cosine and sine shapes, devices and dtypes must match") + if cos.ndim == 2: + return cos.repeat(1, 2).contiguous(), sin.repeat(1, 2).contiguous() + if cos.ndim != 3 or cos.shape[0] != 3 or len(sections) != 3: + raise ValueError("Expected three rotary axes and three mRoPE sections") + offsets = torch.arange(cos.shape[-1], device=cos.device) + height = ((offsets % 3) == 1) & (offsets <= 3 * sections[1]) + width = ((offsets % 3) == 2) & (offsets <= 3 * sections[2]) + combined = [] + for tensor in (cos, sin): + values = torch.where(height, tensor[1], torch.where(width, tensor[2], tensor[0])) + combined.append(values.repeat(1, 2).contiguous()) + return tuple(combined) diff --git a/lightllm/models/qwen4_exp/speculative_cache.py b/lightllm/models/qwen4_exp/speculative_cache.py new file mode 100644 index 0000000000..e1c51570d1 --- /dev/null +++ b/lightllm/models/qwen4_exp/speculative_cache.py @@ -0,0 +1,108 @@ +import torch +import triton +import triton.language as tl + + +@triton.jit +def _stage( + History, + Inputs, + Starts, + Requests, + Pending, + CONTEXT: tl.constexpr, + WIDTH: tl.constexpr, + VERIFY: tl.constexpr, + BLOCK: tl.constexpr, +): + segment = tl.program_id(0) + request = tl.load(Requests + segment).to(tl.int64) + if request >= 0: + start = tl.load(Starts + segment) + count = tl.load(Starts + segment + 1) - start + offset = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + position, column = offset // WIDTH, offset % WIDTH + old = tl.load(History + segment.to(tl.int64) * CONTEXT * WIDTH + offset, position < CONTEXT, other=0) + new = tl.load( + Inputs + (start.to(tl.int64) + position - CONTEXT) * WIDTH + column, + (position >= CONTEXT) & (position < CONTEXT + count), + other=0, + ) + tl.store( + Pending + request * (CONTEXT + VERIFY) * WIDTH + offset, + tl.where(position < CONTEXT, old, new), + (position < CONTEXT + count) & (position < CONTEXT + VERIFY), + ) + + +@triton.jit +def _commit( + Pending, + History, + Requests, + Starts, + Accepted, + HOLD: tl.constexpr, + CONTEXT: tl.constexpr, + WIDTH: tl.constexpr, + VERIFY: tl.constexpr, + BLOCK: tl.constexpr, +): + group = tl.program_id(0) + start = tl.load(Starts + group) + request = tl.load(Requests + start).to(tl.int64) + if request != HOLD: + count = tl.load(Accepted + group) + offset = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + values = tl.load( + Pending + request * (CONTEXT + VERIFY) * WIDTH + count * WIDTH + offset, offset < CONTEXT * WIDTH, other=0 + ) + tl.store(History + request * CONTEXT * WIDTH + offset, values, offset < CONTEXT * WIDTH) + + +class LexicalSpeculativeCache: + def __init__(self, accepted_cache, verify_width): + if verify_width < 1: + raise ValueError("Speculative lexical cache requires positive verify width") + self.verify_width = verify_width + self.fields = [] + for history in (accepted_cache.ngram, accepted_cache.convolution): + requests, context = history.shape[:2] + width = history.shape[2] if history.ndim == 3 else 1 + pending = torch.empty((requests, context + verify_width, width), dtype=history.dtype, device=history.device) + self.fields.append((history, pending, context, width)) + + def stage(self, metadata, batch): + if batch.input_tokens is None or batch.conv_inputs is None: + raise RuntimeError("PLE candidates must be computed before staging") + for (history, pending, context, width), prefix, inputs in zip( + self.fields, (batch.ngram_history, batch.conv_history), (batch.input_tokens, batch.conv_inputs) + ): + if metadata.active_rows.numel(): + _stage[(metadata.active_rows.numel(), triton.cdiv((context + self.verify_width) * width, 256))]( + prefix, + inputs, + metadata.starts, + metadata.active_rows, + pending, + context, + width, + self.verify_width, + 256, + ) + + def commit(self, requests, starts, accepted_lengths, hold_request): + for history, pending, context, width in self.fields: + if starts.numel() and context: + _commit[(starts.numel(), triton.cdiv(context * width, 256))]( + pending, + history, + requests, + starts, + accepted_lengths, + hold_request, + context, + width, + self.verify_width, + 256, + ) diff --git a/lightllm/models/qwen4_exp/state.py b/lightllm/models/qwen4_exp/state.py new file mode 100644 index 0000000000..6800f8a6e6 --- /dev/null +++ b/lightllm/models/qwen4_exp/state.py @@ -0,0 +1,136 @@ +from dataclasses import dataclass +from typing import Optional + +import torch +import triton + +from lightllm.common.qsa.attention import indexed_attention +from lightllm.distributed import all_reduce + + +@dataclass +class QSABatch: + slot_table: torch.Tensor + active_rows: torch.Tensor + key_lengths: torch.Tensor + first_blocks: torch.Tensor + query_rows: torch.Tensor + visible_lengths: torch.Tensor + mem_indices: torch.Tensor + cos: torch.Tensor + sin: torch.Tensor + max_new_blocks: int + + @classmethod + def from_metadata(cls, metadata, slot_table, mem_indices, cos, sin, compress_ratio: int): + if compress_ratio < 1: + raise ValueError("QSA compression ratio must be positive") + return cls( + slot_table, + metadata.active_rows, + metadata.key_lengths, + metadata.first_blocks, + metadata.query_rows, + metadata.visible_lengths, + mem_indices.to(torch.int64), + cos, + sin, + (mem_indices.numel() + compress_ratio - 1) // compress_ratio, + ) + + +@dataclass +class QSAExecutionState: + raw_keys: torch.Tensor + + pooled_keys: torch.Tensor + key_cos: torch.Tensor + key_sin: torch.Tensor + + def attend( + self, + indexer, + hidden, + query, + key, + value, + batch: QSABatch, + *, + decode: bool = False, + attention_splits: Optional[int] = None, + ): + index_query, raw = indexer.project(hidden, batch.cos, batch.sin) + self.raw_keys.index_copy_(0, batch.mem_indices, raw) + indexer.pool( + self.raw_keys, + batch.slot_table, + batch.key_lengths, + self.key_cos, + self.key_sin, + self.pooled_keys, + batch.first_blocks, + batch.max_new_blocks, + request_rows=batch.active_rows, + ) + selected = indexer.select( + index_query, + self.pooled_keys, + batch.query_rows, + batch.visible_lengths, + batch.slot_table, + query_chunk_size=max(64, query.shape[0]) if decode else 1024, + ) + + splits = attention_splits + if splits is None and decode: + splits = min(32, triton.cdiv(selected.shape[1], 64)) + wide_block = query.dtype in (torch.float16, torch.bfloat16) + return indexed_attention( + query.contiguous(), + key, + value, + selected, + splits=splits, + selection_lengths=batch.visible_lengths, + short_rows_direct=decode or attention_splits is not None, + skip_empty_splits=True, + block_size=128 if wide_block else 64, + num_stages=2 if wide_block else None, + ) + + +@dataclass +class PLEBatch: + starts: torch.Tensor + ngram_history: torch.Tensor + conv_history: torch.Tensor + mask: Optional[torch.Tensor] = None + conv_inputs: Optional[torch.Tensor] = None + input_tokens: Optional[torch.Tensor] = None + + def apply(self, residual, tokens, layer_weight, tp_world_size: int, dist_group): + hasher = layer_weight.ngram_hasher.component + if self.mask is not None: + tokens = tokens.masked_fill(~self.mask, hasher.eos_token_id) + indices, _ = hasher.hash(tokens, self.starts, self.ngram_history) + embeddings = layer_weight.ngram_embedding(indices.flatten()).view( + tokens.numel(), hasher.heads * layer_weight.ngram_embedding.dim + ) + if tp_world_size > 1: + all_reduce(embeddings, group=dist_group, async_op=False) + update, conv_inputs = layer_weight.ple.component.forward( + residual, embeddings, self.starts, self.conv_history, self.mask + ) + self.input_tokens, self.conv_inputs = tokens, conv_inputs + return update + + def histories_after(self, layer_weight, accepted_lengths=None): + if self.input_tokens is None or self.conv_inputs is None: + raise RuntimeError("PLE candidates must be computed before requesting histories") + ngram = layer_weight.ngram_hasher.component.history_after( + self.input_tokens, self.starts, self.ngram_history, accepted_lengths + ) + conv = layer_weight.ple.component.history_after( + self.conv_inputs, self.starts, self.conv_history, accepted_lengths + ) + return ngram, conv diff --git a/lightllm/models/qwen4_exp/tokenizer.py b/lightllm/models/qwen4_exp/tokenizer.py new file mode 100644 index 0000000000..807cea9ad5 --- /dev/null +++ b/lightllm/models/qwen4_exp/tokenizer.py @@ -0,0 +1,6 @@ +from lightllm.models.qwen2_vl.model import QWen2VLTokenizer + + +class Qwen4ExpTokenizer(QWen2VLTokenizer): + def count_prompt_tokens(self, text_tokens, image_tokens, image_count, audio_tokens, audio_count): + return text_tokens + image_tokens + audio_tokens diff --git a/lightllm/models/qwen4_exp/triton_kernel/__init__.py b/lightllm/models/qwen4_exp/triton_kernel/__init__.py new file mode 100644 index 0000000000..2ae28399f5 --- /dev/null +++ b/lightllm/models/qwen4_exp/triton_kernel/__init__.py @@ -0,0 +1 @@ +pass diff --git a/lightllm/models/qwen4_exp/triton_kernel/cache_pages.py b/lightllm/models/qwen4_exp/triton_kernel/cache_pages.py new file mode 100644 index 0000000000..c541c1b892 --- /dev/null +++ b/lightllm/models/qwen4_exp/triton_kernel/cache_pages.py @@ -0,0 +1,71 @@ +import triton +import triton.language as tl + + +@triton.jit +def copy_slot_pages( + Slots, + Pages, + MemIndices, + PageIndices, + Readies, + SLOT_STRIDE: tl.constexpr, + LAYER_STRIDE: tl.constexpr, + PAGE_STRIDE: tl.constexpr, + TOKENS: tl.constexpr, + LAYERS: tl.constexpr, + DIM: tl.constexpr, + OFFLOAD: tl.constexpr, + BLOCK: tl.constexpr, +): + page = tl.program_id(0) + page_index = tl.load(PageIndices + page).to(tl.int64) + active = page_index >= 0 + if OFFLOAD: + active = active & ~tl.load(Readies + page).to(tl.int1) + if active: + col = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + mask = col < TOKENS * LAYERS * DIM + token = col // (LAYERS * DIM) + layer = ((col // DIM) % LAYERS).to(tl.int64) + slot = tl.load(MemIndices + page * TOKENS + token, mask, other=-1).to(tl.int64) + mask = mask & (slot >= 0) + live = Slots + slot * SLOT_STRIDE + layer * LAYER_STRIDE + col % DIM + cached = Pages + page_index * PAGE_STRIDE + col + if OFFLOAD: + value = tl.load(live, mask, other=0) + tl.store(cached, value, mask) + else: + value = tl.load(cached, mask, other=0) + tl.store(live, value, mask) + + +@triton.jit +def copy_state_pages( + States, + Pages, + StateIndices, + PageIndices, + Readies, + WIDTH: tl.constexpr, + PAGE_STRIDE: tl.constexpr, + OFFLOAD: tl.constexpr, + BLOCK: tl.constexpr, +): + page = tl.program_id(0) + page_index = tl.load(PageIndices + page).to(tl.int64) + active = page_index >= 0 + if OFFLOAD: + active = active & ~tl.load(Readies + page).to(tl.int1) + if active: + state = tl.load(StateIndices + page).to(tl.int64) + col = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + mask = col < WIDTH + live = States + state * WIDTH + col + cached = Pages + page_index * PAGE_STRIDE + col + if OFFLOAD: + value = tl.load(live, mask, other=0) + tl.store(cached, value, mask) + else: + value = tl.load(cached, mask, other=0) + tl.store(live, value, mask) diff --git a/lightllm/models/qwen4_exp/triton_kernel/gated_norm.py b/lightllm/models/qwen4_exp/triton_kernel/gated_norm.py new file mode 100644 index 0000000000..9f50ce3acc --- /dev/null +++ b/lightllm/models/qwen4_exp/triton_kernel/gated_norm.py @@ -0,0 +1,92 @@ +import torch +import triton +import triton.language as tl + + +@triton.jit +def _gated_norm( + X, + Gate, + Weight, + Out, + input_stride_row, + gate_stride_token, + gate_stride_head, + out_stride_row, + GATE_HEADS: tl.constexpr, + D: tl.constexpr, + EPS: tl.constexpr, + SIGMOID: tl.constexpr, + BLOCK: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + col = tl.arange(0, BLOCK) + x = tl.load(X + row * input_stride_row + col, col < D, other=0).to(tl.float32) + inv = tl.rsqrt(tl.sum(x * x, 0) / D + EPS) + + normalized = (x * inv).to(X.dtype.element_ty).to(tl.float32) + weight = tl.load(Weight + col, col < D, other=0).to(tl.float32) + weighted = (normalized * weight).to(X.dtype.element_ty).to(tl.float32) + gate_token = row // GATE_HEADS + gate_head = row % GATE_HEADS + gate = tl.load( + Gate + gate_token * gate_stride_token + gate_head * gate_stride_head + col, + col < D, + other=0, + ).to(tl.float32) + activation = tl.sigmoid(gate) + if not SIGMOID: + activation = gate * activation + tl.store(Out + row * out_stride_row + col, weighted * activation, col < D) + + +@torch.no_grad() +def gated_norm(input: torch.Tensor, gate: torch.Tensor, weight: torch.Tensor, eps: float, activation: str, out=None): + if activation not in ("sigmoid", "silu"): + raise ValueError("Qwen4Exp GDN output gate must be sigmoid or silu") + if ( + input.ndim != 2 + or input.shape[-1] != weight.numel() + or gate.ndim not in (2, 3) + or gate.shape[-1] != input.shape[-1] + or (gate.ndim == 2 and gate.shape != input.shape) + or (gate.ndim == 3 and gate.shape[0] * gate.shape[1] != input.shape[0]) + or not input.is_contiguous() + or gate.stride(-1) != 1 + ): + raise ValueError("GDN norm requires a contiguous input, a packed gate, and a per-head norm weight") + for tensor in (input, gate, weight): + if not tensor.is_cuda or tensor.dtype != input.dtype or tensor.device != input.device: + raise ValueError("GDN norm tensors must be CUDA tensors with matching device and dtype") + if weight.ndim != 1 or not weight.is_contiguous() or input.ndim < 2 or eps <= 0: + raise ValueError("GDN norm requires a weight vector, token/head rows and positive epsilon") + result = torch.empty_like(input) if out is None else out + if ( + result.shape != input.shape + or result.dtype != input.dtype + or result.device != input.device + or not result.is_contiguous() + ): + raise ValueError("GDN norm output must be contiguous and match the input") + rows = input.shape[0] + gate_heads = 1 if gate.ndim == 2 else gate.shape[1] + gate_stride_token = gate.stride(0) + gate_stride_head = 0 if gate.ndim == 2 else gate.stride(1) + if rows: + _gated_norm[(rows,)]( + input, + gate, + weight, + result, + input.stride(0), + gate_stride_token, + gate_stride_head, + result.stride(0), + GATE_HEADS=gate_heads, + D=input.shape[-1], + EPS=eps, + SIGMOID=activation == "sigmoid", + BLOCK=triton.next_power_of_2(input.shape[-1]), + enable_fp_fusion=False, + ) + return result diff --git a/lightllm/models/qwen4_exp/vision.py b/lightllm/models/qwen4_exp/vision.py new file mode 100644 index 0000000000..93b1d49c19 --- /dev/null +++ b/lightllm/models/qwen4_exp/vision.py @@ -0,0 +1,36 @@ +import torch +import torch.nn.functional as F + +from lightllm.models.qwen3_vl.qwen3_visual import Qwen3VisionTransformerPretrainedModel +from lightllm.models.qwen2_vl.qwen2_visual import VisionRotaryEmbedding + + +class Qwen4ExpVisionTransformer(Qwen3VisionTransformerPretrainedModel): + def __init__(self, kvargs, **config): + config = dict(config) + if config.get("deepstack_visual_indexes"): + raise ValueError("Qwen4Exp vision does not define deepstack outputs") + config["deepstack_visual_indexes"] = [] + rope = config.get("rope_parameters") or {"rope_type": "axial", "rope_theta": 10000.0} + if rope.get("rope_type", "axial") != "axial": + raise ValueError("Qwen4Exp vision requires axial rotary embeddings") + super().__init__(kvargs, **config) + self.rotary_pos_emb = VisionRotaryEmbedding( + self.hidden_size // self.num_heads // 2, theta=rope.get("rope_theta", 10000.0) + ).cuda() + + def fast_pos_embed_interpolate(self, grid_thw): + table = self.pos_embed.weight.float().T.reshape( + 1, self.hidden_size, self.num_grid_per_side, self.num_grid_per_side + ) + merge = self.spatial_merge_size + patches = [] + for frames, height, width in grid_thw.detach().cpu().tolist(): + if min(frames, height, width) < 1 or height % merge or width % merge: + raise ValueError("Vision grids must be positive and divisible by the spatial merge size") + positions = F.interpolate(table, size=(height, width), mode="bilinear", align_corners=True) + positions = positions[0].permute(1, 2, 0) + positions = positions.reshape(height // merge, merge, width // merge, merge, self.hidden_size) + positions = positions.permute(0, 2, 1, 3, 4).reshape(-1, self.hidden_size).repeat(frames, 1) + patches.append(positions.to(self.pos_embed.weight.dtype)) + return torch.cat(patches) diff --git a/lightllm/server/api_cli.py b/lightllm/server/api_cli.py index 41ae769ae6..4bdfbdbf74 100644 --- a/lightllm/server/api_cli.py +++ b/lightllm/server/api_cli.py @@ -1,4 +1,5 @@ import argparse +import json def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: @@ -945,6 +946,11 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: help="""enable torch memory saver, which is used for release_memory and resume_memory during RL training.""", ) parser.add_argument("--enable_weight_cpu_backup", action="store_true", help="""enable weight cpu backup.""") + parser.add_argument( + "--offload_ngram_embedding", + action="store_true", + help="Keep supported n-gram embedding tables in pinned CPU memory; gather requested rows onto CUDA.", + ) parser.add_argument( "--disk_cache_dir", type=str, @@ -984,6 +990,21 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: If radix cache hit rate is low under high load due to limited small-page capacity and LRU eviction, increasing linear_att_cache_size can improve hit rate at the cost of more memory.""", ) + parser.add_argument( + "--ssm_state_mode", + choices=["native", "replay"], + default="native", + help="SSM state storage: native snapshots or deferred GDN/KDA ReplaySSM (FP32/BF16).", + ) + parser.add_argument( + "--ssm_replay_config", + type=json.loads, + default={}, + help='Replay-only JSON object: {"cache_len":16,"compute_mode":"precompute"}. ' + "cache_len: 4/8/16/32/64, at least the verify width; compute_mode: inline/precompute. " + "mtp_step=0 decode uses the same replay history: explicit precompute projects that one token, " + "and the default is inline. MTP defaults to precompute.", + ) parser.add_argument( "--linear_att_ssm_data_type", type=str, diff --git a/lightllm/server/api_lightllm.py b/lightllm/server/api_lightllm.py index b881b8ad1f..1fc24a127b 100644 --- a/lightllm/server/api_lightllm.py +++ b/lightllm/server/api_lightllm.py @@ -30,7 +30,6 @@ async def lightllm_get_score(request: Request, httpserver_manager: HttpServerMan async def lightllm_generate(request: Request, httpserver_manager: HttpServerManager) -> Response: - request_dict = await request.json() prompt = request_dict.pop("inputs") sample_params_dict = request_dict["parameters"] @@ -115,11 +114,10 @@ async def lightllm_generate(request: Request, httpserver_manager: HttpServerMana async def lightllm_generate_stream(request: Request, httpserver_manager: HttpServerManager) -> Response: - request_dict = await request.json() prompt = request_dict.pop("inputs") sample_params_dict = request_dict["parameters"] - _ = sample_params_dict.pop("return_details", False) + return_details = sample_params_dict.pop("return_details", False) _ = sample_params_dict.pop("return_routed_experts", None) sampling_params = SamplingParams() sampling_params.init(tokenizer=httpserver_manager.tokenizer, **sample_params_dict) @@ -155,6 +153,15 @@ async def stream_results() -> AsyncGenerator[bytes, None]: "input_usage": input_usage, } ret["token"]["logprobs"] = metadata["logprobs"] + if return_details: + for name in ( + "prompt_cache_len", + "mtp_accepted_token_num", + "mtp_verify_token_num", + "mtp_verify_step_num", + ): + if name in metadata: + ret["token"][name] = metadata[name] if "prompt_logprobs" in metadata: ret["prompt_logprobs"] = metadata["prompt_logprobs"] ret["prompt_token_ids"] = metadata.get("prompt_token_ids") diff --git a/lightllm/server/api_start.py b/lightllm/server/api_start.py index a42acef884..361d789a64 100644 --- a/lightllm/server/api_start.py +++ b/lightllm/server/api_start.py @@ -43,6 +43,59 @@ def _set_envs_and_config(args: StartArgs): mp.set_start_method("spawn", force=True) +def _validate_and_normalize_ssm_config(args: StartArgs): + assert args.ssm_state_mode in ("native", "replay"), "Invalid ssm_state_mode" + config = args.ssm_replay_config + assert isinstance(config, dict), "ssm_replay_config must be a JSON object" + assert args.ssm_state_mode == "replay" or not config, "ssm_replay_config requires ssm_state_mode=replay" + if args.ssm_state_mode == "replay": + defaults = {"cache_len": 16, "compute_mode": "precompute" if args.mtp_step > 0 else "inline"} + unknown = config.keys() - defaults.keys() + assert not unknown, f"Unknown ssm_replay_config fields: {unknown}" + config = {**defaults, **config} + assert type(config["cache_len"]) is int and config["cache_len"] in ( + 4, + 8, + 16, + 32, + 64, + ), "ReplaySSM capacity must be a power of two and at least 4, up to 64" + assert config["compute_mode"] in ("inline", "precompute"), "Invalid replay compute_mode" + args.ssm_replay_config = config + logger.info(f"SSM mode: replay, config: {config}") + else: + logger.info(f"SSM mode: {args.ssm_state_mode}") + from lightllm.utils.config_utils import get_model_type + + model_type = get_model_type(args.model_dir) + is_gdn_or_kda = model_type in ( + "glm5_next", + "glm5_next_text", + "qwen3_next", + "qwen3_5", + "qwen3_5_moe", + "qwen3_5_text", + "qwen3_5_moe_text", + "qwen4_exp", + "qwen4_exp_text", + "sensenova_loop", + ) + if is_gdn_or_kda: + assert args.mtp_step <= 15, ( + "GDN/KDA requires mtp_step <= 15: " "linear attention MTP state acceptance supports verify width <= 16" + ) + if args.ssm_state_mode == "native": + return + assert args.linear_att_ssm_data_type in ( + "float32", + "bfloat16", + ), "Non-native SSM state modes require FP32 or BF16 state" + assert is_gdn_or_kda, "Non-native SSM state modes currently require a GDN or KDA model" + assert args.mtp_step >= 0, "ssm_state_mode=replay requires mtp_step >= 0" + capacity = config["cache_len"] + assert capacity >= args.mtp_step + 1, "ReplaySSM capacity must cover the verify width" + + def _launch_subprocesses(args: StartArgs): _validate_pd_features(args) _set_envs_and_config(args) @@ -355,6 +408,8 @@ def _launch_subprocesses(args: StartArgs): s3_model_prepare(args.model_dir) + _validate_and_normalize_ssm_config(args) + # 如果args.eos_id 是 None, 从 config.json 中读取 eos_token_id 相关的信息，赋值给 args if args.eos_id is None: from lightllm.utils.config_utils import get_eos_token_ids diff --git a/lightllm/server/core/objs/start_args_type.py b/lightllm/server/core/objs/start_args_type.py index 0d6468f397..0f43364d73 100644 --- a/lightllm/server/core/objs/start_args_type.py +++ b/lightllm/server/core/objs/start_args_type.py @@ -246,6 +246,7 @@ class StartArgs: enable_rl: bool = field(default=False) enable_torch_memory_saver: bool = field(default=False) enable_weight_cpu_backup: bool = field(default=False) + offload_ngram_embedding: bool = field(default=False) hardware_platform: str = field(default="cuda", metadata={"choices": ["cuda", "musa"]}) enable_torch_fallback: bool = field(default=False) enable_triton_fallback: bool = field(default=False) @@ -255,6 +256,8 @@ class StartArgs: weight_version: str = "default" # hybrid attention model (Qwen3Next) + ssm_state_mode: str = "native" + ssm_replay_config: dict[str, object] = field(default_factory=dict) linear_att_hash_page_size: int = field(default=512) linear_att_page_block_num: int = field(default=10000000) disable_linear_att_small_page_cpu_cache: bool = field(default=False) diff --git a/lightllm/server/httpserver/manager.py b/lightllm/server/httpserver/manager.py index 43b06f02ef..541e5b9502 100644 --- a/lightllm/server/httpserver/manager.py +++ b/lightllm/server/httpserver/manager.py @@ -1,3 +1,4 @@ +from lightllm.server.httpserver.slot_availability import RequestSlotAvailability import sys import zmq import zmq.asyncio @@ -56,6 +57,7 @@ def __init__( args: StartArgs, ): self.args: StartArgs = args + self._request_slot_availability = RequestSlotAvailability.from_args(args) ports = get_shm_port_args() context = zmq.asyncio.Context(2) self.send_to_router = context.socket(zmq.PUSH) @@ -344,7 +346,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) @@ -587,6 +588,8 @@ async def _alloc_shm_req_indexes( alloc_deadline = time.monotonic() + alloc_timeout_seconds if alloc_timeout_seconds is not None else None # 空闲槽不足时整组重试，避免多个 n > 1 的请求各占部分槽后互相等待。 + slot_availability = getattr(self, "_request_slot_availability", None) + observed_generation = slot_availability.generation if slot_availability is not None else 0 alloced_req_indexes = await self.shm_req_manager.async_alloc_req_indexes(req_num) # 保持相同的退避起点，仅通过系数让高优先级请求更快地重新尝试获取 shm_req。 sleep_time_factor = 0.2 if high_priority_request else 1 @@ -601,11 +604,20 @@ async def _alloc_shm_req_indexes( f"PD {self.args.run_mode} node is busy: unable to allocate a shm_req object " f"within {alloc_timeout_seconds} seconds" ) - await asyncio.sleep(sleep_time * sleep_time_factor) + if slot_availability is None: + await asyncio.sleep(sleep_time * sleep_time_factor) + else: + await slot_availability.wait(observed_generation, sleep_time * sleep_time_factor) + observed_generation = slot_availability.generation if slot_availability is not None else 0 sleep_time = min(1, sleep_time * 1.1) alloced_req_indexes = await self.shm_req_manager.async_alloc_req_indexes(req_num) return alloced_req_indexes + def _notify_req_slot_release(self): + availability = getattr(self, "_request_slot_availability", None) + if availability is not None: + availability.notify_release() + async def _log_req_header(self, request_headers, group_request_id: int): x_request_id = request_headers.get("X-Request-Id", "") x_session_id = request_headers.get("X-Session-Id", "") @@ -756,7 +768,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 +810,6 @@ async def _wait_to_token_package( req_status: "ReqStatus", request: Request, ): - event = req_status.event unfinished_count = sampling_params.best_of out_token_counter = 0 @@ -997,6 +1007,7 @@ async def recycle_resource_loop(self): logger.debug(f"httpserver release req_id {req.request_id}, index {req.index_in_shm_mem}") await self.shm_req_manager.async_put_back_req_obj(req) await self.shm_req_manager.async_release_req_index(req.index_in_shm_mem) + self._notify_req_slot_release() await self._release_multimodal_resources(req_status.group_req_objs.multimodal_params) # 先保留这个关键得日志，用于方便定位重构中的问题。 diff --git a/lightllm/server/httpserver/slot_availability.py b/lightllm/server/httpserver/slot_availability.py new file mode 100644 index 0000000000..739166fa6d --- /dev/null +++ b/lightllm/server/httpserver/slot_availability.py @@ -0,0 +1,38 @@ +import asyncio + + +class RequestSlotAvailability: + def __init__(self, enabled=True): + self.enabled = enabled + self.generation = 0 + self.event = asyncio.Event() if enabled else None + + @classmethod + def from_args(cls, args): + return cls( + enabled=( + args.run_mode == "normal" + and args.dp == 1 + and args.nnodes == 1 + and not args.enable_rl + and args.httpserver_workers == 1 + ) + ) + + def notify_release(self): + if self.enabled: + self.generation += 1 + self.event.set() + + async def wait(self, observed_generation, timeout): + if not self.enabled: + await asyncio.sleep(timeout) + return False + if observed_generation != self.generation: + return True + self.event.clear() + try: + await asyncio.wait_for(self.event.wait(), timeout=timeout) + return True + except asyncio.TimeoutError: + return False 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 09eb22be36..9c8ba806d6 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -172,6 +172,10 @@ def init_model(self, kvargs): with self.model.torch_memory_saver.region(tag=MemoryTag.KV_CACHE): manager.finish_state_registration() manager.load_full_att_calibration(args=self.model.args) + if self.args.enable_dp_prompt_cache_fetch: + # 所有 rank 完成模型初始化后再注册共享页池，避免与其他 rank 的加载分配重叠。 + torch.cuda.synchronize() + dist.barrier() manager.initialize_storage( self.model.args, token_capacity=self.model.max_total_token_num, mem_fraction=self.model.mem_fraction ) diff --git a/lightllm/server/router/model_infer/mtp_speculative/planner/lightspec.py b/lightllm/server/router/model_infer/mtp_speculative/planner/lightspec.py index b2f3285eee..f2f621b1cf 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/planner/lightspec.py +++ b/lightllm/server/router/model_infer/mtp_speculative/planner/lightspec.py @@ -104,6 +104,10 @@ def plan(self, decode_reqs: List, origin_batch_size: int) -> SpecDecodePlan: max_batch_size = min(origin_batch_size, available_batch_size) all_reqs_have_proposals = req_num_with_proposals == req_num + if not all_reqs_have_proposals: + self.progress_ema_by_config.clear() + self.prefix_survival_by_depth = [None] * self.max_draft_step + if not self.progress_ema_by_config: # Costs come from CUDA Graph capture, while progress requires one # real verification batch. Keep the available proposal intact until diff --git a/lightllm/server/router/req_queue/__init__.py b/lightllm/server/router/req_queue/__init__.py index d34a0b2933..94ff216527 100644 --- a/lightllm/server/router/req_queue/__init__.py +++ b/lightllm/server/router/req_queue/__init__.py @@ -13,6 +13,14 @@ def _get_req_queue_class(args, router, dp_size_in_node: int): if args.diverse_mode: return ChunkedBeamContinuesBatchQueue + from .chunked_prefill.cached_prefix_queue import ( + CachedPrefixChunkedPrefillQueue, + supports_cached_prefix_reservations, + ) + from lightllm.utils.config_utils import get_model_type + + if get_model_type(args.model_dir) in {"qwen4_exp", "qwen4_exp_text"} and supports_cached_prefix_reservations(args): + return CachedPrefixChunkedPrefillQueue if args.disable_chunked_prefill: # 虽然也使用chuncked prefill queue 但是由于 args.chunked_prefill_size = args.max_req_total_len # 所以调度的实际行为类似过去的 continues batch 调度，所以将两种调度的实现统一为一种实现，减少代码重复。 diff --git a/lightllm/server/router/req_queue/cached_prefix.py b/lightllm/server/router/req_queue/cached_prefix.py new file mode 100644 index 0000000000..8cd7aa5ca5 --- /dev/null +++ b/lightllm/server/router/req_queue/cached_prefix.py @@ -0,0 +1,71 @@ +import os + +from .peak_reservation import Reservation, estimate_peak + + +class CachedPrefixReservations: + def __init__(self, args, dp_index): + self.page_tokens = int(args.linear_att_hash_page_size) + name = os.environ.get("LIGHTLLM_UNIQUE_SERVICE_NAME_ID") + self.domain = (name + ":" + str(dp_index) + ":" + str(self.page_tokens)).encode() if name else None + self.memo, self.by_tuple = {}, {} + self.shared_prefix_tokens = 0 + + def reset(self, requests): + active = {int(request.request_id) for request in requests} + self.memo = {identity: value for identity, value in self.memo.items() if identity in active} + self.by_tuple.clear() + + def bind(self, reservation, request): + self.by_tuple[id(reservation)] = self.prefix(request) + return reservation + + def prefix(self, request): + try: + mm = getattr(request, "multimodal_params", None) + if ( + self.domain is None + or mm is None + or mm.images + or mm.audios + or request.is_paused + or request.is_aborted + or request.sample_params.disable_prompt_cache + ): + return () + length = min(int(request.prompt_cache_len), int(request.input_len), int(request.shm_cur_kv_len)) + hashes = request.linear_att_token_hash_list.get_all() if length >= self.page_tokens else [] + count = min(length // self.page_tokens, len(hashes), 32) + if count <= 0: + return () + hashes = tuple(int(value) & ((1 << 64) - 1) for value in hashes[:count]) + identity = (int(request.request_id), count, hashes, int(request.input_len)) + saved = self.memo.get(int(request.request_id)) + if saved is not None and saved[0] == identity: + return saved[1] + if not hasattr(request, "shm_prompt_ids"): + request.link_prompt_ids_shm_array() + ids = request.shm_prompt_ids.arr + if ids is None or ids.ndim != 1 or ids.dtype.kind not in "iu" or len(ids) < count * self.page_tokens: + return () + pages = tuple( + hashes[index].to_bytes(8, "little") + + ids[index * self.page_tokens : (index + 1) * self.page_tokens].astype(" 0 + ) + + +class CachedPrefixChunkedPrefillQueue(ChunkedPrefillQueue): + def __init__(self, args, router, dp_index, dp_size_in_node): + super().__init__(args, router, dp_index, dp_size_in_node) + self.prefix_reservations = CachedPrefixReservations(args, dp_index) + + def _init_cache_list(self, current_batch, is_busy): + requests = ( + [] + if current_batch is None + else [ + request for request in current_batch.reqs if request.sample_params.suggested_dp_index == self.dp_index + ] + ) + self.prefix_reservations.reset(requests) + super()._init_cache_list(current_batch, is_busy) + + def _get_req_reservation(self, request, is_busy): + reservation = super()._get_req_reservation(request, is_busy) + return self.prefix_reservations.bind(reservation, request) + + def _estimate_peak_tokens(self): + return self.prefix_reservations.calculate(self.cache_len_list) diff --git a/lightllm/server/router/req_queue/chunked_prefill/impl.py b/lightllm/server/router/req_queue/chunked_prefill/impl.py index f69ab11c5e..ab8bebc95d 100644 --- a/lightllm/server/router/req_queue/chunked_prefill/impl.py +++ b/lightllm/server/router/req_queue/chunked_prefill/impl.py @@ -8,7 +8,7 @@ class ChunkedPrefillQueue(BaseQueue): def _init_cache_list(self, current_batch: Batch, is_busy): if current_batch is not None: self.cache_len_list = [ - req.get_tuple_tokens(is_busy, self.router.router_statics.ema_req_out_len) + self._get_req_reservation(req, is_busy) for req in current_batch.reqs if req.sample_params.suggested_dp_index == self.dp_index ] @@ -18,17 +18,10 @@ def _init_cache_list(self, current_batch: Batch, is_busy): # @calculate_time(show=True, min_cost_ms=0.1) def _can_add_new_req(self, req: Req, is_busy): - self.cache_len_list.append( - req.get_tuple_tokens(is_busy, self.router.router_statics.ema_req_out_len) - ) # hard to analysis + self.cache_len_list.append(self._get_req_reservation(req, is_busy)) # hard to analysis self.cache_len_list.sort(key=lambda x: -x[1]) - left_out_len_array = np.array([e[1] for e in self.cache_len_list]) - has_run_len_array = np.array([e[0] for e in self.cache_len_list]) - cum_run_len_array = np.cumsum(has_run_len_array) - size_array = np.arange(1, len(self.cache_len_list) + 1, 1) - - need_max_token_num = (left_out_len_array * size_array + cum_run_len_array).max() + need_max_token_num = self._estimate_peak_tokens() ok_token_num = need_max_token_num < self.max_total_tokens ok_req_num = len(self.cache_len_list) <= self.running_max_req_size @@ -79,14 +72,18 @@ def generate_new_batch(self, current_batch: Batch): def _calcu_batch_token_load_batch_not_none(self, current_batch: Batch): is_busy = self.is_busy() self._init_cache_list(current_batch, is_busy) - if len(self.cache_len_list) != 0: - self.cache_len_list.sort(key=lambda x: -x[1]) - left_out_len_array = np.array([e[1] for e in self.cache_len_list]) - has_run_len_array = np.array([e[0] for e in self.cache_len_list]) - cum_run_len_array = np.cumsum(has_run_len_array) - size_array = np.arange(1, len(self.cache_len_list) + 1, 1) - need_max_token_num = (left_out_len_array * size_array + cum_run_len_array).max() - else: - need_max_token_num = 0 + need_max_token_num = self._estimate_peak_tokens() return (need_max_token_num, need_max_token_num / self.max_total_tokens) + + def _get_req_reservation(self, req, is_busy): + return req.get_tuple_tokens(is_busy, self.router.router_statics.ema_req_out_len) + + def _estimate_peak_tokens(self): + if not self.cache_len_list: + return 0 + self.cache_len_list.sort(key=lambda entry: -entry[1]) + remaining = np.array([entry[1] for entry in self.cache_len_list]) + resident = np.array([entry[0] for entry in self.cache_len_list]) + count = np.arange(1, len(self.cache_len_list) + 1) + return (remaining * count + np.cumsum(resident)).max() diff --git a/lightllm/server/router/req_queue/peak_reservation.py b/lightllm/server/router/req_queue/peak_reservation.py new file mode 100644 index 0000000000..9df7ab6144 --- /dev/null +++ b/lightllm/server/router/req_queue/peak_reservation.py @@ -0,0 +1,52 @@ +from dataclasses import dataclass +from typing import Optional, Sequence, Tuple + + +@dataclass(frozen=True) +class Reservation: + base_tokens: int + remaining_tokens: int + prefix_pages: Tuple[bytes, ...] = () + cache_domain: Optional[bytes] = None + + +@dataclass(frozen=True) +class Peak: + legacy_tokens: int + shared_tokens: int + max_prefix_credit: int + + +def estimate_peak(entries: Sequence[Reservation], page_tokens: int) -> Peak: + if type(page_tokens) is not int or page_tokens <= 0: + raise ValueError("Page token count must be a positive integer") + for entry in entries: + if ( + type(entry.base_tokens) is not int + or entry.base_tokens < 0 + or type(entry.remaining_tokens) is not int + or entry.remaining_tokens < 0 + ): + raise ValueError("Reservation counts must be nonnegative integers") + if not isinstance(entry.prefix_pages, tuple) or any(type(page) is not bytes for page in entry.prefix_pages): + raise ValueError("Prefix pages must have exact immutable identities") + if entry.cache_domain is not None and type(entry.cache_domain) is not bytes: + raise ValueError("An explicit shared domain must have immutable identity") + if len(entry.prefix_pages) * page_tokens > entry.base_tokens: + raise ValueError("Known cached pages exceed the resident base") + seen = set() + base = credit = legacy = shared = 0 + for count, entry in enumerate(sorted(entries, key=lambda entry: -entry.remaining_tokens), 1): + base += entry.base_tokens + if entry.cache_domain is not None: + prefix = (entry.cache_domain,) + for page in entry.prefix_pages: + prefix += (page,) + if prefix in seen: + credit += page_tokens + else: + seen.add(prefix) + projected = base + count * entry.remaining_tokens + legacy = max(legacy, projected) + shared = max(shared, projected - credit) + return Peak(legacy, shared, credit) diff --git a/lightllm/server/tokenizer.py b/lightllm/server/tokenizer.py index 12ddbb1c68..1fb836998d 100644 --- a/lightllm/server/tokenizer.py +++ b/lightllm/server/tokenizer.py @@ -124,6 +124,15 @@ def get_tokenizer( tokenizer = QWen3VLTokenizer( tokenizer=tokenizer, image_processor=processor.image_processor, model_cfg=model_cfg ) + elif model_type == "qwen4_exp" and "vision_config" in model_cfg: + from ..models.qwen4_exp.image_processing import Qwen4ExpImageProcessor + from ..models.qwen4_exp.tokenizer import Qwen4ExpTokenizer + + tokenizer = Qwen4ExpTokenizer( + tokenizer=tokenizer, + image_processor=Qwen4ExpImageProcessor.from_pretrained(tokenizer_name), + model_cfg=model_cfg, + ) elif model_type in ["qwen3_5", "sensenova_loop", "qwen3_5_moe"] and "vision_config" in model_cfg: from ..models.qwen3_5.model import QWen3_5Tokenizer from transformers import AutoImageProcessor diff --git a/lightllm/server/visualserver/model_infer/model_rpc.py b/lightllm/server/visualserver/model_infer/model_rpc.py index 6451e70008..b8751db58b 100644 --- a/lightllm/server/visualserver/model_infer/model_rpc.py +++ b/lightllm/server/visualserver/model_infer/model_rpc.py @@ -92,6 +92,10 @@ def exposed_init_model(self, kvargs): self.model = ( Qwen2_5_VisionTransformerPretrainedModel(kvargs, **model_cfg["vision_config"]).eval().bfloat16() ) + elif self.model_type == "qwen4_exp": + from lightllm.models.qwen4_exp.vision import Qwen4ExpVisionTransformer + + self.model = Qwen4ExpVisionTransformer(kvargs, **model_cfg["vision_config"]).eval().bfloat16() elif self.model_type in ["qwen3_vl", "qwen3_vl_moe", "qwen3_5", "sensenova_loop", "qwen3_5_moe"]: self.model = ( Qwen3VisionTransformerPretrainedModel(kvargs, **model_cfg["vision_config"]).eval().bfloat16() diff --git a/lightllm/utils/config_utils.py b/lightllm/utils/config_utils.py index 1b0ce53879..187af4706d 100644 --- a/lightllm/utils/config_utils.py +++ b/lightllm/utils/config_utils.py @@ -216,6 +216,8 @@ def auto_set_fused_shared_experts(args) -> None: "qwen3_5_text", "qwen3_5_moe", "qwen3_5_moe_text", + "qwen4_exp", + "qwen4_exp_text", } if model_type not in supported_model_types: logger.info(f"do not enable fused shared experts: unsupported model_type={model_type}") @@ -288,7 +290,15 @@ def get_eos_token_ids(model_path: str) -> Optional[List[int]]: try: config_json = get_config_json(model_path) model_type = config_json.get("model_type") or config_json.get("text_config", {}).get("model_type") - if model_type in {"qwen3_5", "sensenova_loop", "qwen3_5_text", "qwen3_5_moe", "qwen3_5_moe_text"}: + if model_type in { + "qwen3_5", + "sensenova_loop", + "qwen3_5_text", + "qwen3_5_moe", + "qwen3_5_moe_text", + "qwen4_exp", + "qwen4_exp_text", + }: from transformers import AutoTokenizer eos_token_ids = [] @@ -419,6 +429,8 @@ def has_vision_module(model_path: str) -> bool: return True elif model_type == "neo_chat": return True + elif model_type == "qwen4_exp": + return model_cfg.get("vision_config") is not None elif model_type in ["qwen3_5", "sensenova_loop", "qwen3_5_moe"]: return True else: @@ -464,6 +476,8 @@ def is_linear_att_mixed_model(model_path: str) -> bool: "qwen3_5_moe", "qwen3_5_text", "qwen3_5_moe_text", + "qwen4_exp", + "qwen4_exp_text", "glm5_next", "glm5_next_text", ] @@ -559,6 +573,8 @@ def get_reasoning_parser_for_model(model_path: str) -> Optional[str]: "qwen3_5_moe", "qwen3_5_text", "qwen3_5_moe_text", + "qwen4_exp", + "qwen4_exp_text", ]: return "qwen3" diff --git a/test/benchmark/kernels/benchmark_kda_ssm.py b/test/benchmark/kernels/benchmark_kda_ssm.py new file mode 100644 index 0000000000..f4928a9017 --- /dev/null +++ b/test/benchmark/kernels/benchmark_kda_ssm.py @@ -0,0 +1,162 @@ +"""Measure a full KDA state cycle: all layers, acceptance, folds and export. + +This synthetic kernel benchmark excludes convolutions, model projections and +HTTP scheduling. Use a matched service workload for end-to-end conclusions. +""" + +import argparse +import gc +import json +import statistics +from pathlib import Path + +import torch +import triton + +from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.kda_decode import fused_recurrent_kda + + +def ReplaySSMExecutor(*args, **kwargs): + kda = kwargs.pop("kda", False) + if kda: + from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import _history_owner + else: + from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import _history_owner + + return _history_owner(*args, **kwargs) + + +from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import get_configs, replay_layer + + +def workload(args, batch, width, dtype, mode, acceptance): + layers, heads, dim = args.layers, args.heads, 128 + shape = (1, batch * width, heads, dim) if width > 1 else (batch, 1, heads, dim) + packed = torch.randn(batch * width, heads * dim * 3, dtype=torch.bfloat16, device="cuda") + q, k, v = [x.view(shape) for x in packed.split(heads * dim, -1)] + gate = torch.randn(batch * width, heads * dim, dtype=q.dtype, device="cuda") - 3 + beta = torch.randn(batch * width, heads, dtype=q.dtype, device="cuda") + alog = torch.zeros(heads, device="cuda") + bias = torch.zeros(heads * dim, device="cuda") + reqs = torch.arange(batch, dtype=torch.int32, device="cuda") + cu = torch.arange(batch + 1, dtype=torch.int32, device="cuda") * width if width > 1 else None + counts = reqs % width + 1 if acceptance == "partial" else torch.full_like(reqs, width) + accepted = torch.cat((counts - 1, counts.new_zeros(1))) + snapshot = torch.empty(layers, heads, dim, dim, dtype=dtype, device="cuda") + slots = (batch + 1) * (width if mode == "native" else 1) + state = torch.randn(layers, slots, heads, dim, dim, dtype=dtype, device="cuda") * 0.01 + ssm_backend = None + if mode != "native": + capacity = max(8, triton.next_power_of_2(width)) + ssm_backend = ReplaySSMExecutor(state, capacity, width, q.dtype, kda=True, compute_mode=mode) + if ssm_backend is not None: + ssm_backend._config_is_fixed = True + indexes = reqs[:, None] * width + torch.arange(width, device="cuda", dtype=torch.int32) + if width == 1: + indexes = reqs + bytes_used = ( + sum(x.numel() * x.element_size() for x in vars(ssm_backend).values() if isinstance(x, torch.Tensor)) + if ssm_backend + else state.numel() * state.element_size() + ) + + native_index = 0 if acceptance == "partial" else width - 1 + + def cycle(): + for _ in range(args.rounds): + if ssm_backend is not None: + ssm_backend.prepare_decode(reqs, cu) + for layer in range(layers): + if ssm_backend is None: + fused_recurrent_kda( + q, + k, + v, + gate.view(*shape[:2], -1), + beta.view(*shape[:2], heads), + alog, + bias, + state[layer], + indexes, + cu_seqlens=cu, + num_accepted_tokens=counts if cu is not None else None, + ) + else: + replay_layer(ssm_backend, layer, q, k, v, gate, beta, alog, bias, reqs, cu) + if ssm_backend is not None: + ssm_backend.accept_updates(reqs, accepted) + if isinstance(ssm_backend, ReplaySSMExecutor): + ssm_backend._materialize_accepted_state(reqs[:1], snapshot, snapshot=True) + ssm_backend.merge_accepted_updates(reqs) + else: + snapshot.copy_(state[:, native_index] if ssm_backend is None else state[:, 0]) + + return ssm_backend, cycle, bytes_used + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--batch", type=int, nargs="+", default=[16, 64, 192]) + p.add_argument("--width", type=int, nargs="+", default=[1, 3]) + p.add_argument("--dtype", nargs="+", choices=["float32", "bfloat16"], default=["float32", "bfloat16"]) + p.add_argument("--layers", type=int, default=34) + p.add_argument("--heads", type=int, default=8) + p.add_argument("--rounds", type=int, default=8) + p.add_argument("--repeats", type=int, default=5) + p.add_argument("--tune", action="store_true") + p.add_argument("--output", type=Path, required=True) + args = p.parse_args() + torch.manual_seed(53) + rows = [] + report = dict( + settings={**vars(args), "output": str(args.output)}, + device=torch.cuda.get_device_name(), + torch=torch.__version__, + triton=triton.__version__, + results=rows, + ) + for dtype_name in args.dtype: + for batch in args.batch: + for width in args.width: + for acceptance in ["full", "partial"] if width > 1 else ["full"]: + for mode in ["native", "inline", "precompute"] if width > 1 else ["native", "inline"]: + ssm_backend, cycle, size = workload( + args, batch, width, getattr(torch, dtype_name), mode, acceptance + ) + candidates = [] + if ssm_backend is not None: + for config in get_configs() if args.tune else [{"BV": 32, "num_warps": 2, "num_stages": 1}]: + if mode == "precompute" and config["BV"] < 16: + continue + ssm_backend.run_config = config + cycle() + ms = triton.testing.do_bench_cudagraph(cycle, rep=50) + candidates.append(dict(config=config, ms=ms)) + ssm_backend.run_config = min(candidates, key=lambda r: r["ms"])["config"] + cycle() + samples = [ + triton.testing.do_bench_cudagraph(cycle, rep=100) / args.rounds for _ in range(args.repeats) + ] + row = dict( + batch=batch, + width=width, + dtype=dtype_name, + acceptance=acceptance, + mode=mode, + cache_len=getattr(ssm_backend, "capacity", None), + state_history_bytes=size, + ms_per_round=statistics.median(samples), + samples_ms=samples, + config=dict(ssm_backend.run_config) if ssm_backend and ssm_backend.run_config else None, + candidates=candidates, + ) + rows.append(row) + args.output.write_text(json.dumps(report, indent=2) + "\n") + print(json.dumps(row), flush=True) + del ssm_backend, cycle + gc.collect() + torch.cuda.empty_cache() + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/kernels/benchmark_replayssm_service.py b/test/benchmark/kernels/benchmark_replayssm_service.py new file mode 100644 index 0000000000..a1cee4826a --- /dev/null +++ b/test/benchmark/kernels/benchmark_replayssm_service.py @@ -0,0 +1,119 @@ +"""Compare streamed requests on baseline and optimized services with identical warmup. + +Reports request timings and token IDs for an exact cross-revision comparison. +""" + +import argparse +import concurrent.futures +import json +import statistics +import time +import threading +from pathlib import Path + +import requests +from transformers import AutoTokenizer + +p = argparse.ArgumentParser() +p.add_argument("--output", required=True) +p.add_argument("--batch", type=int, nargs="+", default=[1, 8, 16]) +p.add_argument("--repeats", type=int, default=3) +p.add_argument("--tokens", type=int, default=64) +p.add_argument("--warmups", type=int, default=1) +p.add_argument("--model-dir", required=True) +p.add_argument("--url", default="http://127.0.0.1:18785/generate_stream") +args = p.parse_args() +url = args.url +tok = AutoTokenizer.from_pretrained(args.model_dir, trust_remote_code=True) +prefix = tok.encode( + "Describe how a computer executes a program, from source code to machine instructions. " + "Explain the role of memory, processors, and operating systems. " * 100 +)[:512] + + +def prepare_prompt(index, rep): + suffix = tok.encode( + f"Example number {rep * 1000 + index}. Discuss this scenario in detail. " + + ("An engineer measures performance and validates numerical correctness. " * 100) + )[:512] + return prefix + suffix + + +def stream_request(index, rep, ids, barrier, cohort_start): + payload = { + "inputs": ids, + "parameters": { + "do_sample": False, + "disable_prompt_cache": True, + "ignore_eos": True, + "max_new_tokens": args.tokens, + }, + } + barrier.wait() + start = time.perf_counter() + times = [] + tokens = [] + texts = [] + finished = False + with requests.post(url, json=payload, stream=True, timeout=300) as response: + response.raise_for_status() + # Consume each HTTP chunk; byte-sized iteration can saturate the client CPU. + for line in response.iter_lines(chunk_size=None): + if not line.startswith(b"data:"): + continue + item = json.loads(line[5:]) + assert "error" not in item, item + token = item.get("token", {}) + if token.get("id") is not None: + times.append(time.perf_counter() - cohort_start) + tokens.append(token["id"]) + texts.append(token.get("text", "")) + finished = item.get("finished", finished) + assert finished and len(tokens) == args.tokens, (finished, len(tokens)) + return dict( + index=index, + rep=rep, + prompt_len=len(ids), + token_ids=tokens, + text="".join(texts), + token_times_s=times, + ttft_ms=(times[0] - (start - cohort_start)) * 1000, + tpot_ms=(times[-1] - times[0]) * 1000 / max(1, len(times) - 1), + elapsed_s=time.perf_counter() - start, + ) + + +def run_concurrent_requests(batch, rep): + prompts = [prepare_prompt(i, rep) for i in range(batch)] + barrier = threading.Barrier(batch) + start = time.perf_counter() + with concurrent.futures.ThreadPoolExecutor(max_workers=batch) as pool: + samples = list(pool.map(lambda i: stream_request(i, rep, prompts[i], barrier, start), range(batch))) + elapsed = time.perf_counter() - start + begin = max(r["token_times_s"][0] for r in samples) + 0.5 + end = min(r["token_times_s"][-1] for r in samples) - 0.5 + count = sum(sum(begin <= t <= end for t in r["token_times_s"]) for r in samples) + return dict( + batch=batch, + rep=rep, + elapsed_s=elapsed, + tokens_per_s=sum(len(r["token_ids"]) for r in samples) / elapsed, + steady_tokens_per_s=count / (end - begin) if end > begin else None, + steady_window_s=max(0.0, end - begin), + ttft_ms=statistics.median(r["ttft_ms"] for r in samples), + tpot_ms=statistics.median(r["tpot_ms"] for r in samples), + requests=samples, + ) + + +# Warm every measured graph bucket before timing, with the same requests on both revisions. +for batch in args.batch: + for warm in range(args.warmups): + run_concurrent_requests(batch, 99 + warm) +rows = [] +for batch in args.batch: + for rep in range(args.repeats): + row = run_concurrent_requests(batch, rep) + rows.append(row) + Path(args.output).write_text(json.dumps(rows, indent=2)) + print(json.dumps({k: v for k, v in row.items() if k != "requests"}), flush=True) diff --git a/test/benchmark/qwen38/benchmark_gated_residual.py b/test/benchmark/qwen38/benchmark_gated_residual.py new file mode 100644 index 0000000000..0e3c4fa039 --- /dev/null +++ b/test/benchmark/qwen38/benchmark_gated_residual.py @@ -0,0 +1,93 @@ +import argparse +import json +import os + +import torch +import torch.nn.functional as F +import triton +from triton.testing import do_bench_cudagraph + +from lightllm.common.gated_residual import GatedResidual + + +def eager(component, residual, update, precise_projections=False): + def project(inputs, weight): + if precise_projections: + return F.linear(inputs.double(), weight.double()).to(inputs.dtype) + return F.linear(inputs, weight) + + grouped = residual.float().unflatten(-1, (component.branches, -1)) + normalized = grouped * torch.rsqrt(grouped.square().mean(-1, keepdim=True) + component.eps) + normalized = (normalized.flatten(-2) * (1 + component.norm_weight.float())).to(residual.dtype) + projected = F.silu(project(normalized, component.down_weight) / component.branches) + mix = torch.sigmoid(project(projected, component.up_weight)) + mixed = (normalized * mix).unflatten(-1, (component.branches, -1)).mean(-2) + gates = 2 * torch.sigmoid(project(normalized, component.inject_weight) / component.branches) + return mixed, residual + (update.unsqueeze(-2) * gates.unsqueeze(-1)).flatten(-2) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--tokens", type=int, nargs="+", default=[1, 16, 128, 1024]) + parser.add_argument("--repetitions", type=int, default=3) + args = parser.parse_args() + if not os.environ.get("EXP_ACTIVE"): + parser.error('Wrap this benchmark with exp -m "purpose"') + torch.manual_seed(37) + dim, branches, rank = 2560, 4, 320 + width = dim * branches + + def rand(*shape): + return torch.randn(*shape, device="cuda", dtype=torch.bfloat16) + + component = GatedResidual( + rand(width) * 0.1, + rand(rank, width) / width ** 0.5, + rand(width, rank) / rank ** 0.5, + rand(branches, width) / width ** 0.5, + branches, + ) + print( + json.dumps( + { + "torch": torch.__version__, + "triton": triton.__version__, + "cuda": torch.version.cuda, + "gpu": torch.cuda.get_device_name(), + "dtype": "bfloat16", + "hidden_size": dim, + "branches": branches, + "rank": rank, + "timing": "CUDA graph, median of repetitions", + "validation": "eager equations with FP64 projections quantized at original dtype boundaries", + "eager_timing": "original input-dtype projection; batch-dependent rounding recorded separately", + "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), + } + ) + ) + for tokens in args.tokens: + residual, update = rand(tokens, width), rand(tokens, dim) + + def fused(): + mixed, gates = component.read(residual) + return mixed, component.write(residual, update, gates) + + def reference(): + return eager(component, residual, update) + + actual = fused() + expected = eager(component, residual, update, precise_projections=True) + for value, oracle in zip(actual, expected): + torch.testing.assert_close(value, oracle, atol=2e-3, rtol=torch.finfo(residual.dtype).eps) + results = { + "eager_max_abs": [float((value - old).float().abs().max()) for value, old in zip(actual, reference())] + } + for name, func in [("eager", reference), ("fused", fused)]: + samples = [do_bench_cudagraph(func, rep=100) * 1000 for _ in range(args.repetitions)] + results[name + "_us"] = sorted(samples)[len(samples) // 2] + results[name + "_samples_us"] = samples + print(json.dumps({"tokens": tokens, **results, "speedup": results["eager_us"] / results["fused_us"]})) + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/qwen38/benchmark_lm_head.py b/test/benchmark/qwen38/benchmark_lm_head.py new file mode 100644 index 0000000000..422c9018c2 --- /dev/null +++ b/test/benchmark/qwen38/benchmark_lm_head.py @@ -0,0 +1,111 @@ +import argparse +import importlib.metadata +import json +import os +from pathlib import Path +import subprocess + +import torch +import triton +from safetensors import safe_open + +from check_generation import REVISION, verified_checkpoint + + +@torch.inference_mode() +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model-directory", type=Path, required=True) + parser.add_argument( + "--tp", type=int, default=4, help="Vocabulary partition count; this process runs rank zero only" + ) + parser.add_argument("--rows", type=int, nargs="+", default=[1, 3, 5, 7, 8, 9, 24, 64, 128]) + parser.add_argument("--dtype", choices=["bf16", "fp16"], default="bf16") + parser.add_argument("--samples", type=int, default=3) + parser.add_argument("--rep-ms", type=int, default=40) + args = parser.parse_args() + if not os.environ.get("EXP_ACTIVE"): + raise RuntimeError("Run through exp") + if min(args.rows + [args.tp, args.samples, args.rep_ms]) < 1: + parser.error("Partition count, rows, samples and repeat duration must be positive") + verified = verified_checkpoint(args.model_directory) + checkpoint = json.loads((args.model_directory / "config.json").read_text()) + config = checkpoint.get("text_config", checkpoint) + hidden, vocab = config["hidden_size"], config["vocab_size"] + if args.tp > vocab: + parser.error("Partition count cannot exceed vocabulary size") + + from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.utils.envs_utils import get_env_start_args + + start = StartArgs() + start.hardware_platform, start.data_type, start.tp, start.dp = "cuda", args.dtype, args.tp, 1 + os.environ["LIGHTLLM_START_ARGS"] = json.dumps(vars(start)) + os.environ.update( + LIGHTLLM_DP_WORLD_SIZE=str(args.tp), LIGHTLLM_CURRENT_RANK_IN_DP="0", LIGHTLLM_CURRENT_DEVICE_ID="0" + ) + get_env_start_args.cache_clear() + from lightllm.common.basemodel.layer_weights.meta_weights import LMHeadWeight + from lightllm.models.qwen4_exp.layer_weights.lm_head_weight import Qwen4ExpLMHeadWeight + + dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float16 + baseline = LMHeadWeight(hidden, vocab, "lm_head.weight", dtype) + index = json.loads((args.model_directory / "model.safetensors.index.json").read_text())["weight_map"] + filename = index["lm_head.weight"] + with safe_open(args.model_directory / filename, framework="pt", device="cpu") as source: + baseline.load_hf_weights({"lm_head.weight": source.get_tensor("lm_head.weight")}) + assert baseline.verify_load() + optimized = Qwen4ExpLMHeadWeight(hidden, vocab, "lm_head.weight", dtype, embedding_weight=baseline) + assert optimized.verify_load() and optimized.weight is baseline.weight + torch.manual_seed(3107) + inputs = torch.randn(max(args.rows), hidden, device="cuda", dtype=dtype) + reference = torch.cat([baseline(row[:, None]) for row in inputs], dim=1) + run = Path(os.environ["EXP_RUN_DIR"]) + report = dict( + status="running", + scope="trained rank-zero head component; no distributed serving or vLLM/SGLang performance claim", + revision=REVISION, + configuration=vars(args) | {"model_directory": str(args.model_directory)}, + local_weight_shape=list(baseline.weight.shape), + hardware=subprocess.check_output( + ["nvidia-smi", "--query-gpu=index,name,uuid,memory.used,memory.free", "--format=csv,noheader"], text=True + ), + software={name: importlib.metadata.version(name) for name in ("torch", "triton")}, + cuda=torch.version.cuda, + bf16_reduced_precision=torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction, + source_file=verified["files"][filename], + results=[], + ) + path = run / "results.json" + try: + for rows in args.rows: + tokens = inputs[:rows].T + old, new = baseline(tokens), optimized(tokens) + torch.testing.assert_close(new, reference[:, :rows], atol=0, rtol=0) + timings = {} + for name, weight in (("baseline", baseline), ("optimized", optimized)): + timings[name] = [ + triton.testing.do_bench_cudagraph(lambda: weight(tokens), rep=args.rep_ms) * 1000 + for _ in range(args.samples) + ] + item = dict( + rows=rows, + scalar_matches=True, + baseline_max_abs=float((old - reference[:, :rows]).float().abs().max()), + microseconds=timings, + speedup=sorted(timings["baseline"])[args.samples // 2] + / sorted(timings["optimized"])[args.samples // 2], + ) + report["results"].append(item) + path.write_text(json.dumps(report, indent=2)) + print(json.dumps(item), flush=True) + report["status"] = "passed" + except Exception as error: + report.update(status="failed", error_type=type(error).__name__, error=str(error)) + raise + finally: + path.write_text(json.dumps(report, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/qwen38/benchmark_ngram_hash.py b/test/benchmark/qwen38/benchmark_ngram_hash.py new file mode 100644 index 0000000000..35042cadcb --- /dev/null +++ b/test/benchmark/qwen38/benchmark_ngram_hash.py @@ -0,0 +1,129 @@ +import argparse +import json +import os + +import torch +import triton +from triton.testing import do_bench_cudagraph + +from lightllm.common.ngram_embedding import NGramHasher + + +def eager(hasher, tokens, history): + batch, length = tokens.shape + full = torch.cat((history, tokens.long()), dim=-1) + positions = torch.arange(full.shape[1], device=full.device, dtype=torch.int64) + eos_positions = torch.where(full == hasher.eos_token_id, positions, -1) + inclusive = torch.cummax(eos_positions, dim=1).values + previous_eos = torch.cat((eos_positions.new_full((batch, 1), -1), inclusive[:, :-1]), dim=1) + relative = positions.unsqueeze(0) - previous_eos - 1 + shifted = [full] + for shift in range(1, hasher.context_len + 1): + source = positions - shift + gathered = full.gather(1, source.clamp_min(0).unsqueeze(0).expand(batch, -1)) + valid = (relative >= shift) & (source.unsqueeze(0) >= 0) + shifted.append(torch.where(valid, gathered, hasher.eos_token_id)) + ids = [] + per_order = hasher.heads // hasher.context_len + for order in range(2, hasher.context_len + 2): + mixed = shifted[0] * hasher.multipliers[0] + for index in range(1, order): + mixed = torch.bitwise_xor(mixed, shifted[index] * hasher.multipliers[index]) + begin = (order - 2) * per_order + sizes = hasher.head_vocab_sizes[begin : begin + per_order] + offsets = hasher.head_offsets[begin : begin + per_order] + ids.append(torch.remainder(mixed.unsqueeze(-1), sizes.view(1, 1, -1)) + offsets.view(1, 1, -1)) + return torch.cat(ids, dim=-1)[:, -length:].reshape(-1, hasher.heads), full[:, -hasher.context_len :].contiguous() + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--shapes", nargs="+", default=["1x1", "16x1", "128x1", "1x4096", "16x256"]) + parser.add_argument("--repetitions", type=int, default=3) + args = parser.parse_args() + if not os.environ.get("EXP_ACTIVE"): + parser.error('Wrap this benchmark with exp -m "purpose"') + torch.manual_seed(37) + + sizes = torch.tensor( + [ + 20000003, + 20000023, + 20000033, + 20000047, + 20000059, + 20000063, + 20000069, + 20000077, + 20000081, + 20000093, + 20000107, + 20000147, + 20000153, + 20000159, + 20000161, + 20000171, + ], + device="cuda", + dtype=torch.int64, + ) + offsets = torch.cat([sizes.new_zeros(1), sizes.cumsum(0)[:-1]]) + multipliers = torch.tensor([23703573157769, 20109073645365, 8052911324071], device="cuda", dtype=torch.int64) + hasher = NGramHasher(multipliers, sizes, offsets, 248044) + print( + json.dumps( + { + "torch": torch.__version__, + "triton": triton.__version__, + "cuda": torch.version.cuda, + "gpu": torch.cuda.get_device_name(), + "ngram_size": 3, + "heads_per_ngram": 8, + "multipliers": multipliers.tolist(), + "head_vocab_sizes": sizes.tolist(), + "eos_token_id": hasher.eos_token_id, + "timing": "CUDA graph, median of interleaved repetitions", + "includes_embedding_lookup": False, + "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), + } + ), + flush=True, + ) + for shape in args.shapes: + batch, length = map(int, shape.split("x")) + tokens = torch.randint(0, 248320, (batch, length), dtype=torch.int64, device="cuda") + tokens[:, ::31] = hasher.eos_token_id + starts = torch.arange(batch + 1, device="cuda", dtype=torch.int32) * length + history = hasher.empty_history(batch) + history[:, -1] = 77 + + def fused(): + return hasher.hash(tokens.view(-1), starts, history) + + def reference(): + return eager(hasher, tokens, history) + + for actual, expected in zip(fused(), reference()): + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + samples = {"eager": [], "fused": []} + for repeat in range(args.repetitions): + functions = [("eager", reference), ("fused", fused)] + for name, func in functions[:: 1 if repeat % 2 == 0 else -1]: + samples[name].append(do_bench_cudagraph(func, rep=100) * 1000) + medians = {name + "_us": sorted(values)[len(values) // 2] for name, values in samples.items()} + print( + json.dumps( + { + "batch": batch, + "tokens_per_request": length, + **medians, + "speedup": medians["eager_us"] / medians["fused_us"], + "samples_us": samples, + } + ), + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/qwen38/benchmark_ple.py b/test/benchmark/qwen38/benchmark_ple.py new file mode 100644 index 0000000000..1c98a8481f --- /dev/null +++ b/test/benchmark/qwen38/benchmark_ple.py @@ -0,0 +1,120 @@ +import argparse +import json +import math +import os + +import torch +import torch.nn.functional as F +import triton +from triton.testing import do_bench_cudagraph + +from lightllm.common.ngram_embedding.ple import PLELayer + + +def eager(layer, residual, embeddings, history): + batch, length, width = residual.shape + + def norm(x, weight): + grouped = x.float().unflatten(-1, (layer.branches, -1)) + normalized = grouped * torch.rsqrt(grouped.square().mean(-1, keepdim=True) + layer.eps) + return (normalized.flatten(-2) * (1 + weight.float())).to(x.dtype) + + key = norm(F.linear(embeddings, layer.key_weight), layer.norm_key).unflatten(-1, (layer.branches, -1)) + query = norm(residual, layer.norm_query).unflatten(-1, (layer.branches, -1)) + value = F.linear(embeddings, layer.value_weight) + gate = (key * query).sum(-1, keepdim=True) / math.sqrt(layer.hidden_size) + gate = gate.abs().clamp_min(1e-6).sqrt() * gate.sign() + gated = (torch.sigmoid(gate) * value.unsqueeze(-2)).flatten(-2) + conv_inputs = norm(gated, layer.norm_conv) + full = torch.cat((history, conv_inputs), dim=1) + conv = F.conv1d(full.transpose(1, 2), layer.conv_weight.unsqueeze(1), dilation=layer.dilation, groups=width) + output = gated + F.silu(conv).transpose(1, 2) + return output, full[:, -layer.context_len :].contiguous() + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--shapes", nargs="+", default=["1x1", "16x1", "1x128", "4x128"]) + parser.add_argument("--repetitions", type=int, default=3) + args = parser.parse_args() + if not os.environ.get("EXP_ACTIVE"): + parser.error('Wrap this benchmark with exp -m "purpose"') + torch.manual_seed(137) + dim, branches, embed_dim = 2560, 4, 2560 + width = dim * branches + + def rand(*shape): + return torch.randn(*shape, device="cuda", dtype=torch.bfloat16) + + layer = PLELayer( + rand(width, embed_dim) / embed_dim ** 0.5, + rand(dim, embed_dim) / embed_dim ** 0.5, + rand(width) * 0.1, + rand(width) * 0.1, + rand(width) * 0.1, + rand(width, 4) * 0.1, + branches, + 3, + ) + print( + json.dumps( + { + "torch": torch.__version__, + "triton": triton.__version__, + "cuda": torch.version.cuda, + "gpu": torch.cuda.get_device_name(), + "dtype": "bfloat16", + "hidden_size": dim, + "branches": branches, + "embedding_width": embed_dim, + "kernel_size": 4, + "dilation": 3, + "timing": "CUDA graph, median of interleaved repetitions", + "includes_table_lookup": False, + "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), + } + ), + flush=True, + ) + for shape in args.shapes: + batch, length = map(int, shape.split("x")) + residual = rand(batch, length, width) + embeddings = rand(batch, length, embed_dim) + history = layer.empty_history(batch).normal_() + starts = torch.arange(batch + 1, device="cuda", dtype=torch.int32) * length + + def fused(): + output, conv_inputs = layer.forward( + residual.view(-1, width), embeddings.view(-1, embed_dim), starts, history + ) + return output.view(batch, length, width), layer.history_after(conv_inputs, starts, history) + + def reference(): + return eager(layer, residual, embeddings, history) + + for actual, expected in zip(fused(), reference()): + torch.testing.assert_close( + actual, expected, atol=2 * torch.finfo(residual.dtype).eps, rtol=2 * torch.finfo(residual.dtype).eps + ) + samples = {"eager": [], "fused": []} + for repeat in range(args.repetitions): + functions = [("eager", reference), ("fused", fused)] + for name, func in functions[:: 1 if repeat % 2 == 0 else -1]: + samples[name].append(do_bench_cudagraph(func, rep=100) * 1000) + medians = {name + "_us": sorted(values)[len(values) // 2] for name, values in samples.items()} + print( + json.dumps( + { + "batch": batch, + "tokens_per_request": length, + **medians, + "speedup": medians["eager_us"] / medians["fused_us"], + "samples_us": samples, + } + ), + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/qwen38/benchmark_qsa.py b/test/benchmark/qwen38/benchmark_qsa.py new file mode 100644 index 0000000000..53eab467dd --- /dev/null +++ b/test/benchmark/qwen38/benchmark_qsa.py @@ -0,0 +1,165 @@ +import argparse +import json +import os + +import torch +import torch.nn.functional as F +import triton +from triton.testing import do_bench_cudagraph + +from lightllm.common.qsa.indexer import QSAIndexer +from lightllm.common.qsa.attention import indexed_attention + + +def eager_select(indexer, query, cache, requests, lengths, table): + blocks = cache.shape[1] + count = lengths // indexer.compress_ratio + + scores = torch.matmul(query.float(), cache[0].float().transpose(-1, -2)) + scores = scores.relu().sum(1) / indexer.head_dim ** 0.5 + scores.masked_fill_(torch.arange(blocks, device=query.device)[None, :] >= count[:, None], -float("inf")) + topk = min(blocks, indexer.token_budget // indexer.compress_ratio) + selected = scores.topk(topk, sorted=False).indices + width = indexer.token_budget + indexer.compress_ratio - 1 + col = torch.arange(width, device=query.device)[None, :].expand(query.shape[0], -1) + chosen = count.clamp_max(indexer.token_budget // indexer.compress_ratio) + full = col < chosen[:, None] * indexer.compress_ratio + rank = (col // indexer.compress_ratio).clamp_max(topk - 1) + position = selected.gather(1, rank) * indexer.compress_ratio + col % indexer.compress_ratio + tail_col = col - chosen[:, None] * indexer.compress_ratio + tail = (tail_col >= 0) & (tail_col < (lengths % indexer.compress_ratio)[:, None]) + position = torch.where(full, position, count[:, None] * indexer.compress_ratio + tail_col) + slots = table[requests.long()[:, None], position.clamp(0, table.shape[1] - 1)] + return slots.masked_fill(~(full | tail), -1) + + +def eager_attention(q, k, v, selected, valid_tokens=2048): + from torch.nn.attention import SDPBackend, sdpa_kernel + + valid_slots = selected[:, :valid_tokens] + gathered_k = k[valid_slots].transpose(1, 2) + gathered_v = v[valid_slots].transpose(1, 2) + with sdpa_kernel(SDPBackend.FLASH_ATTENTION): + return F.scaled_dot_product_attention(q.unsqueeze(2), gathered_k, gathered_v, enable_gqa=True).squeeze(2) + + +def measure(reference, fused, repetitions): + samples = {"torch": [], "kernel": []} + for repeat in range(repetitions): + functions = [("torch", reference), ("kernel", fused)] + for name, function in functions[:: 1 if repeat % 2 == 0 else -1]: + samples[name].append(do_bench_cudagraph(function, rep=100) * 1000) + result = {name + "_us": sorted(values)[len(values) // 2] for name, values in samples.items()} + return {**result, "speedup": result["torch_us"] / result["kernel_us"], "samples_us": samples} + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--contexts", nargs="+", type=int, default=[4099, 65539]) + parser.add_argument("--index-rows", nargs="+", type=int, default=[1, 16]) + parser.add_argument("--attention-rows", nargs="+", type=int, default=[1, 16, 128]) + parser.add_argument("--valid-selections", nargs="+", type=int, default=[2048]) + parser.add_argument("--selection-bounds", action="store_true", help="Use native request visibility bounds") + parser.add_argument("--repetitions", type=int, default=3) + args = parser.parse_args() + if not os.environ.get("EXP_ACTIVE"): + parser.error('Wrap this benchmark with exp -m "purpose"') + if any(count < 1 or count > 2048 for count in args.valid_selections): + parser.error("Valid selection counts must lie in [1, 2048]") + torch.manual_seed(537) + + def rand(*shape): + return torch.randn(*shape, device="cuda", dtype=torch.bfloat16) + + indexer = QSAIndexer(rand(640, 2560), rand(128), rand(128)) + print( + json.dumps( + { + "torch": torch.__version__, + "triton": triton.__version__, + "cuda": torch.version.cuda, + "gpu": torch.cuda.get_device_name(), + "dtype": "bfloat16", + "index_heads": 4, + "index_head_dim": 128, + "token_budget": 2048, + "compress_ratio": 4, + "query_heads": 24, + "kv_heads": 2, + "attention_head_dim": 256, + "timing": "CUDA graph, median of interleaved repetitions", + "attention_baseline": "KV gather plus forced Torch FlashAttention, no padding mask", + "selection_baseline": ( + "Torch matmul with shared cached keys, ReLU head reduction, top-k and slot gather" + ), + "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), + "selection_bounds": args.selection_bounds, + } + ), + flush=True, + ) + for context in args.contexts: + cache = rand(1, context // 4, 128) + table = torch.arange(context, device="cuda", dtype=torch.int32).view(1, -1) + for rows in args.index_rows: + query = rand(rows, 4, 128) + requests = torch.zeros(rows, device="cuda", dtype=torch.int32) + lengths = ( + torch.full((rows,), context, device="cuda", dtype=torch.int32) + - torch.arange(rows, device="cuda", dtype=torch.int32) % 7 + ) + + def reference(): + return eager_select(indexer, query, cache, requests, lengths, table) + + def fused(): + return indexer.select(query, cache, requests, lengths, table) + + torch.testing.assert_close(fused().sort(-1).values, reference().sort(-1).values, atol=0, rtol=0) + print( + json.dumps( + { + "component": "cached_block_selection", + "queries": rows, + "context": context, + **measure(reference, fused, args.repetitions), + } + ), + flush=True, + ) + slots = 65539 + key, value = rand(slots, 2, 256), rand(slots, 2, 256) + for rows in args.attention_rows: + query = rand(rows, 24, 256) + choices = torch.stack([torch.randperm(slots, device="cuda")[:2051] for _ in range(rows)]).int() + for count in args.valid_selections: + selected = choices.clone() + selected[:, count:] = -1 + + bound = slots if count == 2048 else count + lengths = torch.full((rows,), bound, device="cuda", dtype=torch.int32) if args.selection_bounds else None + + def reference(): + return eager_attention(query, key, value, selected, count) + + def fused(): + return indexed_attention(query, key, value, selected, selection_lengths=lengths) + + torch.testing.assert_close(fused(), reference(), atol=0.016, rtol=0.016) + print( + json.dumps( + { + "component": "indexed_gqa", + "queries": rows, + "physical_slots": slots, + "selected": count, + "selection_capacity": selected.shape[1], + **measure(reference, fused, args.repetitions), + } + ), + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/qwen38/check_aiperf_timing.py b/test/benchmark/qwen38/check_aiperf_timing.py new file mode 100644 index 0000000000..f282df9374 --- /dev/null +++ b/test/benchmark/qwen38/check_aiperf_timing.py @@ -0,0 +1,60 @@ +import argparse +import hashlib +import json +import os +from pathlib import Path + +from test.benchmark.qwen38.measure_aiperf import audit_records, profiling_records +from test.benchmark.service.stream_timing import summarize_streaming_times + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--replay", type=Path, required=True, help="Completed aiperf-replay.json") + args = parser.parse_args() + if not os.environ.get("EXP_ACTIVE"): + parser.error('Run through exp -m "purpose"') + replay = json.loads(args.replay.read_text()) + if replay["status"] != "aiperf_replay_recorded": + raise ValueError("A completed replay is required") + directory = args.replay.parent + trace_path = directory / "trace.jsonl" + trace = [json.loads(line) for line in trace_path.read_text().splitlines()] + audit = audit_records(directory / "aiperf", trace) + if audit != replay["audit"]: + raise ValueError("Retained replay audit differs from its raw records") + timing = summarize_streaming_times(profiling_records(directory / "aiperf")) + if timing["requests"] != audit["requests"] or timing["empty_text_responses"] != audit["content_errors"]: + raise ValueError("Timing cohort differs from the complete replay") + service = replay["service"] + for key in ("framework", "version", "software", "hardware", "command", "configuration", "model_revision"): + if not service.get(key): + raise ValueError(f"Missing actual service provenance: {key}") + paths = ( + args.replay, + trace_path, + directory / "aiperf/profile_export_raw.jsonl", + Path(__file__), + Path(__file__).parent.parent / "service/stream_timing.py", + ) + report = dict( + status="complete_aiperf_stream_timing_verified", + replay=str(args.replay.resolve()), + service=service, + timing=timing, + aiperf_metrics=audit["metrics"], + hashes={str(path.resolve()): hashlib.sha256(path.read_bytes()).hexdigest() for path in paths}, + comparison_eligible=False, + limitations=( + "Timing only. Complete work and explicit empty-text errors are retained. " + "HTTP completion includes SSE DONE; visible-text boundaries exclude later invisible-token generation. " + "Task quality and matched framework comparisons require separate audits." + ), + ) + run = Path(os.environ["EXP_RUN_DIR"]) + (run / "stream-timing.json").write_text(json.dumps(report, indent=2) + "\n") + print(json.dumps(dict(status=report["status"], timing=timing))) + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/qwen38/check_baseline.py b/test/benchmark/qwen38/check_baseline.py new file mode 100644 index 0000000000..e732a83b9b --- /dev/null +++ b/test/benchmark/qwen38/check_baseline.py @@ -0,0 +1,102 @@ +import argparse +from datetime import datetime, timezone +import hashlib +import importlib.metadata +import inspect +import json +import os +from pathlib import Path +import platform +import subprocess +import sys + + +def check_runtime(framework, model_directory): + report = { + "framework": framework, + "created_at": datetime.now(timezone.utc).isoformat(), + "python": platform.python_version(), + "glibc": platform.libc_ver(), + "executable": sys.executable, + "scope": "Dependency/import/GPU smoke checks only; no serving or performance result", + "software": {}, + "status": "checking", + } + for package in (framework, "torch", "triton", "transformers", "flashinfer-python"): + try: + report["software"][package] = importlib.metadata.version(package) + except importlib.metadata.PackageNotFoundError: + report["software"][package] = None + try: + dependencies = subprocess.run([sys.executable, "-m", "pip", "check"], capture_output=True, text=True) + report["dependencies"] = { + "exit_code": dependencies.returncode, + "output": dependencies.stdout + dependencies.stderr, + } + if dependencies.returncode: + raise RuntimeError("Baseline dependency validation failed") + import torch + + if not torch.cuda.is_available(): + raise RuntimeError("Baseline cannot access CUDA") + report.update(cuda=torch.version.cuda, nccl=torch.cuda.nccl.version(), gpu=torch.cuda.get_device_name()) + a = torch.ones(16, 64, device="cuda", dtype=torch.bfloat16) + b = torch.ones(64, 32, device="cuda", dtype=torch.bfloat16) + torch.testing.assert_close(a @ b, torch.full((16, 32), 64, device="cuda", dtype=torch.bfloat16), atol=0, rtol=0) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = a @ b + a.fill_(2) + graph.replay() + torch.testing.assert_close(captured, torch.full_like(captured, 128), atol=0, rtol=0) + graph.reset() + report["bf16_matmul_and_graph"] = "passed" + architecture = "Qwen4ExpForConditionalGeneration" + if framework == "vllm": + from vllm.config import ModelConfig + from vllm.model_executor.models import ModelRegistry + + config = ModelConfig( + model=str(model_directory.resolve()), skip_tokenizer_init=True, dtype="bfloat16", max_model_len=32 + ) + model_class, resolved_architecture = ModelRegistry.resolve_model_cls(architecture, config) + if resolved_architecture != architecture: + raise RuntimeError("Baseline resolved a different model architecture") + else: + from sglang.srt.models.qwen4_exp import Qwen4ExpForConditionalGeneration + + model_class = Qwen4ExpForConditionalGeneration + source = Path(inspect.getfile(model_class)) + report["model"] = { + "architecture": architecture, + "module": model_class.__module__, + "source": str(source), + "sha256": hashlib.sha256(source.read_bytes()).hexdigest(), + } + report["status"] = "passed" + except Exception as error: + report.update(status="failed", error_type=type(error).__name__, error=str(error)) + return report + + +def main(): + if not os.environ.get("EXP_ACTIVE"): + raise RuntimeError("Run baseline checks through exp") + parser = argparse.ArgumentParser() + parser.add_argument("--framework", choices=("vllm", "sglang"), required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--model-directory", type=Path, default=Path("artifacts/references/qwen4_exp/layer0")) + args = parser.parse_args() + result = check_runtime(args.framework, args.model_directory) + payload = json.dumps(result, indent=2) + "\n" + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(payload) + (Path(os.environ["EXP_RUN_DIR"]) / "baseline-readiness.json").write_text(payload) + print(payload, flush=True) + if result["status"] != "passed": + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/qwen38/check_cold_streams.py b/test/benchmark/qwen38/check_cold_streams.py new file mode 100644 index 0000000000..2c5da4f2d9 --- /dev/null +++ b/test/benchmark/qwen38/check_cold_streams.py @@ -0,0 +1,159 @@ +import argparse +import hashlib +import json +import os +from pathlib import Path +import secrets +import shutil + +from test.benchmark.qwen38.check_generation import REVISION, cases_for, verified_checkpoint +from test.benchmark.service.stream_cohort import StreamingWorkload, run_cohort, summarize_cohort +from test.benchmark.service.token_stream import request_stream + + +def workloads_from_reference(reference, expected_cases): + if ( + reference.get("status") not in ("passed", "recorded") + or reference.get("revision") != REVISION + or reference.get("cases") != expected_cases + ): + raise ValueError("Require a qualified generation reference with the exact pinned inputs") + length = reference.get("output_length") + if type(length) is not int or length < 1: + raise ValueError("Reference output length must be a positive integer") + results = reference.get("results", {}) + if set(results) != set(expected_cases): + raise ValueError("Reference must contain every declared generation case") + workloads = [] + for name, prompt in expected_cases.items(): + ids = results[name].get("token_ids") + if not isinstance(ids, list) or len(ids) != length or any(type(token) is not int or token < 0 for token in ids): + raise ValueError("Reference must contain the complete actual output IDs") + if name.startswith("repeat_"): + original = name.removeprefix("repeat_") + if ( + original not in expected_cases + or prompt != expected_cases[original] + or ids != results[original]["token_ids"] + ): + raise ValueError("Repeated reference prompts and outputs must match") + else: + workloads.append(StreamingWorkload(name, prompt, ids)) + if not workloads: + raise ValueError("Require at least one non-repeated workload") + return workloads + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base-url", required=True) + parser.add_argument("--model-directory", type=Path, required=True) + parser.add_argument("--service-metadata", type=Path, required=True) + parser.add_argument("--reference", type=Path, required=True) + parser.add_argument("--repeats", type=int, default=3) + args = parser.parse_args() + if not os.environ.get("EXP_ACTIVE"): + parser.error('Run through exp -m "purpose"') + if not 1 <= args.repeats <= 20: + parser.error("Repeats must be between 1 and 20") + run = Path(os.environ["EXP_RUN_DIR"]) + shutil.copyfile(__file__, run / "cold-stream-cohorts.py") + verified_checkpoint(args.model_directory) + service = json.loads(args.service_metadata.read_text()) + if service.get("framework") not in ("lightllm", "vllm", "sglang") or any( + not service.get(key) for key in ("version", "command", "hardware", "configuration") + ): + raise ValueError("Record actual service framework, version, command, hardware and configuration") + reference = json.loads(args.reference.read_text()) + lengths = [int(name.removeprefix("long_")) for name in reference.get("cases", {}) if name.startswith("long_")] + workloads = workloads_from_reference(reference, cases_for(args.model_directory, lengths)) + manifest = [dict(name=work.name, prompt_ids=work.prompt_ids, expected_ids=work.expected_ids) for work in workloads] + report = dict( + status="running", + comparison_eligible=False, + scope=__doc__, + reason="Shared hardware diagnostic; exact-output and explicit zero-cache requirements remain mandatory", + service=service, + reference=str(args.reference), + revision=REVISION, + manifest=manifest, + manifest_sha256=hashlib.sha256(json.dumps(manifest, sort_keys=True).encode()).hexdigest(), + cache_policy="native request cache bypass or independent random baseline cache salt on every request", + repeats=args.repeats, + cohorts=[], + ) + output = run / "cold-cohort.json" + + def save(): + output.write_text(json.dumps(report, indent=2)) + + def request(workload, response_path): + options = ( + dict(disable_prompt_cache=True) + if service["framework"] == "lightllm" + else dict(cache_salt=secrets.token_urlsafe(32)) + ) + return request_stream( + args.base_url, + service["framework"], + str(args.model_directory.resolve()), + list(workload.prompt_ids), + len(workload.expected_ids), + 600, + response_path, + **options, + ) + + try: + save() + for repeat in range(args.repeats): + order = [1, 2, 4, 8] + order = order[repeat % 4 :] + order[: repeat % 4] + sequence = workloads[repeat % len(workloads) :] + workloads[: repeat % len(workloads)] + for concurrency in order: + result = run_cohort( + sequence * 4, concurrency, request, run / f"repeat-{repeat}/concurrency-{concurrency}" + ) + result["repeat"] = repeat + transport_complete = all( + row.get("result") is not None + and ( + row["status"] == "passed" + or row.get("error") == "Output IDs differ from the qualified reference" + ) + for row in result["records"] + ) + cached = [ + row.get("result", {}).get("metrics", {}).get("cached_input_tokens") for row in result["records"] + ] + cold = all(type(count) is int and count == 0 for count in cached) + result["zero_cache_observed"] = cold + if not cold: + result.update( + status="failed", + metrics=None, + cache_gate_error="Every request must report exactly zero cached input tokens", + ) + result["unqualified_diagnostic_metrics"] = ( + summarize_cohort([row["result"] for row in result["records"]]) if transport_complete else None + ) + report["cohorts"].append(result) + save() + print( + json.dumps( + dict(repeat=repeat, concurrency=concurrency, status=result["status"], zero_cache_observed=cold) + ), + flush=True, + ) + if result["status"] != "passed": + raise ValueError("Cold request quality, timing or cache gate failed; qualified metrics unavailable") + report["status"] = "passed" + except Exception as error: + report.update(status="failed", error_type=type(error).__name__, error=str(error)) + raise + finally: + save() + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/qwen38/check_cpu_cache.py b/test/benchmark/qwen38/check_cpu_cache.py new file mode 100644 index 0000000000..e5aa24f4dd --- /dev/null +++ b/test/benchmark/qwen38/check_cpu_cache.py @@ -0,0 +1,146 @@ +import argparse +from concurrent.futures import ThreadPoolExecutor +import json +import os +from pathlib import Path + +from check_generation import REVISION, verified_checkpoint +from check_prefix_eviction import cache_hit, eviction_cases, validate_eviction +from check_scheduler import generate + + +def validate_cpu_restore(events, ranks, target_token, initial, restored, minimum_prefix=2048): + if type(ranks) is not int or ranks < 1 or not initial < restored[0] < restored[1]: + raise ValueError("Require valid rank count and ordered restore interval") + fields = {"qsa_raw_keys", "qsa_pooled_keys", "qsa_key_cos", "qsa_key_sin"} + proof = [] + for rank in range(ranks): + local = [event for event in events if event.get("rank") == rank and event.get("first_token") == target_token] + offloads = [ + event + for event in local + if event["event"] == "cpu_offload_queued" + and initial <= event["monotonic"] < restored[0] + and type(event.get("pages")) is int + and event["pages"] > 0 + ] + loads = [ + event + for event in local + if event["event"] == "cpu_load_queued" + and restored[0] <= event["monotonic"] <= restored[1] + and type(event.get("tokens")) is int + and event["tokens"] >= 128 + and type(event.get("kv_length")) is int + and event["kv_length"] >= minimum_prefix + and type(event.get("request_id")) is int + and event["request_id"] >= 0 + and event.get("pinned") is True + ] + complete = [] + for load in loads: + extra = [ + event + for event in local + if event["event"] == "cpu_extra_load_queued" + and event.get("request_id") == load.get("request_id") + and restored[0] <= event["monotonic"] <= load["monotonic"] + and fields <= set(event.get("slot_fields", [])) + and event.get("lexical_layers") + ] + if extra: + complete.append(load) + if not offloads or not complete: + raise ValueError(f"No complete target CPU offload/restore evidence on TP rank {rank}") + proof.append(dict(rank=rank, offloads=len(offloads), restores=len(complete))) + return proof + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base-url", required=True) + parser.add_argument("--model-directory", type=Path, required=True) + parser.add_argument("--service-metadata", type=Path, required=True) + parser.add_argument("--cache-events", type=Path, required=True) + parser.add_argument("--reference", type=Path, required=True) + parser.add_argument("--timeout", type=float, default=600) + args = parser.parse_args() + if not os.environ.get("EXP_ACTIVE"): + parser.error("Run through exp") + verified_checkpoint(args.model_directory) + service = json.loads(args.service_metadata.read_text()) + if ( + service.get("framework") != "lightllm" + or any(not service.get(key) for key in ("version", "command", "hardware", "configuration")) + or service["configuration"].get("cpu_cache") is not True + ): + raise ValueError("Require complete actual native CPU-cache service provenance") + cases = eviction_cases(args.model_directory) + reference = json.loads(args.reference.read_text()) + if ( + reference.get("status") not in ("recorded", "passed") + or reference["revision"] != REVISION + or reference["cases"] != cases + ): + raise ValueError("Require a passing reference with identical checkpoint and pressure inputs") + run = Path(os.environ["EXP_RUN_DIR"]) + report = dict(status="running", scope=__doc__, service=service, revision=REVISION, cases=cases, requests={}) + path = run / "cpu-cache.json" + + def request(label, case): + result = generate(args.base_url, case, [], args.timeout, run / f"{label}.json") + if result["token_ids"] != reference["requests"][label]["token_ids"]: + raise ValueError(f"CPU-cache output differs from qualified reference for {label}") + cache_hit(result) + report["requests"][label] = result + path.write_text(json.dumps(report, indent=2)) + return result + + try: + target = cases["target"] + initial = request("initial", target) + warm = request("warm", target) + if cache_hit(initial) != 0 or cache_hit(warm) < 2048: + raise ValueError("Require fresh target followed by actual warm-prefix reuse") + for name, case in cases.items(): + if name.startswith("small_pressure_"): + request(name, case) + request("after_small_pressure", target) + for name, case in cases.items(): + if name.startswith("pressure_"): + request(name, case) + after = request("after_pressure", target) + if cache_hit(after) < 2048: + raise ValueError("Evicted target did not restore a long cached prefix") + request("rewarm", target) + + with ThreadPoolExecutor(max_workers=4) as executor: + results = list( + executor.map( + lambda index: generate(args.base_url, target, [], args.timeout, run / f"concurrent_{index}.json"), + range(4), + ) + ) + for index, result in enumerate(results): + label = f"concurrent_{index}" + if result["token_ids"] != reference["requests"][label]["token_ids"]: + raise ValueError("Concurrent restored outputs differ from qualified reference") + report["requests"][label] = result + events = [json.loads(line) for line in args.cache_events.read_text().splitlines() if line] + ranks = service["configuration"]["tp"] + report["eviction_proof"] = validate_eviction( + events, ranks, target["inputs"][0], warm["ended"], after["started"] + ) + report["cpu_restore_proof"] = validate_cpu_restore( + events, ranks, target["inputs"][0], initial["started"], (after["started"], after["ended"]) + ) + report["status"] = "passed" + except Exception as error: + report.update(status="failed", error_type=type(error).__name__, error=str(error)) + raise + finally: + path.write_text(json.dumps(report, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/qwen38/check_generation.py b/test/benchmark/qwen38/check_generation.py new file mode 100644 index 0000000000..b0632992b0 --- /dev/null +++ b/test/benchmark/qwen38/check_generation.py @@ -0,0 +1,228 @@ +import argparse +import hashlib +import json +import os +from pathlib import Path +import secrets + +import requests +from tokenizers import Tokenizer + +REVISION = "de4b8e4d43b917e7706784d8bb445c9af86a3540" +TOKENIZER_SHA256 = "0997f410c57a1f4e53b09e4be8f4a172d90edd9564368fb0847030937229b9f3" + + +def verified_checkpoint(root): + acquisition = json.loads((root / "acquisition.json").read_text()) + verified = json.loads((root / "verified_files.json").read_text()) + if acquisition["status"] != "verified" or acquisition["revision"] != REVISION or verified["revision"] != REVISION: + raise ValueError("The complete pinned checkpoint must finish verification before this check") + index = json.loads((root / "model.safetensors.index.json").read_text()) + required = set(index["weight_map"].values()) | {"config.json", "model.safetensors.index.json"} + for name in required: + metadata = verified["files"][name] + if (root / name).stat().st_size != metadata["size"]: + raise ValueError(f"Checkpoint file size changed after verification: {name}") + return verified + + +def cases_for(root, long_lengths=(2049, 4063)): + if not long_lengths or any(type(length) is not int or length < 1 for length in long_lengths): + raise ValueError("Long prompt lengths must be positive integers") + if len(set(long_lengths)) != len(long_lengths): + raise ValueError("Long prompt lengths must be distinct") + path = root / "tokenizer.json" + if hashlib.sha256(path.read_bytes()).hexdigest() != TOKENIZER_SHA256: + raise ValueError("Tokenizer does not match the pinned checkpoint") + tokenizer = Tokenizer.from_file(str(path)) + prompts = { + "english": "The capital of France is", + "chinese": "中国的首都是", + "code_unicode": "def add(a, b):\n # 求和 🐍\n return", + } + cases = {name: tokenizer.encode(prompt, add_special_tokens=False).ids for name, prompt in prompts.items()} + seed = tokenizer.encode("A short factual note about language models and their context. ").ids + for length in long_lengths: + cases[f"long_{length}"] = (seed * (length // len(seed) + 1))[:length] + cases[f"repeat_long_{length}"] = list(cases[f"long_{length}"]) + return cases + + +def generate( + client, base_url, framework, model, tokens, output_length, timeout, response_path=None, *, cold_cache=False +): + if framework == "lightllm": + endpoint = "/generate" + payload = { + "inputs": tokens, + "parameters": { + "do_sample": False, + "ignore_eos": True, + "max_new_tokens": output_length, + "return_details": True, + }, + } + if cold_cache: + payload["parameters"]["disable_prompt_cache"] = True + else: + endpoint = "/v1/completions" + payload = { + "model": model, + "prompt": tokens, + "max_tokens": output_length, + "temperature": 0, + "seed": 3107, + "ignore_eos": True, + "return_token_ids": True, + } + if cold_cache: + payload["cache_salt"] = secrets.token_urlsafe(32) + response = client.post(base_url.rstrip("/") + endpoint, json=payload, timeout=timeout) + if response_path is not None: + response_path.write_text(response.text) + response.raise_for_status() + body = response.json() + if framework == "lightllm": + ids = [token["id"] for token in body["tokens"]] + prompt_length = body["prompt_tokens"] + else: + ids = body["choices"][0]["token_ids"] + prompt_length = body["usage"]["prompt_tokens"] + if not isinstance(ids, list) or len(ids) != output_length or not all(type(token) is int for token in ids): + received = len(ids) if isinstance(ids, list) else type(ids).__name__ + raise ValueError(f"Expected {output_length} integer token IDs, received {received}") + if prompt_length != len(tokens): + raise ValueError("Service prompt token count does not match the shared tokenized input") + if cold_cache: + cached = ( + [token.get("prompt_cache_len") for token in body["tokens"]] + if framework == "lightllm" + else [(body["usage"].get("prompt_tokens_details") or {}).get("cached_tokens")] + ) + if not cached or any(type(count) is not int or count != 0 for count in cached): + raise ValueError("Cold generation must explicitly report zero cached input tokens") + return {"token_ids": ids, "response": body} + + +def main(): + if not os.environ.get("EXP_ACTIVE"): + raise RuntimeError("Run through exp") + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--framework", choices=("lightllm", "vllm", "sglang"), required=True) + parser.add_argument("--base-url", required=True) + parser.add_argument("--model-directory", type=Path, required=True) + parser.add_argument("--service-metadata", type=Path, required=True) + parser.add_argument("--reference", type=Path) + parser.add_argument("--output-length", type=int, default=32) + parser.add_argument("--long-prompt-lengths", type=int, nargs="+", default=[2049, 4063]) + parser.add_argument("--timeout", type=float, default=600) + parser.add_argument( + "--cold-cache", action="store_true", help="Isolate each request and require zero reported cache hits" + ) + parser.add_argument( + "--record-observations", + action="store_true", + help="Retain baseline repeat failures as observations, without qualifying a reference", + ) + args = parser.parse_args() + if args.record_observations and (args.framework == "lightllm" or args.reference or not args.cold_cache): + parser.error("Baseline observations require cold requests and no reference; native qualification stays strict") + if args.output_length < 1: + parser.error("Output length must be positive") + if any(length < 1 for length in args.long_prompt_lengths) or len(set(args.long_prompt_lengths)) != len( + args.long_prompt_lengths + ): + parser.error("Long prompt lengths must be distinct positive integers") + run = Path(os.environ["EXP_RUN_DIR"]) + metadata = json.loads(args.service_metadata.read_text()) + for key in ("framework", "version", "command", "hardware", "configuration"): + if not metadata.get(key): + raise ValueError(f"Service metadata must record {key}") + if metadata["framework"] != args.framework: + raise ValueError("Service metadata names a different framework") + (run / "verified-files.json").write_text(json.dumps(verified_checkpoint(args.model_directory), indent=2)) + cases = cases_for(args.model_directory, args.long_prompt_lengths) + reference = json.loads(args.reference.read_text()) if args.reference else None + if reference and ( + reference["status"] not in ("passed", "recorded") + or reference["revision"] != REVISION + or reference["cases"] != cases + or reference["output_length"] != args.output_length + or reference.get("cache_policy", "service_default") != ("cold" if args.cold_cache else "service_default") + ): + raise ValueError("Reference must use the same revision, exact inputs and output length") + report = { + "status": "running", + "revision": REVISION, + "tokenizer_sha256": TOKENIZER_SHA256, + "service": metadata, + "base_url": args.base_url, + "cases": cases, + "output_length": args.output_length, + "long_prompt_lengths": args.long_prompt_lengths, + "reference": str(args.reference) if args.reference else None, + "cache_policy": "cold" if args.cold_cache else "service_default", + "results": {}, + } + output = run / "generation.json" + client = requests.Session() + client.trust_env = False + active_case = None + try: + for name, tokens in cases.items(): + active_case = name + result = generate( + client, + args.base_url, + args.framework, + str(args.model_directory.resolve()), + tokens, + args.output_length, + args.timeout, + response_path=run / f"response-{name}.txt", + cold_cache=args.cold_cache, + ) + report["results"][name] = result + result["reference_matches"] = ( + result["token_ids"] == reference["results"][name]["token_ids"] if reference else None + ) + output.write_text(json.dumps(report, indent=2)) + print( + json.dumps( + {"case": name, "token_ids": result["token_ids"], "reference_matches": result["reference_matches"]} + ), + flush=True, + ) + active_case = None + results = report["results"] + repeat_matches = all( + results[f"long_{length}"]["token_ids"] == results[f"repeat_long_{length}"]["token_ids"] + for length in args.long_prompt_lengths + ) + report["repeat_consistency"] = dict( + status="passed" if repeat_matches else "failed", + mismatches=[ + length + for length in args.long_prompt_lengths + if results[f"long_{length}"]["token_ids"] != results[f"repeat_long_{length}"]["token_ids"] + ], + ) + if args.record_observations: + report["status"] = "observations_recorded" + report[ + "scope" + ] = "Baseline token/count observations; repeat failures remain explicit, no exact-output qualification" + return + if not repeat_matches or any(item["reference_matches"] is False for item in results.values()): + raise ValueError("Greedy token IDs differ across repeated prompts or reference services") + report["status"] = "passed" if reference else "recorded" + except Exception as error: + report.update(status="failed", failed_case=active_case, error_type=type(error).__name__, error=str(error)) + raise + finally: + output.write_text(json.dumps(report, indent=2)) + client.close() + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/qwen38/check_math_quality.py b/test/benchmark/qwen38/check_math_quality.py new file mode 100644 index 0000000000..be9652f768 --- /dev/null +++ b/test/benchmark/qwen38/check_math_quality.py @@ -0,0 +1,129 @@ +import argparse +import hashlib +import json +import os +from pathlib import Path + +import requests + +from lightllm.server.tokenizer import get_tokenizer +from test.benchmark.qwen38.check_generation import REVISION, TOKENIZER_SHA256, verified_checkpoint +from test.benchmark.qwen38.math_quality import ( + DATASET_REVISION, + DATASET_SHA256, + load_examples, + prepare_cases, + score_answer, + select_examples, +) +from test.benchmark.service.completion_runner import CompletionClient + + +def verified_workload(manifest, root, dataset_path, tokenizer): + if ( + manifest["status"] != "prepared" + or manifest["model_revision"] != REVISION + or manifest["dataset"]["revision"] != DATASET_REVISION + or manifest["dataset"]["sha256"] != DATASET_SHA256 + or manifest["dataset"]["split"] != "test" + or manifest["decoding"] != dict(temperature=0, ignore_eos=False, enable_thinking=False) + or manifest["tokenizer_files"].get("tokenizer.json") != TOKENIZER_SHA256 + ): + raise ValueError("Workload must use pinned model, dataset and greedy non-thinking decoding") + for name in ("tokenizer.json", "tokenizer_config.json", "chat_template.jinja"): + if hashlib.sha256((root / name).read_bytes()).hexdigest() != manifest["tokenizer_files"].get(name): + raise ValueError("Tokenizer files differ from workload preparation") + examples = load_examples(dataset_path) + if len(examples) != manifest["dataset"]["records"]: + raise ValueError("Dataset record count differs") + cases = prepare_cases( + select_examples(examples, manifest["sample_count"], manifest["sample_seed"]), + tokenizer, + manifest["max_output_tokens"], + manifest["context_limit"], + ) + if cases != manifest["cases"]: + raise ValueError("Workload prompts or gold labels differ from pinned preparation") + return cases + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--framework", choices=("lightllm", "vllm", "sglang"), required=True) + parser.add_argument("--base-url", required=True) + parser.add_argument("--model-directory", type=Path, required=True) + parser.add_argument("--service-metadata", type=Path, required=True) + parser.add_argument("--manifest", type=Path, required=True) + parser.add_argument("--dataset", type=Path, required=True) + parser.add_argument("--timeout", type=float, default=600) + parser.add_argument("--concurrency", type=int, default=1) + args = parser.parse_args() + if not 1 <= args.concurrency <= 8: + parser.error("Concurrency must be between one and eight") + if not os.environ.get("EXP_ACTIVE"): + parser.error('Run through exp -m "purpose"') + run = Path(os.environ["EXP_RUN_DIR"]) + verified_checkpoint(args.model_directory) + metadata = json.loads(args.service_metadata.read_text()) + if metadata.get("framework") != args.framework or any( + not metadata.get(key) for key in ("version", "command", "hardware", "configuration") + ): + raise ValueError("Record actual service framework, version, command, hardware and configuration") + manifest = json.loads(args.manifest.read_text()) + tokenizer = get_tokenizer(str(args.model_directory)).tokenizer + cases = verified_workload(manifest, args.model_directory, args.dataset, tokenizer) + config = json.loads((args.model_directory / "generation_config.json").read_text()) + eos = config["eos_token_id"] + eos_ids = sorted(set(eos if isinstance(eos, list) else [eos]) | {tokenizer.eos_token_id}) + report = dict( + status="running", + scope="Pinned non-thinking scored math sample; no performance or broad accuracy claim", + service=metadata, + manifest=manifest, + manifest_sha256=hashlib.sha256(json.dumps(manifest, sort_keys=True).encode()).hexdigest(), + eos_ids=eos_ids, + concurrency=args.concurrency, + results=[], + ) + output = run / "math-quality.json" + client = CompletionClient( + args.framework, + str(args.model_directory.resolve()), + args.base_url, + args.timeout, + eos_ids, + run, + concurrency=args.concurrency, + session_factory=requests.Session, + ) + active_case = None + try: + responses = iter(client.results(cases)) + for case in cases: + active_case = case["dataset_index"] + result = next(responses) + text = tokenizer.decode(result["token_ids"], skip_special_tokens=True) + score = score_answer(text, case["expected_number"]) + score["correct"] = score["correct"] and not result["truncated"] + result.update(dataset_index=active_case, text=text, **score) + report["results"].append(result) + output.write_text(json.dumps(report, indent=2)) + print(json.dumps({key: result[key] for key in ("dataset_index", "correct", "truncated")}), flush=True) + report.update( + status="recorded", + scored=len(cases), + correct=sum(row["correct"] for row in report["results"]), + truncated=sum(row["truncated"] for row in report["results"]), + invalid_final_answer=sum(row["predicted"] is None for row in report["results"]), + ) + report["accuracy"] = report["correct"] / report["scored"] + except Exception as error: + report.update(status="failed", failed_case=active_case, error_type=type(error).__name__, error=str(error)) + raise + finally: + output.write_text(json.dumps(report, indent=2)) + client.close() + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/qwen38/check_multimodal.py b/test/benchmark/qwen38/check_multimodal.py new file mode 100644 index 0000000000..9a2bacd63d --- /dev/null +++ b/test/benchmark/qwen38/check_multimodal.py @@ -0,0 +1,291 @@ +import argparse +import base64 +from concurrent.futures import ThreadPoolExecutor, as_completed +import hashlib +from io import BytesIO +import json +import os +from pathlib import Path +import re +from types import SimpleNamespace + +from PIL import Image +import requests + +from check_generation import REVISION, verified_checkpoint +from check_scheduler import validate_generation + + +def image_fixtures(directory, pressure_count): + directory.mkdir(parents=True, exist_ok=True) + definitions = {"red": ((255, 0, 0), (256, 256)), "blue": ((0, 0, 255), (256, 256))} + definitions["rect_blue"] = ((0, 0, 255), (384, 128)) + definitions.update({f"pressure_{i}": ((255, i + 1, 0), (256, 256)) for i in range(pressure_count)}) + payloads, metadata = {}, {} + for name, (color, size) in definitions.items(): + buffer = BytesIO() + Image.new("RGB", size, color).save(buffer, format="PNG") + data = buffer.getvalue() + (directory / f"{name}.png").write_bytes(data) + payloads[name] = {"type": "base64", "data": base64.b64encode(data).decode("ascii")} + metadata[name] = dict(sha256=hashlib.sha256(data).hexdigest(), width=size[0], height=size[1]) + if len({item["sha256"] for item in metadata.values()}) != len(metadata): + raise ValueError("Cache pressure requires distinct image bytes") + return payloads, metadata + + +def prepare_cases(tokenizer, images): + definitions = { + "red": (["red"], ["red"]), + "blue": (["blue"], ["blue"]), + "rect_blue": (["rect_blue"], ["blue"]), + "pair": (["red", "blue"], ["red", "blue"]), + "reversed_pair": (["blue", "red"], ["blue", "red"]), + } + definitions.update({name: ([name], []) for name in images if name.startswith("pressure_")}) + cases = {} + for name, (names, expected) in definitions.items(): + question = ( + "Name the colors of these images in order. Reply with only the two English color words, " + "separated by a comma." + if len(names) > 1 + else "What color is the image? Reply with one English word." + ) + prompt = tokenizer.tokenizer.apply_chat_template( + [{"role": "user", "content": [{"type": "image"} for _ in names] + [{"type": "text", "text": question}]}], + tokenize=False, + add_generation_prompt=True, + enable_thinking=False, + ) + items = [] + for index, image in enumerate(names): + item = SimpleNamespace( + image_w=images[image]["width"], image_h=images[image]["height"], token_id=100000000 + index * 1024 + ) + item.token_num = tokenizer.get_image_token_length(item) + items.append(item) + expanded = tokenizer.encode(prompt, SimpleNamespace(images=items, audios=[])) + cases[name] = dict( + inputs=prompt, + images=names, + expected_words=expected, + prompt_tokens=len(expanded), + image_tokens=sum(item.token_num for item in items), + max_new_tokens=1 if not expected else 64, + ignore_eos=not expected, + ) + return cases + + +def validate_image_result(body, case, eos_ids, tokenizer): + ids = validate_generation(body, case["prompt_tokens"], case["max_new_tokens"], case["ignore_eos"], eos_ids) + usage = body["input_usage"] + if ( + usage["input_image_tokens"] != case["image_tokens"] + or usage["input_audio_tokens"] != 0 + or usage["input_text_tokens"] != case["prompt_tokens"] - case["image_tokens"] + ): + raise ValueError("Image/text accounting differs from actual expanded input IDs") + answer = tokenizer.decode([token for token in ids if token not in eos_ids], skip_special_tokens=True).strip() + if case["expected_words"] and re.findall(r"[a-z]+", answer.lower()) != case["expected_words"]: + raise ValueError(f"Image color/order response differs from the fixture: {answer!r}") + return dict(token_ids=ids, answer=answer, response=body) + + +def verify_reallocation(events, image_md5): + allocations = [ + item + for event in events + if event["event"] == "alloc" + for item in event["records"] + if item["md5"].startswith(image_md5 + "_") + ] + if not allocations: + raise ValueError("No actual image-cache allocation was observed") + first, last = allocations[0], allocations[-1] + evicted = {event["record"]["id"] for event in events if event["event"] == "evicted"} + embedded = {item["id"] for event in events if event["event"] == "embedded" for item in event["records"]} + if first["id"] == last["id"] or first["token_id"] == last["token_id"] or first["id"] not in evicted: + raise ValueError("The original image was not demonstrably evicted and assigned new virtual IDs") + if not {first["id"], last["id"]}.issubset(embedded): + raise ValueError("Both original and reallocated images require actual embedding-completion evidence") + return dict(original=first, reallocated=last, original_evicted=True, both_embedded=True) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base-url", required=True) + parser.add_argument("--model-directory", type=Path, required=True) + parser.add_argument("--service-metadata", type=Path, required=True) + parser.add_argument("--reference", type=Path) + parser.add_argument("--cache-events", type=Path) + parser.add_argument("--prepare-only", action="store_true") + parser.add_argument("--timeout", type=float, default=600) + args = parser.parse_args() + if not os.environ.get("EXP_ACTIVE"): + parser.error("Run through exp") + verified_checkpoint(args.model_directory) + service = json.loads(args.service_metadata.read_text()) + if service.get("framework") != "lightllm" or any( + not service.get(key) for key in ("version", "command", "hardware", "configuration") + ): + raise ValueError("Require actual native service provenance") + command = service["command"] + capacity = int(command[command.index("--cache_capacity") + 1]) + if not 3 <= capacity <= 16: + raise ValueError("This bounded two-image workload requires explicit cache capacity between 3 and 16") + run = Path(os.environ["EXP_RUN_DIR"]) + payloads, images = image_fixtures(run / "images", capacity + 1) + from lightllm.server.tokenizer import get_tokenizer + + tokenizer = get_tokenizer(str(args.model_directory)) + cases = prepare_cases(tokenizer, images) + source = { + name: hashlib.sha256((args.model_directory / name).read_bytes()).hexdigest() + for name in ( + "config.json", + "tokenizer.json", + "tokenizer_config.json", + "preprocessor_config.json", + "generation_config.json", + ) + } + generation_eos = json.loads((args.model_directory / "generation_config.json").read_text())["eos_token_id"] + eos_ids = sorted( + set(generation_eos if isinstance(generation_eos, list) else [generation_eos]) | {tokenizer.eos_token_id} + ) + reference = json.loads(args.reference.read_text()) if args.reference else None + if reference and ( + reference["status"] not in ("passed", "recorded") + or reference["revision"] != REVISION + or reference["source"] != source + or reference["cases"] != cases + or reference["images"] != images + or reference["eos_ids"] != eos_ids + ): + raise ValueError("Require a passing reference with identical images, prompt template, processor and EOS") + report = dict( + status="running", + scope=__doc__, + revision=REVISION, + service=service, + source=source, + reference=str(args.reference) if args.reference else None, + images=images, + cases=cases, + eos_ids=eos_ids, + cache_capacity=capacity, + serial={}, + repeats=[], + pressure=[], + waves=[], + ) + path = run / "multimodal.json" + active = None + + def save(): + path.write_text(json.dumps(report, indent=2)) + + if args.prepare_only: + report["status"] = "prepared" + save() + return + + def request(name, label, cold=False): + case = cases[name] + params = dict( + do_sample=False, + return_details=True, + print_eos_token=True, + disable_prompt_cache=cold, + max_new_tokens=case["max_new_tokens"], + ignore_eos=case["ignore_eos"], + ) + with requests.Session() as client: + client.trust_env = False + response = client.post( + args.base_url.rstrip("/") + "/generate", + json=dict( + inputs=case["inputs"], + parameters=params, + multimodal_params=dict(images=[payloads[n] for n in case["images"]]), + ), + timeout=args.timeout, + ) + (run / f"{label}.txt").write_text(response.text) + response.raise_for_status() + return validate_image_result(response.json(), case, eos_ids, tokenizer.tokenizer) + + def compare(name, result): + if result["token_ids"] != report["serial"][name]["token_ids"]: + raise ValueError(f"Repeated/concurrent image IDs differ: {name}") + + try: + save() + for name in ("red", "blue", "rect_blue", "pair", "reversed_pair"): + active = {"phase": "serial", "case": name} + case = cases[name] + with requests.Session() as client: + client.trust_env = False + response = client.post( + args.base_url.rstrip("/") + "/tokens", + json=dict( + text=case["inputs"], multimodal_params=dict(images=[payloads[n] for n in case["images"]]) + ), + timeout=args.timeout, + ) + (run / f"count-{name}.txt").write_text(response.text) + response.raise_for_status() + if response.json()["ntokens"] != case["prompt_tokens"]: + raise ValueError("Public image token count differs from actual expansion") + result = request(name, f"serial-{name}") + report["serial"][name] = result + save() + if reference and result["token_ids"] != reference["serial"][name]["token_ids"]: + raise ValueError(f"Ordinary/MTP image output IDs differ: {name}") + for cold in (False, True): + for name in ("red", "rect_blue", "pair"): + active = dict(phase="repeat", case=name, cold=cold) + result = request(name, f"repeat-{int(cold)}-{name}", cold) + report["repeats"].append(dict(case=name, cold=cold, result=result)) + save() + compare(name, result) + for index in range(capacity + 1): + name = f"pressure_{index}" + active = dict(phase="cache_pressure", case=name) + report["pressure"].append(request(name, name)) + save() + active = dict(phase="after_cache_pressure", case="red") + report["after_cache_pressure"] = request("red", "after-cache-pressure") + save() + compare("red", report["after_cache_pressure"]) + for concurrency in (2, 4): + active = dict(phase="concurrent", concurrency=concurrency) + wave = dict(concurrency=concurrency, results={}) + report["waves"].append(wave) + with ThreadPoolExecutor(max_workers=concurrency) as executor: + pending = { + executor.submit(request, name, f"wave-{concurrency}-{index}"): (index, name) + for index, name in enumerate(["red", "blue", "pair", "reversed_pair"] * 2) + } + for future in as_completed(pending): + index, name = pending[future] + result = future.result() + wave["results"][str(index)] = dict(case=name, result=result) + save() + compare(name, result) + if args.cache_events: + events = [json.loads(line) for line in args.cache_events.read_text().splitlines()] + digest = hashlib.md5((run / "images/red.png").read_bytes()).hexdigest() + report["cache_reallocation"] = verify_reallocation(events, digest) + report["cache_events"] = str(args.cache_events) + report["status"] = "passed" if reference else "recorded" + except Exception as error: + report.update(status="failed", failed_request=active, error_type=type(error).__name__, error=str(error)) + raise + finally: + save() + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/qwen38/check_prefix_eviction.py b/test/benchmark/qwen38/check_prefix_eviction.py new file mode 100644 index 0000000000..e604bfbaf4 --- /dev/null +++ b/test/benchmark/qwen38/check_prefix_eviction.py @@ -0,0 +1,174 @@ +import argparse +from concurrent.futures import ThreadPoolExecutor +import json +import os +from pathlib import Path + +from check_generation import REVISION, cases_for, verified_checkpoint +from check_scheduler import generate + + +def eviction_cases(root, pressure_count=12, prompt_length=2369): + if pressure_count < 2 or not 2049 <= prompt_length <= 4063: + raise ValueError("Eviction workload requires distinct long prefixes") + seed = cases_for(root)["long_4063"] + cases = { + "target" + if index == 0 + else f"pressure_{index - 1}": dict( + inputs=[1000 + index, *seed[: prompt_length - 1]], + max_new_tokens=32 if index == 0 else 8, + ignore_eos=True, + ) + for index in range(pressure_count + 1) + } + cases.update( + { + f"small_pressure_{index}": dict(inputs=[2000 + index, *seed[:64]], max_new_tokens=8, ignore_eos=True) + for index in range(20) + } + ) + return cases + + +def cache_hit(result): + value = result["response"]["tokens"][0]["prompt_cache_len"] + if type(value) is not int or not 0 <= value < result["response"]["prompt_tokens"]: + raise ValueError("Invalid actual prompt-cache hit length") + return value + + +def validate_eviction(events, ranks, target_token, started, ended): + if type(ranks) is not int or ranks < 1 or not started < ended: + raise ValueError("Require valid TP ranks and pressure interval") + during = [event for event in events if started <= event["monotonic"] <= ended] + proof = [] + for rank in range(ranks): + local = [event for event in during if event.get("rank") == rank] + nodes = [ + event + for event in local + if event["event"] == "node_evicted" + and event["first_token"] == target_token + and event.get("big_state") is not None + ] + small = [ + event + for event in local + if event["event"] == "small_state_evicted" + and event["first_token"] == target_token + and event.get("small_state") is not None + ] + freed = [event for event in local if event["event"] == "state_freed"] + + def observed_free(event, kind): + return any( + item["pool"] == event[f"{kind}_pool"] + and event[f"{kind}_state"] in item["indexes"] + and item["monotonic"] <= event["monotonic"] + for item in freed + ) + + if not any(observed_free(event, "big") for event in nodes): + raise ValueError(f"No observed target KV/big-state eviction on TP rank {rank}") + if not any(observed_free(event, "small") for event in small): + raise ValueError(f"No observed target small-state eviction on TP rank {rank}") + proof.append(dict(rank=rank, target_nodes=len(nodes), target_small_states=len(small))) + return proof + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base-url", required=True) + parser.add_argument("--model-directory", type=Path, required=True) + parser.add_argument("--service-metadata", type=Path, required=True) + parser.add_argument("--cache-events", type=Path, required=True) + parser.add_argument("--reference", type=Path) + parser.add_argument("--timeout", type=float, default=600) + args = parser.parse_args() + if not os.environ.get("EXP_ACTIVE"): + parser.error("Run through exp") + verified_checkpoint(args.model_directory) + service = json.loads(args.service_metadata.read_text()) + if service.get("framework") != "lightllm" or any( + not service.get(key) for key in ("version", "command", "hardware", "configuration") + ): + raise ValueError("Record complete actual native service provenance") + cases = eviction_cases(args.model_directory) + reference = json.loads(args.reference.read_text()) if args.reference else None + if reference and ( + reference.get("status") not in ("recorded", "passed") + or reference["revision"] != REVISION + or reference["cases"] != cases + ): + raise ValueError("Require a passing reference with identical checkpoint and eviction inputs") + run = Path(os.environ["EXP_RUN_DIR"]) + report = dict( + status="running", + scope=__doc__, + revision=REVISION, + service=service, + cases=cases, + reference=str(args.reference) if args.reference else None, + cache_events=str(args.cache_events), + requests={}, + ) + path = run / "prefix-eviction.json" + + def request(label, case): + result = generate(args.base_url, case, [], args.timeout, run / f"{label}.json") + if reference and result["token_ids"] != reference["requests"][label]["token_ids"]: + raise ValueError(f"Output differs from ordinary reference for {label}") + return result + + def save(label, result): + report["requests"][label] = result + path.write_text(json.dumps(report, indent=2)) + + try: + target = cases["target"] + save("initial", request("initial", target)) + save("warm", request("warm", target)) + initial = report["requests"]["initial"] + warm = report["requests"]["warm"] + if cache_hit(initial) != 0 or cache_hit(warm) < 2048 or initial["token_ids"] != warm["token_ids"]: + raise ValueError("Require a fresh target and exact warm-prefix reuse before pressure") + for name, case in cases.items(): + if name.startswith("small_pressure_"): + save(name, request(name, case)) + save("after_small_pressure", request("after_small_pressure", target)) + small = report["requests"]["after_small_pressure"] + if not 0 < cache_hit(small) < cache_hit(warm) or small["token_ids"] != initial["token_ids"]: + raise ValueError("Small-state pressure must fall back to a shorter big-state prefix with exact output") + for name, case in cases.items(): + if name.startswith("pressure_"): + save(name, request(name, case)) + save("after_pressure", request("after_pressure", target)) + after = report["requests"]["after_pressure"] + if cache_hit(after) >= cache_hit(warm) or after["token_ids"] != initial["token_ids"]: + raise ValueError("Target must lose its warm prefix and preserve exact output after pressure") + save("rewarm", request("rewarm", target)) + if cache_hit(report["requests"]["rewarm"]) < 2048: + raise ValueError("Reallocated target prefix did not become reusable") + with ThreadPoolExecutor(max_workers=4) as executor: + pending = { + f"concurrent_{index}": executor.submit(request, f"concurrent_{index}", target) for index in range(4) + } + for label, future in pending.items(): + save(label, future.result()) + if any(report["requests"][label]["token_ids"] != initial["token_ids"] for label in ["rewarm", *pending]): + raise ValueError("Reallocated target outputs differ under reuse/concurrency") + events = [json.loads(line) for line in args.cache_events.read_text().splitlines() if line] + report["eviction_proof"] = validate_eviction( + events, service["configuration"]["tp"], target["inputs"][0], warm["ended"], after["started"] + ) + report["status"] = "passed" if reference else "recorded" + except Exception as error: + report.update(status="failed", error_type=type(error).__name__, error=str(error)) + raise + finally: + path.write_text(json.dumps(report, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/qwen38/check_retrieval.py b/test/benchmark/qwen38/check_retrieval.py new file mode 100644 index 0000000000..b9ef80fa9d --- /dev/null +++ b/test/benchmark/qwen38/check_retrieval.py @@ -0,0 +1,114 @@ +import argparse +import hashlib +import json +import os +from pathlib import Path +import secrets +import shutil +import time + +import requests + +from lightllm.server.tokenizer import get_tokenizer +from test.benchmark.qwen38.check_generation import verified_checkpoint +from test.benchmark.qwen38.retrieval_quality import load_workload, score_records +from test.benchmark.service.completion import completion_request, completion_tokens + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--framework", choices=("lightllm", "vllm", "sglang"), required=True) + parser.add_argument("--base-url", required=True) + parser.add_argument("--model-directory", type=Path, required=True) + parser.add_argument("--service-metadata", type=Path, required=True) + parser.add_argument("--manifest", type=Path, required=True) + parser.add_argument("--timeout", type=float, default=1200) + args = parser.parse_args() + if not os.environ.get("EXP_ACTIVE"): + parser.error('Run through exp -m "purpose"') + run = Path(os.environ["EXP_RUN_DIR"]) + shutil.copyfile(__file__, run / "client.py") + verified_checkpoint(args.model_directory) + manifest = load_workload(args.manifest, args.model_directory) + service = json.loads(args.service_metadata.read_text()) + if service.get("framework") != args.framework or any( + not service.get(key) for key in ("version", "command", "hardware", "configuration") + ): + raise ValueError("Record actual service framework, version, command, hardware and configuration") + configuration = service["configuration"] + context = configuration.get("context_limit", configuration.get("effective_token_budget", 0)) + if type(context) is not int or context < manifest["context_limit"]: + raise ValueError("Service must declare the actual full context budget") + tokenizer = get_tokenizer(str(args.model_directory)).tokenizer + config = json.loads((args.model_directory / "generation_config.json").read_text()) + eos = config["eos_token_id"] + eos_ids = sorted(set(eos if isinstance(eos, list) else [eos]) | {tokenizer.eos_token_id}) + report = dict( + status="running", + scope=__doc__, + service=service, + manifest_sha256=hashlib.sha256(json.dumps(manifest, sort_keys=True).encode()).hexdigest(), + results={}, + ) + output = run / "long-retrieval.json" + client = requests.Session() + client.trust_env = False + active_case = None + try: + for case in manifest["cases"]: + active_case = case["name"] + endpoint, payload = completion_request( + args.framework, str(args.model_directory.resolve()), case["prompt_ids"], case["max_output_tokens"] + ) + if args.framework == "lightllm": + payload["parameters"]["disable_prompt_cache"] = True + else: + payload["cache_salt"] = secrets.token_urlsafe(32) + started = time.monotonic() + response = client.post(args.base_url.rstrip("/") + endpoint, json=payload, timeout=args.timeout) + ended = time.monotonic() + (run / f"{active_case}.json").write_text(response.text) + response.raise_for_status() + body = response.json() + result = completion_tokens(body, args.framework, case["prompt_ids"], case["max_output_tokens"], eos_ids) + if args.framework == "lightllm": + cached = body["tokens"][0]["prompt_cache_len"] + else: + cached = (body["usage"].get("prompt_tokens_details") or {}).get("cached_tokens") + if type(cached) is not int or cached != 0: + raise ValueError("Every context case must report exactly zero cached prompt tokens") + text = tokenizer.decode(result["token_ids"], skip_special_tokens=True) + score = score_records(text, case["expected_records"]) + score["correct"] = score["correct"] and not result["truncated"] + result.update( + text=text, + prompt_tokens=len(case["prompt_ids"]), + cached_input_tokens=cached, + observed_request_s=ended - started, + **score, + ) + report["results"][active_case] = result + output.write_text(json.dumps(report, indent=2)) + print( + json.dumps( + dict( + case=active_case, + **{key: result[key] for key in ("prompt_tokens", "correct", "records", "truncated")}, + ) + ), + flush=True, + ) + report.update(cases=len(manifest["cases"]), correct=sum(row["correct"] for row in report["results"].values())) + if report["cases"] != 4 or report["correct"] != report["cases"]: + raise ValueError("Some declared-context retrieval answers differ from the pinned known records") + report["status"] = "passed" + except Exception as error: + report.update(status="failed", failed_case=active_case, error_type=type(error).__name__, error=str(error)) + raise + finally: + output.write_text(json.dumps(report, indent=2)) + client.close() + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/qwen38/check_scheduler.py b/test/benchmark/qwen38/check_scheduler.py new file mode 100644 index 0000000000..37b17b129d --- /dev/null +++ b/test/benchmark/qwen38/check_scheduler.py @@ -0,0 +1,215 @@ +import argparse +from concurrent.futures import ThreadPoolExecutor, as_completed +import hashlib +import json +import os +from pathlib import Path +import threading +import time + +import requests +from transformers import AutoTokenizer + +from check_generation import REVISION, TOKENIZER_SHA256, cases_for, verified_checkpoint + + +def scheduler_cases(root): + raw = cases_for(root) + tokenizer = AutoTokenizer.from_pretrained(root, local_files_only=True) + result = { + name: {"inputs": raw[name], "max_new_tokens": length, "ignore_eos": True} + for name, length in ( + ("english", 32), + ("chinese", 31), + ("code_unicode", 17), + ("long_2049", 32), + ("long_4063", 32), + ) + } + result["eos_history"] = { + "inputs": [248044] + raw["english"] + [248046] + raw["chinese"], + "max_new_tokens": 19, + "ignore_eos": True, + } + for name, question in ( + ("chat_ok", "Reply with exactly OK."), + ("chat_number", "What is 2+2? Reply only with the number."), + ): + result[name] = { + "inputs": tokenizer.apply_chat_template( + [{"role": "user", "content": question}], + tokenize=True, + add_generation_prompt=True, + enable_thinking=False, + ), + "max_new_tokens": 96, + "ignore_eos": False, + } + generation_eos = json.loads((root / "generation_config.json").read_text())["eos_token_id"] + eos = set(generation_eos if isinstance(generation_eos, list) else [generation_eos]) + eos.add(tokenizer.eos_token_id) + return result, sorted(eos) + + +def generate(base_url, case, eos_ids, timeout, response_path, cold=False): + parameters = {key: value for key, value in case.items() if key != "inputs"} + parameters.update(do_sample=False, return_details=True, print_eos_token=True, disable_prompt_cache=cold) + started = time.monotonic() + with requests.Session() as client: + client.trust_env = False + response = client.post( + base_url.rstrip("/") + "/generate", + json={"inputs": case["inputs"], "parameters": parameters}, + timeout=timeout, + ) + ended = time.monotonic() + response_path.write_text(response.text) + response.raise_for_status() + body = response.json() + ids = validate_generation(body, len(case["inputs"]), case["max_new_tokens"], case["ignore_eos"], eos_ids) + return {"token_ids": ids, "response": body, "started": started, "ended": ended} + + +def validate_generation(body, prompt_tokens, limit, ignore_eos, eos_ids): + ids = [token["id"] for token in body["tokens"]] + if not ids or any(type(token) is not int or token < 0 for token in ids) or len(ids) > limit: + raise ValueError("Invalid output token IDs or output length") + if body["prompt_tokens"] != prompt_tokens or body["count_output_tokens"] != len(ids): + raise ValueError("Service prompt/output counts do not match exact token IDs") + if ignore_eos: + if len(ids) != limit or body["finish_reason"] != "length": + raise ValueError("Fixed-length request did not reach its requested length") + elif ids[-1] not in eos_ids or any(token in eos_ids for token in ids[:-1]) or body["finish_reason"] != "stop": + raise ValueError("EOS request must stop at its first declared EOS token") + return ids + + +def monitor(base_url, stop, summary, path): + with requests.Session() as client: + client.trust_env = False + while not stop.is_set(): + try: + response = client.get(base_url.rstrip("/") + "/metrics", timeout=5) + response.raise_for_status() + path.write_text(response.text) + for line in response.text.splitlines(): + if line.startswith(("lightllm_batch_current_size", "lightllm_num_running_reqs")): + summary["maximum_router_batch"] = max(summary["maximum_router_batch"], float(line.split()[-1])) + summary["samples"] += 1 + except Exception as error: + summary["last_error"] = str(error) + stop.wait(0.02) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base-url", required=True) + parser.add_argument("--model-directory", type=Path, required=True) + parser.add_argument("--service-metadata", type=Path, required=True) + parser.add_argument("--reference", type=Path) + parser.add_argument("--concurrency", type=int, nargs="+", default=[2, 4, 8]) + parser.add_argument("--timeout", type=float, default=600) + args = parser.parse_args() + if not os.environ.get("EXP_ACTIVE"): + parser.error('Run through exp -m "purpose"') + if not args.concurrency or min(args.concurrency) < 2 or max(args.concurrency) > 8: + parser.error("Concurrent waves must use between two and eight requests") + run = Path(os.environ["EXP_RUN_DIR"]) + verified_checkpoint(args.model_directory) + metadata = json.loads(args.service_metadata.read_text()) + if metadata.get("framework") != "lightllm" or any( + not metadata.get(key) for key in ("version", "command", "hardware", "configuration") + ): + raise ValueError("Record the actual native service version, command, hardware and configuration") + cases, eos_ids = scheduler_cases(args.model_directory) + source = { + name: hashlib.sha256((args.model_directory / name).read_bytes()).hexdigest() + for name in ("tokenizer.json", "tokenizer_config.json", "generation_config.json") + } + reference = json.loads(args.reference.read_text()) if args.reference else None + if reference and ( + reference.get("status") not in ("passed", "recorded") + or reference["revision"] != REVISION + or reference["cases"] != cases + or reference["source"] != source + or reference["eos_ids"] != eos_ids + ): + raise ValueError("Reference must pass with identical checkpoint, tokenizer, inputs and EOS contract") + report = dict( + status="running", + scope="full trained scheduler/EOS/cache correctness; no performance claim", + revision=REVISION, + tokenizer_sha256=TOKENIZER_SHA256, + source=source, + service=metadata, + cases=cases, + eos_ids=eos_ids, + reference=str(args.reference) if args.reference else None, + serial={}, + waves=[], + mismatches=[], + ) + output = run / "scheduler.json" + stop = threading.Event() + metrics = {"maximum_router_batch": 0, "samples": 0} + thread = threading.Thread(target=monitor, args=(args.base_url, stop, metrics, run / "metrics.txt"), daemon=True) + thread.start() + active = None + try: + for name, case in cases.items(): + active = {"wave": "serial", "case": name} + result = generate(args.base_url, case, eos_ids, args.timeout, run / f"serial-{name}.txt") + report["serial"][name] = result + if reference and result["token_ids"] != reference["serial"][name]["token_ids"]: + report["mismatches"].append({"wave": "serial_reference", "case": name}) + output.write_text(json.dumps(report, indent=2)) + names = list(cases) + plans = [(count, False) for count in args.concurrency] + [(8, True), (8, False)] + for index, (count, cold) in enumerate(plans): + order = names[index:] + names[:index] + if index == len(plans) - 1: + order = list(reversed(order)) + wave = {"concurrency": count, "disable_prompt_cache": cold, "results": {}} + report["waves"].append(wave) + with ThreadPoolExecutor(max_workers=count) as executor: + pending = { + executor.submit( + generate, + args.base_url, + cases[name], + eos_ids, + args.timeout, + run / f"wave-{index}-{name}.txt", + cold, + ): name + for name in order + } + for future in as_completed(pending): + name = pending[future] + active = {"wave": index, "case": name} + result = future.result() + wave["results"][name] = result + if result["token_ids"] != report["serial"][name]["token_ids"]: + report["mismatches"].append({"wave": index, "case": name}) + output.write_text(json.dumps(report, indent=2)) + print( + json.dumps({"wave": index, "concurrency": count, "cold": cold, "mismatches": report["mismatches"]}), + flush=True, + ) + if report["mismatches"]: + raise ValueError("Serial/concurrent or ordinary/speculative token IDs differ") + if metrics["maximum_router_batch"] < 2: + raise ValueError("No actual multi-request router batch observed; packed serving remains unverified") + report["status"] = "passed" if reference else "recorded" + except Exception as error: + report.update(status="failed", failed_request=active, error_type=type(error).__name__, error=str(error)) + raise + finally: + stop.set() + thread.join(timeout=6) + report["router_metrics"] = metrics + output.write_text(json.dumps(report, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/qwen38/check_streaming.py b/test/benchmark/qwen38/check_streaming.py new file mode 100644 index 0000000000..196d75edfb --- /dev/null +++ b/test/benchmark/qwen38/check_streaming.py @@ -0,0 +1,76 @@ +import argparse +import hashlib +import json +import os +from pathlib import Path + +from check_generation import REVISION, verified_checkpoint +from test.benchmark.service.token_stream import request_stream + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--framework", choices=("lightllm", "vllm", "sglang"), required=True) + parser.add_argument("--base-url", required=True) + parser.add_argument("--model-directory", type=Path, required=True) + parser.add_argument("--service-metadata", type=Path, required=True) + parser.add_argument("--reference", type=Path, required=True) + parser.add_argument("--cases", nargs="+", default=["english", "chinese", "long_2049"]) + parser.add_argument("--timeout", type=float, default=600) + args = parser.parse_args() + if not os.environ.get("EXP_ACTIVE"): + parser.error("Run through exp") + verified_checkpoint(args.model_directory) + reference = json.loads(args.reference.read_text()) + service = json.loads(args.service_metadata.read_text()) + if service.get("framework") != args.framework or any( + not service.get(key) for key in ("version", "command", "hardware", "configuration") + ): + raise ValueError("Record actual service provenance") + if reference["status"] not in ("passed", "recorded") or reference["revision"] != REVISION: + raise ValueError("Require a passing pinned scheduler reference") + for name, expected in reference["source"].items(): + if hashlib.sha256((args.model_directory / name).read_bytes()).hexdigest() != expected: + raise ValueError(f"Checkpoint/tokenizer changed after reference: {name}") + run = Path(os.environ["EXP_RUN_DIR"]) + report = dict( + status="running", + scope="real streaming protocol/greedy parity; no performance comparison", + revision=REVISION, + service=service, + reference=str(args.reference), + results={}, + ) + path = run / "streaming.json" + active = None + try: + for name in args.cases: + active = name + case = reference["cases"][name] + if not case["ignore_eos"]: + raise ValueError("Streaming fixed-length gate requires EOS ignored") + result = request_stream( + args.base_url, + args.framework, + str(args.model_directory.resolve()), + case["inputs"], + case["max_new_tokens"], + args.timeout, + run / f"response-{name}.sse", + ) + result["reference_matches"] = result["token_ids"] == reference["serial"][name]["token_ids"] + report["results"][name] = result + path.write_text(json.dumps(report, indent=2)) + print(json.dumps({"case": name, **result}), flush=True) + if not all(result["reference_matches"] for result in report["results"].values()): + raise ValueError("Streamed token IDs differ from the exact scheduler reference") + report["status"] = "passed" + except Exception as error: + report.update(status="failed", failed_case=active, error_type=type(error).__name__, error=str(error)) + raise + finally: + path.write_text(json.dumps(report, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/qwen38/generation_observations.py b/test/benchmark/qwen38/generation_observations.py new file mode 100644 index 0000000000..cebf1130b3 --- /dev/null +++ b/test/benchmark/qwen38/generation_observations.py @@ -0,0 +1,48 @@ +from test.benchmark.qwen38.check_generation import REVISION +from test.benchmark.service.stream_cohort import StreamingWorkload + + +def workloads_from_observations(report, expected_cases): + if ( + report.get("status") != "observations_recorded" + or report.get("service", {}).get("framework") not in ("vllm", "sglang") + or report.get("revision") != REVISION + or report.get("cases") != expected_cases + or report.get("reference") is not None + or report.get("cache_policy") != "cold" + ): + raise ValueError("Require explicit cold baseline observations with pinned inputs") + results, length = report.get("results", {}), report.get("output_length") + if type(length) is not int or length < 1 or set(results) != set(expected_cases): + raise ValueError("Require complete baseline observations and a positive output length") + workloads, differences = [], [] + for name, prompt in expected_cases.items(): + result = results[name] + ids, body = result["token_ids"], result["response"] + usage, choices = body["usage"], body["choices"] + cached = (usage.get("prompt_tokens_details") or {}).get("cached_tokens") + if ( + len(choices) != 1 + or choices[0]["token_ids"] != ids + or usage["prompt_tokens"] != len(prompt) + or usage["completion_tokens"] != length + or len(ids) != length + or type(cached) is not int + or cached != 0 + ): + raise ValueError("Baseline observations must preserve raw IDs, counts and explicit zero cache") + work = StreamingWorkload(name, prompt, ids) + if name.startswith("repeat_"): + original = name.removeprefix("repeat_") + if original not in expected_cases or prompt != expected_cases[original]: + raise ValueError("Repeat observations require identical pinned inputs") + if ids != results[original]["token_ids"]: + differences.append(int(original.removeprefix("long_"))) + else: + workloads.append(work) + quality = dict(status="failed" if differences else "passed", mismatches=differences) + if quality != report.get("repeat_consistency"): + raise ValueError("Baseline repeat-consistency summary differs from actual observations") + if not workloads: + raise ValueError("Require a non-repeated workload") + return workloads, quality diff --git a/test/benchmark/qwen38/math_quality.py b/test/benchmark/qwen38/math_quality.py new file mode 100644 index 0000000000..fd7bbc8c8d --- /dev/null +++ b/test/benchmark/qwen38/math_quality.py @@ -0,0 +1,127 @@ +from dataclasses import dataclass +from decimal import Decimal +import hashlib +import json +import random +import re + +DATASET_REVISION = "3101c7d5072418e28b9008a6636bde82a006892c" +DATASET_SHA256 = "3730d312f6e3440559ace48831e51066acaca737f6eabec99bccb9e4b3c39d14" +NUMBER = re.compile(r"[+-]?(?:\d{1,3}(?:,\d{3})+|\d+)(?:\.\d+)?") +INSTRUCTION = ( + "Solve the following math problem. Show your reasoning briefly, then end with a final line " + "exactly in the format #### . Use a plain decimal number.\n\nProblem:\n" +) + + +def normalize_number(value): + if not isinstance(value, str) or not NUMBER.fullmatch(value.strip()): + return None + number = Decimal(value.strip().replace(",", "")) + if number == 0: + return "0" + text = format(number, "f") + return text.rstrip("0").rstrip(".") if "." in text else text + + +def extract_final_number(completion): + if not isinstance(completion, str): + return None + lines = completion.rstrip().splitlines() + if not lines or sum(line.strip().startswith("####") for line in lines) != 1: + return None + match = re.fullmatch(r"####[ \t]+(.+)", lines[-1].strip()) + return normalize_number(match[1]) if match else None + + +def score_answer(completion, expected): + gold = normalize_number(expected) + if gold is None: + raise ValueError("Expected answer must be a finite plain decimal number") + predicted = extract_final_number(completion) + return dict(expected=gold, predicted=predicted, correct=predicted is not None and predicted == gold) + + +def compare_outputs(reference, actual): + digest = reference.get("manifest_sha256") + if ( + reference.get("status") != "recorded" + or actual.get("status") != "recorded" + or not isinstance(digest, str) + or not re.fullmatch(r"[0-9a-f]{64}", digest) + or actual.get("manifest_sha256") != digest + ): + raise ValueError("Require complete scored runs with identical workload manifests") + + def rows(report): + values = report["results"] + if type(report.get("scored")) is not int or report["scored"] < 1 or len(values) != report["scored"]: + raise ValueError("Scored count differs from completed cases") + indices = [row["dataset_index"] for row in values] + if any(type(index) is not int or index < 0 for index in indices) or len(set(indices)) != len(indices): + raise ValueError("Require distinct actual dataset indices") + return {row["dataset_index"]: row for row in values} + + expected, observed = rows(reference), rows(actual) + if set(expected) != set(observed): + raise ValueError("Scored runs contain different selected cases") + fields = ("token_ids", "finish_reason", "truncated", "correct", "predicted") + mismatches = [ + index for index in sorted(expected) if any(expected[index][key] != observed[index][key] for key in fields) + ] + return dict(status="passed" if not mismatches else "failed", cases=len(expected), mismatches=mismatches) + + +@dataclass(frozen=True) +class MathExample: + index: int + question: str + expected: str + + +def load_examples(path): + payload = path.read_bytes() + if hashlib.sha256(payload).hexdigest() != DATASET_SHA256: + raise ValueError("GSM8K test data differs from the pinned source") + examples = [] + for index, line in enumerate(payload.splitlines()): + row = json.loads(line) + question = row.get("question") + answer = extract_final_number(row.get("answer")) + if not isinstance(question, str) or not question.strip() or answer is None: + raise ValueError("Dataset requires a question and unambiguous final numeric answer") + examples.append(MathExample(index, question, answer)) + if not examples or len({example.question for example in examples}) != len(examples): + raise ValueError("Dataset requires distinct questions") + return tuple(examples) + + +def select_examples(examples, limit, seed): + if type(limit) is not int or not 1 <= limit <= len(examples) or type(seed) is not int: + raise ValueError("Require a valid sample count and integer seed") + return tuple(examples[index] for index in random.Random(seed).sample(range(len(examples)), limit)) + + +def prepare_cases(examples, tokenizer, max_output_tokens=512, context_limit=4096): + if any(type(value) is not int or value < 1 for value in (max_output_tokens, context_limit)): + raise ValueError("Output and context limits must be positive integers") + cases = [] + for example in examples: + messages = [{"role": "user", "content": INSTRUCTION + example.question}] + ids = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, enable_thinking=False) + if not isinstance(ids, list) or not ids or any(type(token) is not int or token < 0 for token in ids): + raise ValueError("Tokenizer must return actual nonnegative integer prompt IDs") + if len(ids) + max_output_tokens > context_limit: + raise ValueError("Math example exceeds the declared serving context") + cases.append( + dict( + dataset_index=example.index, + question_sha256=hashlib.sha256(example.question.encode()).hexdigest(), + prompt_ids=ids, + expected_number=example.expected, + max_output_tokens=max_output_tokens, + ) + ) + if not cases or len({case["dataset_index"] for case in cases}) != len(cases): + raise ValueError("Require distinct selected examples") + return cases diff --git a/test/benchmark/qwen38/measure_aiperf.py b/test/benchmark/qwen38/measure_aiperf.py new file mode 100644 index 0000000000..fc61f16102 --- /dev/null +++ b/test/benchmark/qwen38/measure_aiperf.py @@ -0,0 +1,310 @@ +import argparse +import hashlib +import json +import math +import os +from pathlib import Path +import shutil +import statistics +import subprocess + +import requests + + +def tokenizer_cache(model, directory, revision): + snapshot = directory / "models--Qwen--Qwen3.8-Flash-Next" / "snapshots" / revision + snapshot.mkdir(parents=True, exist_ok=True) + hashes = {} + for name in ( + "config.json", + "tokenizer.json", + "tokenizer_config.json", + "special_tokens_map.json", + "added_tokens.json", + ): + original = model.resolve() / name + if not original.exists(): + if name in ("config.json", "tokenizer.json", "tokenizer_config.json"): + raise ValueError(f"Missing pinned tokenizer file: {name}") + continue + destination = snapshot / name + if not destination.exists(): + destination.symlink_to(original) + if destination.read_bytes() != original.read_bytes(): + raise ValueError("Cached tokenizer differs from the pinned checkpoint") + hashes[name] = hashlib.sha256(original.read_bytes()).hexdigest() + return hashes + + +def profiling_records(directory): + records = [json.loads(line) for line in (directory / "profile_export_raw.jsonl").read_text().splitlines()] + records = [row for row in records if row["metadata"]["benchmark_phase"] == "profiling"] + + def index(row): + identity = row["metadata"]["conversation_id"] + if not identity.startswith("session_") or not identity.removeprefix("session_").isdigit(): + raise ValueError("Unexpected source conversation identity") + return int(identity.removeprefix("session_")) + + records.sort(key=index) + if [index(row) for row in records] != list(range(len(records))): + raise ValueError("Missing or duplicated source conversations") + return records + + +def audit_records(directory, trace): + summary = json.loads((directory / "profile_export_aiperf.json").read_text()) + records = profiling_records(directory) + if len(records) != len(trace): + raise ValueError("Incomplete replay") + if sorted(row["metadata"]["session_num"] for row in records) != list(range(len(trace))): + raise ValueError("Missing, duplicated or reordered dataset sessions") + requests_for_comparison, starts, cache_counts, empty_content = [], [], [], [] + for index, (record, expected) in enumerate(zip(records, trace)): + payload, metadata = record["payload"], record["metadata"] + if record["status"] != 200 or metadata.get("was_cancelled"): + raise ValueError("Failed or cancelled request") + if payload["max_tokens"] != expected["output_length"] or not payload.get("ignore_eos"): + raise ValueError("Replay changed the specified output work") + usage, done, finished, has_text = None, False, False, False + for response in record["responses"]: + for packet in response.get("packets", []): + if packet.get("name") != "data": + continue + if packet["value"] == "[DONE]": + done = True + continue + data = json.loads(packet["value"]) + usage = data.get("usage") or usage + finished |= any(choice.get("finish_reason") == "length" for choice in data.get("choices", [])) + has_text |= any(bool(choice.get("text")) for choice in data.get("choices", [])) + if not done or not finished or not usage: + raise ValueError("Missing stream terminator, length finish or actual usage") + if type(usage.get("completion_tokens")) is not int or usage["completion_tokens"] != expected["output_length"]: + raise ValueError("Actual completion count differs from the trace") + if ( + type(usage.get("prompt_tokens")) is not int + or not 0 < usage["prompt_tokens"] + usage["completion_tokens"] <= 4096 + ): + raise ValueError("Actual request exceeds the shared context limit") + cache_counts.append((usage.get("prompt_tokens_details") or {}).get("cached_tokens")) + if not has_text: + empty_content.append(index) + requests_for_comparison.append( + dict( + prompt=payload["prompt"], + max_tokens=payload["max_tokens"], + temperature=payload["temperature"], + ignore_eos=payload["ignore_eos"], + ) + ) + starts.append(metadata["request_start_ns"]) + + errors = summary.get("error_summary", []) + for error in errors: + detail = error["error_details"] + if ( + detail["type"] != "InvalidInferenceResultError" + or "No responses with actual content" not in detail["message"] + ): + raise ValueError("AIPerf reported a transport, protocol or other request error") + if sum(error["count"] for error in errors) != len(empty_content): + raise ValueError("AIPerf errors do not match the complete empty-text responses") + if summary["request_count"]["avg"] != len(trace) - len(empty_content): + raise ValueError("AIPerf successful request count mismatch") + successful_tokens = sum(row["output_length"] for index, row in enumerate(trace) if index not in empty_content) + if summary["output_token_count"]["sum"] != successful_tokens: + raise ValueError("AIPerf successful token count mismatch") + + lag = [(start - starts[0]) / 1e6 - (row["timestamp"] - trace[0]["timestamp"]) for start, row in zip(starts, trace)] + required = ( + "request_throughput", + "output_token_throughput", + "time_to_first_token", + "inter_token_latency", + "request_latency", + ) + for name in required: + if not math.isfinite(summary[name]["avg"]): + raise ValueError(f"Invalid AIPerf metric: {name}") + return dict( + requests=len(records), + output_tokens=sum(row["output_length"] for row in trace), + content_errors=len(empty_content), + error_rate=len(empty_content) / len(trace), + empty_content_sessions=empty_content, + successful_output_tokens=successful_tokens, + metric_scope="AIPerf successful content-bearing responses; empty-text responses remain errors", + request_manifest_sha256=hashlib.sha256( + json.dumps(requests_for_comparison, sort_keys=True).encode() + ).hexdigest(), + cached_tokens=cache_counts, + relative_schedule_lag_ms=dict(median=statistics.median(lag), max=max(lag), min=min(lag)), + metrics={name: summary[name] for name in required}, + ) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--service-metadata", type=Path, required=True) + parser.add_argument("--trace-directory", type=Path, required=True) + parser.add_argument("--model-directory", type=Path, required=True) + parser.add_argument("--aiperf", type=Path, required=True) + parser.add_argument("--speed", type=float, default=1.0) + parser.add_argument( + "--fresh-owned-service", + action="store_true", + help="Native replay in a newly owned service with no prior workload requests", + ) + args = parser.parse_args() + if not os.environ.get("EXP_ACTIVE"): + parser.error('Run through exp -m "purpose"') + if not math.isfinite(args.speed) or args.speed <= 0: + parser.error("Speed must be finite and positive") + service = json.loads(args.service_metadata.read_text()) + for key in ( + "framework", + "version", + "software", + "command", + "hardware", + "configuration", + "base_url", + "model_revision", + ): + if not service.get(key): + raise ValueError(f"Missing actual service provenance: {key}") + if service["status"] not in ("ready", "measuring_aiperf"): + raise ValueError("AIPerf requires an exclusively owned, ready service") + manifest = json.loads((args.trace_directory / "manifest.json").read_text()) + raw = (args.trace_directory / "replay.jsonl").read_bytes() + if hashlib.sha256(raw).hexdigest() != manifest["trace_sha256"]: + raise ValueError("Trace checksum mismatch") + trace = [json.loads(line) for line in raw.splitlines()] + if len(trace) != manifest["requests"] or manifest["context_length"] != 4096: + raise ValueError("Trace manifest mismatch") + run = Path(os.environ["EXP_RUN_DIR"]) + shutil.copyfile(__file__, run / Path(__file__).name) + version = subprocess.check_output([str(args.aiperf), "--version"], text=True).strip() + if version != "0.9.0": + raise ValueError("Comparison is pinned to AIPerf 0.9.0") + for row in trace: + row["timestamp"] /= args.speed + trace_path = run / "trace.jsonl" + trace_path.write_text("".join(json.dumps(row) + "\n" for row in trace)) + report = dict( + status="running", + service=service, + aiperf_version=version, + manifest=manifest, + speed=args.speed, + concurrency_ceiling=8, + cache_policy="flush before each replay; prefix reuse within trace enabled", + precision_scope="Performance only; reuse independent model task-quality results", + comparison_eligible=False, + ) + + def save(): + (run / "aiperf-replay.json").write_text(json.dumps(report, indent=2) + "\n") + + try: + save() + cache = args.trace_directory.resolve() / "tokenizer-cache" + report["tokenizer_sha256"] = tokenizer_cache(args.model_directory, cache, service["model_revision"]) + environment = dict( + os.environ, + HF_HUB_CACHE=str(cache), + HUGGINGFACE_HUB_CACHE=str(cache), + HF_HUB_OFFLINE="1", + TRANSFORMERS_OFFLINE="1", + ) + framework = service["framework"] + if framework not in ("lightllm", "vllm", "sglang"): + raise ValueError("Unsupported framework") + reset = "/reset_prefix_cache" if framework == "vllm" else "/flush_cache" + if args.fresh_owned_service: + import psutil + + if framework != "lightllm" or not service["configuration"].get("fresh_aiperf_replay"): + raise ValueError("Fresh-instance cache isolation is restricted to the native replay owner") + if not service.get("fresh_service_no_quality_requests") or service.get("checks"): + raise ValueError("Fresh replay requires a service with no preceding workload requests") + owner = psutil.Process(service["launcher_pid"]) + allowed = {"qwen38-h100-native-fresh-service.py", "qwen38-h100-aiperf-profile.py"} + if owner.create_time() != service["launcher_created"] or not any( + Path(value).name in allowed for value in owner.cmdline() + ): + raise ValueError("Fresh replay owner identity mismatch") + report[ + "cache_policy" + ] = "New owned native instance per replay; no prior workload requests; prefix reuse within trace enabled" + report["cache_reset"] = dict( + method="fresh_owned_service", container=service["container"]["name"], owner_pid=owner.pid + ) + save() + else: + with requests.Session() as session: + session.trust_env = False + response = session.post(service["base_url"] + reset, timeout=120) + report["cache_reset"] = dict(status=response.status_code, body=response.text) + save() + response.raise_for_status() + if response.text.strip().lower() == "false": + raise ValueError("Server refused cache reset") + if framework == "vllm" and response.json().get("success") is not True: + raise ValueError("vLLM did not confirm prefix-cache reset") + command = [ + str(args.aiperf), + "profile", + "--model", + str(args.model_directory.resolve()), + "--tokenizer", + "Qwen/Qwen3.8-Flash-Next", + "--tokenizer-revision", + service["model_revision"], + "--url", + service["base_url"], + "--endpoint-type", + "completions", + "--streaming", + "--input-file", + str(trace_path), + "--custom-dataset-type", + "mooncake_trace", + "--fixed-schedule", + "--fixed-schedule-auto-offset", + "--concurrency", + "8", + "--request-count", + str(len(trace)), + "--random-seed", + "42", + "--use-server-token-count", + "--export-level", + "raw", + "--artifact-dir", + str(run / "aiperf"), + "--ui", + "none", + "--workers-max", + "4", + "--no-auto-plot", + "--request-timeout-seconds", + "600", + "--extra-inputs", + json.dumps(dict(temperature=0, ignore_eos=True)), + ] + report["command"] = command + save() + subprocess.run(command, check=True, env=environment) + report["audit"] = audit_records(run / "aiperf", trace) + report["status"] = "aiperf_replay_recorded" + except BaseException as error: + report.update(status="failed", error_type=type(error).__name__, error=str(error)) + raise + finally: + save() + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/qwen38/measure_fixed_work.py b/test/benchmark/qwen38/measure_fixed_work.py new file mode 100644 index 0000000000..c072a790e2 --- /dev/null +++ b/test/benchmark/qwen38/measure_fixed_work.py @@ -0,0 +1,152 @@ +import argparse +import hashlib +import json +import os +from pathlib import Path +import secrets +import shutil + +from test.benchmark.qwen38.check_cold_streams import workloads_from_reference +from test.benchmark.qwen38.check_generation import cases_for, verified_checkpoint +from test.benchmark.qwen38.generation_observations import workloads_from_observations +from test.benchmark.service.stream_cohort import run_cohort, summarize_cohort +from test.benchmark.service.token_stream import request_stream + + +def summarize_fixed_work(cohort, workloads): + expected = {work.name: work for work in workloads} + results = [] + for record in cohort["records"]: + if record["status"] != "passed" and record.get("error") != "Output IDs differ from the qualified reference": + raise ValueError("Incomplete or invalid streaming request") + work = expected[record["workload"]] + result = record["result"] + metrics = result["metrics"] + if metrics["input_tokens"] != len(work.prompt_ids) or len(result["token_ids"]) != len(work.expected_ids): + raise ValueError("Actual input/output lengths differ from the workload") + cached = metrics.get("cached_input_tokens") + if type(cached) is not int or cached != 0: + raise ValueError("Every measured request must explicitly report zero cached input tokens") + results.append(result) + if len(results) != len(workloads): + raise ValueError("Missing workload requests") + return summarize_cohort(results) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--service-metadata", type=Path, required=True) + parser.add_argument("--reference", type=Path, required=True) + parser.add_argument("--model-directory", type=Path, required=True) + parser.add_argument("--repeats", type=int, default=3) + parser.add_argument("--concurrency", type=int, nargs="+", default=[1, 2, 4, 8]) + parser.add_argument( + "--observed-generation", + action="store_true", + help="Measure count-valid baseline observations; retain repeat-quality failures", + ) + args = parser.parse_args() + if not os.environ.get("EXP_ACTIVE"): + parser.error('Run through exp -m "purpose"') + if not 2 <= args.repeats <= 20 or any(c not in (1, 2, 4, 8) for c in args.concurrency): + parser.error("Require 2–20 repeats and concurrency in 1/2/4/8") + service = json.loads(args.service_metadata.read_text()) + for key in ("framework", "version", "software", "command", "hardware", "configuration", "base_url"): + if not service.get(key): + raise ValueError(f"Missing actual service provenance: {key}") + if service["framework"] not in ("lightllm", "vllm", "sglang"): + raise ValueError("Unsupported framework") + verified_checkpoint(args.model_directory) + reference = json.loads(args.reference.read_text()) + lengths = [int(name.removeprefix("long_")) for name in reference["cases"] if name.startswith("long_")] + expected_cases = cases_for(args.model_directory, lengths) + if args.observed_generation: + if ( + service["framework"] not in ("vllm", "sglang") + or reference.get("service", {}).get("framework") != service["framework"] + ): + raise ValueError("Observation timing is restricted to matching baseline frameworks") + workloads, reference_quality = workloads_from_observations(reference, expected_cases) + else: + workloads = workloads_from_reference(reference, expected_cases) + reference_quality = dict(status="qualified_generation") + run = Path(os.environ["EXP_RUN_DIR"]) + shutil.copyfile(__file__, run / Path(__file__).name) + manifest = [dict(name=w.name, prompt_ids=w.prompt_ids, output_length=len(w.expected_ids)) for w in workloads] + report = dict( + status="running", + scope=__doc__, + service=service, + manifest=manifest, + manifest_sha256=hashlib.sha256(json.dumps(manifest, sort_keys=True).encode()).hexdigest(), + reference_sha256=hashlib.sha256(args.reference.read_bytes()).hexdigest(), + reference_kind="observations" if args.observed_generation else "qualified_generation", + reference_quality=reference_quality, + service_sha256=hashlib.sha256(args.service_metadata.read_bytes()).hexdigest(), + repeats=args.repeats, + cohorts=[], + warmups=[], + comparison_eligible=False, + source_sha256=hashlib.sha256(Path(__file__).read_bytes()).hexdigest(), + ) + + def save(): + (run / "fixed-work.json").write_text(json.dumps(report, indent=2)) + + def request(work, path): + options = ( + dict(disable_prompt_cache=True) + if service["framework"] == "lightllm" + else dict(cache_salt=secrets.token_urlsafe(32)) + ) + return request_stream( + service["base_url"], + service["framework"], + str(args.model_directory.resolve()), + list(work.prompt_ids), + len(work.expected_ids), + 600, + path, + **options, + ) + + try: + save() + + for concurrency in args.concurrency: + cohort = run_cohort(workloads * 4, concurrency, request, run / f"warmup-{concurrency}") + report["warmups"].append(cohort) + summarize_fixed_work(cohort, workloads * 4) + save() + for repeat in range(args.repeats): + order = ( + args.concurrency[repeat % len(args.concurrency) :] + args.concurrency[: repeat % len(args.concurrency)] + ) + sequence = workloads[repeat % len(workloads) :] + workloads[: repeat % len(workloads)] + for concurrency in order: + cohort = run_cohort(sequence * 4, concurrency, request, run / f"repeat-{repeat}-c{concurrency}") + cohort["repeat"] = repeat + report["cohorts"].append(cohort) + cohort["fixed_work_metrics"] = summarize_fixed_work(cohort, sequence * 4) + save() + print( + json.dumps( + dict( + repeat=repeat, + concurrency=concurrency, + exact_status=cohort["status"], + output_tokens_per_s=cohort["fixed_work_metrics"]["output_tokens_per_s"], + ) + ), + flush=True, + ) + report["status"] = "fixed_work_measured" + except BaseException as error: + report.update(status="failed", error_type=type(error).__name__, error=str(error)) + raise + finally: + save() + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/qwen38/prepare_aiperf_trace.py b/test/benchmark/qwen38/prepare_aiperf_trace.py new file mode 100644 index 0000000000..da8e294f2f --- /dev/null +++ b/test/benchmark/qwen38/prepare_aiperf_trace.py @@ -0,0 +1,72 @@ +import argparse +import hashlib +import json +from pathlib import Path +from urllib.request import urlopen + + +REVISION = "3cca71daccf2a7afb8fe3f0295358f70e3a69fdb" +URL = f"https://raw.githubusercontent.com/kvcache-ai/Mooncake/{REVISION}/FAST25-release/traces/conversation_trace.jsonl" +SHA256 = "b8cbb061a85206d729d91cdc2981f43c9e0d99209dce588d3af5f7934408b9df" + + +def select_requests(rows, count, context_length): + if count < 1 or context_length < 2: + raise ValueError("Positive count and context length >= 2 are required") + eligible = [] + previous = -1 + for index, row in enumerate(rows): + if row["timestamp"] < previous: + raise ValueError("Trace timestamps must be nondecreasing") + previous = row["timestamp"] + if ( + row["input_length"] > 0 + and row["output_length"] > 0 + and row["input_length"] + row["output_length"] <= context_length + ): + eligible.append((index, row)) + if len(eligible) < count: + raise ValueError("Not enough requests fit the shared context limit") + return eligible[:count], len(eligible) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output-directory", type=Path, required=True) + parser.add_argument("--source", type=Path) + parser.add_argument("--count", type=int, default=256) + parser.add_argument("--context-length", type=int, default=4096) + args = parser.parse_args() + raw = args.source.read_bytes() if args.source else urlopen(URL, timeout=60).read() + if hashlib.sha256(raw).hexdigest() != SHA256: + raise ValueError("Public trace differs from the pinned source") + rows = [json.loads(line) for line in raw.splitlines() if line.strip()] + selected, eligible = select_requests(rows, args.count, args.context_length) + directory = args.output_directory + directory.mkdir(parents=True, exist_ok=True) + (directory / "conversation_trace.original.jsonl").write_bytes(raw) + trace = directory / "replay.jsonl" + trace.write_text("".join(json.dumps(row, separators=(",", ":")) + "\n" for _, row in selected)) + report = dict( + source_url=URL, + revision=REVISION, + source_sha256=SHA256, + source_requests=len(rows), + eligible_requests=eligible, + selected_indices=[index for index, _ in selected], + requests=len(selected), + context_length=args.context_length, + selection="First eligible requests in source order; no input/output clipping", + prompt_content="Anonymized lengths and 512-token prefix hashes; AIPerf synthesizes text", + timing="Original millisecond timestamps, normalized by AIPerf; speed multipliers reported separately", + trace_sha256=hashlib.sha256(trace.read_bytes()).hexdigest(), + schedule_span_ms=selected[-1][1]["timestamp"] - selected[0][1]["timestamp"], + input_tokens=sum(row["input_length"] for _, row in selected), + output_tokens=sum(row["output_length"] for _, row in selected), + ) + (directory / "manifest.json").write_text(json.dumps(report, indent=2) + "\n") + print(json.dumps(report, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/qwen38/retrieval_quality.py b/test/benchmark/qwen38/retrieval_quality.py new file mode 100644 index 0000000000..1206964ff1 --- /dev/null +++ b/test/benchmark/qwen38/retrieval_quality.py @@ -0,0 +1,51 @@ +import hashlib +import json +import re + +from test.benchmark.qwen38.check_generation import REVISION + +MANIFEST_SHA256 = "3214c9f09a0eb3bae49061d4658daeaea251b4befcfa5214278d796e5762d7cd" +LABELS = ("ALPHA", "BETA", "GAMMA") +ANSWER = re.compile(r"\s*ALPHA\s*=\s*(\d+)\s*;\s*BETA\s*=\s*(\d+)\s*;\s*GAMMA\s*=\s*(\d+)\s*") + + +def score_records(text, expected): + if set(expected) != set(LABELS) or any( + not isinstance(value, str) or not re.fullmatch(r"[0-9]+", value) for value in expected.values() + ): + raise ValueError("Require the three declared record labels and exact digit strings") + match = ANSWER.fullmatch(text) if isinstance(text, str) else None + predicted = dict(zip(LABELS, match.groups())) if match else None + records = {label: predicted is not None and predicted[label] == expected[label] for label in LABELS} + return dict(predicted=predicted, expected=expected, records=records, correct=all(records.values())) + + +def load_workload(path, model_directory): + manifest = json.loads(path.read_text()) + if hashlib.sha256(json.dumps(manifest, sort_keys=True).encode()).hexdigest() != MANIFEST_SHA256: + raise ValueError("Retrieval workload differs from pinned preparation") + config = json.loads((model_directory / "config.json").read_text()) + context = config.get("text_config", config)["max_position_embeddings"] + if ( + manifest["status"] != "prepared" + or manifest["model_revision"] != REVISION + or manifest["context_limit"] != context + or manifest["decoding"] != dict(temperature=0, ignore_eos=False, enable_thinking=False) + ): + raise ValueError("Retrieval workload model, context or decoding differs") + for name in ("tokenizer.json", "tokenizer_config.json", "chat_template.jinja"): + if hashlib.sha256((model_directory / name).read_bytes()).hexdigest() != manifest["tokenizer_files"].get(name): + raise ValueError("Tokenizer differs from retrieval preparation") + for case in manifest["cases"]: + ids, limit = case["prompt_ids"], case["max_output_tokens"] + if ( + not isinstance(ids, list) + or not ids + or any(type(token) is not int or token < 0 for token in ids) + or type(limit) is not int + or limit < 1 + or len(ids) + limit > context + ): + raise ValueError("Retrieval inputs exceed context or lack exact token IDs") + score_records("", case["expected_records"]) + return manifest diff --git a/test/benchmark/qwen38/score_reference.py b/test/benchmark/qwen38/score_reference.py new file mode 100644 index 0000000000..a42d9b66c2 --- /dev/null +++ b/test/benchmark/qwen38/score_reference.py @@ -0,0 +1,208 @@ +import argparse +import hashlib +import importlib.metadata +import importlib.util +import json +import os +from pathlib import Path +import subprocess +import sys +import time + +import torch +from safetensors import safe_open + +from check_generation import REVISION, cases_for, verified_checkpoint + +SOURCE_SHA256 = "87c27c88424cea73d5d8d9179f76a77a97e6e4249ce43ecff79a8b255fab0504" + + +def load_source(path): + if hashlib.sha256(path.read_bytes()).hexdigest() != SOURCE_SHA256: + raise ValueError("Reference equations do not match the pinned source") + name = "transformers.models.qwen4_exp.modeling_qwen4_exp" + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +def verify_embedding_rows(model, root): + index = json.loads((root / "model.safetensors.index.json").read_text())["weight_map"] + weight = model.model.language_model.layers[1].ple.ple_embedding.ngram_embedding.weight + if weight.is_meta or tuple(weight.shape) != (320001536, 160) or weight.dtype != torch.bfloat16: + raise ValueError("Reference lexical table has the wrong shape, residency or dtype") + checked = 0 + for shard in range(128): + key = f"model.language_model.layers.1.ple.ple_embedding.ngram_embedding.shard_{shard}.weight" + with safe_open(root / index[key], framework="pt", device="cpu") as source: + tensor = source.get_slice(key) + rows, width = tensor.get_shape() + if (rows, width) != (2500012, 160): + raise ValueError(f"Unexpected source shard shape: {key}") + for row in (0, rows // 2, rows - 1): + actual = weight[shard * rows + row].cpu() + torch.testing.assert_close(actual, tensor[row : row + 1][0], atol=0, rtol=0) + checked += 1 + return {"shards": 128, "exact_rows": checked, "device": str(weight.device)} + + +def forward_reference(model, inputs, cache, length, chunk_size=None): + count = inputs.shape[1] + if count < 1 or length < count or (chunk_size is not None and chunk_size < 1): + raise ValueError("Reference forward requires nonempty inputs and valid lengths") + chunk_size = count if chunk_size is None else chunk_size + for start in range(0, count, chunk_size): + stop = min(start + chunk_size, count) + absolute_start, absolute_stop = length - count + start, length - count + stop + output = model( + input_ids=inputs[:, start:stop], + position_ids=torch.arange(absolute_start, absolute_stop, device=inputs.device).unsqueeze(0), + attention_mask=torch.ones((1, absolute_stop), device=inputs.device, dtype=torch.long), + past_key_values=cache, + use_cache=True, + logits_to_keep=1, + ) + cache = output.past_key_values + return output + + +@torch.inference_mode() +def score_case(model, prompt, result, directory, case, prefill_chunk_size=None): + tokens = result["token_ids"] + native = result["response"]["tokens"] + if [item["id"] for item in native] != tokens: + raise ValueError("Native probability records do not match generated IDs") + device = model.model.language_model.embed_tokens.weight.device + inputs = torch.tensor([prompt], device=device, dtype=torch.long) + cache = None + rows, scores = [], [] + for step, target in enumerate(tokens): + length = len(prompt) + step + output = forward_reference(model, inputs, cache, length, prefill_chunk_size) + cache = output.past_key_values + logits = output.logits[0, -1].float().cpu() + if not bool(torch.isfinite(logits).all()): + raise ValueError(f"Nonfinite reference logits at {case} step {step}") + scores.append(logits) + reference_logprob = float(logits[target] - logits.logsumexp(0)) + native_logprob = float(native[step]["logprob"]) + greedy = int(logits.argmax()) + rows.append( + dict( + step=step, + token_id=target, + reference_greedy_id=greedy, + reference_logprob=reference_logprob, + native_logprob=native_logprob, + logprob_delta=native_logprob - reference_logprob, + reference_greedy_gap=float(logits[greedy] - logits[target]), + ) + ) + inputs = torch.tensor([[target]], device=device, dtype=torch.long) + torch.save(torch.stack(scores), directory / f"{case}-scores.pt") + return dict( + steps=rows, + greedy_matches=sum(row["token_id"] == row["reference_greedy_id"] for row in rows), + reference_mean_nll=-sum(row["reference_logprob"] for row in rows) / len(rows), + native_mean_nll=-sum(row["native_logprob"] for row in rows) / len(rows), + mean_absolute_logprob_delta=sum(abs(row["logprob_delta"]) for row in rows) / len(rows), + maximum_absolute_logprob_delta=max(abs(row["logprob_delta"]) for row in rows), + ) + + +def main(): + if not os.environ.get("EXP_ACTIVE") or os.environ.get("USE_HUB_KERNELS") != "NO": + raise RuntimeError("Run through exp with USE_HUB_KERNELS=NO") + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model-directory", type=Path, required=True) + parser.add_argument("--source", type=Path, required=True) + parser.add_argument("--device-plan", type=Path, required=True) + parser.add_argument("--generation", type=Path, required=True) + parser.add_argument("--cases", nargs="+", default=["english", "chinese", "code_unicode"]) + parser.add_argument("--prefill-chunk-size", type=int) + args = parser.parse_args() + if args.prefill_chunk_size is not None and args.prefill_chunk_size < 1: + parser.error("Prefill chunk size must be positive") + root, run = args.model_directory, Path(os.environ["EXP_RUN_DIR"]) + verified_checkpoint(root) + trace = json.loads(args.generation.read_text()) + cases = cases_for(root) + if trace["revision"] != REVISION or trace["cases"] != cases or trace["service"]["framework"] != "lightllm": + raise ValueError("Generation must contain native outputs for the pinned shared inputs") + for case in args.cases: + if not trace["results"][case]["token_ids"]: + raise ValueError(f"Missing generated IDs: {case}") + plan = json.loads(args.device_plan.read_text()) + device_map = plan["device_map"] + if not device_map or not all(type(device) is int for device in device_map.values()): + raise ValueError("This reference diagnostic requires a complete GPU placement") + module = load_source(args.source) + report = dict( + status="loading", + scope="identical native token histories; diagnostic only, no quality parity or performance claim", + revision=REVISION, + source_sha256=SOURCE_SHA256, + source=str(args.source.resolve()), + generation=str(args.generation.resolve()), + device_plan=plan, + command=sys.argv, + cuda_visible_devices=os.environ.get("CUDA_VISIBLE_DEVICES"), + software={name: importlib.metadata.version(name) for name in ("torch", "transformers", "accelerate")}, + hardware=subprocess.check_output( + ["nvidia-smi", "--query-gpu=index,name,uuid,memory.used,memory.free", "--format=csv,noheader"], text=True + ), + configuration=dict( + dtype="bfloat16", + attention="sdpa", + experts="eager", + use_hub_kernels=False, + allow_tf32=torch.backends.cuda.matmul.allow_tf32, + bf16_reduced_precision=torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction, + deterministic_algorithms=torch.are_deterministic_algorithms_enabled(), + prefill_chunk_size=args.prefill_chunk_size, + ), + results={}, + ) + path = run / "reference-scores.json" + path.write_text(json.dumps(report, indent=2)) + started = time.monotonic() + try: + model, loading = module.Qwen4ExpForConditionalGeneration.from_pretrained( + str(root), + dtype=torch.bfloat16, + device_map=device_map, + local_files_only=True, + use_safetensors=True, + attn_implementation="sdpa", + experts_implementation="eager", + output_loading_info=True, + ) + model.eval() + (run / "loading-info.json").write_text(json.dumps(loading, indent=2, default=str)) + if loading.get("missing_keys") or loading.get("mismatched_keys"): + raise ValueError("Reference checkpoint loading is incomplete") + if any(not key.startswith("mtp.") for key in loading.get("unexpected_keys", [])): + raise ValueError("Unexpected non-MTP reference weights") + report["lexical_row_verification"] = verify_embedding_rows(model, root) + report.update(status="scoring", load_seconds=time.monotonic() - started) + path.write_text(json.dumps(report, indent=2)) + for case in args.cases: + result = score_case(model, cases[case], trace["results"][case], run, case, args.prefill_chunk_size) + report["results"][case] = result + path.write_text(json.dumps(report, indent=2)) + print( + json.dumps({"case": case, **{key: value for key, value in result.items() if key != "steps"}}), + flush=True, + ) + report["status"] = "recorded" + except Exception as error: + report.update(status="failed", error_type=type(error).__name__, error=str(error)) + raise + finally: + path.write_text(json.dumps(report, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/service/cache_capacity.py b/test/benchmark/service/cache_capacity.py new file mode 100644 index 0000000000..0fb53a7578 --- /dev/null +++ b/test/benchmark/service/cache_capacity.py @@ -0,0 +1,50 @@ +import math + +from prometheus_client.parser import text_string_to_metric_families + + +def plan_vllm_capacity(metrics_text, model_length, target_tokens): + if any(type(value) is not int or value < 1 for value in (model_length, target_tokens)): + raise ValueError("Require positive integer context and token budgets") + if target_tokens % model_length: + raise ValueError("Token budget must contain whole full-context requests") + samples = [ + sample + for family in text_string_to_metric_families(metrics_text) + for sample in family.samples + if sample.name == "vllm:cache_config_info" + ] + if len(samples) != 1 or samples[0].value != 1: + raise ValueError("Require exactly one initialized cache-info sample") + fields = samples[0].labels + try: + blocks = int(fields["num_gpu_blocks"]) + concurrency = float(fields["kv_cache_max_concurrency"]) + reported_tokens = int(fields["kv_cache_size_tokens"]) + except (KeyError, TypeError, ValueError) as error: + raise ValueError("Initialized block count and group-aware capacity are required") from error + if blocks < 1 or not math.isfinite(concurrency) or concurrency <= 0 or reported_tokens < 1: + raise ValueError("Runtime cache geometry must be finite and positive") + if reported_tokens != int(concurrency * model_length): + raise ValueError("Reported capacity does not match the recorded context length") + if target_tokens > reported_tokens: + raise ValueError("Capacity planning only shrinks the observed initialized pool") + ratio = blocks / concurrency + per_request = round(ratio) + if per_request < 1 or not math.isclose(ratio, per_request, rel_tol=1e-12, abs_tol=1e-9): + raise ValueError("Cannot infer an integer device-pool block budget per request") + return dict( + status="planned", + scope="Logical single-engine budget inferred from initialized metrics; verify again after launch", + model_length=model_length, + target_tokens=target_tokens, + observed_tokens=reported_tokens, + observed_num_gpu_blocks=blocks, + observed_max_concurrency=concurrency, + blocks_per_full_request=per_request, + num_gpu_blocks_override=per_request * (target_tokens // model_length), + cache_geometry={ + key: fields.get(key) + for key in ("block_size", "mamba_block_size", "mamba_cache_mode", "mamba_ssm_cache_dtype") + }, + ) diff --git a/test/benchmark/service/completion.py b/test/benchmark/service/completion.py new file mode 100644 index 0000000000..9cb190a53c --- /dev/null +++ b/test/benchmark/service/completion.py @@ -0,0 +1,61 @@ +def completion_request(framework, model, prompt_ids, limit): + if not prompt_ids or any(type(token) is not int or token < 0 for token in prompt_ids): + raise ValueError("Require nonnegative integer prompt IDs") + if type(limit) is not int or limit < 1: + raise ValueError("Require a positive integer output limit") + if framework == "lightllm": + return "/generate", { + "inputs": prompt_ids, + "parameters": dict( + do_sample=False, + ignore_eos=False, + max_new_tokens=limit, + return_details=True, + print_eos_token=True, + ), + } + if framework in ("vllm", "sglang"): + return "/v1/completions", dict( + model=model, + prompt=prompt_ids, + max_tokens=limit, + temperature=0, + seed=3107, + ignore_eos=False, + return_token_ids=True, + ) + raise ValueError("Unsupported framework") + + +def completion_tokens(body, framework, prompt_ids, limit, eos_ids): + if framework == "lightllm": + ids = [token["id"] for token in body["tokens"]] + prompt_count, output_count = body["prompt_tokens"], body["count_output_tokens"] + reason = body["finish_reason"] + elif framework in ("vllm", "sglang"): + if len(body["choices"]) != 1 or body["choices"][0].get("index") != 0: + raise ValueError("Require exactly one completion choice") + choice = body["choices"][0] + if choice.get("prompt_token_ids", prompt_ids) != prompt_ids: + raise ValueError("Service changed the exact prompt IDs") + ids, reason = choice["token_ids"], choice["finish_reason"] + prompt_count, output_count = body["usage"]["prompt_tokens"], body["usage"]["completion_tokens"] + else: + raise ValueError("Unsupported framework") + if not isinstance(ids, list) or not ids or any(type(token) is not int or token < 0 for token in ids): + raise ValueError("Require actual nonnegative output token IDs") + if ( + type(prompt_count) is not int + or type(output_count) is not int + or prompt_count != len(prompt_ids) + or output_count != len(ids) + or len(ids) > limit + ): + raise ValueError("Service counts differ from exact prompt/output IDs") + if reason not in ("stop", "length") or any(token in eos_ids for token in ids[:-1]): + raise ValueError("Invalid termination reason or output beyond first EOS") + if reason == "length" and (len(ids) != limit or ids[-1] in eos_ids): + raise ValueError("Length termination must exhaust the budget without EOS") + if framework == "lightllm" and reason == "stop" and ids[-1] not in eos_ids: + raise ValueError("Native natural termination must expose its declared EOS") + return dict(token_ids=ids, finish_reason=reason, truncated=reason == "length") diff --git a/test/benchmark/service/completion_runner.py b/test/benchmark/service/completion_runner.py new file mode 100644 index 0000000000..95a7a56cdd --- /dev/null +++ b/test/benchmark/service/completion_runner.py @@ -0,0 +1,63 @@ +from concurrent.futures import ThreadPoolExecutor +import threading + +import requests + +from .completion import completion_request, completion_tokens + + +class CompletionClient: + def __init__( + self, + framework, + model, + base_url, + timeout, + eos_ids, + raw_directory, + *, + concurrency=1, + session_factory=requests.Session, + ): + if type(concurrency) is not int or not 1 <= concurrency <= 8: + raise ValueError("Completion concurrency must be between one and eight") + self.framework, self.model, self.base_url = framework, model, base_url + self.timeout, self.eos_ids, self.raw_directory = timeout, eos_ids, raw_directory + self.concurrency, self.session_factory = concurrency, session_factory + self.local = threading.local() + self.sessions, self.lock, self.executor = [], threading.Lock(), None + + def _complete(self, case): + index = case["dataset_index"] + if type(index) is not int or index < 0: + raise ValueError("Require a nonnegative integer dataset index") + if not hasattr(self.local, "session"): + session = self.session_factory() + session.trust_env = False + self.local.session = session + with self.lock: + self.sessions.append(session) + endpoint, payload = completion_request( + self.framework, self.model, case["prompt_ids"], case["max_output_tokens"] + ) + response = self.local.session.post(self.base_url.rstrip("/") + endpoint, json=payload, timeout=self.timeout) + (self.raw_directory / f"response-{index}.txt").write_text(response.text) + response.raise_for_status() + return completion_tokens( + response.json(), self.framework, case["prompt_ids"], case["max_output_tokens"], self.eos_ids + ) + + def results(self, cases): + if self.concurrency == 1: + return map(self._complete, cases) + if self.executor is not None: + raise RuntimeError("A completion client accepts one concurrent workload") + self.executor = ThreadPoolExecutor(max_workers=self.concurrency) + + return self.executor.map(self._complete, cases) + + def close(self): + if self.executor is not None: + self.executor.shutdown(wait=True, cancel_futures=True) + for session in self.sessions: + session.close() diff --git a/test/benchmark/service/stream_cohort.py b/test/benchmark/service/stream_cohort.py new file mode 100644 index 0000000000..9ebc722ddc --- /dev/null +++ b/test/benchmark/service/stream_cohort.py @@ -0,0 +1,166 @@ +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +import math +from pathlib import Path +from queue import Empty, SimpleQueue +from statistics import fmean +import threading + + +@dataclass(frozen=True) +class StreamingWorkload: + name: str + prompt_ids: tuple[int, ...] + expected_ids: tuple[int, ...] + + def __post_init__(self): + if not isinstance(self.name, str) or not self.name: + raise ValueError("Workload requires a name") + for field in ("prompt_ids", "expected_ids"): + values = getattr(self, field) + if not isinstance(values, (list, tuple)) or not values or any(type(x) is not int or x < 0 for x in values): + raise ValueError("Workload requires exact nonnegative prompt and reference token IDs") + object.__setattr__(self, field, tuple(values)) + + +def distribution(values): + if not values: + return None + if any(not math.isfinite(value) or value < 0 for value in values): + raise ValueError("Latency samples must be finite and nonnegative") + ordered = sorted(values) + result = dict(count=len(values), min=ordered[0], max=ordered[-1], mean=fmean(values)) + for percentile in (50, 90, 95, 99): + position = (len(ordered) - 1) * percentile / 100 + left, right = math.floor(position), math.ceil(position) + result[f"p{percentile}"] = ordered[left] + (ordered[right] - ordered[left]) * (position - left) + return result + + +def _validate_result(result): + metrics, timeline = result["metrics"], result["timeline"] + started, first, last, ended = (timeline[key] for key in ("started", "first_token", "last_token", "ended")) + if not all(math.isfinite(value) for value in (started, first, last, ended)): + raise ValueError("Request timeline must be finite") + if not started <= first <= last <= ended or ended <= started: + raise ValueError("Request timeline is incomplete or nonmonotonic") + ids = result["token_ids"] + if not ids or any(type(token) is not int or token < 0 for token in ids): + raise ValueError("Result requires actual integer output IDs") + if type(metrics["input_tokens"]) is not int or metrics["input_tokens"] < 1: + raise ValueError("Result requires the actual input token count") + if type(metrics["output_tokens"]) is not int or metrics["output_tokens"] != len(ids): + raise ValueError("Output count differs from actual token IDs") + gaps = metrics["receiver_itl_s"] + if len(gaps) != len(ids) - 1 or any(not math.isfinite(gap) or gap < 0 for gap in gaps): + raise ValueError("Receiver gaps must account for every output position") + expected = dict(ttft_s=first - started, e2e_s=ended - started) + if len(ids) > 1: + expected["tpot_s"] = (last - first) / (len(ids) - 1) + elif metrics["tpot_s"] is not None: + raise ValueError("A single output token has no TPOT") + for key, value in expected.items(): + if not math.isclose(metrics[key], value, rel_tol=1e-9, abs_tol=1e-9): + raise ValueError(f"{key} differs from the captured request timeline") + if not math.isclose(sum(gaps), last - first, rel_tol=1e-9, abs_tol=1e-9): + raise ValueError("Receiver gaps differ from the token arrival interval") + + +def summarize_cohort(results): + if not results: + raise ValueError("An empty cohort cannot be measured") + for result in results: + _validate_result(result) + elapsed = max(item["timeline"]["ended"] for item in results) - min(item["timeline"]["started"] for item in results) + input_tokens = sum(item["metrics"]["input_tokens"] for item in results) + output_tokens = sum(len(item["token_ids"]) for item in results) + cache_counts = [item["metrics"].get("cached_input_tokens") for item in results] + if any( + count is not None and (type(count) is not int or not 0 <= count <= item["metrics"]["input_tokens"]) + for count, item in zip(cache_counts, results) + ): + raise ValueError("Invalid cached input count") + cached_tokens = sum(cache_counts) if all(count is not None for count in cache_counts) else None + events = sorted( + [(item["timeline"]["started"], 1) for item in results] + [(item["timeline"]["ended"], -1) for item in results] + ) + active = maximum = 0 + for _, delta in events: + active += delta + maximum = max(maximum, active) + gaps = [gap for item in results for gap in item["metrics"]["receiver_itl_s"]] + return dict( + requests=len(results), + input_tokens=input_tokens, + output_tokens=output_tokens, + cache_reported_requests=sum(count is not None for count in cache_counts), + cached_input_tokens=cached_tokens, + uncached_input_tokens=input_tokens - cached_tokens if cached_tokens is not None else None, + elapsed_s=elapsed, + requests_per_s=len(results) / elapsed, + input_tokens_per_s=input_tokens / elapsed, + output_tokens_per_s=output_tokens / elapsed, + input_token_definition="requested prompt IDs; includes cached tokens", + maximum_client_inflight=maximum, + ttft_s=distribution([item["metrics"]["ttft_s"] for item in results]), + tpot_s=distribution([item["metrics"]["tpot_s"] for item in results if item["metrics"]["tpot_s"] is not None]), + e2e_s=distribution([item["metrics"]["e2e_s"] for item in results]), + receiver_itl_s=distribution(gaps), + receiver_zero_gap_fraction=sum(gap == 0 for gap in gaps) / len(gaps) if gaps else None, + itl_definition="client arrival gaps; tokens delivered together share a timestamp", + percentile_definition="linear interpolation between ordered observations", + ) + + +def run_cohort(workloads, concurrency, request, response_directory: Path): + workloads = tuple(workloads) + if type(concurrency) is not int or not 1 <= concurrency <= len(workloads): + raise ValueError("Concurrency must be positive and no larger than the request count") + if any(not isinstance(item, StreamingWorkload) for item in workloads): + raise ValueError("Every request requires a qualified streaming workload") + response_directory.mkdir(parents=True, exist_ok=True) + pending = SimpleQueue() + for index, workload in enumerate(workloads): + pending.put((index, workload)) + records = [None] * len(workloads) + ready = threading.Barrier(concurrency + 1) + release = threading.Event() + + def worker(): + ready.wait(timeout=30) + release.wait() + while True: + try: + index, workload = pending.get_nowait() + except Empty: + return + path = response_directory / f"{index:06d}.sse" + record = dict(index=index, workload=workload.name, response=str(path), status="running") + try: + result = request(workload, path) + record["result"] = result + _validate_result(result) + if result["metrics"]["input_tokens"] != len(workload.prompt_ids): + raise ValueError("Service input length differs from the qualified workload") + if tuple(result["token_ids"]) != workload.expected_ids: + raise ValueError("Output IDs differ from the qualified reference") + record["status"] = "passed" + except Exception as error: + record.update(status="failed", error_type=type(error).__name__, error=str(error)) + records[index] = record + + with ThreadPoolExecutor(max_workers=concurrency) as executor: + futures = [executor.submit(worker) for _ in range(concurrency)] + try: + ready.wait(timeout=30) + finally: + release.set() + for future in futures: + future.result() + passed = all(item["status"] == "passed" for item in records) + return dict( + status="passed" if passed else "failed", + concurrency=concurrency, + records=records, + metrics=summarize_cohort([item["result"] for item in records]) if passed else None, + ) diff --git a/test/benchmark/service/stream_timing.py b/test/benchmark/service/stream_timing.py new file mode 100644 index 0000000000..3bce3d7254 --- /dev/null +++ b/test/benchmark/service/stream_timing.py @@ -0,0 +1,80 @@ +import json +import math + + +def _stats(values): + if not values: + return None + ordered = sorted(values) + + def percentile(fraction): + index = (len(ordered) - 1) * fraction + lower, upper = math.floor(index), math.ceil(index) + return ordered[lower] + (ordered[upper] - ordered[lower]) * (index - lower) + + return dict(count=len(ordered), p50=percentile(0.5), p95=percentile(0.95), p99=percentile(0.99), max=ordered[-1]) + + +def _nanoseconds(value): + if type(value) is not int or value < 0: + raise ValueError("Expected a nonnegative integer timestamp in nanoseconds") + return value + + +def streaming_completion_times(record): + metadata = record["metadata"] + if record["status"] != 200 or metadata.get("was_cancelled"): + raise ValueError("Cannot measure a failed or cancelled streaming response") + start = _nanoseconds(metadata["request_start_ns"]) + end = _nanoseconds(metadata["request_end_ns"]) + origin = _nanoseconds(record["start_perf_ns"]) + if end < start: + raise ValueError("HTTP end precedes request start") + text_times, done_times = [], [] + previous = origin + for response in record["responses"]: + timestamp = _nanoseconds(response["perf_ns"]) + if timestamp < previous: + raise ValueError("Nonmonotonic streaming packet timestamps") + previous = timestamp + for packet in response.get("packets", []): + if packet.get("name") != "data": + continue + if done_times: + raise ValueError("Unexpected SSE data after DONE") + if packet["value"] == "[DONE]": + done_times.append(timestamp) + continue + payload = json.loads(packet["value"]) + if any(choice.get("text") for choice in payload.get("choices", [])): + text_times.append(timestamp) + if len(done_times) != 1: + raise ValueError("Expected one complete SSE DONE marker") + + if abs((done_times[0] - origin) - (end - start)) > 1_000_000: + raise ValueError("HTTP completion and SSE DONE elapsed times disagree") + completion = (end - start) / 1e6 + first = (text_times[0] - origin) / 1e6 if text_times else None + last = (text_times[-1] - origin) / 1e6 if text_times else None + return dict( + http_completion_ms=completion, + first_visible_text_ms=first, + last_visible_text_ms=last, + done_after_last_visible_text_ms=(done_times[0] - text_times[-1]) / 1e6 if text_times else None, + ) + + +def summarize_streaming_times(records): + rows = [streaming_completion_times(record) for record in records] + if not rows: + raise ValueError("No streaming records to summarize") + content = [row for row in rows if row["last_visible_text_ms"] is not None] + return dict( + requests=len(rows), + empty_text_responses=len(rows) - len(content), + http_completion_all_requests_ms=_stats([row["http_completion_ms"] for row in rows]), + http_completion_content_requests_ms=_stats([row["http_completion_ms"] for row in content]), + first_visible_text_ms=_stats([row["first_visible_text_ms"] for row in content]), + last_visible_text_ms=_stats([row["last_visible_text_ms"] for row in content]), + done_after_last_visible_text_ms=_stats([row["done_after_last_visible_text_ms"] for row in content]), + ) diff --git a/test/benchmark/service/token_stream.py b/test/benchmark/service/token_stream.py new file mode 100644 index 0000000000..fc63f66afd --- /dev/null +++ b/test/benchmark/service/token_stream.py @@ -0,0 +1,237 @@ +import codecs +from dataclasses import dataclass, field +import json +import math + + +class SSEDecoder: + def __init__(self): + self.decoder = codecs.getincrementaldecoder("utf-8")() + self.buffer = "" + self.data = [] + + def feed(self, chunk: bytes) -> list[str]: + self.buffer += self.decoder.decode(chunk) + events = [] + while "\n" in self.buffer: + line, self.buffer = self.buffer.split("\n", 1) + line = line.removesuffix("\r") + if not line: + if self.data: + events.append("\n".join(self.data)) + self.data.clear() + elif line.startswith("data:"): + self.data.append(line[5:].removeprefix(" ")) + return events + + def finish(self): + self.buffer += self.decoder.decode(b"", final=True) + if self.buffer.strip() or self.data: + raise ValueError("Truncated SSE event at EOF") + + +@dataclass +class TokenStream: + framework: str + prompt_ids: list[int] + output_length: int + started: float + token_ids: list[int] = field(default_factory=list, init=False) + arrivals: list[float] = field(default_factory=list, init=False) + finished: bool = field(default=False, init=False) + done: bool = field(default=False, init=False) + events: int = field(default=0, init=False) + last_received: float = field(default=0.0, init=False, repr=False) + cached_input_tokens: int | None = field(default=None, init=False) + + def __post_init__(self): + if self.framework not in ("lightllm", "vllm", "sglang") or self.output_length < 1: + raise ValueError("Expected a supported framework and positive output length") + if not self.prompt_ids or any(type(token) is not int or token < 0 for token in self.prompt_ids): + raise ValueError("Prompt must contain exact integer token IDs") + self.last_received = self.started + if not math.isfinite(self.started): + raise ValueError("Start timestamp must be finite") + + def record(self, payload: str, received: float): + if not math.isfinite(received) or received < self.last_received: + raise ValueError("Token arrival timestamps must be finite and monotonic") + self.last_received = received + if payload == "[DONE]": + if not self.finished: + raise ValueError("Stream ended before the terminal generation event") + self.done = True + return + if self.done: + raise ValueError("Received data after DONE") + body = json.loads(payload) + if body.get("error"): + raise ValueError(f"Service stream error: {body['error']}") + self.events += 1 + if self.framework == "lightllm": + token = body["token"] + if token["prompt_tokens"] != len(self.prompt_ids): + raise ValueError("Stream prompt count differs from actual request IDs") + if token["count_output_tokens"] != len(self.token_ids) + 1: + raise ValueError("Native stream skipped or duplicated an output position") + ids = [token["id"]] + self._record_cache_count(token.get("prompt_cache_len")) + reason = body.get("finish_reason") if body.get("finished") is True else None + else: + choices = body.get("choices", []) + if len(choices) > 1 or (choices and choices[0].get("index") != 0): + raise ValueError("Benchmark requires one completion choice") + choice = choices[0] if choices else {} + ids = choice.get("token_ids") + if ids is None: + ids = [] + if choice.get("text") and not ids: + raise ValueError("Nonempty text event omitted explicit token IDs") + prompt = choice.get("prompt_token_ids") + if prompt is not None and prompt != self.prompt_ids: + raise ValueError("Service changed the exact tokenized prompt") + reason = choice.get("finish_reason") + if not isinstance(ids, list) or any(type(token) is not int or token < 0 for token in ids): + raise ValueError("Stream must expose integer token IDs") + if ids and self.finished: + raise ValueError("Received output tokens after terminal generation") + self.token_ids.extend(ids) + self.arrivals.extend([received] * len(ids)) + if len(self.token_ids) > self.output_length: + raise ValueError("Service exceeded the requested output length") + if reason is not None: + if reason != "length" or len(self.token_ids) != self.output_length: + raise ValueError("Fixed-length stream did not finish at its requested token count") + self.finished = True + usage = body.get("usage") + if usage is not None: + self._record_cache_count((usage.get("prompt_tokens_details") or {}).get("cached_tokens")) + if usage is not None and ( + usage["prompt_tokens"] != len(self.prompt_ids) or usage["completion_tokens"] != len(self.token_ids) + ): + raise ValueError("Stream usage does not match observed exact token counts") + + def _record_cache_count(self, count): + if count is None: + return + if type(count) is not int or not 0 <= count <= len(self.prompt_ids): + raise ValueError("Cached input count must be an integer within the exact prompt") + if self.cached_input_tokens is not None and count != self.cached_input_tokens: + raise ValueError("Cached input count changed within one request") + self.cached_input_tokens = count + + def summarize(self, ended: float) -> dict: + if not self.finished or len(self.token_ids) != self.output_length: + raise ValueError("Incomplete generation cannot be included in benchmark results") + if self.framework != "lightllm" and not self.done: + raise ValueError("OpenAI-compatible stream omitted DONE") + if not math.isfinite(ended) or ended < self.last_received: + raise ValueError("Invalid request completion timestamp") + gaps = [right - left for left, right in zip(self.arrivals, self.arrivals[1:])] + return { + "input_tokens": len(self.prompt_ids), + "cached_input_tokens": self.cached_input_tokens, + "uncached_input_tokens": ( + len(self.prompt_ids) - self.cached_input_tokens if self.cached_input_tokens is not None else None + ), + "output_tokens": len(self.token_ids), + "events": self.events, + "ttft_s": self.arrivals[0] - self.started, + "tpot_s": sum(gaps) / len(gaps) if gaps else None, + "e2e_s": ended - self.started, + "receiver_itl_s": gaps, + "itl_definition": "client arrival gaps; IDs delivered in one event share a timestamp", + } + + +def request_stream( + base_url, + framework, + model, + prompt_ids, + output_length, + timeout, + response_path, + *, + disable_prompt_cache=False, + cache_salt=None, +): + import time + import requests + + if disable_prompt_cache and framework != "lightllm": + raise ValueError("Per-request prompt-cache disabling is only implemented for native requests") + if cache_salt is not None: + if framework not in ("vllm", "sglang"): + raise ValueError("Cache salts are only implemented for baseline requests") + if ( + not isinstance(cache_salt, str) + or not 1 <= len(cache_salt) <= 128 + or any(char in cache_salt for char in "@/\\\x00") + ): + raise ValueError("Cache salt must contain 1--128 characters without '@', '/', backslash or NUL") + if framework == "lightllm": + endpoint = "/generate_stream" + payload = { + "inputs": prompt_ids, + "parameters": { + "do_sample": False, + "ignore_eos": True, + "max_new_tokens": output_length, + "return_details": True, + }, + } + if disable_prompt_cache: + payload["parameters"]["disable_prompt_cache"] = True + elif framework in ("vllm", "sglang"): + endpoint = "/v1/completions" + payload = { + "model": model, + "prompt": prompt_ids, + "temperature": 0, + "seed": 3107, + "ignore_eos": True, + "max_tokens": output_length, + "stream": True, + "stream_options": {"include_usage": True}, + "return_token_ids": True, + } + if cache_salt is not None: + payload["cache_salt"] = cache_salt + else: + raise ValueError("Unsupported serving framework") + raw = [] + decoder = SSEDecoder() + try: + with requests.Session() as client: + client.trust_env = False + started = time.perf_counter() + stream = TokenStream(framework, prompt_ids, output_length, started) + with client.post(base_url.rstrip("/") + endpoint, json=payload, stream=True, timeout=timeout) as response: + if response.status_code != 200: + raw.append(response.content) + response.raise_for_status() + if not response.headers.get("content-type", "").startswith("text/event-stream"): + raw.append(response.content) + raise ValueError("Expected an SSE response") + if "chunked" not in response.headers.get("transfer-encoding", "").lower(): + raise ValueError("Token timing requires HTTP/1.1 chunked transfer") + for chunk in response.iter_content(chunk_size=None): + received = time.perf_counter() + raw.append(chunk) + for event in decoder.feed(chunk): + stream.record(event, received) + ended = time.perf_counter() + decoder.finish() + return { + "token_ids": stream.token_ids, + "metrics": stream.summarize(ended), + "timeline": { + "started": started, + "first_token": stream.arrivals[0], + "last_token": stream.arrivals[-1], + "ended": ended, + }, + } + finally: + response_path.write_bytes(b"".join(raw)) diff --git a/unit_tests/benchmark/qwen38/test_aiperf_replay.py b/unit_tests/benchmark/qwen38/test_aiperf_replay.py new file mode 100644 index 0000000000..2217ea9732 --- /dev/null +++ b/unit_tests/benchmark/qwen38/test_aiperf_replay.py @@ -0,0 +1,139 @@ +import copy +import json + +import pytest + +from test.benchmark.qwen38.measure_aiperf import audit_records +from test.benchmark.qwen38.prepare_aiperf_trace import select_requests + + +def test_trace_selection_preserves_bursts_prefixes_and_lengths(): + rows = [ + dict(timestamp=0, input_length=4000, output_length=200, hash_ids=[1]), + dict(timestamp=10, input_length=20, output_length=3, hash_ids=[2]), + dict(timestamp=10, input_length=30, output_length=4, hash_ids=[2, 3]), + dict(timestamp=100, input_length=10, output_length=5, hash_ids=[4]), + ] + selected, count = select_requests(rows, 2, 4096) + assert count == 3 and selected == [(1, rows[1]), (2, rows[2])] + with pytest.raises(ValueError, match="Not enough"): + select_requests(rows, 4, 4096) + with pytest.raises(ValueError, match="nondecreasing"): + select_requests(list(reversed(rows)), 1, 4096) + + +def fixture_record(): + return dict( + metadata=dict( + benchmark_phase="profiling", + conversation_id="session_000000", + session_num=0, + request_start_ns=1000, + was_cancelled=False, + ), + status=200, + payload=dict(prompt=["hello"], max_tokens=2, temperature=0, ignore_eos=True), + responses=[ + dict( + packets=[ + dict( + name="data", + value=json.dumps( + dict( + choices=[dict(text="hello", finish_reason="length")], + usage=dict(prompt_tokens=1, completion_tokens=2), + ) + ), + ), + dict(name="data", value="[DONE]"), + ] + ) + ], + ) + + +def write_artifacts(path, records): + metrics = ( + "request_throughput", + "output_token_throughput", + "time_to_first_token", + "inter_token_latency", + "request_latency", + ) + summary = dict(error_summary=[], **{name: dict(avg=1) for name in metrics}) + summary.update(request_count=dict(avg=len(records)), output_token_count=dict(sum=2 * len(records))) + (path / "profile_export_aiperf.json").write_text(json.dumps(summary)) + (path / "profile_export_raw.jsonl").write_text("".join(json.dumps(row) + "\n" for row in records)) + + +def test_raw_audit_accepts_complete_stream_and_excludes_model_alias(tmp_path): + row = fixture_record() + write_artifacts(tmp_path, [row]) + trace = [dict(timestamp=0, output_length=2)] + result = audit_records(tmp_path, trace) + row["payload"]["model"] = "different-service-alias" + write_artifacts(tmp_path, [row]) + assert audit_records(tmp_path, trace)["request_manifest_sha256"] == result["request_manifest_sha256"] + + +@pytest.mark.parametrize("failure", ["missing", "duplicate", "cancelled", "http", "no_done", "count"]) +def test_raw_audit_rejects_invalid_work(tmp_path, failure): + row = fixture_record() + rows = [row] + if failure == "missing": + rows = [] + elif failure == "duplicate": + rows.append(copy.deepcopy(row)) + elif failure == "cancelled": + row["metadata"]["was_cancelled"] = True + elif failure == "http": + row["status"] = 500 + elif failure == "no_done": + row["responses"][0]["packets"].pop() + elif failure == "count": + packet = row["responses"][0]["packets"][0] + body = json.loads(packet["value"]) + body["usage"]["completion_tokens"] = 1 + packet["value"] = json.dumps(body) + write_artifacts(tmp_path, rows) + with pytest.raises(ValueError): + audit_records(tmp_path, [dict(timestamp=0, output_length=2)]) + + +@pytest.mark.parametrize("error_type", ["InvalidInferenceResultError", "TimeoutError"]) +def test_complete_empty_text_stays_an_explicit_error(tmp_path, error_type): + empty, valid = fixture_record(), fixture_record() + valid["metadata"].update(session_num=1, conversation_id="session_000001", request_start_ns=2000) + packet = empty["responses"][0]["packets"][0] + body = json.loads(packet["value"]) + body["choices"][0]["text"] = "" + packet["value"] = json.dumps(body) + write_artifacts(tmp_path, [empty, valid]) + path = tmp_path / "profile_export_aiperf.json" + summary = json.loads(path.read_text()) + summary.update( + request_count=dict(avg=1), + output_token_count=dict(sum=2), + error_summary=[dict(count=1, error_details=dict(type=error_type, message="No responses with actual content"))], + ) + path.write_text(json.dumps(summary)) + trace = [dict(timestamp=0, output_length=2), dict(timestamp=1, output_length=2)] + if error_type == "TimeoutError": + with pytest.raises(ValueError, match="other request error"): + audit_records(tmp_path, trace) + else: + result = audit_records(tmp_path, trace) + assert result["error_rate"] == 0.5 and result["empty_content_sessions"] == [0] + assert result["output_tokens"] == 4 and result["successful_output_tokens"] == 2 + + +def test_concurrent_dispatch_order_does_not_reorder_source_requests(tmp_path): + first, second = fixture_record(), fixture_record() + second["metadata"].update(session_num=1, conversation_id="session_000001", request_start_ns=2000) + second["payload"]["prompt"] = ["different"] + trace = [dict(timestamp=0, output_length=2), dict(timestamp=1, output_length=2)] + write_artifacts(tmp_path, [first, second]) + reference = audit_records(tmp_path, trace) + first["metadata"]["session_num"], second["metadata"]["session_num"] = 1, 0 + write_artifacts(tmp_path, [second, first]) + assert audit_records(tmp_path, trace) == reference diff --git a/unit_tests/benchmark/qwen38/test_check_cold_streams.py b/unit_tests/benchmark/qwen38/test_check_cold_streams.py new file mode 100644 index 0000000000..abd432cfc1 --- /dev/null +++ b/unit_tests/benchmark/qwen38/test_check_cold_streams.py @@ -0,0 +1,61 @@ +from copy import deepcopy + +import pytest + +from test.benchmark.qwen38.check_cold_streams import workloads_from_reference +from test.benchmark.qwen38.check_generation import REVISION + + +def reference(): + return dict( + status="recorded", + revision=REVISION, + output_length=2, + cases=dict(english=[1, 2], long_3=[1, 2, 3], repeat_long_3=[1, 2, 3]), + results={name: dict(token_ids=[7, 8]) for name in ("english", "long_3", "repeat_long_3")}, + ) + + +def test_qualified_workloads_keep_order_and_exclude_repeat_controls(): + evidence = reference() + workloads = workloads_from_reference(evidence, evidence["cases"]) + assert [(work.name, work.prompt_ids, work.expected_ids) for work in workloads] == [ + ("english", (1, 2), (7, 8)), + ("long_3", (1, 2, 3), (7, 8)), + ] + + +@pytest.mark.parametrize( + "field,value", [("status", "failed"), ("revision", "other"), ("output_length", True), ("output_length", 0)] +) +def test_failed_or_unpinned_reference_cannot_authorize_timing(field, value): + evidence = reference() + expected = deepcopy(evidence["cases"]) + evidence[field] = value + with pytest.raises(ValueError): + workloads_from_reference(evidence, expected) + + +@pytest.mark.parametrize( + "change", ["prompt", "missing_case", "extra_case", "truncated", "invalid_id", "repeat_prompt", "repeat_output"] +) +def test_incomplete_or_edited_evidence_cannot_authorize_timing(change): + evidence = reference() + expected = deepcopy(evidence["cases"]) + if change == "prompt": + evidence["cases"]["english"] = [99, 2] + elif change == "missing_case": + del evidence["results"]["long_3"] + elif change == "extra_case": + evidence["results"]["other"] = dict(token_ids=[7, 8]) + elif change == "truncated": + evidence["results"]["long_3"]["token_ids"] = [7] + elif change == "invalid_id": + evidence["results"]["long_3"]["token_ids"] = [7, True] + elif change == "repeat_prompt": + expected["repeat_long_3"] = [2, 2, 3] + evidence["cases"]["repeat_long_3"] = list(expected["repeat_long_3"]) + elif change == "repeat_output": + evidence["results"]["repeat_long_3"]["token_ids"] = [7, 9] + with pytest.raises(ValueError): + workloads_from_reference(evidence, expected) diff --git a/unit_tests/benchmark/qwen38/test_check_cpu_cache.py b/unit_tests/benchmark/qwen38/test_check_cpu_cache.py new file mode 100644 index 0000000000..cff922e327 --- /dev/null +++ b/unit_tests/benchmark/qwen38/test_check_cpu_cache.py @@ -0,0 +1,66 @@ +import copy +import importlib +from pathlib import Path + +import pytest + + +@pytest.fixture +def checker(monkeypatch): + monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[3] / "test/benchmark/qwen38")) + return importlib.import_module("check_cpu_cache") + + +def events(): + result = [] + for rank in range(4): + common = dict(rank=rank, first_token=1000, request_id=80) + result.extend( + [ + dict(common, event="cpu_offload_queued", monotonic=1, pages=10), + dict( + common, + event="cpu_extra_load_queued", + monotonic=4.5, + lexical_layers=[1], + slot_fields=["qsa_raw_keys", "qsa_pooled_keys", "qsa_key_cos", "qsa_key_sin"], + ), + dict(common, event="cpu_load_queued", monotonic=5, tokens=2368, kv_length=2368, pinned=True), + ] + ) + return result + + +def test_complete_target_cpu_operations_required_on_every_rank(checker): + proof = checker.validate_cpu_restore(events(), 4, 1000, 0.5, (4, 6)) + assert proof == [dict(rank=rank, offloads=1, restores=1) for rank in range(4)] + + +@pytest.mark.parametrize("kind", ["cpu_offload_queued", "cpu_load_queued", "cpu_extra_load_queued"]) +def test_partial_transfer_evidence_cannot_pass(checker, kind): + incomplete = [event for event in events() if not (event["rank"] == 3 and event["event"] == kind)] + with pytest.raises(ValueError, match="TP rank 3"): + checker.validate_cpu_restore(incomplete, 4, 1000, 0.5, (4, 6)) + + +@pytest.mark.parametrize( + "kind, field, value", + [ + ("cpu_load_queued", "pinned", False), + ("cpu_load_queued", "tokens", 0), + ("cpu_load_queued", "kv_length", 64), + ("cpu_load_queued", "monotonic", 3), + ("cpu_load_queued", "request_id", None), + ("cpu_extra_load_queued", "request_id", 81), + ("cpu_extra_load_queued", "slot_fields", ["qsa_raw_keys"]), + ("cpu_extra_load_queued", "lexical_layers", []), + ("cpu_offload_queued", "first_token", 1001), + ], +) +def test_stale_wrong_request_or_incomplete_model_fields_cannot_pass(checker, kind, field, value): + altered = copy.deepcopy(events()) + for event in altered: + if event["rank"] == 0 and event["event"] == kind: + event[field] = value + with pytest.raises(ValueError, match="TP rank 0"): + checker.validate_cpu_restore(altered, 4, 1000, 0.5, (4, 6)) diff --git a/unit_tests/benchmark/qwen38/test_check_generation.py b/unit_tests/benchmark/qwen38/test_check_generation.py new file mode 100644 index 0000000000..7e9f38ce5c --- /dev/null +++ b/unit_tests/benchmark/qwen38/test_check_generation.py @@ -0,0 +1,95 @@ +import hashlib +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from test.benchmark.qwen38 import check_generation + + +@pytest.fixture +def tokenizer_root(tmp_path, monkeypatch): + payload = b"pinned tokenizer fixture" + (tmp_path / "tokenizer.json").write_bytes(payload) + monkeypatch.setattr(check_generation, "TOKENIZER_SHA256", hashlib.sha256(payload).hexdigest()) + + class Tokenizer: + def encode(self, prompt, **kwargs): + return SimpleNamespace(ids=[7, 8, 9]) + + monkeypatch.setattr(check_generation.Tokenizer, "from_file", lambda path: Tokenizer()) + return tmp_path + + +def test_default_workload_preserves_prior_long_cases(tokenizer_root): + cases = check_generation.cases_for(tokenizer_root) + assert list(cases) == [ + "english", + "chinese", + "code_unicode", + "long_2049", + "repeat_long_2049", + "long_4063", + "repeat_long_4063", + ] + assert len(cases["long_2049"]) == 2049 + assert len(cases["long_4063"]) == 4063 + + +def test_custom_lengths_keep_exact_repeated_inputs_without_aliasing(tokenizer_root): + cases = check_generation.cases_for(tokenizer_root, (16, 31)) + assert len(cases["long_16"]) == 16 and len(cases["long_31"]) == 31 + assert cases["long_16"] == cases["repeat_long_16"] + assert cases["long_31"] == cases["repeat_long_31"] + cases["repeat_long_16"][0] = 100 + assert cases["long_16"][0] == 7 + assert "long_2049" not in cases + + +@pytest.mark.parametrize("lengths", [(), (0,), (-1,), (True,), (2.0,), (16, 16)]) +def test_invalid_lengths_fail_before_loading_tokenizer(lengths): + with pytest.raises(ValueError, match="Long prompt lengths"): + check_generation.cases_for(Path("missing-tokenizer"), lengths) + + +def test_custom_lengths_still_require_pinned_tokenizer(tokenizer_root): + (tokenizer_root / "tokenizer.json").write_bytes(b"changed") + with pytest.raises(ValueError, match="pinned checkpoint"): + check_generation.cases_for(tokenizer_root, (31,)) + + +@pytest.mark.parametrize("framework", ["lightllm", "vllm", "sglang"]) +@pytest.mark.parametrize("cached", [0, 1, None, False]) +def test_cold_generation_isolates_requests_and_requires_explicit_zero_cache(framework, cached, tmp_path): + body = ( + dict(prompt_tokens=2, tokens=[dict(id=7, prompt_cache_len=cached), dict(id=8, prompt_cache_len=cached)]) + if framework == "lightllm" + else dict( + choices=[dict(token_ids=[7, 8])], + usage=dict(prompt_tokens=2, prompt_tokens_details=dict(cached_tokens=cached)), + ) + ) + sent = [] + serialized = json.dumps(body) + + def post(url, json, timeout): + sent.append(json) + return SimpleNamespace(text=serialized, json=lambda: body, raise_for_status=lambda: None) + + client = SimpleNamespace(post=post) + for index in range(2): + raw = tmp_path / f"response-{index}.json" + arguments = (client, "http://service", framework, "model", [1, 2], 2, 10) + if type(cached) is int and cached == 0: + result = check_generation.generate(*arguments, response_path=raw, cold_cache=True) + assert result["token_ids"] == [7, 8] + else: + with pytest.raises(ValueError, match="zero cached"): + check_generation.generate(*arguments, response_path=raw, cold_cache=True) + assert json.loads(raw.read_text()) == body + if framework == "lightllm": + assert all(payload["parameters"]["disable_prompt_cache"] is True for payload in sent) + else: + assert all(isinstance(payload["cache_salt"], str) and payload["cache_salt"] for payload in sent) + assert sent[0]["cache_salt"] != sent[1]["cache_salt"] diff --git a/unit_tests/benchmark/qwen38/test_check_math_quality.py b/unit_tests/benchmark/qwen38/test_check_math_quality.py new file mode 100644 index 0000000000..5cb9681f71 --- /dev/null +++ b/unit_tests/benchmark/qwen38/test_check_math_quality.py @@ -0,0 +1,133 @@ +from copy import deepcopy +import hashlib +import json +import sys +from types import SimpleNamespace + +import pytest + +from test.benchmark.qwen38 import check_math_quality as check +from test.benchmark.qwen38.math_quality import MathExample + + +@pytest.fixture +def workload(tmp_path, monkeypatch): + files = {} + for name in ("tokenizer.json", "tokenizer_config.json", "chat_template.jinja"): + payload = name.encode() + (tmp_path / name).write_bytes(payload) + files[name] = hashlib.sha256(payload).hexdigest() + monkeypatch.setattr(check, "TOKENIZER_SHA256", files["tokenizer.json"]) + examples = (MathExample(0, "What is two plus two?", "4"),) + monkeypatch.setattr(check, "load_examples", lambda path: examples) + + class Tokenizer: + def apply_chat_template(self, *args, **kwargs): + return [1, 2] + + tokenizer = Tokenizer() + manifest = dict( + status="prepared", + model_revision=check.REVISION, + dataset=dict(revision=check.DATASET_REVISION, sha256=check.DATASET_SHA256, split="test", records=1), + tokenizer_files=files, + decoding=dict(temperature=0, ignore_eos=False, enable_thinking=False), + sample_count=1, + sample_seed=38, + max_output_tokens=3, + context_limit=5, + cases=check.prepare_cases(examples, tokenizer, 3, 5), + ) + return manifest, tmp_path, tokenizer + + +def test_verifies_reconstructed_prompts_and_gold_labels(workload): + manifest, root, tokenizer = workload + assert check.verified_workload(manifest, root, root / "data", tokenizer) == manifest["cases"] + for field, changed in [("expected_number", "999"), ("prompt_ids", [1, 3]), ("dataset_index", 1)]: + forged = deepcopy(manifest) + forged["cases"][0][field] = changed + with pytest.raises(ValueError, match="prompts or gold"): + check.verified_workload(forged, root, root / "data", tokenizer) + + +def test_rejects_changed_tokenizer_and_wrong_protocol(workload): + manifest, root, tokenizer = workload + forged = deepcopy(manifest) + forged["decoding"]["ignore_eos"] = True + with pytest.raises(ValueError, match="greedy non-thinking"): + check.verified_workload(forged, root, root / "data", tokenizer) + (root / "chat_template.jinja").write_bytes(b"changed") + with pytest.raises(ValueError, match="Tokenizer files"): + check.verified_workload(manifest, root, root / "data", tokenizer) + + +@pytest.mark.parametrize("bad_count", [False, True]) +def test_client_keeps_wrong_and_truncated_answers_in_denominator(tmp_path, monkeypatch, bad_count): + metadata = dict( + framework="lightllm", version="fixture", command=["fixture"], hardware="fixture", configuration={"tp": 1} + ) + (tmp_path / "service.json").write_text(json.dumps(metadata)) + (tmp_path / "manifest.json").write_text("{}") + (tmp_path / "generation_config.json").write_text(json.dumps(dict(eos_token_id=[99]))) + cases = [dict(dataset_index=i, prompt_ids=[1, 2], max_output_tokens=3, expected_number="4") for i in range(3)] + + class Tokenizer: + eos_token_id = 99 + + def decode(self, ids, **kwargs): + return {11: "#### 4", 12: "#### 5", 13: "#### 4"}[ids[0]] + + class Session: + index = 0 + + def post(self, url, *, json, timeout): + assert "expected_number" not in json and json["inputs"] == [1, 2] + ids, reason = [([11, 99], "stop"), ([12, 99], "stop"), ([13, 14, 15], "length")][self.index] + self.index += 1 + body = dict( + tokens=[dict(id=token) for token in ids], + prompt_tokens=2, + count_output_tokens=1 if bad_count else len(ids), + finish_reason=reason, + ) + return SimpleNamespace(text="fixture", raise_for_status=lambda: None, json=lambda: body) + + def close(self): + pass + + monkeypatch.setattr(check, "get_tokenizer", lambda path: SimpleNamespace(tokenizer=Tokenizer())) + monkeypatch.setattr(check, "verified_checkpoint", lambda path: None) + monkeypatch.setattr(check, "verified_workload", lambda *args: cases) + monkeypatch.setattr(check.requests, "Session", Session) + monkeypatch.setenv("EXP_ACTIVE", "1") + monkeypatch.setenv("EXP_RUN_DIR", str(tmp_path)) + monkeypatch.setattr( + sys, + "argv", + [ + "check_math_quality.py", + "--framework", + "lightllm", + "--base-url", + "http://fixture", + "--model-directory", + str(tmp_path), + "--service-metadata", + str(tmp_path / "service.json"), + "--manifest", + str(tmp_path / "manifest.json"), + "--dataset", + str(tmp_path / "data"), + ], + ) + if bad_count: + with pytest.raises(ValueError, match="counts"): + check.main() + report = json.loads((tmp_path / "math-quality.json").read_text()) + assert report["status"] == "failed" and "accuracy" not in report + else: + check.main() + report = json.loads((tmp_path / "math-quality.json").read_text()) + assert report["status"] == "recorded" + assert (report["scored"], report["correct"], report["truncated"], report["accuracy"]) == (3, 1, 1, 1 / 3) diff --git a/unit_tests/benchmark/qwen38/test_check_multimodal.py b/unit_tests/benchmark/qwen38/test_check_multimodal.py new file mode 100644 index 0000000000..928a72549c --- /dev/null +++ b/unit_tests/benchmark/qwen38/test_check_multimodal.py @@ -0,0 +1,98 @@ +import importlib +from pathlib import Path +from types import SimpleNamespace + +import pytest + + +@pytest.fixture +def checker(monkeypatch): + monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[3] / "test/benchmark/qwen38")) + return importlib.import_module("check_multimodal") + + +def fixture(): + case = dict(prompt_tokens=90, image_tokens=64, max_new_tokens=64, ignore_eos=False, expected_words=["red"]) + body = dict( + prompt_tokens=90, + count_output_tokens=2, + finish_reason="stop", + tokens=[dict(id=7), dict(id=248046)], + input_usage=dict(input_text_tokens=26, input_image_tokens=64, input_audio_tokens=0), + ) + tokenizer = SimpleNamespace(decode=lambda ids, **_: " ".join({7: "Red", 8: "blue"}[value] for value in ids)) + return case, body, tokenizer + + +def test_image_count_eos_and_semantic_validation(checker): + case, body, tokenizer = fixture() + result = checker.validate_image_result(body, case, [248044, 248046], tokenizer) + assert result["token_ids"] == [7, 248046] and result["answer"] == "Red" + + +@pytest.mark.parametrize( + "field,value", [("input_text_tokens", 27), ("input_image_tokens", 63), ("input_audio_tokens", 1)] +) +def test_incorrect_modality_counts_cannot_pass(checker, field, value): + case, body, tokenizer = fixture() + body["input_usage"][field] = value + with pytest.raises(ValueError, match="accounting"): + checker.validate_image_result(body, case, [248046], tokenizer) + + +def test_wrong_color_and_swapped_image_order_cannot_pass(checker): + case, body, tokenizer = fixture() + body["tokens"][0]["id"] = 8 + with pytest.raises(ValueError, match="color/order"): + checker.validate_image_result(body, case, [248046], tokenizer) + case["expected_words"] = ["red", "blue"] + body["tokens"] = [dict(id=8), dict(id=7), dict(id=248046)] + body["count_output_tokens"] = 3 + with pytest.raises(ValueError, match="color/order"): + checker.validate_image_result(body, case, [248046], tokenizer) + + +def test_pressure_request_has_explicit_fixed_length_contract(checker): + case, body, tokenizer = fixture() + case.update(expected_words=[], max_new_tokens=1, ignore_eos=True) + body.update(tokens=[dict(id=7)], count_output_tokens=1, finish_reason="length") + assert checker.validate_image_result(body, case, [248046], tokenizer)["token_ids"] == [7] + + +def cache_events(): + first = dict(id=1, token_id=100000000, token_num=64, md5="digest_0") + last = dict(id=2, token_id=100000384, token_num=64, md5="digest_0") + return [ + dict(event="alloc", records=[first]), + dict(event="embedded", records=[first]), + dict(event="evicted", record=first), + dict(event="alloc", records=[last]), + dict(event="embedded", records=[last]), + ] + + +def test_reallocation_requires_new_virtual_ids_eviction_and_embedding_completion(checker): + proof = checker.verify_reallocation(cache_events(), "digest") + assert proof["original_evicted"] and proof["both_embedded"] + assert proof["original"]["token_id"] != proof["reallocated"]["token_id"] + + +@pytest.mark.parametrize("missing", ["alloc", "evicted", "embedded"]) +def test_incomplete_cache_evidence_cannot_pass(checker, missing): + events = [item for item in cache_events() if item["event"] != missing] + with pytest.raises(ValueError): + checker.verify_reallocation(events, "digest") + + +def test_cache_reallocation_must_change_virtual_range(checker): + events = cache_events() + events[3]["records"][0]["token_id"] = events[0]["records"][0]["token_id"] + with pytest.raises(ValueError, match="new virtual IDs"): + checker.verify_reallocation(events, "digest") + + +def test_fixture_hashes_distinguish_cache_pressure_images(checker, tmp_path): + payload, images = checker.image_fixtures(tmp_path, 5) + assert len(payload) == len(images) == 8 + assert len({image["sha256"] for image in images.values()}) == 8 + assert images["rect_blue"]["width"] == 384 and images["rect_blue"]["height"] == 128 diff --git a/unit_tests/benchmark/qwen38/test_check_prefix_eviction.py b/unit_tests/benchmark/qwen38/test_check_prefix_eviction.py new file mode 100644 index 0000000000..e5b4e80f53 --- /dev/null +++ b/unit_tests/benchmark/qwen38/test_check_prefix_eviction.py @@ -0,0 +1,86 @@ +import copy +import importlib +from pathlib import Path + +import pytest + + +@pytest.fixture +def checker(monkeypatch): + monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[3] / "test/benchmark/qwen38")) + return importlib.import_module("check_prefix_eviction") + + +def events(): + result = [] + for rank in range(4): + result.extend( + [ + dict(event="state_freed", rank=rank, monotonic=2.0, pool=10, indexes=[3]), + dict(event="node_evicted", rank=rank, monotonic=2.1, first_token=1000, big_pool=10, big_state=3), + dict(event="state_freed", rank=rank, monotonic=3.0, pool=20, indexes=[4]), + dict( + event="small_state_evicted", + rank=rank, + monotonic=3.1, + first_token=1000, + small_pool=20, + small_state=4, + ), + ] + ) + return result + + +def test_actual_target_big_and_small_state_releases_required_on_every_rank(checker): + proof = checker.validate_eviction(events(), 4, 1000, 1, 4) + assert {item["rank"] for item in proof} == {0, 1, 2, 3} + + +@pytest.mark.parametrize("missing", ["node_evicted", "small_state_evicted", "state_freed"]) +def test_partial_eviction_evidence_cannot_pass(checker, missing): + incomplete = [event for event in events() if event["event"] != missing] + with pytest.raises(ValueError, match="No observed target"): + checker.validate_eviction(incomplete, 4, 1000, 1, 4) + + +def test_one_rank_trace_cannot_prove_four_rank_eviction(checker): + with pytest.raises(ValueError, match="rank 1"): + checker.validate_eviction([event for event in events() if event["rank"] == 0], 4, 1000, 1, 4) + + +@pytest.mark.parametrize("change", ["other_target", "other_pool", "other_index", "late_free", "outside_interval"]) +def test_unrelated_or_out_of_order_events_cannot_prove_target_eviction(checker, change): + changed = copy.deepcopy(events()) + for event in changed: + if change == "other_target" and "first_token" in event: + event["first_token"] = 1001 + elif change == "other_pool" and event["event"] == "state_freed": + event["pool"] += 50 + elif change == "other_index" and event["event"] == "state_freed": + event["indexes"] = [900] + elif change == "late_free" and event["event"] == "state_freed": + event["monotonic"] = 3.9 + elif change == "outside_interval": + event["monotonic"] += 10 + with pytest.raises(ValueError, match="No observed target"): + checker.validate_eviction(changed, 4, 1000, 1, 4) + + +@pytest.mark.parametrize("value", [-1, True, 2305, 4.0]) +def test_invalid_cache_hit_accounting_rejected(checker, value): + result = dict(response=dict(prompt_tokens=2305, tokens=[dict(prompt_cache_len=value)])) + with pytest.raises(ValueError, match="cache hit"): + checker.cache_hit(result) + + +def test_pressure_prefixes_diverge_before_first_hash_page(checker, monkeypatch): + monkeypatch.setattr(checker, "cases_for", lambda _: {"long_4063": [9] * 4063}) + cases = checker.eviction_cases(Path("unused")) + assert len({tuple(case["inputs"][:64]) for case in cases.values()}) == 33 + assert (len(cases["target"]["inputs"]) - 1) % 256 == 64 + assert all( + len(case["inputs"]) == (65 if name.startswith("small_pressure_") else 2369) + and len(case["inputs"]) + case["max_new_tokens"] <= 4096 + for name, case in cases.items() + ) diff --git a/unit_tests/benchmark/qwen38/test_check_scheduler.py b/unit_tests/benchmark/qwen38/test_check_scheduler.py new file mode 100644 index 0000000000..3ff7319ec7 --- /dev/null +++ b/unit_tests/benchmark/qwen38/test_check_scheduler.py @@ -0,0 +1,121 @@ +import importlib +import json +from pathlib import Path +import threading + +import pytest +import requests + + +@pytest.fixture +def client_module(monkeypatch): + root = Path(__file__).resolve().parents[3] + monkeypatch.syspath_prepend(str(root / "test/benchmark/qwen38")) + return importlib.import_module("check_scheduler") + + +def response_for(monkeypatch, module, body, http_error=False): + class Response: + text = json.dumps(body) + + def raise_for_status(self): + if http_error: + raise requests.HTTPError("503") + + def json(self): + return body + + class Session: + def __enter__(self): + return self + + def __exit__(self, *args): + pass + + def post(self, url, json, timeout): + assert json["parameters"]["do_sample"] is False + assert json["parameters"]["return_details"] is True + assert json["parameters"]["print_eos_token"] is True + return Response() + + monkeypatch.setattr(module.requests, "Session", Session) + + +@pytest.mark.parametrize( + "ids,ignore_eos,reason", [([3, 4], True, "length"), ([248046], False, "stop"), ([3, 248044], False, "stop")] +) +def test_valid_termination_preserves_raw_response(client_module, monkeypatch, tmp_path, ids, ignore_eos, reason): + body = { + "tokens": [{"id": token} for token in ids], + "prompt_tokens": 3, + "count_output_tokens": len(ids), + "finish_reason": reason, + } + response_for(monkeypatch, client_module, body) + case = {"inputs": [7, 8, 9], "max_new_tokens": 2, "ignore_eos": ignore_eos} + path = tmp_path / "response.txt" + result = client_module.generate("http://localhost", case, [248044, 248046], 10, path) + assert result["token_ids"] == ids and json.loads(path.read_text()) == body + + +@pytest.mark.parametrize( + "ids,ignore_eos,reason,prompt,count", + [ + ([], True, "length", 3, 0), + ([True, 4], True, "length", 3, 2), + ([-1, 4], True, "length", 3, 2), + ([3, 4, 5], True, "length", 3, 3), + ([3], True, "length", 3, 1), + ([3, 4], True, "stop", 3, 2), + ([3, 4], True, "length", 2, 2), + ([3, 4], True, "length", 3, 1), + ([3], False, "length", 3, 1), + ([248046, 3], False, "stop", 3, 2), + ([3, 248046], False, "length", 3, 2), + ], +) +def test_invalid_response_cannot_pass(client_module, monkeypatch, tmp_path, ids, ignore_eos, reason, prompt, count): + body = { + "tokens": [{"id": token} for token in ids], + "prompt_tokens": prompt, + "count_output_tokens": count, + "finish_reason": reason, + } + response_for(monkeypatch, client_module, body) + case = {"inputs": [7, 8, 9], "max_new_tokens": 2, "ignore_eos": ignore_eos} + path = tmp_path / "response.txt" + with pytest.raises(ValueError): + client_module.generate("http://localhost", case, [248044, 248046], 10, path) + assert json.loads(path.read_text()) == body + + +def test_http_error_is_saved(client_module, monkeypatch, tmp_path): + body = {"error": "service unavailable"} + response_for(monkeypatch, client_module, body, http_error=True) + path = tmp_path / "response.txt" + with pytest.raises(requests.HTTPError): + client_module.generate("http://localhost", {"inputs": [7]}, [248046], 10, path) + assert json.loads(path.read_text()) == body + + +def test_monitor_records_server_batch_gauges(client_module, monkeypatch, tmp_path): + stop = threading.Event() + payload = '# a comment\nlightllm_num_running_reqs{model_name="test"} 4.0\nlightllm_batch_current_size 3.0\n' + + class Session: + def __enter__(self): + return self + + def __exit__(self, *args): + pass + + def get(self, *args, **kwargs): + stop.set() + return type("Response", (), {"text": payload, "raise_for_status": lambda self: None})() + + monkeypatch.setattr(client_module.requests, "Session", Session) + summary = {"maximum_router_batch": 0, "samples": 0} + path = tmp_path / "metrics.txt" + client_module.monitor("http://localhost", stop, summary, path) + assert summary == {"maximum_router_batch": 4.0, "samples": 1} + assert path.read_text() == payload diff --git a/unit_tests/benchmark/qwen38/test_generation_observations.py b/unit_tests/benchmark/qwen38/test_generation_observations.py new file mode 100644 index 0000000000..ca27acc960 --- /dev/null +++ b/unit_tests/benchmark/qwen38/test_generation_observations.py @@ -0,0 +1,62 @@ +import copy + +import pytest + +from test.benchmark.qwen38.check_cold_streams import workloads_from_reference +from test.benchmark.qwen38.check_generation import REVISION +from test.benchmark.qwen38.generation_observations import workloads_from_observations + + +def observations(): + cases = {"long_16": [1, 2], "repeat_long_16": [1, 2]} + results = {} + for name, ids in [("long_16", [7, 8]), ("repeat_long_16", [7, 9])]: + results[name] = dict( + token_ids=ids, + response=dict( + choices=[dict(token_ids=ids)], + usage=dict(prompt_tokens=2, completion_tokens=2, prompt_tokens_details=dict(cached_tokens=0)), + ), + ) + return dict( + status="observations_recorded", + service=dict(framework="vllm"), + revision=REVISION, + cases=cases, + reference=None, + cache_policy="cold", + results=results, + output_length=2, + repeat_consistency=dict(status="failed", mismatches=[16]), + ) + + +def test_observations_retain_repeat_failure_and_cannot_qualify_a_reference(): + report = observations() + before = copy.deepcopy(report) + workloads, quality = workloads_from_observations(report, report["cases"]) + assert report == before + assert len(workloads) == 1 and workloads[0].expected_ids == (7, 8) + assert quality == dict(status="failed", mismatches=[16]) + with pytest.raises(ValueError, match="qualified"): + workloads_from_reference(report, report["cases"]) + + +@pytest.mark.parametrize("change", ["native", "cached", "count", "edited_ids", "false_pass", "incomplete"]) +def test_observations_reject_invalid_protocol_or_hidden_quality_failure(change): + report = observations() + result = report["results"]["long_16"] + if change == "native": + report["service"]["framework"] = "lightllm" + elif change == "cached": + result["response"]["usage"]["prompt_tokens_details"]["cached_tokens"] = 1 + elif change == "count": + result["response"]["usage"]["completion_tokens"] = 1 + elif change == "edited_ids": + result["response"]["choices"][0]["token_ids"] = [7, 10] + elif change == "false_pass": + report["repeat_consistency"] = dict(status="passed", mismatches=[]) + elif change == "incomplete": + del report["results"]["repeat_long_16"] + with pytest.raises(ValueError): + workloads_from_observations(report, report["cases"]) diff --git a/unit_tests/benchmark/qwen38/test_math_quality.py b/unit_tests/benchmark/qwen38/test_math_quality.py new file mode 100644 index 0000000000..1bf77ab80b --- /dev/null +++ b/unit_tests/benchmark/qwen38/test_math_quality.py @@ -0,0 +1,145 @@ +import hashlib +import json +from copy import deepcopy + +import pytest + +from test.benchmark.qwen38 import math_quality as math + + +@pytest.mark.parametrize( + "value, expected", + [("1,234.00", "1234"), ("+001.2500", "1.25"), ("-0.00", "0"), ("-42", "-42"), ("0.0010", "0.001")], +) +def test_decimal_normalization_is_exact(value, expected): + assert math.normalize_number(value) == expected + large = "123456789012345678901234567890123456789" + assert math.normalize_number(large) == large + + +@pytest.mark.parametrize("value", ["1,23", "1.2.3", "1/2", "1e3", "NaN", "inf", "", None]) +def test_invalid_plain_numbers_are_not_guessed(value): + assert math.normalize_number(value) is None + + +def test_score_uses_final_numeric_answer_only(): + assert math.score_answer("Work: 10 - 4 = 6.\n#### 6.00\n", "6")["correct"] + assert not math.score_answer("Work: 6 + 1 = 7.\n#### 7", "6")["correct"] + with pytest.raises(ValueError, match="Expected answer"): + math.score_answer("#### 6", "unknown") + + +@pytest.mark.parametrize("text", ["The answer is 6", "#### 6\nMore work", "#### 7\n#### 6", "#### $6", "#### NaN"]) +def test_missing_ambiguous_or_nonfinal_answers_are_invalid(text): + assert math.extract_final_number(text) is None + + +def test_loader_pins_bytes_and_rejects_duplicate_questions(tmp_path, monkeypatch): + path = tmp_path / "test.jsonl" + rows = [ + dict(question="Question one", answer="A calculation.\n#### 4"), + dict(question="Question two", answer="#### 5"), + ] + payload = "\n".join(json.dumps(row) for row in rows).encode() + path.write_bytes(payload) + monkeypatch.setattr(math, "DATASET_SHA256", hashlib.sha256(payload).hexdigest()) + examples = math.load_examples(path) + assert [(row.index, row.expected) for row in examples] == [(0, "4"), (1, "5")] + path.write_bytes(payload + b"\n") + with pytest.raises(ValueError, match="pinned source"): + math.load_examples(path) + duplicate = "\n".join(json.dumps(rows[0]) for _ in range(2)).encode() + path.write_bytes(duplicate) + monkeypatch.setattr(math, "DATASET_SHA256", hashlib.sha256(duplicate).hexdigest()) + with pytest.raises(ValueError, match="distinct questions"): + math.load_examples(path) + + +def test_sampling_is_reproducible_and_without_replacement(): + examples = tuple(math.MathExample(i, f"Question {i}", "4") for i in range(20)) + selected = math.select_examples(examples, 8, 38) + assert selected == math.select_examples(examples, 8, 38) + assert len({row.index for row in selected}) == 8 + + +@pytest.mark.parametrize("limit, seed", [(0, 1), (2, 1), (True, 1), (1, True), (1, 1.0)]) +def test_invalid_selection_is_rejected(limit, seed): + with pytest.raises(ValueError): + math.select_examples((math.MathExample(0, "Question", "4"),), limit, seed) + + +def test_prompts_do_not_include_gold_solutions_and_preserve_token_budget(): + class Tokenizer: + messages = [] + + def apply_chat_template(self, messages, **kwargs): + self.messages.append(messages) + assert kwargs == dict(tokenize=True, add_generation_prompt=True, enable_thinking=False) + return [11, 12] + + tokenizer = Tokenizer() + first = math.prepare_cases([math.MathExample(0, "What is two plus two?", "4")], tokenizer, 3, 5) + second = math.prepare_cases([math.MathExample(0, "What is two plus two?", "999")], tokenizer, 3, 5) + assert tokenizer.messages[0] == tokenizer.messages[1] + assert first[0]["prompt_ids"] == second[0]["prompt_ids"] == [11, 12] + assert first[0]["expected_number"] == "4" and second[0]["expected_number"] == "999" + with pytest.raises(ValueError, match="serving context"): + math.prepare_cases([math.MathExample(0, "Question", "4")], tokenizer, 4, 5) + + +@pytest.mark.parametrize("ids", [[], [True], [-1], [1.0], "tokens"]) +def test_invalid_tokenizer_payloads_are_rejected(ids): + class Tokenizer: + def apply_chat_template(self, *args, **kwargs): + return ids + + with pytest.raises(ValueError, match="actual nonnegative"): + math.prepare_cases([math.MathExample(0, "Question", "4")], Tokenizer()) + + +def scored_report(): + return dict( + status="recorded", + manifest_sha256="a" * 64, + scored=2, + results=[ + dict( + dataset_index=1, token_ids=[11, 99], finish_reason="stop", truncated=False, correct=True, predicted="4" + ), + dict( + dataset_index=2, + token_ids=[12, 13], + finish_reason="length", + truncated=True, + correct=False, + predicted=None, + ), + ], + ) + + +def test_output_comparison_keeps_wrong_and_truncated_cases(): + reference = scored_report() + assert math.compare_outputs(reference, deepcopy(reference)) == dict(status="passed", cases=2, mismatches=[]) + for field, value in [("token_ids", [12, 14]), ("finish_reason", "stop"), ("truncated", False), ("correct", True)]: + actual = deepcopy(reference) + actual["results"][1][field] = value + assert math.compare_outputs(reference, actual) == dict(status="failed", cases=2, mismatches=[2]) + + +@pytest.mark.parametrize("field,value", [("status", "failed"), ("manifest_sha256", "b" * 64), ("scored", 1)]) +def test_output_comparison_rejects_incomplete_or_different_workloads(field, value): + reference = scored_report() + actual = deepcopy(reference) + actual[field] = value + with pytest.raises(ValueError): + math.compare_outputs(reference, actual) + + +def test_output_comparison_rejects_duplicate_or_different_cases(): + reference = scored_report() + actual = deepcopy(reference) + for index in (1, 3): + actual["results"][1]["dataset_index"] = index + with pytest.raises(ValueError): + math.compare_outputs(reference, actual) diff --git a/unit_tests/benchmark/qwen38/test_measure_fixed_work.py b/unit_tests/benchmark/qwen38/test_measure_fixed_work.py new file mode 100644 index 0000000000..95d6cc3ca9 --- /dev/null +++ b/unit_tests/benchmark/qwen38/test_measure_fixed_work.py @@ -0,0 +1,55 @@ +from copy import deepcopy + +import pytest + +from test.benchmark.qwen38.measure_fixed_work import summarize_fixed_work +from test.benchmark.service.stream_cohort import StreamingWorkload, run_cohort +from unit_tests.benchmark.service.test_stream_cohort import observation + + +def fixture(): + work = StreamingWorkload("prompt", [12, 13], [20, 21]) + result = observation(1, [2, 3], 4, ids=[22, 23]) + result["metrics"]["cached_input_tokens"] = 0 + return work, result + + +def test_id_mismatch_keeps_exact_gate_failed_but_measures_equal_work(tmp_path): + work, result = fixture() + cohort = run_cohort([work], 1, lambda *_: result, tmp_path) + before = deepcopy(cohort) + metrics = summarize_fixed_work(cohort, [work]) + assert metrics["output_tokens_per_s"] == pytest.approx(2 / 3) + assert cohort == before + assert cohort["status"] == "failed" and cohort["metrics"] is None + + +@pytest.mark.parametrize("cached", [None, True, -1, 1]) +def test_unknown_or_nonzero_cache_cannot_produce_timing(tmp_path, cached): + work, result = fixture() + result["metrics"]["cached_input_tokens"] = cached + cohort = run_cohort([work], 1, lambda *_: result, tmp_path) + with pytest.raises(ValueError): + summarize_fixed_work(cohort, [work]) + + +@pytest.mark.parametrize("fault", ["transport", "input_length", "output_length", "timeline", "missing"]) +def test_invalid_work_never_produces_timing(tmp_path, fault): + work, result = fixture() + if fault == "input_length": + result["metrics"]["input_tokens"] = 3 + elif fault == "output_length": + result["token_ids"] = [22] + elif fault == "timeline": + result["timeline"]["ended"] = 0 + + def request(*_): + if fault == "transport": + raise TimeoutError("timed out") + return result + + cohort = run_cohort([work], 1, request, tmp_path) + if fault == "missing": + cohort["records"] = [] + with pytest.raises(ValueError): + summarize_fixed_work(cohort, [work]) diff --git a/unit_tests/benchmark/qwen38/test_retrieval_quality.py b/unit_tests/benchmark/qwen38/test_retrieval_quality.py new file mode 100644 index 0000000000..9311c87d8f --- /dev/null +++ b/unit_tests/benchmark/qwen38/test_retrieval_quality.py @@ -0,0 +1,62 @@ +import hashlib +import json + +import pytest + +from test.benchmark.qwen38 import retrieval_quality as quality + +EXPECTED = dict(ALPHA="00123", BETA="456", GAMMA="789") + + +def test_exact_digit_strings_and_all_records_are_required(): + assert quality.score_records(" ALPHA = 00123;\n BETA=456; GAMMA=789\n", EXPECTED)["correct"] + result = quality.score_records("ALPHA=123; BETA=456; GAMMA=789", EXPECTED) + assert not result["correct"] and result["records"] == dict(ALPHA=False, BETA=True, GAMMA=True) + + +@pytest.mark.parametrize( + "text", + [ + None, + "", + "ALPHA=00123; BETA=456", + "ALPHA=00123; BETA=456; GAMMA=789.\n", + "ALPHA=00123; BETA=456; GAMMA=789; ALPHA=123", + "Here: ALPHA=00123; BETA=456; GAMMA=789", + ], +) +def test_incomplete_or_ambiguous_answers_cannot_pass(text): + result = quality.score_records(text, EXPECTED) + assert result["predicted"] is None and not result["correct"] + + +def test_manifest_and_actual_tokenizer_are_pinned(tmp_path, monkeypatch): + (tmp_path / "config.json").write_text(json.dumps(dict(text_config=dict(max_position_embeddings=10)))) + files = {} + for name in ("tokenizer.json", "tokenizer_config.json", "chat_template.jinja"): + payload = name.encode() + (tmp_path / name).write_bytes(payload) + files[name] = hashlib.sha256(payload).hexdigest() + manifest = dict( + status="prepared", + model_revision=quality.REVISION, + context_limit=10, + decoding=dict(temperature=0, ignore_eos=False, enable_thinking=False), + tokenizer_files=files, + cases=[dict(prompt_ids=[1, 2], max_output_tokens=8, expected_records=EXPECTED)], + ) + path = tmp_path / "manifest.json" + path.write_text(json.dumps(manifest)) + monkeypatch.setattr( + quality, "MANIFEST_SHA256", hashlib.sha256(json.dumps(manifest, sort_keys=True).encode()).hexdigest() + ) + assert quality.load_workload(path, tmp_path) == manifest + manifest["cases"][0]["expected_records"]["ALPHA"] = "999" + path.write_text(json.dumps(manifest)) + with pytest.raises(ValueError, match="pinned preparation"): + quality.load_workload(path, tmp_path) + manifest["cases"][0]["expected_records"]["ALPHA"] = "00123" + path.write_text(json.dumps(manifest)) + (tmp_path / "chat_template.jinja").write_bytes(b"changed") + with pytest.raises(ValueError, match="Tokenizer"): + quality.load_workload(path, tmp_path) diff --git a/unit_tests/benchmark/qwen38/test_score_reference.py b/unit_tests/benchmark/qwen38/test_score_reference.py new file mode 100644 index 0000000000..d265be3b48 --- /dev/null +++ b/unit_tests/benchmark/qwen38/test_score_reference.py @@ -0,0 +1,49 @@ +import importlib +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + + +@pytest.fixture +def scorer(monkeypatch): + monkeypatch.syspath_prepend(str(Path(__file__).resolve().parents[3] / "test/benchmark/qwen38")) + return importlib.import_module("score_reference") + + +class RecordingModel: + def __init__(self): + self.calls = [] + + def __call__(self, **kwargs): + prior = kwargs["past_key_values"] or () + ids = tuple(kwargs["input_ids"][0].tolist()) + positions = kwargs["position_ids"][0].tolist() + assert positions == list(range(len(prior), len(prior) + len(ids))) + assert kwargs["attention_mask"].tolist() == [[1] * (len(prior) + len(ids))] + assert kwargs["use_cache"] and kwargs["logits_to_keep"] == 1 + self.calls.append(kwargs) + history = prior + ids + return SimpleNamespace(past_key_values=history, logits=torch.tensor([[sum(history)]])) + + +@pytest.mark.parametrize("chunk_size,expected_calls", [(None, 1), (1, 9), (4, 3), (16, 1)]) +def test_prefill_chunking_preserves_history_and_next_decode_position(scorer, chunk_size, expected_calls): + model = RecordingModel() + inputs = torch.arange(9).unsqueeze(0) + output = scorer.forward_reference(model, inputs, None, 9, chunk_size) + assert len(model.calls) == expected_calls + assert output.past_key_values == tuple(range(9)) + assert output.logits.item() == 36 + output = scorer.forward_reference(model, torch.tensor([[17]]), output.past_key_values, 10, chunk_size) + assert output.past_key_values == (*range(9), 17) + assert output.logits.item() == 53 + + +@pytest.mark.parametrize("count,length,chunk_size", [(0, 0, None), (2, 1, 1), (2, 2, 0)]) +def test_invalid_forward_lengths_are_rejected_before_call(scorer, count, length, chunk_size): + model = RecordingModel() + with pytest.raises(ValueError, match="valid lengths"): + scorer.forward_reference(model, torch.zeros((1, count)), None, length, chunk_size) + assert not model.calls diff --git a/unit_tests/benchmark/service/test_cache_capacity.py b/unit_tests/benchmark/service/test_cache_capacity.py new file mode 100644 index 0000000000..01bc5cceab --- /dev/null +++ b/unit_tests/benchmark/service/test_cache_capacity.py @@ -0,0 +1,62 @@ +import json + +import pytest + +from test.benchmark.service.cache_capacity import plan_vllm_capacity + + +def metrics(**overrides): + fields = dict( + engine="0", + num_gpu_blocks="12345", + kv_cache_max_concurrency=str(12345 / 258), + kv_cache_size_tokens=str(int(12345 / 258 * 4096)), + block_size="16", + mamba_block_size="256", + mamba_cache_mode="align", + mamba_ssm_cache_dtype="float32", + ) + fields.update(overrides) + labels = ",".join(f"{key}={json.dumps(value)}" for key, value in fields.items()) + return f"# TYPE vllm:cache_config_info gauge\nvllm:cache_config_info{{{labels}}} 1\n" + + +def test_hybrid_capacity_does_not_use_attention_block_size_alone(): + plan = plan_vllm_capacity(metrics(), 4096, 16384) + assert plan["blocks_per_full_request"] == 258 + assert plan["num_gpu_blocks_override"] == 1032 + assert plan["cache_geometry"]["mamba_cache_mode"] == "align" + + +def test_unverified_pool_growth_is_not_planned(): + with pytest.raises(ValueError, match="only shrinks"): + plan_vllm_capacity(metrics(), 4096, 196608) + + +@pytest.mark.parametrize("text", ["", metrics() + metrics(engine="1")]) +def test_missing_or_ambiguous_engines_cannot_yield_capacity_plan(text): + with pytest.raises(ValueError, match="exactly one"): + plan_vllm_capacity(text, 4096, 16384) + + +@pytest.mark.parametrize( + "fields", + [ + dict(num_gpu_blocks="None"), + dict(num_gpu_blocks="0"), + dict(kv_cache_max_concurrency="nan"), + dict(kv_cache_max_concurrency="inf"), + dict(kv_cache_size_tokens="9"), + dict(num_gpu_blocks="12346"), + dict(kv_cache_max_concurrency="0"), + ], +) +def test_uninitialized_or_inconsistent_runtime_metrics_are_rejected(fields): + with pytest.raises(ValueError): + plan_vllm_capacity(metrics(**fields), 4096, 16384) + + +@pytest.mark.parametrize("length,budget", [(True, 16384), (0, 16384), (4096, -1), (4096, 16385)]) +def test_invalid_or_partial_context_budgets_are_rejected(length, budget): + with pytest.raises(ValueError): + plan_vllm_capacity(metrics(), length, budget) diff --git a/unit_tests/benchmark/service/test_completion.py b/unit_tests/benchmark/service/test_completion.py new file mode 100644 index 0000000000..6e568ef264 --- /dev/null +++ b/unit_tests/benchmark/service/test_completion.py @@ -0,0 +1,88 @@ +from copy import deepcopy + +import pytest + +from test.benchmark.service.completion import completion_request, completion_tokens + + +def native(ids=(11, 99), reason="stop"): + return dict( + tokens=[dict(id=token) for token in ids], + prompt_tokens=2, + count_output_tokens=len(ids), + finish_reason=reason, + ) + + +def baseline(ids=(11,), reason="stop"): + return dict( + choices=[dict(index=0, token_ids=list(ids), finish_reason=reason)], + usage=dict(prompt_tokens=2, completion_tokens=len(ids)), + ) + + +@pytest.mark.parametrize("framework", ["lightllm", "vllm", "sglang"]) +def test_request_sends_only_prompt_and_decoding_options(framework): + endpoint, payload = completion_request(framework, "model", [1, 2], 3) + if framework == "lightllm": + assert endpoint == "/generate" and payload["inputs"] == [1, 2] + assert payload["parameters"]["print_eos_token"] + assert payload["parameters"]["ignore_eos"] is False + else: + assert endpoint == "/v1/completions" and payload["prompt"] == [1, 2] + assert payload["ignore_eos"] is False and payload["return_token_ids"] + assert "expected_number" not in payload + + +@pytest.mark.parametrize("framework, factory", [("lightllm", native), ("vllm", baseline), ("sglang", baseline)]) +def test_natural_eos_and_budget_exhaustion_remain_distinct(framework, factory): + result = completion_tokens(factory(), framework, [1, 2], 3, [99]) + assert result["finish_reason"] == "stop" and not result["truncated"] + truncated = completion_tokens(factory((11, 12, 13), "length"), framework, [1, 2], 3, [99]) + assert truncated["token_ids"] == [11, 12, 13] and truncated["truncated"] + + +@pytest.mark.parametrize("framework", ["vllm", "sglang"]) +def test_baseline_may_include_or_omit_terminal_eos(framework): + for ids in [(11,), (11, 99)]: + assert completion_tokens(baseline(ids), framework, [1, 2], 3, [99])["token_ids"] == list(ids) + + +@pytest.mark.parametrize("framework, factory", [("lightllm", native), ("vllm", baseline), ("sglang", baseline)]) +@pytest.mark.parametrize("ids, reason", [((99, 11), "stop"), ((11,), "length"), ((11, 12, 99), "length")]) +def test_invalid_termination_cannot_be_scored(framework, factory, ids, reason): + with pytest.raises(ValueError): + completion_tokens(factory(ids, reason), framework, [1, 2], 3, [99]) + + +def test_native_stop_requires_observed_eos(): + with pytest.raises(ValueError, match="declared EOS"): + completion_tokens(native((11,)), "lightllm", [1, 2], 3, [99]) + + +@pytest.mark.parametrize("ids", [[], [True], [-1], [1.0], [11, 12, 13, 99]]) +def test_invalid_output_ids_and_over_budget_are_rejected(ids): + with pytest.raises(ValueError): + completion_tokens(baseline(ids), "vllm", [1, 2], 3, [99]) + + +def test_usage_and_prompt_identity_are_checked(): + original = baseline() + for field, value in [("prompt_tokens", 3), ("completion_tokens", 2), ("completion_tokens", True)]: + body = deepcopy(original) + body["usage"][field] = value + with pytest.raises(ValueError, match="counts"): + completion_tokens(body, "vllm", [1, 2], 3, [99]) + body = deepcopy(original) + body["choices"][0]["prompt_token_ids"] = [1, 3] + with pytest.raises(ValueError, match="exact prompt"): + completion_tokens(body, "vllm", [1, 2], 3, [99]) + body["choices"] *= 2 + with pytest.raises(ValueError, match="one completion"): + completion_tokens(body, "vllm", [1, 2], 3, [99]) + + +@pytest.mark.parametrize("prompt, limit", [([], 3), ([True], 3), ([-1], 3), ([1], True), ([1], 0)]) +def test_bad_requests_fail_before_transport(prompt, limit): + with pytest.raises(ValueError): + completion_request("lightllm", "model", prompt, limit) diff --git a/unit_tests/benchmark/service/test_completion_runner.py b/unit_tests/benchmark/service/test_completion_runner.py new file mode 100644 index 0000000000..9509ca39ca --- /dev/null +++ b/unit_tests/benchmark/service/test_completion_runner.py @@ -0,0 +1,111 @@ +from contextlib import contextmanager +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +import json +import threading + +import pytest +import requests + +from test.benchmark.service.completion_runner import CompletionClient + + +@contextmanager +def service(*, overlap=False, bad_count=False): + barrier, second_done = threading.Barrier(2), threading.Event() + finished = [] + + class Handler(BaseHTTPRequestHandler): + def log_message(self, *args): + pass + + def do_POST(self): + payload = json.loads(self.rfile.read(int(self.headers["Content-Length"]))) + assert "expected_number" not in payload + index = payload["inputs"][0] + if overlap: + barrier.wait(timeout=5) + if index == 1: + assert second_done.wait(timeout=5) + body = dict( + tokens=[dict(id=10 + index), dict(id=99)], + prompt_tokens=1, + count_output_tokens=1 if bad_count and index == 2 else 2, + finish_reason="stop", + ) + data = json.dumps(body).encode() + self.send_response(200) + self.send_header("Content-Length", str(len(data))) + self.end_headers() + self.wfile.write(data) + finished.append(index) + if index == 2: + second_done.set() + + server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield f"http://127.0.0.1:{server.server_port}", finished + finally: + server.shutdown() + server.server_close() + thread.join(timeout=5) + + +def cases(): + return [ + dict(dataset_index=17, prompt_ids=[1], max_output_tokens=3, expected_number="secret"), + dict(dataset_index=3, prompt_ids=[2], max_output_tokens=3, expected_number="secret"), + ] + + +def test_actual_concurrent_http_keeps_order_and_sessions_separate(tmp_path): + sessions = [] + + class Session(requests.Session): + def __init__(self): + super().__init__() + self.owner, self.was_closed = threading.get_ident(), False + sessions.append(self) + + def post(self, *args, **kwargs): + assert self.owner == threading.get_ident() and not self.trust_env + return super().post(*args, **kwargs) + + def close(self): + self.was_closed = True + super().close() + + with service(overlap=True) as (url, finished): + client = CompletionClient( + "lightllm", + "model", + url, + 10, + [99], + tmp_path, + concurrency=2, + session_factory=Session, + ) + try: + results = list(client.results(cases())) + finally: + client.close() + assert finished == [2, 1] + assert [x["token_ids"] for x in results] == [[11, 99], [12, 99]] + assert len(sessions) == 2 and all(x.was_closed for x in sessions) + assert {x.name for x in tmp_path.glob("response-*.txt")} == {"response-17.txt", "response-3.txt"} + + +@pytest.mark.parametrize("concurrency", [1, 2]) +def test_count_failure_is_not_dropped_and_raw_response_survives(tmp_path, concurrency): + with service(bad_count=True) as (url, finished): + client = CompletionClient("lightllm", "model", url, 10, [99], tmp_path, concurrency=concurrency) + try: + results = iter(client.results(cases())) + assert next(results)["token_ids"] == [11, 99] + with pytest.raises(ValueError, match="counts"): + next(results) + finally: + client.close() + assert json.loads((tmp_path / "response-3.txt").read_text())["count_output_tokens"] == 1 diff --git a/unit_tests/benchmark/service/test_stream_cohort.py b/unit_tests/benchmark/service/test_stream_cohort.py new file mode 100644 index 0000000000..48ce7d4a08 --- /dev/null +++ b/unit_tests/benchmark/service/test_stream_cohort.py @@ -0,0 +1,185 @@ +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +import json +import threading + +import pytest + +from test.benchmark.service.stream_cohort import StreamingWorkload, distribution, run_cohort, summarize_cohort +from test.benchmark.service.token_stream import request_stream + + +def observation(started, arrivals, ended, input_tokens=2, ids=None): + gaps = [right - left for left, right in zip(arrivals, arrivals[1:])] + return dict( + token_ids=list(range(len(arrivals))) if ids is None else ids, + metrics=dict( + input_tokens=input_tokens, + output_tokens=len(arrivals), + ttft_s=arrivals[0] - started, + tpot_s=sum(gaps) / len(gaps) if gaps else None, + e2e_s=ended - started, + receiver_itl_s=gaps, + ), + timeline=dict(started=started, first_token=arrivals[0], last_token=arrivals[-1], ended=ended), + ) + + +def test_overlapping_requests_use_elapsed_time_and_actual_token_counts(): + results = [observation(1, [1.5, 2.5, 4.5], 5), observation(2, [2.7], 4, input_tokens=4)] + metrics = summarize_cohort(results) + assert metrics["elapsed_s"] == 4 + assert metrics["output_tokens"] == 4 and metrics["output_tokens_per_s"] == 1 + assert metrics["input_tokens_per_s"] == 1.5 + assert metrics["maximum_client_inflight"] == 2 + assert metrics["ttft_s"]["p50"] == pytest.approx(0.6) + assert metrics["e2e_s"]["p50"] == 3 + assert metrics["tpot_s"]["count"] == 1 + assert metrics["receiver_itl_s"]["p50"] == 1.5 + + +def test_cache_aggregates_require_complete_observability(): + results = [observation(1, [1.5], 2), observation(2, [2.5], 3)] + results[0]["metrics"]["cached_input_tokens"] = 1 + partial = summarize_cohort(results) + assert partial["cache_reported_requests"] == 1 + assert partial["cached_input_tokens"] is None and partial["uncached_input_tokens"] is None + results[1]["metrics"]["cached_input_tokens"] = 0 + complete = summarize_cohort(results) + assert complete["cache_reported_requests"] == 2 + assert complete["cached_input_tokens"] == 1 and complete["uncached_input_tokens"] == 3 + + +def test_bursts_and_single_tokens_keep_their_metric_meaning(): + metrics = summarize_cohort([observation(1, [2, 2, 3], 4)]) + assert metrics["receiver_zero_gap_fraction"] == 0.5 + metrics = summarize_cohort([observation(1, [2], 4)]) + assert metrics["tpot_s"] is None and metrics["receiver_itl_s"] is None + assert metrics["receiver_zero_gap_fraction"] is None + + +def test_percentiles_are_interpolated_and_invalid_samples_rejected(): + metrics = distribution([4, 2, 1, 3]) + assert metrics["p50"] == 2.5 and metrics["p95"] == pytest.approx(3.85) + assert distribution([]) is None + for values in ([float("nan")], [float("inf")], [-1]): + with pytest.raises(ValueError): + distribution(values) + + +@pytest.mark.parametrize( + "field,value", + [ + ("started", float("nan")), + ("started", 3), + ("ended", 1), + ("first_token", 4), + ("last_token", 1), + ], +) +def test_invalid_timelines_cannot_produce_throughput(field, value): + result = observation(1, [2, 3], 4) + result["timeline"][field] = value + with pytest.raises(ValueError): + summarize_cohort([result]) + + +@pytest.mark.parametrize( + "field,value", + [("output_tokens", 99), ("output_tokens", True), ("input_tokens", 0), ("receiver_itl_s", []), ("ttft_s", 0)], +) +def test_inconsistent_counts_or_timing_cannot_pass(field, value): + result = observation(1, [2, 3], 4) + result["metrics"][field] = value + with pytest.raises(ValueError): + summarize_cohort([result]) + + +def test_reference_mismatch_invalidates_entire_cohort_and_keeps_evidence(tmp_path): + tasks = [StreamingWorkload("ok", [11, 12], [7, 8]), StreamingWorkload("bad", [11, 12], [7, 8])] + + def request(task, path): + ids = [7, 8] if task.name == "ok" else [7, 9] + path.write_text("raw response") + return observation(1, [2, 3], 4, ids=ids) + + report = run_cohort(tasks, 2, request, tmp_path) + assert report["status"] == "failed" and report["metrics"] is None + assert [record["status"] for record in report["records"]] == ["passed", "failed"] + assert report["records"][1]["result"]["token_ids"] == [7, 9] + assert (tmp_path / "000001.sse").read_text() == "raw response" + + +def test_transport_failure_cannot_be_hidden_by_successful_requests(tmp_path): + task = StreamingWorkload("error", [11], [7]) + + def request(task, path): + path.write_text("partial response") + raise TimeoutError("stream stalled") + + report = run_cohort([task], 1, request, tmp_path) + assert report["status"] == "failed" and report["metrics"] is None + assert report["records"][0]["error_type"] == "TimeoutError" + + +@pytest.mark.parametrize("concurrency", [0, -1, 2, True]) +def test_invalid_concurrency_is_rejected(concurrency, tmp_path): + with pytest.raises(ValueError): + run_cohort([StreamingWorkload("x", [1], [2])], concurrency, lambda *_: None, tmp_path) + + +def test_real_chunked_sse_closed_loop_cohort(tmp_path): + pair = threading.Barrier(2) + + class Handler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def log_message(self, *_): + pass + + def do_POST(self): + body = json.loads(self.rfile.read(int(self.headers["Content-Length"]))) + assert body["prompt"] == [11, 12] and body["return_token_ids"] is True + pair.wait(timeout=5) + self.send_response(200) + self.send_header("content-type", "text/event-stream") + self.send_header("transfer-encoding", "chunked") + self.end_headers() + events = [ + {"choices": [{"index": 0, "token_ids": [7, 8]}]}, + {"choices": [{"index": 0, "token_ids": [9], "finish_reason": "length"}]}, + {"choices": [], "usage": {"prompt_tokens": 2, "completion_tokens": 3}}, + ] + for payload in [*(json.dumps(event) for event in events), "[DONE]"]: + chunk = ("data: " + payload + "\n\n").encode() + self.wfile.write(f"{len(chunk):x}\r\n".encode() + chunk + b"\r\n") + self.wfile.flush() + self.wfile.write(b"0\r\n\r\n") + self.wfile.flush() + + server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + task = StreamingWorkload("fixture", [11, 12], [7, 8, 9]) + + def request(task, path): + return request_stream( + f"http://127.0.0.1:{server.server_port}", + "vllm", + "fixture", + list(task.prompt_ids), + len(task.expected_ids), + 5, + path, + ) + + try: + report = run_cohort([task] * 4, 2, request, tmp_path) + finally: + server.shutdown() + server.server_close() + thread.join(timeout=5) + assert report["status"] == "passed" + assert report["metrics"]["requests"] == 4 + assert report["metrics"]["output_tokens"] == 12 + assert report["metrics"]["maximum_client_inflight"] == 2 + assert len(list(tmp_path.glob("*.sse"))) == 4 diff --git a/unit_tests/benchmark/service/test_stream_timing.py b/unit_tests/benchmark/service/test_stream_timing.py new file mode 100644 index 0000000000..bca4c755e3 --- /dev/null +++ b/unit_tests/benchmark/service/test_stream_timing.py @@ -0,0 +1,79 @@ +import copy +import json + +import pytest + +from test.benchmark.service.stream_timing import streaming_completion_times, summarize_streaming_times + + +def stream(text="hello"): + origin = 10_000_000_000 + return dict( + status=200, + start_perf_ns=origin, + metadata=dict(request_start_ns=1_000_000, request_end_ns=3_001_000_000, was_cancelled=False), + responses=[ + dict( + perf_ns=origin + 500_000_000, + packets=[dict(name="data", value=json.dumps(dict(choices=[dict(text=text)])))], + ), + dict( + perf_ns=origin + 1_000_000_000, + packets=[dict(name="data", value=json.dumps(dict(choices=[dict(text=text)])))], + ), + dict( + perf_ns=origin + 2_000_000_000, + packets=[dict(name="data", value=json.dumps(dict(choices=[dict(text="")])))], + ), + dict(perf_ns=origin + 3_000_000_000, packets=[dict(name="data", value="[DONE]")]), + ], + ) + + +def test_invisible_tokens_do_not_shorten_http_latency(): + assert streaming_completion_times(stream()) == dict( + http_completion_ms=3000, + first_visible_text_ms=500, + last_visible_text_ms=1000, + done_after_last_visible_text_ms=2000, + ) + + +def test_empty_text_retains_http_timing_without_invented_text_latency(): + empty = stream("") + empty["metadata"]["request_end_ns"] += 1_000_000_000 + empty["responses"][-1]["perf_ns"] += 1_000_000_000 + result = summarize_streaming_times([stream(), empty]) + assert result["requests"] == 2 and result["empty_text_responses"] == 1 + assert result["http_completion_all_requests_ms"]["p95"] == 3950 + assert result["http_completion_content_requests_ms"]["p95"] == 3000 + assert result["last_visible_text_ms"]["p95"] == 1000 + assert summarize_streaming_times([empty])["last_visible_text_ms"] is None + + +@pytest.mark.parametrize( + "failure", + ["no_done", "duplicate_done", "after_done", "out_of_order", "clock", "cancelled", "http", "float_timestamp"], +) +def test_rejects_invalid_stream_boundaries(failure): + record = stream() + if failure == "no_done": + record["responses"].pop() + elif failure == "duplicate_done": + record["responses"].append(copy.deepcopy(record["responses"][-1])) + elif failure == "after_done": + response = copy.deepcopy(record["responses"][0]) + response["perf_ns"] = record["responses"][-1]["perf_ns"] + record["responses"].append(response) + elif failure == "out_of_order": + record["responses"][1]["perf_ns"] = record["responses"][0]["perf_ns"] - 1 + elif failure == "clock": + record["metadata"]["request_end_ns"] += record["start_perf_ns"] + elif failure == "cancelled": + record["metadata"]["was_cancelled"] = True + elif failure == "http": + record["status"] = 500 + elif failure == "float_timestamp": + record["responses"][0]["perf_ns"] = float(record["responses"][0]["perf_ns"]) + with pytest.raises(ValueError): + streaming_completion_times(record) diff --git a/unit_tests/benchmark/service/test_token_stream.py b/unit_tests/benchmark/service/test_token_stream.py new file mode 100644 index 0000000000..4ef3e19069 --- /dev/null +++ b/unit_tests/benchmark/service/test_token_stream.py @@ -0,0 +1,287 @@ +import json + +import pytest + +from test.benchmark.service.token_stream import SSEDecoder, TokenStream + + +def test_sse_handles_split_utf8_crlf_and_multiline_json(): + decoder = SSEDecoder() + raw = ': heartbeat\r\ndata: {"text": "中",\r\ndata: "token_ids": [7, 8]}\r\n\r\ndata: [DONE]\n\n'.encode() + events = [] + for byte in raw: + events.extend(decoder.feed(bytes([byte]))) + decoder.finish() + assert json.loads(events[0]) == {"text": "中", "token_ids": [7, 8]} + assert events[1] == "[DONE]" + + +@pytest.mark.parametrize("raw", [b'data: {"x": 1}', b'data: {"x": 1}\n', b'data: "\xe4']) +def test_truncated_sse_cannot_pass(raw): + decoder = SSEDecoder() + decoder.feed(raw) + with pytest.raises((ValueError, UnicodeDecodeError)): + decoder.finish() + + +@pytest.mark.parametrize("framework", ["vllm", "sglang"]) +def test_burst_token_count_and_arrival_metrics(framework): + stream = TokenStream(framework, [11, 12], 3, 1.0) + stream.record(json.dumps({"choices": [{"index": 0, "text": "ab", "token_ids": [7, 8]}]}), 1.5) + stream.record( + json.dumps({"choices": [{"index": 0, "text": "c", "token_ids": [9], "finish_reason": "length"}]}), 1.8 + ) + stream.record(json.dumps({"choices": [], "usage": {"prompt_tokens": 2, "completion_tokens": 3}}), 1.9) + stream.record("[DONE]", 1.9) + result = stream.summarize(2.0) + assert result["output_tokens"] == 3 and result["events"] == 3 + assert result["ttft_s"] == 0.5 + assert result["tpot_s"] == pytest.approx(0.15) + assert result["receiver_itl_s"] == pytest.approx([0, 0.3]) + + +def test_native_empty_text_is_still_an_output_token(): + stream = TokenStream("lightllm", [11], 1, 1.0) + stream.record( + json.dumps( + { + "token": {"id": 7, "text": "", "prompt_tokens": 1, "count_output_tokens": 1}, + "finished": True, + "finish_reason": "length", + } + ), + 1.2, + ) + assert stream.summarize(1.3)["output_tokens"] == 1 + + +@pytest.mark.parametrize("framework", ["lightllm", "vllm", "sglang"]) +@pytest.mark.parametrize("cached", [None, 0, 1, 2]) +def test_optional_cache_counts_preserve_unknown_and_zero(framework, cached): + stream = TokenStream(framework, [11, 12], 1, 1.0) + if framework == "lightllm": + body = dict(token=dict(id=7, prompt_tokens=2, count_output_tokens=1), finished=True, finish_reason="length") + if cached is not None: + body["token"]["prompt_cache_len"] = cached + else: + body = dict( + choices=[dict(index=0, token_ids=[7], finish_reason="length")], + usage=dict(prompt_tokens=2, completion_tokens=1), + ) + if cached is not None: + body["usage"]["prompt_tokens_details"] = dict(cached_tokens=cached) + stream.record(json.dumps(body), 1.2) + if framework != "lightllm": + stream.record("[DONE]", 1.3) + result = stream.summarize(1.4) + assert result["cached_input_tokens"] == cached + assert result["uncached_input_tokens"] == (None if cached is None else 2 - cached) + + +@pytest.mark.parametrize("cached", [-1, True, 3, 1.0]) +def test_invalid_cached_usage_cannot_pass(cached): + stream = TokenStream("vllm", [11, 12], 1, 1.0) + with pytest.raises(ValueError, match="Cached input"): + stream.record( + json.dumps( + dict( + choices=[], + usage=dict(prompt_tokens=2, completion_tokens=0, prompt_tokens_details=dict(cached_tokens=cached)), + ) + ), + 1.1, + ) + + +def test_cached_usage_cannot_change_within_request(): + stream = TokenStream("vllm", [11, 12], 1, 1.0) + stream.record( + json.dumps( + dict( + choices=[], + usage=dict(prompt_tokens=2, completion_tokens=0, prompt_tokens_details=dict(cached_tokens=1)), + ) + ), + 1.1, + ) + with pytest.raises(ValueError, match="changed"): + stream.record( + json.dumps( + dict( + choices=[], + usage=dict(prompt_tokens=2, completion_tokens=0, prompt_tokens_details=dict(cached_tokens=0)), + ) + ), + 1.2, + ) + + +@pytest.mark.parametrize("cold", [False, True]) +def test_native_transport_requests_observable_details_and_explicit_cold_policy(monkeypatch, tmp_path, cold): + import requests + from test.benchmark.service.token_stream import request_stream + + class Response: + status_code = 200 + headers = {"content-type": "text/event-stream", "transfer-encoding": "chunked"} + + def __enter__(self): + return self + + def __exit__(self, *args): + pass + + def iter_content(self, chunk_size): + body = dict( + token=dict(id=7, prompt_tokens=2, count_output_tokens=1, prompt_cache_len=0 if cold else 1), + finished=True, + finish_reason="length", + ) + yield ("data:" + json.dumps(body) + "\n\n").encode() + + class Session: + def __enter__(self): + return self + + def __exit__(self, *args): + pass + + def post(self, url, json, **kwargs): + assert json["parameters"]["return_details"] is True + assert json["parameters"].get("disable_prompt_cache", False) is cold + return Response() + + monkeypatch.setattr(requests, "Session", Session) + result = request_stream( + "http://localhost", "lightllm", "model", [11, 12], 1, 10, tmp_path / "response.sse", disable_prompt_cache=cold + ) + assert result["metrics"]["cached_input_tokens"] == (0 if cold else 1) + + +def test_baseline_request_cannot_silently_ignore_native_cache_policy(tmp_path): + from test.benchmark.service.token_stream import request_stream + + with pytest.raises(ValueError, match="only implemented for native"): + request_stream( + "http://localhost", "vllm", "model", [11], 1, 10, tmp_path / "response.sse", disable_prompt_cache=True + ) + + +@pytest.mark.parametrize( + "body", + [ + {"choices": [{"index": 0, "text": "x"}]}, + {"choices": [{"index": 0, "token_ids": [True]}]}, + {"choices": [{"index": 0, "token_ids": [7, 8, 9]}]}, + {"choices": [{"index": 0, "token_ids": [7], "finish_reason": "stop"}]}, + {"choices": [{"index": 0, "prompt_token_ids": [99]}]}, + {"choices": [], "usage": {"prompt_tokens": 99, "completion_tokens": 0}}, + {"error": {"message": "aborted"}}, + {"choices": [{"index": 0, "token_ids": [-1]}]}, + {"choices": [{"index": 0, "token_ids": 0}]}, + ], +) +def test_invalid_stream_events_are_rejected(body): + stream = TokenStream("vllm", [11], 2, 1.0) + with pytest.raises(ValueError): + stream.record(json.dumps(body), 1.1) + + +def test_truncated_generation_and_missing_done_cannot_be_reported(): + stream = TokenStream("vllm", [11], 1, 1.0) + with pytest.raises(ValueError): + stream.record("[DONE]", 1.1) + with pytest.raises(ValueError): + stream.summarize(1.2) + stream.record(json.dumps({"choices": [{"index": 0, "token_ids": [7], "finish_reason": "length"}]}), 1.2) + with pytest.raises(ValueError): + stream.summarize(1.3) + with pytest.raises(ValueError): + stream.record(json.dumps({"choices": [{"index": 0, "token_ids": [8]}]}), 1.3) + + +def test_control_events_also_require_monotonic_timestamps(): + stream = TokenStream("vllm", [11], 1, 1.0) + stream.record(json.dumps({"choices": []}), 1.5) + with pytest.raises(ValueError): + stream.record(json.dumps({"choices": []}), 1.4) + stream.record(json.dumps({"choices": [{"index": 0, "token_ids": [7], "finish_reason": "length"}]}), 1.6) + stream.record("[DONE]", 1.7) + with pytest.raises(ValueError): + stream.summarize(1.65) + + +@pytest.mark.parametrize("framework", ["vllm", "sglang"]) +@pytest.mark.parametrize("cache_salt", [None, "isolated-prefix", "x" * 128]) +def test_transport_uses_chunk_boundaries_and_saves_raw_after_timing(monkeypatch, tmp_path, framework, cache_salt): + import requests + import time + from test.benchmark.service.token_stream import request_stream + + bodies = [ + {"choices": [{"index": 0, "text": "ab", "token_ids": [7, 8]}]}, + {"choices": [{"index": 0, "token_ids": [9], "finish_reason": "length"}]}, + {"choices": [], "usage": {"prompt_tokens": 1, "completion_tokens": 3}}, + ] + chunks = [("data: " + json.dumps(body) + "\n\n").encode() for body in bodies] + chunks.append(b"data: [DONE]\n\n") + + class Response: + status_code = 200 + headers = {"content-type": "text/event-stream", "transfer-encoding": "chunked"} + + def __enter__(self): + return self + + def __exit__(self, *args): + pass + + def iter_content(self, chunk_size): + assert chunk_size is None + yield from chunks + + class Session: + def __enter__(self): + return self + + def __exit__(self, *args): + pass + + def post(self, url, json, **kwargs): + assert json["prompt"] == [11] and json["return_token_ids"] is True + assert json["stream_options"]["include_usage"] is True + if cache_salt is None: + assert "cache_salt" not in json + else: + assert json["cache_salt"] == cache_salt + return Response() + + monkeypatch.setattr(requests, "Session", Session) + timestamps = iter([1.0, 1.5, 1.8, 1.9, 1.9, 2.0]) + monkeypatch.setattr(time, "perf_counter", lambda: next(timestamps)) + path = tmp_path / "response.sse" + result = request_stream("http://localhost", framework, "model", [11], 3, 10, path, cache_salt=cache_salt) + assert result["token_ids"] == [7, 8, 9] + assert result["metrics"]["ttft_s"] == 0.5 + assert result["metrics"]["e2e_s"] == 1.0 + assert result["timeline"] == {"started": 1.0, "first_token": 1.5, "last_token": 1.8, "ended": 2.0} + assert path.read_bytes() == b"".join(chunks) + + +@pytest.mark.parametrize("cache_salt", ["", "x" * 129, 1, True, ["salt"], "x@y", "x/y", "x\\y", "x\x00y"]) +def test_invalid_cache_salt_is_rejected_before_transport(tmp_path, cache_salt): + from test.benchmark.service.token_stream import request_stream + + with pytest.raises(ValueError, match="Cache salt must"): + request_stream( + "http://localhost", "vllm", "model", [11], 1, 10, tmp_path / "response.sse", cache_salt=cache_salt + ) + + +def test_native_cache_salt_cannot_be_silently_ignored(tmp_path): + from test.benchmark.service.token_stream import request_stream + + with pytest.raises(ValueError, match="only implemented for baseline"): + request_stream( + "http://localhost", "lightllm", "model", [11], 1, 10, tmp_path / "response.sse", cache_salt="salt" + ) diff --git a/unit_tests/common/att_resources/model_helpers.py b/unit_tests/common/att_resources/model_helpers.py index 10c52c8cea..86e62cb671 100644 --- a/unit_tests/common/att_resources/model_helpers.py +++ b/unit_tests/common/att_resources/model_helpers.py @@ -253,6 +253,8 @@ def build_hybrid_resources( 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 @@ -282,6 +284,9 @@ def build_hybrid_resources( 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, diff --git a/unit_tests/common/att_resources/test_hybrid_layer_mapping.py b/unit_tests/common/att_resources/test_hybrid_layer_mapping.py index 15d29bd85d..fb9ce7ba5b 100644 --- a/unit_tests/common/att_resources/test_hybrid_layer_mapping.py +++ b/unit_tests/common/att_resources/test_hybrid_layer_mapping.py @@ -9,6 +9,7 @@ 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", @@ -41,6 +42,7 @@ def test_hybrid_layer_mapping_and_mtp_states( expected_indices: tuple[int, ...], linear_layers: int, draft_layers: int, + mode: str, ) -> None: model_config = dict( num_hidden_layers=len(layer_types), @@ -56,7 +58,13 @@ def test_hybrid_layer_mapping_and_mtp_states( full_attention_interval=4, layer_types=layer_types, ) - args = StartArgs(data_type="bfloat16", linear_att_ssm_data_type="float32") + 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( @@ -70,6 +78,11 @@ def test_hybrid_layer_mapping_and_mtp_states( 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) diff --git a/unit_tests/common/att_resources/test_page_pool_storage.py b/unit_tests/common/att_resources/test_page_pool_storage.py index 032f7674eb..5eb5ba52cc 100644 --- a/unit_tests/common/att_resources/test_page_pool_storage.py +++ b/unit_tests/common/att_resources/test_page_pool_storage.py @@ -58,9 +58,9 @@ def test_shared_registration_lives_until_last_view( 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) + def record_unregister(addresses: tuple[int, ...], device_id: int) -> None: + unregistered.extend(addresses) + unregister(addresses=addresses, device_id=device_id) monkeypatch.setattr(utils, "_unregister_page_pool_storage", record_unregister) tensor = utils.create_page_pool_tensor(size=(4096,), dtype=torch.float32) @@ -75,6 +75,92 @@ def record_unregister(ptr: int, device_id: int) -> None: assert unregistered == [ptr] +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +def test_chunked_shared_pages_keep_dma_boundaries_and_lifetime(cpu_args, monkeypatch, dtype) -> None: + cpu_args.enable_dp_prompt_cache_fetch = True + monkeypatch.setattr(utils, "_HOST_REGISTRATION_BYTES", 16384) + runtime = torch.cuda.cudart() + registered, released = [], [] + + class RecordingRuntime: + def cudaHostRegister(self, address, size, flags): + status = runtime.cudaHostRegister(address, size, flags) + if int(status) == 0: + registered.append((address, size)) + return status + + def cudaHostUnregister(self, address): + released.append(address) + return runtime.cudaHostUnregister(address) + + monkeypatch.setattr(torch.cuda, "cudart", lambda: RecordingRuntime()) + tensor = utils.create_page_pool_tensor(size=(7, 3, 512), dtype=dtype) + cell_bytes = tensor[0].numel() * tensor.element_size() + assert len(registered) > 1 + assert all((address - tensor.data_ptr()) % cell_bytes == 0 for address, _ in registered) + assert all(size <= 16384 for _, size in registered) + destination = torch.empty((3, 5, 512), dtype=dtype, device="cuda") + for index in range(tensor.shape[0]): + tensor[index].fill_(index + 1) + destination[:, 1].copy_(tensor[index], non_blocking=True) + torch.cuda.synchronize() + assert torch.equal(destination[:, 1].cpu(), tensor[index]) + + view = tensor[-1] + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + destination[:, 1].copy_(view, non_blocking=True) + view.fill_(17) + del tensor + gc.collect() + assert not released + assert view.is_shared() and view.is_pinned() + graph.replay() + torch.cuda.synchronize() + assert torch.equal(destination[:, 1].cpu(), view) + graph.reset() + del view + gc.collect() + assert released == [address for address, _ in reversed(registered)] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_private_storage_keeps_whole_tensor_dma(monkeypatch) -> None: + monkeypatch.setattr(utils, "_HOST_REGISTRATION_BYTES", 16384) + tensor = torch.full((7, 3, 512), 7, dtype=torch.float32) + utils.register_page_pool_tensor(tensor) + destination = tensor.to("cuda", non_blocking=True) + torch.cuda.synchronize() + assert torch.equal(destination.cpu(), tensor) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_shared_pool_partial_registration_failure_rolls_back(cpu_args, monkeypatch) -> None: + cpu_args.enable_dp_prompt_cache_fetch = True + monkeypatch.setattr(utils, "_HOST_REGISTRATION_BYTES", 16384) + runtime = torch.cuda.cudart() + registered, released = [], [] + + class FailingRuntime: + def cudaHostRegister(self, address, size, flags): + if registered: + return 1 + status = runtime.cudaHostRegister(address, size, flags) + if int(status) == 0: + registered.append(address) + return status + + def cudaHostUnregister(self, address): + released.append(address) + return runtime.cudaHostUnregister(address) + + monkeypatch.setattr(torch.cuda, "cudart", lambda: FailingRuntime()) + with pytest.raises(torch.cuda.CudaError): + utils.create_page_pool_tensor(size=(7, 3, 512), dtype=torch.float32) + assert registered and released == registered + + @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 diff --git a/unit_tests/common/basemodel/attention/linear/test_gdn.py b/unit_tests/common/basemodel/attention/linear/test_gdn.py index 8ec77c6f4d..22b561b822 100644 --- a/unit_tests/common/basemodel/attention/linear/test_gdn.py +++ b/unit_tests/common/basemodel/attention/linear/test_gdn.py @@ -74,6 +74,7 @@ def test_decode_state_initializes_normal_layout(): 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 + assert state.b1_mtp_cu_q_seq_len_int64 is None @pytest.mark.parametrize("role", ["main", "draft0"]) @@ -91,6 +92,7 @@ def test_decode_state_initializes_fixed_mtp_layout(role: str): 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.b1_mtp_cu_q_seq_len_int64, torch.tensor([0, 3, 6], dtype=torch.int64)) 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( @@ -124,6 +126,7 @@ def build_params(**kwargs): state.init_state() assert state.b1_mtp_cu_q_seq_len is expected_cu_q_seq_len + torch.testing.assert_close(state.b1_mtp_cu_q_seq_len_int64, expected_cu_q_seq_len.to(torch.int64)) 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( @@ -140,6 +143,89 @@ def build_params(**kwargs): assert build_calls[0]["hold_req_id"] == 5 +def test_mtp_recurrent_layers_reuse_converted_offsets(monkeypatch): + import lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_mtp as conv_mtp + + state, infer_state = _create_decode_state( + draft_step=2, + dynamic_layout=False, + b_req_idx=torch.tensor([0, 0, 0], dtype=torch.int32), + req_to_mtp_state_index=torch.tensor([0, 0], dtype=torch.int32), + ) + state.init_state() + state.backend.activation = "silu" + state.backend._rearrange_mixed_qkv = lambda mixed, decode: (mixed, mixed, mixed) + offsets = [] + monkeypatch.setattr(conv_mtp, "causal_conv1d_update", lambda mixed, *args, **kwargs: mixed) + + def recurrent(**kwargs): + offsets.append(kwargs["cu_seqlens"]) + return kwargs["q"], None + + monkeypatch.setattr(gdn, "mtp_fused_recurrent_gated_delta_rule", recurrent) + weight = SimpleNamespace( + linear_conv1d=SimpleNamespace(mm_param=SimpleNamespace(weight=None), bias=None), + linear_A_log=SimpleNamespace(weight=None), + linear_dt_bias=SimpleNamespace(weight=None), + ) + tensor = torch.zeros((3, 1)) + for _ in range(2): + state._gdn_mtp_kernel(tensor, tensor, tensor, tensor, tensor, state.infer_state, weight) + assert len(offsets) == 2 + assert all(value is state.b1_mtp_cu_q_seq_len_int64 for value in offsets) + torch.testing.assert_close(offsets[0], torch.tensor([0, 3], dtype=torch.int64)) + + +@pytest.mark.parametrize( + "device", + ["cpu", pytest.param("cuda", marks=pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required"))], +) +def test_graph_state_copy_refreshes_cached_mtp_offsets(monkeypatch, device): + layouts = [ + torch.tensor([0, 2, 3, 3], dtype=torch.int32, device=device), + torch.tensor([0, 1, 3, 3], dtype=torch.int32, device=device), + ] + current = [0] + + def build_params(**kwargs): + return ( + layouts[current[0]], + torch.tensor([0, 1, 2], dtype=torch.int32, device=device), + torch.ones(3, dtype=torch.int32, device=device), + ) + + monkeypatch.setattr(gdn, "build_dynamic_mtp_linear_att_state_params", build_params) + + def create(): + state, infer_state = _create_decode_state( + draft_step=2, + dynamic_layout=True, + b_req_idx=torch.tensor([0, 0, 1] if current[0] == 0 else [0, 1, 1], dtype=torch.int32, device=device), + req_to_mtp_state_index=torch.zeros(3, dtype=torch.int32, device=device), + ) + state.init_state() + return state + + captured = create() + cached = captured.b1_mtp_cu_q_seq_len_int64 + graph = None + if device == "cuda": + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + observed = cached.clone() + current[0] = 1 + fresh = create() + captured.copy_for_decode_cuda_graph(fresh) + assert captured.b1_mtp_cu_q_seq_len_int64 is cached + expected = torch.tensor([0, 1, 3, 3], dtype=torch.int64, device=device) + torch.testing.assert_close(cached, expected) + assert captured.b1_mtp_cu_q_seq_len.dtype == torch.int32 + if graph is not None: + graph.replay() + torch.testing.assert_close(observed, expected) + graph.reset() + + def test_linear_backend_is_abstract_and_triton_supplies_chunk_kernel(): assert inspect.isabstract(gdn.LinearAttBackend) backend = object.__new__(TritonLinearAttBackend) @@ -177,6 +263,34 @@ def test_missing_linear_backend_arg_uses_auto_selection(monkeypatch, auto_linear 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_qwen35_linear_prefill_att_backend_class(index=1) + decode_backend = linear_create_utils.get_qwen35_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_flashqla_validation_loads_linear_cache_config(monkeypatch): linear_config = SimpleNamespace( tp_world_size=2, 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..fb5168f661 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_acceptance_fusion.py @@ -0,0 +1,243 @@ +import pytest +import torch +import triton + +from lightllm.common.basemodel.triton_kernel.mtp_utils import linear_att_mtp_state_index_update as update + + +def ReplaySSMExecutor(*args, **kwargs): + kda = kwargs.pop("kda", False) + if kda: + from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import _history_owner + else: + from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import _history_owner + + return _history_owner(*args, **kwargs) + + +def ssm_snapshot(backend, req): + return backend.accepted_ssm(req) + + +from types import SimpleNamespace +from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState +from lightllm.common.basemodel.triton_kernel.linear_att.ssm_acceptance import commit_ssm_prefix +from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import replay_layer + + +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 + ssm_backend.capacity = max(4, triton.next_power_of_2(width)) + 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) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("replay", [False, True]) +def test_request_counts_share_compiled_kernel(monkeypatch, replay): + 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 + indices = torch.zeros(513, dtype=torch.int32, device="cuda") + cursors = torch.zeros_like(indices) if replay else None + 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) + if replay: + cursors.zero_() + commit_ssm_prefix(indices, reqs[:count], lengths[:count], cursors) + assert torch.equal(indices[:count], lengths[:count] - 1) + assert (indices[count:] == -1).all() + if replay: + 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) + m._ssm_executor.accept_updates(args[1][args[0].long()], m.mtp_state_index) + + +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()] + m._ssm_executor.commit_accepted(m.mtp_state_index, reqs, lengths) + + +@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 + + class _IndexOnly: + def commit_accepted(self, indices, req_indices, accepted_lengths): + commit_ssm_prefix(indices, req_indices, accepted_lengths) + + m._ssm_executor = _IndexOnly() + accept_fused(m, args) + assert m.mtp_state_index.tolist() == expected_index + + 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, + run_config={"BV": 32, "num_warps": 1, "num_stages": 1}, + ) + 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(1, 12, 2, 32, dtype=torch.bfloat16, device="cuda") + k = torch.randn_like(q) + v = torch.randn(1, 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] + ssm_backend.prepare_decode(reqs, cu) + outputs = [replay_layer(ssm_backend, layer, q, k, v, a, b, alog, alog, reqs, cu)[0] 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.merge_accepted_updates(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.clear_history(1) + manager._ssm_executor.ssm_state[:, 1].zero_() + manager.mtp_state_index[1] = 0 + for ssm_backend in ssm_backends: + ssm_backend.merge_accepted_updates(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..9021f50195 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 @@ -6,6 +6,31 @@ from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_mtp import causal_conv1d_update +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("activation", [None, "silu"]) +def test_low_precision_products_accumulate_in_fp32(dtype, activation): + torch.manual_seed(2907) + x = torch.randn(6, 512, device="cuda", dtype=dtype) * 0.3 + weight = torch.randn(512, 4, device="cuda", dtype=dtype) * 0.2 + prefix = torch.randn(2, 512, 3, device="cuda", dtype=dtype) * 0.3 + full = torch.cat([prefix, x.view(2, 3, 512).transpose(1, 2)], -1).float() + expected = torch.nn.functional.conv1d(full, weight.float()[:, None], groups=512) + if activation == "silu": + expected = torch.nn.functional.silu(expected) + expected = expected.transpose(1, 2).reshape(6, 512).to(dtype) + actual = causal_conv1d_update( + x.clone(), + torch.nn.functional.pad(prefix, (0, 2)), + weight, + mtp_step=2, + conv_state_indices=torch.tensor([0, 1], device="cuda", dtype=torch.int32), + num_accepted_tokens=torch.ones(2, device="cuda", dtype=torch.int32), + query_start_loc=torch.tensor([0, 3, 6], device="cuda", dtype=torch.int32), + activation=activation, + ) + torch.testing.assert_close(actual, expected, atol=1e-7, rtol=1e-6) + + def causal_conv1d_ref( x: torch.Tensor, conv_state: torch.Tensor, diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_fused_recurrent_packed_decode.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_fused_recurrent_packed_decode.py new file mode 100644 index 0000000000..aae63dd572 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_fused_recurrent_packed_decode.py @@ -0,0 +1,62 @@ +import pytest +import torch + +from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.fused_recurrent import ( + fused_recurrent_gated_delta_rule, + fused_recurrent_gated_delta_rule_packed_decode, +) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_packed_decode_matches_materialized_decode(): + torch.manual_seed(17) + batch, key_heads, value_heads = 2, 2, 4 + key_dim = value_dim = 128 + device = "cuda" + dtype = torch.bfloat16 + + mixed_qkv = torch.randn( + batch, + 2 * key_heads * key_dim + value_heads * value_dim, + device=device, + dtype=dtype, + ) + a_raw = torch.randn(batch, value_heads, device=device, dtype=dtype) + b_raw = torch.randn_like(a_raw) + a_log = torch.randn(value_heads, device=device, dtype=torch.float32) + dt_bias = torch.randn(value_heads, device=device, dtype=torch.float32) + state = torch.randn(batch + 1, value_heads, key_dim, value_dim, device=device, dtype=dtype) + state_reference = state.clone() + state_packed = state.clone() + state_indices = torch.arange(1, batch + 1, device=device, dtype=torch.int32) + + q = mixed_qkv[:, : key_heads * key_dim].view(batch, 1, key_heads, key_dim) + k = mixed_qkv[:, key_heads * key_dim : 2 * key_heads * key_dim].view(batch, 1, key_heads, key_dim) + v = mixed_qkv[:, 2 * key_heads * key_dim :].view(batch, 1, value_heads, value_dim) + + reference, _ = fused_recurrent_gated_delta_rule( + q, + k, + v, + initial_state=state_reference, + ssm_state_indices=state_indices, + use_qk_l2norm_in_kernel=True, + A_log=a_log, + dt_bias=dt_bias, + a_raw=a_raw, + b_raw=b_raw, + ) + packed, _ = fused_recurrent_gated_delta_rule_packed_decode( + mixed_qkv, + a_raw, + b_raw, + a_log, + dt_bias, + key_dim ** -0.5, + state_packed, + state_indices, + use_qk_l2norm_in_kernel=True, + ) + + torch.testing.assert_close(packed, reference, atol=0.02, rtol=0.02) + torch.testing.assert_close(state_packed, state_reference, atol=0.02, rtol=0.02) 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 index 11035c83a3..b280f4fabb 100644 --- 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 @@ -163,5 +163,57 @@ def test_variable_length_cu_seqlens(seed): ) +@pytest.mark.parametrize("state_dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("steps", [1, 2, 3, 4]) +def test_match_decode_preserves_every_sequential_snapshot(state_dtype, steps): + torch.manual_seed(3107) + batch, heads, value_heads, width = 2, 4, 12, 128 + count = batch * steps + q = torch.randn(batch, steps, heads, width, device="cuda", dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn(batch, steps, value_heads, width, device="cuda", dtype=torch.bfloat16) + a = torch.randn(batch, steps, value_heads, device="cuda", dtype=torch.bfloat16) + b = torch.randn_like(a) + indices = torch.arange(count, device="cuda", dtype=torch.int32).view(batch, steps) + accepted = torch.tensor([1, steps], device="cuda", dtype=torch.int32) + initial = torch.randn(count, value_heads, width, width, device="cuda", dtype=state_dtype) * 0.1 + state = initial[indices.gather(1, accepted.long()[:, None] - 1).flatten().long()].clone() + expected_states = initial.clone() + gates = dict(A_log=torch.zeros(value_heads, device="cuda"), dt_bias=torch.zeros(value_heads, device="cuda")) + outputs = [] + for step in range(steps): + output, _ = fused_recurrent_gated_delta_rule( + q=q[:, step : step + 1], + k=k[:, step : step + 1], + v=v[:, step : step + 1], + initial_state=state, + inplace_final_state=True, + ssm_state_indices=torch.arange(batch, device="cuda", dtype=torch.int32), + use_qk_l2norm_in_kernel=True, + a_raw=a[:, step], + b_raw=b[:, step], + **gates, + ) + outputs.append(output) + expected_states.index_copy_(0, indices[:, step].long(), state) + actual, actual_states = mtp_fused_recurrent_gated_delta_rule( + q=q.reshape(1, count, heads, width), + k=k.reshape(1, count, heads, width), + v=v.reshape(1, count, value_heads, width), + initial_state=initial.clone(), + cu_seqlens=torch.arange(batch + 1, device="cuda", dtype=torch.int64) * steps, + ssm_state_indices=indices, + ssm_state_write_indices=indices, + num_accepted_tokens=accepted, + a_raw=a.reshape(count, value_heads), + b_raw=b.reshape(count, value_heads), + run_config={"BV": 32, "num_warps": 4, "num_stages": 1}, + **gates, + ) + expected = torch.cat(outputs, dim=1).reshape_as(actual) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + torch.testing.assert_close(actual_states, expected_states, atol=0, rtol=0) + + 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.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py new file mode 100644 index 0000000000..22235598ed --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm.py @@ -0,0 +1,776 @@ +import pytest +import torch + + +def ReplaySSMExecutor(*args, **kwargs): + kda = kwargs.pop("kda", False) + if kda: + from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import _history_owner + else: + from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import _history_owner + + return _history_owner(*args, **kwargs) + + +def ssm_snapshot(backend, req): + return backend.accepted_ssm(req) + + +from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import get_configs, replay_layer + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) +@pytest.mark.parametrize("flush", [False, True], ids=["compress", "export"]) +@pytest.mark.parametrize("capture", [False, True], ids=["eager", "graph"]) +def test_multi_layer_fold_keeps_cursor_until_all_tiles_finish(kda, flush, capture): + state = torch.ones(36, 9, 6, 128, 128, device="cuda") + backend = ReplaySSMExecutor( + state, + capacity=16, + verify_width=4, + kda=kda, + num_key_heads=2, + run_config={"BV": 32, "num_warps": 4, "num_stages": 1}, + ) + backend.raw_keys.zero_() + backend.raw_values.zero_() + backend.gates.fill_(-0.1) + backend.betas.zero_() + backend.cursors[:8] = 13 + reqs = torch.arange(8, device="cuda", dtype=torch.int32) + expected = torch.full_like(state[:, :8], torch.exp(torch.tensor(-1.3)).item()) + targets = [torch.empty_like(state[:, 0]) for _ in range(8)] if flush else None + + def fold(): + if flush: + backend.export_accepted_state(list(range(8)), targets, flush=True) + else: + backend.maybe_compress_prefix(reqs) + + if capture: + fold() + state.fill_(1) + backend.cursors[:8] = 13 + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + fold() + state.fill_(1) + backend.cursors[:8] = 13 + graph.replay() + else: + fold() + if flush: + torch.testing.assert_close(torch.stack(targets, dim=1), expected, atol=1e-6, rtol=1e-6) + torch.testing.assert_close(state[:, :8], expected, atol=1e-6, rtol=1e-6) + assert not backend.cursors.any() + + +def reference(q, k, v, a, b, a_log, bias, state, kda=False, gate_lower_bound=-5.0): + q, k, v = q.float(), k.float(), v.float() + 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() + q = q.repeat_interleave(v.shape[-2] // q.shape[-2], -2) + k = k.repeat_interleave(v.shape[-2] // k.shape[-2], -2) + if kda: + g = gate_lower_bound * torch.sigmoid(a_log.exp()[:, None] * (a.float().view_as(k) + bias.view_as(k))) + state = state * g.exp()[..., None] + else: + g = -a_log.exp() * torch.nn.functional.softplus(a.float() + bias) + state = state * g.exp()[..., None, None] + d = b.float().sigmoid()[..., None] * (v - torch.einsum("...kv,...k->...v", state, k)) + state = state + k[..., None] * d[..., None, :] + return torch.einsum("...kv,...k->...v", state, q), state + + +@pytest.mark.parametrize("capacity", [4, 8]) +@pytest.mark.parametrize("bv", [8, 32], ids=["small_bv_fallback", "precomputed"]) +@pytest.mark.parametrize("state_dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("kda", [False, True], ids=["gdn", "kda"]) +@pytest.mark.parametrize("activation_dtype", [torch.float16, torch.bfloat16]) +def test_fp16_correction_history_preserves_finite_recurrence(capacity, bv, state_dtype, kda, activation_dtype): + state = torch.zeros(1, 2, 1, 128, 128, dtype=state_dtype, device="cuda") + ssm_backend = ReplaySSMExecutor( + state, + capacity, + 2, + activation_dtype, + num_key_heads=1, + kda=kda, + compute_mode="precompute", + run_config={"BV": bv, "num_warps": 1}, + ) + reqs = torch.tensor([0, 1], dtype=torch.int32, device="cuda") + cu = torch.tensor([0, 2, 4], dtype=torch.int32, device="cuda") + accepted = torch.tensor([1, 0], dtype=torch.int32, device="cuda") + q = torch.zeros(1, 4, 1, 128, dtype=activation_dtype, device="cuda") + q[0, :2, 0, 0] = 1 + k = torch.zeros_like(q) + k[0, 0, 0, 0] = -0.25 + k[0, 1, 0, 0] = 1 + v = torch.zeros_like(q) + v[0, :2, 0, 0] = 60000 + a = torch.full((4, 128 if kda else 1), -20, dtype=activation_dtype, device="cuda") + b = torch.full((4, 1), 20, dtype=activation_dtype, device="cuda") + alog = torch.zeros(1, device="cuda") + bias = torch.zeros(128 if kda else 1, device="cuda") + expected_state = state[0, 0].float().clone() + for round_idx in range(2): + ssm_backend.prepare_decode(reqs, cu) + # The second round must consume the first round's accepted history. + # For capacity 4 the next span exactly fills it, so no fold should hide + # an overflowed history correction by reconstructing from raw inputs. + assert ssm_backend.cursors[0].item() == 2 * round_idx + assert not state.any() + actual = replay_layer(ssm_backend, 0, q, k, v, a, b, alog, bias, reqs, cu)[0] + expected = [] + for token in range(2): + output, expected_state = reference( + q[0, token], k[0, token], v[0, token], a[token], b[token], alog, bias, expected_state, kda=kda + ) + expected.append(output) + assert torch.isfinite(expected_state).all() + assert torch.isfinite(actual).all() + torch.testing.assert_close( + actual[0, :2].float(), torch.stack(expected).to(actual.dtype).float(), atol=4, rtol=1e-3 + ) + ssm_backend.accept_updates(reqs, accepted) + # The correction is larger than FP16 can represent even though the + # recurrent state, inputs, and attention outputs are all finite. + assert ssm_backend.deltas[0, 0, 0, 1, 0].item() > torch.finfo(torch.float16).max + assert torch.isfinite(ssm_backend.deltas[0, 0, :, :4]).all() + assert ssm_backend.cursors[1].item() == 0 + 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, + run_config={"BV": 32, "num_warps": 1, "num_stages": 1}, + ) + 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 = torch.empty_like(expected) + + backend.export_accepted_state([0], [target], flush=True) + 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 = torch.empty_like(expected) + backend.export_accepted_state([0], [empty_history_target], flush=True) + 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)]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_replay_acceptance_flush_and_graph( + width, + dims, + capacity, + dtype, + run_config=None, + compute_mode="inline", + activation_dtype=torch.bfloat16, + gate_shift=-3, + head_geometry=(2, 4), + kda=False, + gate_lower_bound=-5.0, +): + torch.manual_seed(21) + kdim, vdim = dims + h, hv = head_geometry + slots = 5 + state = (torch.randn(2, slots, hv, kdim, vdim, device="cuda") * 0.01).to(dtype) + ssm_backend = ReplaySSMExecutor( + state, + capacity, + width, + num_key_heads=h, + kda=kda, + gate_lower_bound=gate_lower_bound, + compute_mode=compute_mode, + run_config=run_config, + activation_dtype=activation_dtype, + ) + expected = state.float().clone() + pending = [0, 0] + state_atol = 2e-5 if dtype == torch.float32 else 0.002 + # Independent FP32 reductions can straddle a BF16 rounding boundary. + # Allow one relative BF16 spacing as well as accumulated near-zero error. + state_rtol = 2e-4 if dtype == torch.float32 else torch.finfo(dtype).eps + reqs = torch.tensor([2, 0, slots - 1], device="cuda", dtype=torch.int32) + cu = torch.tensor([0, width, 2 * width - 1, 2 * width - 1], device="cuda", dtype=torch.int32) + if width == 1: + cu = torch.tensor([0, 1, 2, 2], device="cuda", dtype=torch.int32) + tokens = 3 * width + q = torch.randn(1, tokens, h, kdim, device="cuda", dtype=activation_dtype) + k = torch.randn_like(q) + v = torch.randn(1, tokens, hv, vdim, device="cuda", dtype=activation_dtype) + a = torch.randn(tokens, hv * kdim if kda else hv, device="cuda", dtype=activation_dtype) + gate_shift + b = torch.randn(tokens, hv, device="cuda", dtype=activation_dtype) + alog = torch.randn(hv, device="cuda") * 0.1 + bias = torch.randn(hv * kdim if kda else hv, device="cuda") * 0.1 + accepted = torch.zeros(slots, dtype=torch.int32, device="cuda") + + def step(): + ssm_backend.prepare_decode(reqs, cu) + out = [replay_layer(ssm_backend, layer, q, k, v, a, b, alog, bias, reqs, cu)[0] for layer in range(2)] + ssm_backend.accept_updates(reqs, accepted) + return out + + # Compile on disposable state, then capture with the HOLD slot only. + saved = state.clone() + step() + state.copy_(saved) + ssm_backend.cursors.zero_() + reqs.fill_(slots - 1) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + outputs = step() + reqs.copy_(torch.tensor([2, 0, slots - 1], device="cuda", dtype=torch.int32)) + for iteration in range(3 * capacity + 5): + counts = [iteration % width + 1, min((iteration + 1) % width + 1, max(width - 1, 1))] + accepted[2] = counts[0] - 1 + accepted[0] = counts[1] - 1 + for seq, req in enumerate([2, 0]): + query_len = width if seq == 0 else max(width - 1, 1) + if pending[seq] + query_len > capacity: + ssm_backend.merge_accepted_updates(torch.tensor([req], dtype=torch.int32, device="cuda")) + expected[:, req] = expected[:, req].to(dtype).float() + pending[seq] = 0 + pending[seq] += counts[seq] + graph.replay() + for layer in range(2): + for seq, req in enumerate([2, 0]): + start = seq * width + length = width if seq == 0 else max(width - 1, 1) + cur = expected[layer, req].clone() + for j in range(length): + t = start + j + out, cur = reference( + q[0, t], + k[0, t], + v[0, t], + a[t], + b[t], + alog, + bias, + cur, + kda=kda, + gate_lower_bound=gate_lower_bound, + ) + torch.testing.assert_close(outputs[layer][0, t].float(), out, atol=0.006, rtol=0.02) + if j + 1 == counts[seq]: + expected[layer, req].copy_(cur) + if iteration == 2 * capacity + 1: + ssm_backend.merge_accepted_updates(reqs[:2]) + expected = expected.to(dtype).float() + pending = [0, 0] + torch.testing.assert_close(state[:, :4].float(), expected[:, :4], atol=state_atol, rtol=state_rtol) + ssm_backend.merge_accepted_updates(reqs[:2]) + 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}, + 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} + ssm_backends = [ + ReplaySSMExecutor( + state.clone(), + 16 if width == 1 else 8, + width, + dtype, + num_key_heads=2, + compute_mode=mode, + run_config=config, + ) + 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((1, tokens, 2, 128), device="cuda", dtype=dtype) + k = torch.randn_like(q) + v = torch.randn((1, 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.merge_accepted_updates(overflow_reqs) + for ssm_backend in ssm_backends: + ssm_backend.prepare_decode(reqs, cu) + outputs.append(replay_layer(ssm_backend, 0, q, k, v, a, b, alog, alog, reqs, cu)[0]) + ssm_backend.accept_updates(reqs, accepted) + 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.merge_accepted_updates(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.merge_accepted_updates(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 = [] + 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)[0]) + 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, + run_config=config, + ) + for mode in ("inline", "precompute") + ] + cu = None if width == 1 else 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((1, width, heads, 128), device="cuda", dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn((1, 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)[0] 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: + backend.accept_updates(reqs, accepted, width=width) + backend.maybe_compress_prefix(reqs) + for backend in backends: + backend.merge_accepted_updates(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("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(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, compute_mode=compute_mode, run_config={"BV": 32, "num_warps": 1} + ) + 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(1, 8, 2, 128, dtype=dtype, device="cuda") + k = torch.randn_like(q) + v = torch.randn(1, 8, 4, 128, dtype=dtype, device="cuda") + a = torch.full((8, 4), -3.0, dtype=dtype, device="cuda") + b = torch.randn_like(a) + alog = torch.zeros(4, device="cuda") + + def forward(): + return [replay_layer(ssm_backend, layer, q, k, v, a, b, alog, alog, reqs, cu)[0] for layer in range(2)] + + def step(): + positions = ssm_backend.prepare_decode(reqs, cu) + out = forward() + ssm_backend.accept_updates(reqs, accepted) + 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.merge_accepted_updates(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.prepare_decode(reqs, cu)) + graph.replay() + if not capture_metadata: + ssm_backend.accept_updates(reqs, accepted) + 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[0, t], k[0, t], v[0, t], a[t], b[t], alog, alog, current) + torch.testing.assert_close(outputs[layer][0, 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.merge_accepted_updates(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.bounded_att import ConvSSMStateConfig + from unit_tests.common.test_linear_state import make_state + + torch.manual_seed(42) + state_rtol = 2e-5 if dtype == torch.float32 else torch.finfo(dtype).eps + config = ConvSSMStateConfig( + logical_layer_ids=(0, 1), + global_k_heads=2, + global_v_heads=2, + key_dim=32, + value_dim=64, + conv_width=3, + conv_dtype=torch.bfloat16, + ssm_dtype=dtype, + ) + state, _, layout = make_state(mode="replay", config=config) + state.initialize_page_pools(big_page_capacity=2, small_page_capacity=2) + state.mtp_state_index[1] = 1 + state.conv_state.normal_() + state.ssm_state.zero_() + state.raw_keys.normal_() + state.raw_values.normal_() + state.gates.uniform_(-0.1, -0.01) + state.betas.uniform_(0.1, 0.9) + active_before = state.ssm_state[:, 1].clone() + conv_expected = state.conv_state[:, 1, :, 1:4].clone() + block = torch.empty((1, layout.byte_end), dtype=torch.uint8, device="cuda") + kind = "big" if save_big else "small" + pool = state.big_pages_pool if save_big else state.small_pages_pool + for count in (2, 1): + state.cursors[1] = count + expected = active_before.float().clone() + for i in range(count): + key = state.raw_keys[:, 1, :, i].float() + key /= (key.square().sum(-1, keepdim=True) + 1e-6).sqrt() + expected *= state.gates[:, 1, :, i].exp()[..., None, None] + delta = state.betas[:, 1, :, i, None] * ( + state.raw_values[:, 1, :, i].float() - torch.einsum("lhkv,lhk->lhv", expected, key) + ) + expected += key[..., None] * delta[..., None, :] + materialized = ssm_snapshot(state, 1) + torch.testing.assert_close(materialized, expected.to(dtype), rtol=state_rtol, atol=2e-5) + state.save_bounded_att_runtime_to_pages([1], [8], kind, [0]) + state.write_to_pd_page(1, block, 0, 8, layout) + torch.cuda.synchronize() + torch.testing.assert_close(pool.ssm_state[0], materialized.cpu(), rtol=0, atol=0) + torch.testing.assert_close(pool.conv_state[0], conv_expected.cpu(), rtol=0, atol=0) + torch.testing.assert_close(state.ssm_state[:, 1], active_before, rtol=0, atol=0) + assert state.cursors[1].item() == count + for restore in ( + lambda: state.restore_bounded_att_runtime_from_pages([2], [8], kind, [0]), + lambda: state.read_from_pd_page(block, 2, 0, 8, layout), + ): + state.cursors[2] = 3 + state.mtp_state_index[2] = 2 + restore() + assert state.cursors[2].item() == state.mtp_state_index[2].item() == 0 + torch.testing.assert_close(state.ssm_state[:, 2], materialized, rtol=0, atol=0) + torch.testing.assert_close(state.conv_state[:, 2, :, :3], conv_expected, rtol=0, atol=0) + state.reset_bounded_att_state([2]) + assert not state.ssm_state[:, 2].any() + 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, run_config={"BV": 32, "num_warps": 1}) + 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.clear_history(reqs[folding:live]) + 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.merge_accepted_updates(reqs[:live]) + backend.prepare_decode(reqs, cu) + actual = replay_layer( + backend, + 0, + q.reshape(1, -1, 1, 32), + k.reshape(1, -1, 1, 32), + v.reshape(1, -1, 1, 32), + a.flatten(0, 1), + b.flatten(0, 1), + alog, + bias, + reqs, + cu, + )[0].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) + backend.commit_accepted(indices, reqs[:live], lengths) + backend.merge_accepted_updates(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]) 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..7261cf1adb --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_replayssm_execution.py @@ -0,0 +1,473 @@ +import pytest +import torch + + +def ReplaySSMExecutor(*args, **kwargs): + kda = kwargs.pop("kda", False) + if kda: + from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import _history_owner + else: + from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import _history_owner + + return _history_owner(*args, **kwargs) + + +from lightllm.common.basemodel.triton_kernel.linear_att.ssm_autotune import replay_layer +from lightllm.common.basemodel.triton_kernel.linear_att.mtp_state_params import ( + build_dynamic_mtp_linear_att_state_params, +) +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 import replayssm + + kernels = ("_decode_replay", "_advance_cursor") + hashes = {name: set() for name in kernels} + launches = [] + for name in kernels: + kernel = getattr(replayssm, name) + + 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, run_config={"BV": 32, "num_warps": 1}) + 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() + backend.prepare_decode(ids) + out = replay_layer(backend, 0, q[:count], k[:count], v[:count], a[:count], beta[:count], alog, bias, ids)[0] + backend.accept_single_token(ids) + 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.merge_accepted_updates(live) + captured.merge_accepted_updates(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", + run_config=config, + ) + 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(1, 12, 2, 128, device="cuda", dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn(1, 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.prepare_decode(reqs, cu) + + 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)[0] for layer in range(3)] + if index == 0: + ssm_backend.accept_updates(reqs, accepted) + else: + manager._ssm_executor.commit_accepted( + manager.mtp_state_index, + rows[starts.long()], + accepted[rows[starts.long()]] + 1, + ) + 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.merge_accepted_updates(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, + run_config={"BV": 32, "num_warps": 1, "num_stages": 1}, + ) + 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") + + def forward(ssm_backend): + return [replay_layer(ssm_backend, layer, q, k, v, a, beta, alog, bias, reqs)[0] for layer in range(3)] + + def warm(ssm_backend): + ssm_backend.prepare_decode(reqs) + out = forward(ssm_backend) + ssm_backend.accept_updates(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.prepare_decode(reqs) + out = forward(ssm_backend) + if capture_metadata: + ssm_backend.accept_updates(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.merge_accepted_updates(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.prepare_decode(reqs)) + graph.replay() + if not capture_metadata: + ssm_backend.accept_updates(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} + key_heads = 12 if kda else 4 + ssm_backends = [ + ReplaySSMExecutor(state, 8, 4, torch.bfloat16, num_key_heads=key_heads, kda=kda, run_config=config) + for state in [dense, sparse] + ] + q = torch.randn(1, 8, key_heads, 128, device="cuda", dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn(1, 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.merge_accepted_updates(reqs.new_tensor(overflow)) + pos = ssm_backend.prepare_decode(reqs, cu) + out = [replay_layer(ssm_backend, layer, q, k, v, a, beta, alog, bias, reqs, cu)[0] for layer in range(48)] + ssm_backend.accept_updates(reqs, accepted) + 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", run_config=config) + 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(1, batch * 4, 2, 128, device="cuda", dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn(1, 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(): + ssm_backend.prepare_decode(ids, offsets) + out = [ + replay_layer(ssm_backend, layer, q, k, v, a, beta, alog, bias, ids, offsets)[0] for layer in range(2) + ] + ssm_backend.accept_updates(ids, accepted) + 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.merge_accepted_updates(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, + run_config={"BV": 32, "num_warps": 1, "num_stages": 1}, + ) + reqs = torch.tensor([0, 1], device="cuda", dtype=torch.int32) + 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) + backend.merge_accepted_updates(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) + backend.accept_single_token(reqs) + backend.maybe_compress_prefix(reqs) + 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 new file mode 100644 index 0000000000..726eec31c4 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_ssm_autotune_memory.py @@ -0,0 +1,82 @@ +import pytest +import torch + +from lightllm.common.basemodel.triton_kernel.linear_att import ssm_autotune as tuning + + +def ReplaySSMExecutor(*args, **kwargs): + kda = kwargs.pop("kda", False) + if kda: + from lightllm.common.attention_state.bounded_att.kda.replay_kda_state import _history_owner + else: + from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import _history_owner + + return _history_owner(*args, **kwargs) + + +@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((1, tokens, 2, 32) if width > 1 else (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 if width > 1 else None + _, kwargs = tuning.rebuild_inputs(ssm_backend, "replay", q, q, q, gates, gates, a_log, a_log, cu_seqlens=cu) + scratch = kwargs["ssm_backend"] + 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 if width > 1 else None + if cu is not None: + q = q.unsqueeze(0) + inputs = dict( + ssm_backend=ssm_backend, mode="replay", q=q, k=q, v=q, a=gates, b=gates, a_log=a_log, bias=bias, cu_seqlens=cu + ) + 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.select_config + 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.resolve_run_config(**inputs) + 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) + + reqs = torch.arange(3, device="cuda", dtype=torch.int32) + ssm_backend.prepare_decode(reqs, cu) + tuning.replay_layer(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/test_gated_rmsnorm_strided_gate.py b/unit_tests/common/basemodel/triton_kernel/test_gated_rmsnorm_strided_gate.py new file mode 100644 index 0000000000..64215880e1 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/test_gated_rmsnorm_strided_gate.py @@ -0,0 +1,41 @@ +import pytest +import torch + +from lightllm.common.basemodel.triton_kernel.norm.gated_rmsnorm import gated_rmsnorm_forward + + +@pytest.mark.parametrize("norm_before_gate", [True, False]) +def test_gated_rmsnorm_accepts_strided_3d_gate(norm_before_gate): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for gated RMSNorm test") + + torch.manual_seed(123) + tokens, heads, head_dim = 5, 4, 128 + packed_dim = heads * head_dim + 256 + x = torch.randn((tokens * heads, head_dim), device="cuda", dtype=torch.bfloat16) + weight = torch.randn((head_dim,), device="cuda", dtype=torch.bfloat16) + packed = torch.randn((tokens, packed_dim), device="cuda", dtype=torch.bfloat16) + z = packed[:, 128 : 128 + heads * head_dim].view(tokens, heads, head_dim) + assert not z.is_contiguous() + + run_config = {"BLOCK_N": head_dim, "num_warps": 1} + expected = gated_rmsnorm_forward( + x=x, + weight=weight, + bias=None, + eps=1e-6, + z=z.contiguous().view(-1, head_dim), + norm_before_gate=norm_before_gate, + run_config=run_config, + ) + actual = gated_rmsnorm_forward( + x=x, + weight=weight, + bias=None, + eps=1e-6, + z=z, + norm_before_gate=norm_before_gate, + run_config=run_config, + ) + + torch.testing.assert_close(actual, expected, rtol=0, atol=0) diff --git a/unit_tests/common/basemodel/triton_kernel/test_scalar_linear.py b/unit_tests/common/basemodel/triton_kernel/test_scalar_linear.py new file mode 100644 index 0000000000..2cbd3559d6 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/test_scalar_linear.py @@ -0,0 +1,70 @@ +import pytest +import torch + +from lightllm.common.basemodel.triton_kernel.scalar_linear import scalar_linear +from lightllm.common.quantization.scalar_projection import StableScalarProjection +from lightllm.common.quantization.quantize_method import WeightPack + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("hidden", [65, 2560]) +def test_reference_batching_strides_and_graph(dtype, hidden): + torch.manual_seed(2319) + input = torch.randn((336, hidden * 2), device="cuda", dtype=dtype)[1:335, ::2] + weight = torch.randn((1, hidden * 2), device="cuda", dtype=dtype)[:, ::2] + out = torch.empty((334, 3), device="cuda", dtype=dtype)[:, 1:2] + expected = (input.double() * weight.double()).sum(-1, keepdim=True).to(dtype) + assert scalar_linear(input, weight, out=out) is out + torch.testing.assert_close(out, expected, atol=0, rtol=0) + for size in (1, 3, 8, 24, 64, 128, 256, 334): + pieces = [scalar_linear(chunk, weight) for chunk in input.split(size)] + assert torch.equal(torch.cat(pieces), out) + permutation = torch.randperm(334, device="cuda") + assert torch.equal(scalar_linear(input[permutation], weight), out[permutation]) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + scalar_linear(input, weight, out=out) + input.mul_(0.5) + graph.replay() + torch.cuda.synchronize() + expected = (input.double() * weight.double()).sum(-1, keepdim=True).to(dtype) + torch.testing.assert_close(out, expected, atol=0, rtol=0) + + +def test_small_terms_survive_before_final_dtype_cast(): + input = torch.tensor([[1, 2 ** -8, 2 ** -24, 2 ** -24]], device="cuda", dtype=torch.bfloat16) + weight = torch.ones_like(input) + actual = scalar_linear(input, weight) + torch.testing.assert_close(actual, torch.full_like(actual, 1.0078125), atol=0, rtol=0) + + +def test_dispatch_preserves_weight_layout_output_and_fallbacks(monkeypatch): + monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", str(torch.cuda.current_device())) + method = StableScalarProjection() + input = torch.randn((7, 65), device="cuda", dtype=torch.bfloat16) + weight = torch.randn((1, 65), device="cuda", dtype=input.dtype) + pack = WeightPack(weight=weight, weight_scale=None, weight_zero_point=None) + out = torch.empty((7, 1), device="cuda", dtype=input.dtype) + assert method.method_name == "none" + assert method.apply(input, pack, out=out, use_custom_tensor_mananger=False) is out + expected = (input.double() @ weight.double().t()).to(input.dtype) + torch.testing.assert_close(out, expected, atol=0, rtol=0) + torch.testing.assert_close(method.apply(input, pack), expected, atol=0, rtol=0) + + for outputs, bias in ((2, None), (1, torch.ones((1,), device="cuda", dtype=input.dtype))): + weight = torch.randn((outputs, 65), device="cuda", dtype=input.dtype) + pack = WeightPack(weight=weight, weight_scale=None, weight_zero_point=None) + actual = method.apply(input, pack, bias=bias, use_custom_tensor_mananger=False) + expected = torch.mm(input, weight.t()) if bias is None else torch.addmm(bias, input, weight.t()) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + +def test_empty_input_and_invalid_output(): + input = torch.empty((0, 65), device="cuda", dtype=torch.bfloat16) + weight = torch.ones((1, 65), device="cuda", dtype=input.dtype) + assert scalar_linear(input, weight).shape == (0, 1) + with pytest.raises(ValueError, match="output must match"): + scalar_linear(input, weight, out=torch.empty((0, 2), device="cuda", dtype=input.dtype)) diff --git a/unit_tests/common/gated_residual/test_gated_residual.py b/unit_tests/common/gated_residual/test_gated_residual.py new file mode 100644 index 0000000000..4833299460 --- /dev/null +++ b/unit_tests/common/gated_residual/test_gated_residual.py @@ -0,0 +1,178 @@ +import pytest +import torch +import torch.nn.functional as F + +from lightllm.common.gated_residual import GatedResidual + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +def make_component(dim, branches, dtype, combine=True): + width, rank = dim * branches, 32 + + def rand(*shape): + return torch.randn(*shape, device="cuda", dtype=dtype) + + return GatedResidual( + rand(width) * 0.1, + rand(rank, width) / width ** 0.5, + rand(width, rank) / rank ** 0.5, + rand(branches, width) / width ** 0.5 if combine else None, + branches, + ) + + +def reference_read(component, x): + grouped = x.float().unflatten(-1, (component.branches, -1)) + normalized = grouped * torch.rsqrt(grouped.square().mean(-1, keepdim=True) + component.eps) + normalized = (normalized.flatten(-2) * (1 + component.norm_weight.float())).to(x.dtype) + lowrank = F.silu(F.linear(normalized, component.down_weight) / component.branches) + gates = torch.sigmoid(F.linear(lowrank, component.up_weight)) + mixed = (normalized * gates).unflatten(-1, (component.branches, -1)).mean(-2) + if component.inject_weight is None: + return mixed + injection = 2 * torch.sigmoid(F.linear(normalized, component.inject_weight) / component.branches) + return mixed, injection + + +def reference_write(residual, update, gates): + return residual + (update.unsqueeze(-2) * gates.unsqueeze(-1)).flatten(-2) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("tokens,dim,branches", [(1, 65, 3), (17, 128, 4), (2, 2560, 4)]) +@pytest.mark.parametrize("combine", [True, False]) +def test_reference_parity(dtype, tokens, dim, branches, combine): + torch.manual_seed(37) + component = make_component(dim, branches, dtype, combine) + residual = torch.randn(1, tokens, dim * branches, device="cuda", dtype=dtype) + original = residual.clone() + actual, expected = component.read(residual), reference_read(component, residual) + + tolerance = max(1e-6, torch.finfo(dtype).eps) + if combine: + for result, ref in zip(actual, expected): + torch.testing.assert_close(result, ref, atol=tolerance, rtol=tolerance) + update = torch.randn(1, tokens, dim, device="cuda", dtype=dtype) + written = component.write(residual, update, actual[1]) + + torch.testing.assert_close(written, reference_write(residual, update, actual[1]), atol=0, rtol=0) + else: + torch.testing.assert_close(actual, expected, atol=tolerance, rtol=tolerance) + torch.testing.assert_close(residual, original, atol=0, rtol=0) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("dim,branches", [(65, 3), (2560, 4)]) +def test_injection_gate_projection_is_independent_of_batch(dtype, dim, branches): + from lightllm.common.gated_residual.kernels import injection_gates + + torch.manual_seed(3107) + x = torch.randn(9, dim * branches, device="cuda", dtype=dtype) + weight = torch.randn(branches, dim * branches, device="cuda", dtype=dtype) * 0.01 + actual = injection_gates(x, weight, branches) + singles = torch.cat([injection_gates(row[None], weight, branches) for row in x]) + chunks = torch.cat([injection_gates(x[start : start + 3], weight, branches) for start in range(0, len(x), 3)]) + torch.testing.assert_close(actual, singles, atol=0, rtol=0) + torch.testing.assert_close(actual, chunks, atol=0, rtol=0) + expected = 2 * torch.sigmoid((x.double() @ weight.double().T).to(dtype) / branches) + tolerance = max(1e-6, torch.finfo(dtype).eps) + torch.testing.assert_close(actual, expected, atol=tolerance, rtol=tolerance) + + +def test_saturated_mix_and_write(): + from lightllm.common.gated_residual.kernels import mix_streams + + component = make_component(65, 3, torch.bfloat16) + x = torch.randn(2, 195, device="cuda", dtype=torch.bfloat16) + logits = torch.linspace(-100, 100, 195, device="cuda", dtype=x.dtype).expand(2, -1).contiguous() + expected = (x * torch.sigmoid(logits)).view(2, 3, 65).mean(1) + torch.testing.assert_close(mix_streams(x, logits, 3), expected, atol=0, rtol=0) + gates = torch.tensor([[0, 1, 2], [2, 0, 1]], device="cuda", dtype=x.dtype) + update = torch.randn(2, 65, device="cuda", dtype=x.dtype) + torch.testing.assert_close(component.write(x, update, gates), reference_write(x, update, gates), atol=0, rtol=0) + + +def test_checkpoint_binding_and_graph_replay(): + component = make_component(128, 4, torch.bfloat16) + prefix = "model.language_model.layers.0.attn_hyper_connection" + weights = { + f"{prefix}.hc_norm.weight": component.norm_weight, + f"{prefix}.input_mix_weight_down.weight": component.down_weight, + f"{prefix}.input_mix_weight_up.weight": component.up_weight, + f"{prefix}.block_inject_weight.weight": component.inject_weight, + } + bound = GatedResidual.from_checkpoint(weights, prefix, branches=4) + assert bound.down_weight is component.down_weight + final = GatedResidual.from_checkpoint(weights, prefix, branches=4, use_combine=False) + residual = torch.randn(3, 512, device="cuda", dtype=torch.bfloat16) + update = torch.randn(3, 128, device="cuda", dtype=residual.dtype) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + _, gates = bound.read(residual) + bound.write(residual, update, gates) + final.read(residual) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + mixed, gates = bound.read(residual) + output = bound.write(residual, update, gates) + final_output = final.read(output) + for _ in range(2): + residual.normal_() + graph.replay() + expected_mixed, expected_gates = reference_read(bound, residual) + expected_output = reference_write(residual, update, expected_gates) + torch.testing.assert_close(mixed, expected_mixed, atol=2e-3, rtol=2e-3) + torch.testing.assert_close(output, expected_output, atol=2e-3, rtol=2e-3) + torch.testing.assert_close(final_output, reference_read(final, expected_output), atol=2e-3, rtol=2e-3) + + +def test_empty_and_invalid_inputs(): + component = make_component(65, 3, torch.bfloat16) + x = torch.empty(0, 195, device="cuda", dtype=torch.bfloat16) + mixed, gates = component.read(x) + assert component.write(x, mixed, gates).shape == x.shape + with pytest.raises(ValueError, match="last dimension"): + component.read(torch.empty(2, 194, device="cuda", dtype=x.dtype)) + with pytest.raises(ValueError, match="contiguous"): + component.read(torch.empty(2, 390, device="cuda", dtype=x.dtype)[:, ::2]) + with pytest.raises(ValueError, match="matching token"): + component.write(x, mixed, torch.empty(1, 3, device="cuda", dtype=x.dtype)) + with pytest.raises(ValueError, match="share"): + component.read(x.float()) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("dim,branches,rank", [(65, 3, 9), (2560, 4, 320)]) +def test_mixer_is_batch_invariant_and_matches_fp64(dtype, dim, branches, rank): + from lightllm.common.gated_residual.kernels import grouped_norm + + torch.manual_seed(701) + width = dim * branches + + def rand(*shape): + return torch.randn(*shape, device="cuda", dtype=dtype) + + component = GatedResidual( + rand(width) * 0.1, + rand(rank, width) / width ** 0.5, + rand(width, rank) / rank ** 0.5, + rand(branches, width) / width ** 0.5, + branches, + ) + residual = rand(129, width) + expected = [component.read(row[None, :]) for row in residual] + scalar = tuple(torch.cat([item[index] for item in expected]) for index in range(2)) + for count in (3, 8, 16, 24, 64, 101, 128, 129): + outputs = [component.read(residual[start : start + count]) for start in range(0, 129, count)] + for index in range(2): + torch.testing.assert_close(torch.cat([item[index] for item in outputs]), scalar[index], atol=0, rtol=0) + normalized = grouped_norm(residual, component.norm_weight, branches, component.eps) + projected = F.linear(normalized.double(), component.down_weight.double()).to(dtype) + lowrank = F.silu(projected / branches) + logits = F.linear(lowrank.double(), component.up_weight.double()).to(dtype) + oracle = (normalized * torch.sigmoid(logits)).unflatten(-1, (branches, dim)).mean(-2) + torch.testing.assert_close(scalar[0], oracle, atol=2e-3, rtol=torch.finfo(dtype).eps) diff --git a/unit_tests/common/ngram_embedding/test_hash.py b/unit_tests/common/ngram_embedding/test_hash.py new file mode 100644 index 0000000000..ad217acf17 --- /dev/null +++ b/unit_tests/common/ngram_embedding/test_hash.py @@ -0,0 +1,206 @@ +import pytest +import torch + +from lightllm.common.ngram_embedding import NGramHasher + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +EOS = 248044 + + +def make_hasher(ngram_size=3, heads_per_ngram=8, overflow=False): + heads = (ngram_size - 1) * heads_per_ngram + multipliers = [39719048695937, 29866592791859, 5701618528319, 123456789123, 9876543212345][:ngram_size] + if overflow: + multipliers = [(1 << 62) + i * 141 + 1 for i in range(ngram_size)] + sizes = torch.arange(heads, device="cuda", dtype=torch.int64) * 2 + 20000003 + offsets = torch.cat([torch.zeros(1, device="cuda", dtype=torch.int64), sizes.cumsum(0)[:-1]]) + return NGramHasher(torch.tensor(multipliers, device="cuda", dtype=torch.int64), sizes, offsets, EOS) + + +def pack(sequences, dtype=torch.int64): + lengths = torch.tensor([0] + [len(seq) for seq in sequences], device="cuda", dtype=torch.int32) + tokens = torch.tensor([token for seq in sequences for token in seq], device="cuda", dtype=dtype) + return tokens, lengths.cumsum(0).to(torch.int32) + + +def reference(hasher, sequences, history): + rows, next_histories = [], [] + context = hasher.context_len + multipliers = hasher.multipliers.cpu() + sizes = hasher.head_vocab_sizes.tolist() + offsets = hasher.head_offsets.tolist() + per_order = hasher.heads // context + for request, sequence in enumerate(sequences): + full = history[request].tolist() + sequence + for index in range(context, len(full)): + shifted = [full[index]] + for shift in range(1, context + 1): + preceding = full[index - shift : index] + shifted.append(EOS if EOS in preceding else full[index - shift]) + row = [] + for order in range(2, context + 2): + mixed = torch.tensor(shifted[0], dtype=torch.int64) * multipliers[0] + for pos in range(1, order): + mixed ^= torch.tensor(shifted[pos], dtype=torch.int64) * multipliers[pos] + begin = (order - 2) * per_order + for head in range(begin, begin + per_order): + row.append((int(mixed) % sizes[head]) + offsets[head]) + rows.append(row) + next_histories.append(full[-context:]) + indices = torch.tensor(rows, device="cuda", dtype=torch.int64).reshape(-1, hasher.heads) + next_history = torch.tensor(next_histories, device="cuda", dtype=torch.int64).reshape(-1, context) + return indices, next_history + + +@pytest.mark.parametrize("ngram_size,heads_per_ngram", [(2, 1), (3, 8), (4, 3), (5, 2)]) +@pytest.mark.parametrize("dtype", [torch.int32, torch.int64]) +@pytest.mark.parametrize("overflow", [False, True]) +def test_hash_matches_eos_segment_reference(ngram_size, heads_per_ngram, dtype, overflow): + hasher = make_hasher(ngram_size, heads_per_ngram, overflow) + sequences = [[], [7], [EOS], [11, EOS, 13, 14, EOS, EOS, 15, 16, 17], [], list(range(130))] + history = hasher.empty_history(len(sequences)) + history[0, -1] = 777 + history[3, -1] = 999 + if hasher.context_len > 1: + history[3, -2] = EOS + before = history.clone() + tokens, starts = pack(sequences, dtype) + actual = hasher.hash(tokens, starts, history) + expected = reference(hasher, sequences, history) + for result, ref in zip(actual, expected): + torch.testing.assert_close(result, ref, atol=0, rtol=0) + torch.testing.assert_close(history, before, atol=0, rtol=0) + assert torch.all(actual[0] >= hasher.head_offsets) + assert torch.all(actual[0] < hasher.head_offsets + hasher.head_vocab_sizes) + + +def test_chunked_prefill_and_decode_match_full_sequence(): + hasher = make_hasher() + sequences = [[1, 2, EOS, 3, 4, 5, EOS, 6], [11, EOS, 12, 13, EOS, 14], [EOS, 31]] + original = hasher.empty_history(len(sequences)) + expected, full_history = hasher.hash(*pack(sequences), original) + history = original + per_request = [[] for _ in sequences] + consumed = [0] * len(sequences) + for counts in [(1, 0, 1), (2, 3, 1), (1, 1, 0), (4, 2, 0)]: + chunks = [sequence[start : start + count] for sequence, start, count in zip(sequences, consumed, counts)] + tokens, starts = pack(chunks) + result, history = hasher.hash(tokens, starts, history) + for request, count in enumerate(counts): + per_request[request].append(result[int(starts[request]) : int(starts[request + 1])]) + consumed[request] += count + combined = torch.cat([torch.cat(chunks) for chunks in per_request]) + torch.testing.assert_close(combined, expected, atol=0, rtol=0) + torch.testing.assert_close(history, full_history, atol=0, rtol=0) + + +def test_speculative_rejection_and_cache_snapshot(): + hasher = make_hasher() + sequences = [[1, 2, 3, EOS], [11, EOS, 12], [21, 22], []] + tokens, starts = pack(sequences) + initial = hasher.empty_history(len(sequences)) + initial[:, -1] = 99 + snapshot = initial.clone() + + hasher.hash(tokens, starts, initial) + accepted = torch.tensor([1, 2, 0, 0], device="cuda", dtype=torch.int32) + committed = hasher.history_after(tokens, starts, initial, accepted) + expected = reference(hasher, [seq[:n] for seq, n in zip(sequences, [1, 2, 0, 0])], snapshot)[1] + torch.testing.assert_close(committed, expected, atol=0, rtol=0) + torch.testing.assert_close(initial, snapshot, atol=0, rtol=0) + + followup = [[7], [8], [9], [10]] + result = hasher.hash(*pack(followup), committed) + ref = reference(hasher, followup, expected) + for actual, expected in zip(result, ref): + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + order = torch.tensor([2, 0, 2], device="cuda") + histories = snapshot.index_select(0, order) + sequences = [[31], [32, EOS, 33], [34]] + for actual, expected in zip(hasher.hash(*pack(sequences), histories), reference(hasher, sequences, histories)): + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + +def test_checkpoint_and_cuda_graph_replay(): + source = make_hasher() + prefix = "model.language_model.layers.1.ple.ple_embedding" + buffers = { + f"{prefix}.layer_multipliers": source.multipliers, + f"{prefix}.ngram_heads_vocab_sizes": source.head_vocab_sizes, + f"{prefix}.ngram_heads_offsets": source.head_offsets, + } + hasher = NGramHasher.from_checkpoint(buffers, prefix, EOS) + assert hasher.multipliers is source.multipliers + history = hasher.empty_history(3) + tokens, starts = pack([[1, 2, 3], [11, 12], [21, 22]]) + accepted = torch.tensor([3, 1, 2], dtype=torch.int32, device="cuda") + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + hasher.hash(tokens, starts, history) + hasher.history_after(tokens, starts, history, accepted) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + indices, next_history = hasher.hash(tokens, starts, history) + committed = hasher.history_after(tokens, starts, history, accepted) + for sequences, counts in [ + ([[EOS, 2], [], [11, 12, 13, 14, 15]], [1, 0, 4]), + ([[], [31, EOS, 32], [41, 42, 43, 44]], [0, 2, 0]), + ]: + updated_tokens, updated_starts = pack(sequences) + tokens.copy_(updated_tokens) + starts.copy_(updated_starts) + accepted.copy_(torch.tensor(counts, device="cuda", dtype=torch.int32)) + history[:, -1] = 151 + graph.replay() + ref_indices, ref_history = reference(hasher, sequences, history) + torch.testing.assert_close(indices, ref_indices, atol=0, rtol=0) + torch.testing.assert_close(next_history, ref_history, atol=0, rtol=0) + ref_committed = reference(hasher, [seq[:n] for seq, n in zip(sequences, counts)], history)[1] + torch.testing.assert_close(committed, ref_committed, atol=0, rtol=0) + + +def test_empty_batches_and_bad_layouts(): + hasher = make_hasher() + for sequences in [[], [[], [], []]]: + history = hasher.empty_history(len(sequences)) + indices, next_history = hasher.hash(*pack(sequences), history) + assert indices.shape == (0, hasher.heads) + torch.testing.assert_close(next_history, history, atol=0, rtol=0) + tokens, starts = pack([[1, 2, 3]]) + history = hasher.empty_history(1) + with pytest.raises(ValueError, match="packed"): + hasher.hash(tokens.view(1, -1), starts, history) + with pytest.raises(ValueError, match="Starts"): + hasher.hash(tokens, starts.long(), history) + with pytest.raises(ValueError, match="History"): + hasher.hash(tokens, starts, history.int()) + with pytest.raises(ValueError, match="contiguous"): + hasher.hash(tokens, starts, torch.empty(1, 4, device="cuda", dtype=torch.int64)[:, ::2]) + with pytest.raises(ValueError, match="Accepted"): + hasher.history_after(tokens, starts, history, torch.ones(1, device="cuda", dtype=torch.int64)) + + +@pytest.mark.parametrize("invalid", ["multipliers", "empty_heads", "dtype", "zero_size", "overlap"]) +def test_invalid_checkpoint_buffers(invalid): + from dataclasses import replace + + hasher = make_hasher() + with pytest.raises(ValueError): + if invalid == "multipliers": + replace(hasher, multipliers=hasher.multipliers[:1]) + elif invalid == "empty_heads": + replace(hasher, head_vocab_sizes=hasher.head_vocab_sizes[:0]) + elif invalid == "dtype": + replace(hasher, multipliers=hasher.multipliers.int()) + elif invalid == "zero_size": + sizes = hasher.head_vocab_sizes.clone() + sizes[0] = 0 + replace(hasher, head_vocab_sizes=sizes) + else: + offsets = hasher.head_offsets.clone() + offsets[1] = 0 + replace(hasher, head_offsets=offsets) diff --git a/unit_tests/common/ngram_embedding/test_pinned_embedding.py b/unit_tests/common/ngram_embedding/test_pinned_embedding.py new file mode 100644 index 0000000000..57baa1744f --- /dev/null +++ b/unit_tests/common/ngram_embedding/test_pinned_embedding.py @@ -0,0 +1,88 @@ +import importlib +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.basemodel.layer_weights.meta_weights.pinned_embedding_weight import PinnedRowShardedEmbeddingWeight +from lightllm.models.qwen4_exp import pinned_memory + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +@pytest.fixture(autouse=True) +def runtime(monkeypatch): + base = importlib.import_module("lightllm.common.basemodel.layer_weights.meta_weights.base_weight") + platform = importlib.import_module("lightllm.common.basemodel.layer_weights.meta_weights.platform_op") + monkeypatch.setattr(base, "get_current_device_id", torch.cuda.current_device) + monkeypatch.setattr( + platform, + "get_env_start_args", + lambda: SimpleNamespace(hardware_platform="cuda", enable_torch_fallback=False, enable_triton_fallback=False), + ) + monkeypatch.setattr(pinned_memory, "_LARGE_BYTES", 1) + monkeypatch.setattr(pinned_memory, "_CHUNK_BYTES", 65536) + + +def weight(dtype, rank=0, world=1): + return PinnedRowShardedEmbeddingWeight( + 160, 1024, "table.weight", dtype, 4, cpu_allocator=pinned_memory.pinned_zeros, tp_rank=rank, tp_world_size=world + ) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +def test_streamed_host_shards_and_vocab_parallel_gather(dtype): + table = (torch.arange(1024 * 160).reshape(1024, 160) % 127 - 64).to(dtype) + ids = torch.tensor([0, 255, 256, 340, 341, 511, 512, 681, 682, 1023, -1, 1024], device="cuda") + outputs = [] + for rank in range(3): + value = weight(dtype, rank, 3) + address = value.weight.data_ptr() + assert not value.verify_load() and value.weight.device.type == "cpu" and value.weight.is_pinned() + assert len(value.weight._qwen4_pinned_allocation.addresses) > 1 + for index in (2, 0, 3, 1): + value.load_hf_weights({f"table.shard_{index}.weight": table[index * 256 : (index + 1) * 256]}) + assert value.verify_load() and value.weight.data_ptr() == address + torch.testing.assert_close(value.weight, table[value.tp_vocab_start_id : value.tp_vocab_end_id], atol=0, rtol=0) + output = value(ids) + assert output.is_cuda + outputs.append(output) + expected = table.cuda().index_select(0, ids.clamp(0, 1023)) + expected[(ids < 0) | (ids >= 1024)] = 0 + torch.testing.assert_close(torch.stack(outputs).sum(0), expected, atol=0, rtol=0) + + +def test_device_checkpoint_staging_and_dynamic_graph_replay(): + table = (torch.arange(1024 * 160).reshape(1024, 160) % 127).to(torch.bfloat16).cuda() + value = weight(torch.bfloat16) + value.load_hf_weights({"table.weight": table}) + assert value.verify_load() + ids = torch.tensor([0, 191, 192, 511, 512, 1023], device="cuda") + output = torch.empty((ids.numel(), 160), dtype=table.dtype, device="cuda") + value(ids, out=output) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + value(ids, out=output) + ids.copy_(torch.tensor([1023, 512, 511, 192, 191, 0], device="cuda")) + graph.replay() + torch.testing.assert_close(output, table.index_select(0, ids), atol=0, rtol=0) + graph.reset() + + +def test_rejects_pageable_allocator_and_wrong_output(): + with pytest.raises(ValueError, match="pinned CPU"): + PinnedRowShardedEmbeddingWeight( + 160, + 1024, + "table.weight", + torch.bfloat16, + 4, + cpu_allocator=lambda shape, **kwargs: torch.zeros(shape, dtype=kwargs["dtype"]), + tp_rank=0, + tp_world_size=1, + ) + value = weight(torch.bfloat16) + with pytest.raises(ValueError, match="CUDA vector"): + value(torch.tensor([0])) + with pytest.raises(ValueError, match="output"): + value(torch.tensor([0], device="cuda"), out=torch.empty((1, 160), device="cuda", dtype=torch.float32)) diff --git a/unit_tests/common/ngram_embedding/test_ple.py b/unit_tests/common/ngram_embedding/test_ple.py new file mode 100644 index 0000000000..677a923734 --- /dev/null +++ b/unit_tests/common/ngram_embedding/test_ple.py @@ -0,0 +1,259 @@ +import math +from dataclasses import replace + +import pytest +import torch +import torch.nn.functional as F + +from lightllm.common.ngram_embedding.ple import PLELayer + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +def make_layer(dim=65, branches=3, embed_dim=32, kernel=4, dilation=3, dtype=torch.bfloat16): + def rand(*shape): + return torch.randn(*shape, device="cuda", dtype=dtype) + + width = dim * branches + return PLELayer( + rand(width, embed_dim) / embed_dim ** 0.5, + rand(dim, embed_dim) / embed_dim ** 0.5, + rand(width) * 0.1, + rand(width) * 0.1, + rand(width) * 0.1, + rand(width, kernel) * 0.1, + branches, + dilation, + ) + + +def norm(x, weight, branches, eps): + grouped = x.float().unflatten(-1, (branches, -1)) + grouped = grouped * torch.rsqrt(grouped.square().mean(-1, keepdim=True) + eps) + return (grouped.flatten(-2) * (1 + weight.float())).to(x.dtype) + + +def reference_values(layer, residual, embeddings, mask=None): + key = norm(F.linear(embeddings, layer.key_weight), layer.norm_key, layer.branches, layer.eps) + query = norm(residual, layer.norm_query, layer.branches, layer.eps) + key, query = [x.unflatten(-1, (layer.branches, layer.hidden_size)) for x in [key, query]] + value = F.linear(embeddings, layer.value_weight) + gate = (key * query).sum(-1, keepdim=True) / math.sqrt(layer.hidden_size) + gate = gate.abs().clamp_min(1e-6).sqrt() * gate.sign() + gated = (torch.sigmoid(gate) * value.unsqueeze(-2)).flatten(-2) + normalized = norm(gated, layer.norm_conv, layer.branches, layer.eps) + if mask is not None: + gated = gated * mask[:, None] + normalized = normalized * mask[:, None] + return gated, normalized + + +def reference_conv(layer, values, gated, lengths, history): + result, next_state = [], [] + offset = 0 + for request, count in enumerate(lengths): + current = values[offset : offset + count] + full = torch.cat([history[request], current]) + if count: + conv = ( + F.conv1d( + full.T.unsqueeze(0), + layer.conv_weight.unsqueeze(1), + dilation=layer.dilation, + groups=values.shape[1], + ) + .squeeze(0) + .T + ) + result.append(gated[offset : offset + count] + F.silu(conv)) + next_state.append(full[-layer.context_len :] if layer.context_len else full[:0]) + offset += count + output = torch.cat(result) if result else values.new_empty((0, values.shape[1])) + state = torch.stack(next_state) if next_state else history.clone() + return output, state + + +def starts_for(lengths): + return torch.tensor([0] + list(lengths), device="cuda", dtype=torch.int32).cumsum(0).to(torch.int32) + + +def assert_close(actual, expected, dtype): + tolerance = max(2e-6, torch.finfo(dtype).eps * 2) + torch.testing.assert_close(actual, expected, atol=tolerance, rtol=tolerance) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("dim,branches,kernel,dilation", [(65, 3, 4, 3), (128, 4, 3, 2), (32, 2, 1, 1)]) +@pytest.mark.parametrize("masked", [False, True]) +def test_ple_reference_equations(dtype, dim, branches, kernel, dilation, masked): + torch.manual_seed(137) + layer = make_layer(dim, branches, kernel=kernel, dilation=dilation, dtype=dtype) + lengths = [0, 1, 2, 13, 0, 7] + count = sum(lengths) + residual = torch.randn(count, dim * branches, device="cuda", dtype=dtype) + embeddings = torch.randn(count, 32, device="cuda", dtype=dtype) + mask = (torch.arange(count, device="cuda") % 3 != 0) if masked else None + history = layer.empty_history(len(lengths)).normal_() + original = history.clone() + actual, conv_inputs = layer.forward(residual, embeddings, starts_for(lengths), history, mask) + gated, expected_conv_inputs = reference_values(layer, residual, embeddings, mask) + expected, expected_state = reference_conv(layer, expected_conv_inputs, gated, lengths, history) + assert_close(conv_inputs, expected_conv_inputs, dtype) + assert_close(actual, expected, dtype) + assert_close(layer.history_after(conv_inputs, starts_for(lengths), history), expected_state, dtype) + torch.testing.assert_close(history, original, atol=0, rtol=0) + + +def test_production_hidden_width(): + torch.manual_seed(137) + layer = make_layer(dim=2560, branches=4, embed_dim=2560) + residual = torch.randn(2, 10240, device="cuda", dtype=torch.bfloat16) + embeddings = torch.randn(2, 2560, device="cuda", dtype=torch.bfloat16) + history = layer.empty_history(2) + output, inputs = layer.forward(residual, embeddings, starts_for([1, 1]), history) + gated, expected_inputs = reference_values(layer, residual, embeddings) + expected, _ = reference_conv(layer, expected_inputs, gated, [1, 1], history) + assert_close(inputs, expected_inputs, residual.dtype) + assert_close(output, expected, residual.dtype) + + +def test_chunking_speculative_acceptance_and_request_isolation(): + torch.manual_seed(137) + layer = make_layer(dim=32, branches=2, dtype=torch.float32) + lengths = [8, 6, 2] + starts = starts_for(lengths) + residual = torch.randn(sum(lengths), 64, device="cuda") + embeddings = torch.randn(sum(lengths), 32, device="cuda") + history = layer.empty_history(3).normal_() + full, full_inputs = layer.forward(residual, embeddings, starts, history) + full_state = layer.history_after(full_inputs, starts, history) + state = history.clone() + consumed = [0] * 3 + per_request = [[] for _ in range(3)] + for counts in [(1, 0, 1), (2, 3, 1), (1, 1, 0), (4, 2, 0)]: + slices = [slice(int(starts[i]) + consumed[i], int(starts[i]) + consumed[i] + n) for i, n in enumerate(counts)] + chunks = [torch.cat([tensor[s] for s in slices]) for tensor in [residual, embeddings]] + chunk_starts = starts_for(counts) + output, conv_inputs = layer.forward(*chunks, chunk_starts, state) + state = layer.history_after(conv_inputs, chunk_starts, state) + for i, n in enumerate(counts): + per_request[i].append(output[int(chunk_starts[i]) : int(chunk_starts[i + 1])]) + consumed[i] += n + assert_close(torch.cat([torch.cat(parts) for parts in per_request]), full, residual.dtype) + assert_close(state, full_state, residual.dtype) + accepted = torch.tensor([1, 4, 0], device="cuda", dtype=torch.int32) + committed = layer.history_after(full_inputs, starts, history, accepted) + expected = torch.stack( + [ + torch.cat([history[i], full_inputs[int(starts[i]) : int(starts[i]) + n]])[-layer.context_len :] + for i, n in enumerate([1, 4, 0]) + ] + ) + torch.testing.assert_close(committed, expected, atol=0, rtol=0) + + order = torch.tensor([2, 0, 2], device="cuda") + reordered = committed.index_select(0, order) + new_residual, new_embeddings = residual[:3], embeddings[:3] + actual, _ = layer.forward(new_residual, new_embeddings, starts_for([1, 1, 1]), reordered) + gated, normalized = reference_values(layer, new_residual, new_embeddings) + expected, _ = reference_conv(layer, normalized, gated, [1, 1, 1], reordered) + assert_close(actual, expected, residual.dtype) + + +def test_checkpoint_binding_and_graph_replay(): + layer = make_layer(dim=32, branches=2) + prefix = "model.language_model.layers.1.ple" + weights = { + f"{prefix}.{name}": tensor + for name, tensor in [ + ("key_proj.weight", layer.key_weight), + ("value_proj.weight", layer.value_weight), + ("norm_key.weight", layer.norm_key), + ("norm_query.weight", layer.norm_query), + ("norm_conv.weight", layer.norm_conv), + ("conv1d.weight", layer.conv_weight.unsqueeze(1)), + ] + } + bound = PLELayer.from_checkpoint(weights, prefix, branches=2, dilation=3) + assert bound.key_weight is layer.key_weight + assert bound.conv_weight.data_ptr() == layer.conv_weight.data_ptr() + residual = torch.randn(7, 64, device="cuda", dtype=torch.bfloat16) + embeddings = torch.randn(7, 32, device="cuda", dtype=torch.bfloat16) + starts = starts_for([3, 2, 2]) + accepted = torch.tensor([3, 1, 2], device="cuda", dtype=torch.int32) + mask = torch.ones(7, device="cuda", dtype=torch.bool) + history = layer.empty_history(3) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + _, conv_inputs = bound.forward(residual, embeddings, starts, history, mask) + bound.history_after(conv_inputs, starts, history, accepted) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output, conv_inputs = bound.forward(residual, embeddings, starts, history, mask) + next_state = bound.history_after(conv_inputs, starts, history, accepted) + for lengths, counts in [([2, 0, 5], [1, 0, 4]), ([0, 3, 4], [0, 2, 0])]: + residual.normal_() + embeddings.normal_() + history.normal_() + starts.copy_(starts_for(lengths)) + accepted.copy_(torch.tensor(counts, device="cuda", dtype=torch.int32)) + mask[::2] = False + graph.replay() + gated, expected_inputs = reference_values(bound, residual, embeddings, mask) + expected, _ = reference_conv(bound, expected_inputs, gated, lengths, history) + assert_close(output, expected, residual.dtype) + expected_state = [] + for i, count in enumerate(counts): + expected_state.append( + torch.cat([history[i], expected_inputs[int(starts[i]) : int(starts[i]) + count]])[-bound.context_len :] + ) + assert_close(next_state, torch.stack(expected_state), residual.dtype) + + +def test_empty_batches_and_invalid_inputs(): + layer = make_layer(dim=32, branches=2) + for lengths in [[], [0, 0, 0]]: + residual = layer.key_weight.new_empty((0, 64)) + embeddings = layer.key_weight.new_empty((0, 32)) + history = layer.empty_history(len(lengths)) + output, conv_inputs = layer.forward(residual, embeddings, starts_for(lengths), history) + assert output.shape == (0, 64) + torch.testing.assert_close( + layer.history_after(conv_inputs, starts_for(lengths), history), history, atol=0, rtol=0 + ) + with pytest.raises(ValueError, match="projection"): + replace(layer, value_weight=layer.value_weight[:1]) + with pytest.raises(ValueError, match="convolution"): + replace(layer, conv_weight=layer.conv_weight[:, :0]) + residual = layer.key_weight.new_empty((1, 64)) + embeddings = layer.key_weight.new_empty((1, 32)) + with pytest.raises(ValueError, match="history"): + layer.forward(residual, embeddings, starts_for([1]), layer.empty_history(2)) + with pytest.raises(ValueError, match="mask"): + layer.forward(residual, embeddings, starts_for([1]), layer.empty_history(1), torch.ones(1, device="cuda")) + with pytest.raises(ValueError, match="accepted"): + layer.history_after(residual, starts_for([1]), layer.empty_history(1), torch.ones(1, device="cuda")) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +def test_gate_zero_near_zero_and_signed_saturation(dtype): + from lightllm.common.ngram_embedding.ple_kernels import _gate + + dots = torch.tensor([0, 1e-8, -1e-8, 1e-5, -1e-5, 100, -100], device="cuda", dtype=dtype) + key = torch.zeros(7, 4, device="cuda", dtype=dtype) + query = torch.zeros_like(key) + key[:, 0] = dots * 2 + query[:, 0] = 1 + value = torch.tensor([[-2, -1, 1, 2]], device="cuda", dtype=dtype).expand(7, -1).contiguous() + gate = (key * query).sum(-1, keepdim=True) / 2 + gate = gate.abs().clamp_min(1e-6).sqrt() * gate.sign() + expected = torch.sigmoid(gate) * value + actual = torch.empty_like(key) + _gate[(7, 1)](key, query, value, None, actual, 4, 1, False, 4, enable_fp_fusion=False) + if dtype == torch.float32: + torch.testing.assert_close(actual, expected, atol=5e-7, rtol=5e-7) + else: + torch.testing.assert_close(actual, expected, atol=0, rtol=0) diff --git a/unit_tests/common/ngram_embedding/test_row_sharded_embedding.py b/unit_tests/common/ngram_embedding/test_row_sharded_embedding.py new file mode 100644 index 0000000000..cacb940aa4 --- /dev/null +++ b/unit_tests/common/ngram_embedding/test_row_sharded_embedding.py @@ -0,0 +1,161 @@ +from types import SimpleNamespace +import importlib + +import pytest +import torch + +from lightllm.common.basemodel.layer_weights.meta_weights.row_sharded_embedding_weight import RowShardedEmbeddingWeight + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +PREFIX = "model.language_model.layers.1.ple.ple_embedding.ngram_embedding" + + +@pytest.fixture(autouse=True) +def runtime_configuration(monkeypatch): + base = importlib.import_module("lightllm.common.basemodel.layer_weights.meta_weights.base_weight") + platform = importlib.import_module("lightllm.common.basemodel.layer_weights.meta_weights.platform_op") + monkeypatch.setattr(base, "get_current_device_id", lambda: torch.cuda.current_device()) + monkeypatch.setattr( + platform, + "get_env_start_args", + lambda: SimpleNamespace( + hardware_platform="cuda", + enable_torch_fallback=False, + enable_triton_fallback=False, + ), + ) + + +def make_weight(rank=0, world=1): + return RowShardedEmbeddingWeight(5, 24, f"{PREFIX}.weight", torch.bfloat16, 4, rank, world) + + +def shards(table): + return [{f"{PREFIX}.shard_{i}.weight": part} for i, part in enumerate(table.chunk(4))] + + +@pytest.mark.parametrize("world", [1, 2, 3, 8]) +def test_streamed_shards_and_vocab_parallel_lookup(world): + table = torch.arange(120, dtype=torch.bfloat16).view(24, 5) + checkpoint = shards(table) + ids = torch.tensor([0, 5, 6, 7, 8, 11, 12, 15, 17, 18, 23], device="cuda", dtype=torch.int64) + outputs = [] + for rank in range(world): + weight = make_weight(rank, world) + assert not weight.verify_load() + required, loaded = set(), set() + for i in range(4): + if max(i * 6, weight.tp_vocab_start_id) < min((i + 1) * 6, weight.tp_vocab_end_id): + required.add(i) + for index in [2, 0, 3, 1]: + weight.load_hf_weights(checkpoint[index]) + loaded.add(index) + assert weight.verify_load() == (loaded >= required) + torch.testing.assert_close( + weight.weight.cpu(), table[weight.tp_vocab_start_id : weight.tp_vocab_end_id], atol=0, rtol=0 + ) + outputs.append(weight(ids)) + + torch.testing.assert_close( + torch.stack(outputs).sum(0), torch.nn.functional.embedding(ids, table.cuda()), atol=0, rtol=0 + ) + + +def test_missing_shards_bad_shapes_and_other_prefixes(): + weight = make_weight() + table = torch.arange(120, dtype=torch.bfloat16).view(24, 5) + checkpoint = shards(table) + for index in [0, 1, 3]: + weight.load_hf_weights(checkpoint[index]) + assert not weight.verify_load() + weight.load_hf_weights({"other.shard_2.weight": table[:6]}) + assert not weight.verify_load() + with pytest.raises(ValueError, match="shard 2"): + weight.load_hf_weights({f"{PREFIX}.shard_2.weight": table[:5]}) + assert not weight.verify_load() + weight.load_hf_weights(checkpoint[2]) + assert weight.verify_load() + + weight.load_hf_weights({"other.weight": table}) + assert weight.verify_load() + + +def test_full_checkpoint_compatibility_and_reject_mixed_formats(): + table = torch.arange(120, dtype=torch.bfloat16).view(24, 5) + weight = make_weight(1, 3) + weight.load_hf_weights({f"{PREFIX}.weight": table}) + assert weight.verify_load() + weight.load_hf_weights({"other.weight": table}) + assert weight.verify_load() + torch.testing.assert_close(weight.weight.cpu(), table[8:16], atol=0, rtol=0) + with pytest.raises(ValueError, match="mix"): + weight.load_hf_weights(shards(table)[0]) + partial = make_weight() + partial.load_hf_weights(shards(table)[1]) + with pytest.raises(ValueError, match="mix"): + partial.load_hf_weights({f"{PREFIX}.weight": table}) + with pytest.raises(ValueError, match="dimensions"): + make_weight().load_hf_weights({f"{PREFIX}.weight": table[:5]}) + + +def test_row_shard_configuration_and_graph_replay(): + with pytest.raises(ValueError, match="divisible"): + RowShardedEmbeddingWeight(5, 23, f"{PREFIX}.weight", torch.bfloat16, 4, 0, 1) + table = torch.arange(120, dtype=torch.bfloat16).view(24, 5) + weight = make_weight() + for part in shards(table): + weight.load_hf_weights(part) + ids = torch.tensor([0, 6, 12, 23], device="cuda", dtype=torch.int64) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + weight(ids) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = weight(ids) + for values in [[23, 12, 6, 0], [1, 7, 13, 19]]: + ids.copy_(torch.tensor(values, device="cuda")) + graph.replay() + torch.testing.assert_close(output.cpu(), table[values], atol=0, rtol=0) + + +def test_packed_hash_lookup_and_ple_composition(): + from lightllm.common.ngram_embedding import NGramHasher + from lightllm.common.ngram_embedding.ple import PLELayer + + torch.manual_seed(137) + sizes = torch.tensor([2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53], device="cuda", dtype=torch.int64) + offsets = torch.cat([sizes.new_zeros(1), sizes.cumsum(0)[:-1]]) + hasher = NGramHasher( + torch.tensor([23703573157769, 20109073645365, 8052911324071], device="cuda"), sizes, offsets, 248044 + ) + tokens = torch.tensor([1, 2, 248044, 3, 4, 5, 6, 7], device="cuda", dtype=torch.int64) + starts = torch.tensor([0, 3, 4, 4, 8], device="cuda", dtype=torch.int32) + indices, _ = hasher.hash(tokens, starts, hasher.empty_history(4)) + + table = torch.randn(512, 4, dtype=torch.bfloat16) + partition_outputs = [] + for rank in range(3): + weight = RowShardedEmbeddingWeight(4, 512, f"{PREFIX}.weight", torch.bfloat16, 128, rank, 3) + for index, shard in reversed(list(enumerate(table.chunk(128)))): + weight.load_hf_weights({f"{PREFIX}.shard_{index}.weight": shard}) + assert weight.verify_load() + partition_outputs.append(weight(indices.flatten())) + embeddings = torch.stack(partition_outputs).sum(0).view(tokens.numel(), 64) + expected_embeddings = torch.nn.functional.embedding(indices, table.cuda()).flatten(-2) + torch.testing.assert_close(embeddings, expected_embeddings, atol=0, rtol=0) + + def rand(*shape): + return torch.randn(*shape, device="cuda", dtype=torch.bfloat16) + + layer = PLELayer( + rand(64, 64) / 8, rand(32, 64) / 8, rand(64) * 0.1, rand(64) * 0.1, rand(64) * 0.1, rand(64, 4) * 0.1, 2, 3 + ) + history = layer.empty_history(4) + residual = rand(tokens.numel(), 64) + output, conv_inputs = layer.forward(residual, embeddings, starts, history) + expected_output, expected_inputs = layer.forward(residual, expected_embeddings, starts, history) + torch.testing.assert_close(output, expected_output, atol=0, rtol=0) + torch.testing.assert_close(conv_inputs, expected_inputs, atol=0, rtol=0) diff --git a/unit_tests/common/qsa/test_attention.py b/unit_tests/common/qsa/test_attention.py new file mode 100644 index 0000000000..e935056012 --- /dev/null +++ b/unit_tests/common/qsa/test_attention.py @@ -0,0 +1,133 @@ +import pytest +import torch +import torch.nn.functional as F + +from lightllm.common.qsa.attention import indexed_attention + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +def reference(query, key, value, indices, scale=None): + selected_k = key[indices.clamp_min(0)].transpose(1, 2) + selected_v = value[indices.clamp_min(0)].transpose(1, 2) + mask = (indices >= 0)[:, None, None, :] + return F.scaled_dot_product_attention( + query.unsqueeze(2), selected_k, selected_v, attn_mask=mask, scale=scale, enable_gqa=True + ).squeeze(2) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize( + "rows,qheads,kvheads,dim,vdim,count", [(3, 6, 2, 32, 16, 79), (1, 24, 2, 256, 256, 2051), (5, 3, 1, 65, 33, 11)] +) +@pytest.mark.parametrize("splits", [1, 4, None]) +@pytest.mark.parametrize("bounded", [False, True]) +def test_sparse_gqa_matches_sdpa(dtype, rows, qheads, kvheads, dim, vdim, count, splits, bounded): + torch.manual_seed(437) + slots = count + 31 + q = torch.randn(rows, qheads, dim, device="cuda", dtype=dtype) + k = torch.randn(slots, kvheads, dim, device="cuda", dtype=dtype) + v = torch.randn(slots, kvheads, vdim, device="cuda", dtype=dtype) + indices = torch.stack([torch.randperm(slots, device="cuda")[:count] for _ in range(rows)]).int() + indices[:, ::7] = -1 + lengths = torch.full((rows,), count + 17, device="cuda", dtype=torch.int32) if bounded else None + if bounded: + indices[0, count // 2 :] = -1 + lengths[0] = count // 2 + if rows > 1: + indices[-1].fill_(-1) + if bounded: + lengths[-1] = -1 + before = (q.clone(), k.clone(), v.clone(), indices.clone()) + actual = indexed_attention(q, k, v, indices, splits=splits, selection_lengths=lengths) + expected = reference(q, k, v, indices) + tolerance = 2e-5 if dtype == torch.float32 else torch.finfo(dtype).eps * 2 + torch.testing.assert_close(actual, expected, atol=tolerance, rtol=tolerance) + for tensor, original in zip((q, k, v, indices), before): + torch.testing.assert_close(tensor, original, atol=0, rtol=0) + if rows > 1: + assert torch.count_nonzero(actual[-1]) == 0 + + +@pytest.mark.parametrize("bounded", [False, True]) +def test_attention_cuda_graph_replay_and_custom_scale(bounded): + q = torch.randn(2, 6, 32, device="cuda", dtype=torch.bfloat16) + k = torch.randn(128, 2, 32, device="cuda", dtype=q.dtype) + v = torch.randn(128, 2, 16, device="cuda", dtype=q.dtype) + indices = torch.arange(79, device="cuda", dtype=torch.int32).expand(2, -1).contiguous() + lengths = torch.full((2,), 79, device="cuda", dtype=torch.int32) if bounded else None + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + indexed_attention(q, k, v, indices, scale=0.17, selection_lengths=lengths) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = indexed_attention(q, k, v, indices, scale=0.17, selection_lengths=lengths) + for bounds in ([13, 0], [128, 63]): + q.normal_() + indices.copy_(torch.stack([torch.randperm(128, device="cuda")[:79] for _ in range(2)]).int()) + indices[:, ::3] = -1 + if bounded: + lengths.copy_(torch.tensor(bounds, device="cuda", dtype=torch.int32)) + indices.masked_fill_(torch.arange(79, device="cuda")[None, :] >= lengths[:, None], -1) + graph.replay() + torch.testing.assert_close(output, reference(q, k, v, indices, scale=0.17), atol=0.016, rtol=0.016) + + +def test_empty_selection_and_invalid_shapes(): + q = torch.randn(2, 6, 32, device="cuda", dtype=torch.bfloat16) + k = torch.randn(10, 2, 32, device="cuda", dtype=q.dtype) + v = torch.randn(10, 2, 16, device="cuda", dtype=q.dtype) + ids = torch.empty(2, 0, device="cuda", dtype=torch.int32) + assert torch.count_nonzero(indexed_attention(q, k, v, ids)) == 0 + assert indexed_attention(q[:0], k, v, ids[:0]).shape == (0, 6, 16) + with pytest.raises(ValueError, match="divisible"): + indexed_attention(q[:, :5].contiguous(), k, v, ids) + with pytest.raises(ValueError, match="int32"): + indexed_attention(q, k, v, ids.long()) + with pytest.raises(ValueError, match="Selection lengths"): + indexed_attention(q, k, v, ids, selection_lengths=torch.zeros(2, device="cuda", dtype=torch.int64)) + + +def test_native_interleaved_kv_views_match_contiguous_cache(): + q = torch.randn(3, 24, 256, device="cuda", dtype=torch.bfloat16) + packed = torch.randn(97, 4, 256, device="cuda", dtype=q.dtype) + k, v = packed[:, :2], packed[:, 2:] + assert not k.is_contiguous() and not v.is_contiguous() + indices = torch.stack([torch.randperm(97, device="cuda")[:79] for _ in range(3)]).int() + indices[:, ::7] = -1 + actual = indexed_attention(q, k, v, indices) + expected = indexed_attention(q, k.contiguous(), v.contiguous(), indices) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +def test_direct_short_rows_preserve_fixed_partition_and_graph_replay(dtype): + torch.manual_seed(701) + rows, slots, width = 5, 4096, 2051 + q = torch.randn(rows, 6, 256, device="cuda", dtype=dtype) + k = torch.randn(slots, 1, 256, device="cuda", dtype=dtype) + v = torch.randn_like(k) + ids = torch.stack([torch.randperm(slots, device="cuda")[:width] for _ in range(rows)]).int() + lengths = torch.tensor([0, 16, 64, 65, 2048], device="cuda", dtype=torch.int32) + + def compute(direct): + return indexed_attention(q, k, v, ids, splits=32, selection_lengths=lengths, short_rows_direct=direct) + + torch.testing.assert_close(compute(True), compute(False), atol=0, rtol=0) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + compute(True) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = compute(True) + for bounds in ([2048, 65, 64, 0, 16], [64, 2048, 0, 65, 63]): + lengths.copy_(torch.tensor(bounds, device="cuda", dtype=torch.int32)) + q.normal_() + graph.replay() + torch.testing.assert_close(actual, compute(False), atol=0, rtol=0) diff --git a/unit_tests/common/qsa/test_flashinfer_topk.py b/unit_tests/common/qsa/test_flashinfer_topk.py new file mode 100644 index 0000000000..5bc432abf1 --- /dev/null +++ b/unit_tests/common/qsa/test_flashinfer_topk.py @@ -0,0 +1,84 @@ +from dataclasses import replace + +import pytest +import torch + +from lightllm.common.qsa.indexer import QSAIndexer +from lightllm.common.qsa.topk import flashinfer_topk_indices, flashinfer_topk_supported + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +@pytest.fixture(autouse=True) +def require_ragged_topk(): + if not flashinfer_topk_supported(torch.device("cuda", 0), 65542, 512): + pytest.skip("FlashInfer graph-safe TopK on Hopper is required") + + +def test_ragged_graph_uses_visible_lengths_and_chunk_offset(): + rows, columns, k, first = 17, 65542, 512, 5 + scores = torch.randn(rows, columns, device="cuda") + visible = torch.zeros(first + rows + 3, dtype=torch.int32, device="cuda") + lengths = visible[first : first + rows] + limits = torch.tensor([0, 3, 2047, 2048, 2051, 65536, 262144], dtype=torch.int32, device="cuda") + cols = torch.arange(columns, device="cuda") + + def select(): + return flashinfer_topk_indices(scores, k, visible_lengths=visible, first=first) + + select() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + selected = select() + for step in range(12): + lengths.copy_(limits[(torch.arange(rows, device="cuda") + step) % limits.numel()]) + scores.normal_() + if step % 3 == 0: + scores.round_() + scores.masked_fill_(cols[None, :] >= lengths[:, None] // 4, 1e9) + scores.masked_fill_((lengths <= k * 4)[:, None], float("nan")) + graph.replay() + for row, complete in enumerate((lengths // 4).tolist()): + actual = selected[row] + if complete <= k: + expected = torch.arange(k, dtype=torch.int32, device="cuda") + expected[expected >= complete] = -1 + assert torch.equal(actual, expected) + else: + assert bool(((actual >= 0) & (actual < complete)).all()) + assert torch.unique(actual).numel() == k + expected = torch.topk(scores[row, :complete], k).values + values = scores[row, actual.long()].sort(descending=True).values + torch.testing.assert_close(values, expected, atol=0, rtol=0) + + +@pytest.mark.parametrize("chunk_size", [3, 64]) +def test_indexer_ragged_topk_preserves_slots_across_graph_transitions(chunk_size): + rows, capacity = 17, 262169 + norm = torch.ones(128, dtype=torch.bfloat16, device="cuda") + original = QSAIndexer(torch.randn(640, 16, dtype=torch.bfloat16, device="cuda"), norm, norm) + optimized = replace(original, use_cuda_topk=True) + cache = torch.zeros(capacity + 3, 128, dtype=torch.bfloat16, device="cuda") + # Binary keys and power-of-two queries give unique scores at the selection boundary. + bits = torch.arange(17, device="cuda") + cache[:, :17] = ((torch.arange(cache.shape[0], device="cuda")[:, None] // 4 >> bits) & 1).to(cache.dtype) + query = torch.zeros(rows, 4, 128, dtype=cache.dtype, device="cuda") + query[:, :, :17] = (2.0 ** bits).to(cache.dtype) + table = torch.arange(capacity, dtype=torch.int32, device="cuda").expand(2, -1).contiguous() + requests = torch.arange(rows, dtype=torch.int32, device="cuda") // 4 % 2 + lengths = torch.full((rows,), 262144, dtype=torch.int32, device="cuda") + data = query, cache, requests, lengths, table + original.select(*data, query_chunk_size=chunk_size) + optimized.select(*data, query_chunk_size=chunk_size) + before, after = torch.cuda.CUDAGraph(), torch.cuda.CUDAGraph() + with torch.cuda.graph(before): + expected = original.select(*data, query_chunk_size=chunk_size) + with torch.cuda.graph(after): + actual = optimized.select(*data, query_chunk_size=chunk_size) + limits = torch.tensor([0, 3, 2047, 2048, 2051, 65536, 262144], dtype=torch.int32, device="cuda") + for step in range(12): + lengths.copy_(limits[(torch.arange(rows, device="cuda") + step) % limits.numel()]) + requests.copy_((torch.arange(rows, device="cuda") + step) % 2) + before.replay() + after.replay() + assert torch.equal(actual.sort(dim=-1).values, expected.sort(dim=-1).values) diff --git a/unit_tests/common/qsa/test_indexer.py b/unit_tests/common/qsa/test_indexer.py new file mode 100644 index 0000000000..d82dee83bf --- /dev/null +++ b/unit_tests/common/qsa/test_indexer.py @@ -0,0 +1,419 @@ +import math + +import pytest +import torch +import torch.nn.functional as F +import triton + +from lightllm.common.qsa.indexer import QSAIndexer +from lightllm.common.qsa.indexer_kernels import _scores + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("minimum", [0, 2]) +def test_grouped_scores_keep_chunk_offsets_short_rows_and_mixed_requests(dtype, minimum): + query = torch.randn(14, 4, 128, device="cuda", dtype=dtype) + cache = torch.randn(204, 128, device="cuda", dtype=dtype) + table = torch.randperm(204, device="cuda").int().view(3, 68) + requests = torch.tensor([0] * 9 + [1, 1, 2, 2, 2], device="cuda", dtype=torch.int32) + lengths = torch.tensor([0] * 5 + [0, 3, 8, 16, 7, 9, 12, 17, 11], device="cuda", dtype=torch.int32) + outputs = [torch.full((9, 17), 42.0, device="cuda") for _ in range(2)] + + def run(): + for group, output in zip((1, 4), outputs): + block = 64 if group == 1 else 128 + _scores[(triton.cdiv(9, group), 1)]( + query, + cache, + requests, + lengths, + table, + output, + 5, + 4, + 128, + 17, + 68, + True, + 4, + 16, + 128, + block, + minimum, + "ieee" if dtype == torch.float32 else "tf32", + GROUP_SIZE=group, + TOKENS=9, + ) + + run() + torch.testing.assert_close(outputs[1], outputs[0], atol=1e-5, rtol=1e-5) + if minimum: + assert torch.all(outputs[1][:2] == 42) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + lengths.zero_() + query.normal_() + graph.replay() + torch.testing.assert_close(outputs[1], outputs[0], atol=1e-5, rtol=1e-5) + + +def make_indexer(dtype=torch.bfloat16, dim=32, heads=4, budget=8, ratio=4): + projection = torch.randn((heads + 1) * dim, 32, device="cuda", dtype=dtype) / math.sqrt(32) + return QSAIndexer( + projection, + torch.randn(dim, device="cuda", dtype=dtype) * 0.1, + torch.randn(dim, device="cuda", dtype=dtype) * 0.1, + heads, + budget, + ratio, + ) + + +def rotary(rows, width, dtype): + angles = torch.randn(rows, width // 2, device="cuda").repeat(1, 2) + return angles.cos().to(dtype), angles.sin().to(dtype) + + +def norm(x, weight, eps): + normalized = x.float() * torch.rsqrt(x.float().square().mean(-1, keepdim=True) + eps) + return (normalized * (1 + weight.float())).to(x.dtype) + + +def rotate(x, cos, sin): + width = cos.shape[-1] + first, second = x[..., :width].chunk(2, -1) + return torch.cat([x[..., :width] * cos + torch.cat([-second, first], -1) * sin, x[..., width:]], -1) + + +def reference_cache(indexer, raw, table, lengths, cos, sin): + result = indexer.empty_cache(table.shape[0], table.shape[1]).fill_(float("nan")) + for row, length in enumerate(lengths): + count = length // indexer.compress_ratio + if not count: + continue + slots = table[row, : count * indexer.compress_ratio].view(count, indexer.compress_ratio).long() + pooled = raw[slots].float().mean(1).to(raw.dtype) + normalized = norm(pooled, indexer.key_norm, indexer.eps) + result[row, :count] = rotate(normalized, cos[slots[:, 0]], sin[slots[:, 0]]) + return result + + +def reference_selection(indexer, query, cache, requests, lengths, table): + rows = [] + for row, (request, length) in enumerate(zip(requests.tolist(), lengths.tolist())): + count = length // indexer.compress_ratio + if count: + scores = (query[row].float() @ cache[request, :count].float().T).relu().sum(0) / math.sqrt(indexer.head_dim) + blocks = scores.topk(min(count, indexer.token_budget // indexer.compress_ratio)).indices + positions = ( + blocks[:, None] * indexer.compress_ratio + torch.arange(indexer.compress_ratio, device=query.device) + ).flatten() + else: + positions = torch.empty(0, device=query.device, dtype=torch.int64) + positions = torch.cat([positions, torch.arange(count * indexer.compress_ratio, length, device=query.device)]) + selected = table[request, positions] + rows.append( + F.pad(selected, (0, indexer.token_budget + indexer.compress_ratio - 1 - selected.numel()), value=-1) + ) + return torch.stack(rows) if rows else table.new_empty((0, indexer.token_budget + indexer.compress_ratio - 1)) + + +def assert_sets(actual, expected): + torch.testing.assert_close(actual.sort(-1).values, expected.sort(-1).values, atol=0, rtol=0) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("dim,rotary_dim,heads,ratio", [(32, 16, 4, 4), (65, 32, 3, 3), (128, 64, 4, 4)]) +def test_projection_pooling_and_causal_selection(dtype, dim, rotary_dim, heads, ratio): + torch.manual_seed(337) + indexer = make_indexer(dtype, dim, heads, budget=2 * ratio, ratio=ratio) + hidden = torch.randn(10, 32, device="cuda", dtype=dtype) + qcos, qsin = rotary(10, rotary_dim, dtype) + query, raw_new = indexer.project(hidden, qcos, qsin) + projected = F.linear(hidden, indexer.projection) + expected_q = projected[:, : heads * dim].reshape(10, heads, dim) + expected_q = rotate(norm(expected_q, indexer.query_norm, indexer.eps), qcos[:, None], qsin[:, None]) + tolerance = max(2e-6, torch.finfo(dtype).eps * 2) + torch.testing.assert_close(query, expected_q, atol=tolerance, rtol=tolerance) + torch.testing.assert_close(raw_new, projected[:, heads * dim :], atol=0, rtol=0) + assert not raw_new.is_contiguous() + raw_copy = torch.empty_like(raw_new) + raw_copy.index_copy_(0, torch.arange(raw_new.shape[0], device="cuda"), raw_new) + torch.testing.assert_close(raw_copy, raw_new.contiguous(), atol=0, rtol=0) + raw = torch.randn(100, dim, device="cuda", dtype=dtype) + cos, sin = rotary(100, rotary_dim, dtype) + table = torch.randperm(100, device="cuda")[:87].view(3, 29).int() + lengths = [23, 17, 0] + cache = indexer.empty_cache(3, 29).fill_(float("nan")) + raw_before = raw.clone() + indexer.pool( + raw, + table, + torch.tensor(lengths, device="cuda", dtype=torch.int32), + cos, + sin, + cache, + torch.zeros(3, device="cuda", dtype=torch.int32), + 29 // ratio, + ) + expected_cache = reference_cache(indexer, raw, table, lengths, cos, sin) + torch.testing.assert_close(cache, expected_cache, atol=tolerance, rtol=tolerance, equal_nan=True) + torch.testing.assert_close(raw, raw_before, atol=0, rtol=0) + requests = torch.tensor([0, 0, 1, 0, 1, 1, 0, 1, 1, 0], device="cuda", dtype=torch.int32) + visible = torch.tensor([0, 1, 3, 4, 5, 8, 9, 15, 17, 23], device="cuda", dtype=torch.int32) + actual = indexer.select(query, cache, requests, visible, table, query_chunk_size=3) + + expected = reference_selection(indexer, query, cache, requests, visible, table) + assert_sets(actual, expected) + for row, (request, length) in enumerate(zip(requests.tolist(), visible.tolist())): + allowed = set(table[request, :length].tolist()) + assert set(actual[row][actual[row] >= 0].tolist()) <= allowed + + +def test_incremental_pool_cache_and_cuda_graph_replay(): + torch.manual_seed(337) + indexer = make_indexer() + raw = torch.randn(48, 32, device="cuda", dtype=torch.bfloat16) + cos, sin = rotary(48, 16, raw.dtype) + table = torch.randperm(48, device="cuda")[:40].view(2, 20).int() + cache = indexer.empty_cache(2, 20).fill_(float("nan")) + lengths = torch.tensor([7, 5], device="cuda", dtype=torch.int32) + first = torch.zeros_like(lengths) + indexer.pool(raw, table, lengths, cos, sin, cache, first, 2) + original = cache[:, :1].clone() + hidden = torch.randn(2, 32, device="cuda", dtype=raw.dtype) + qcos, qsin = rotary(2, 16, raw.dtype) + requests = torch.tensor([0, 1], device="cuda", dtype=torch.int32) + first.fill_(1) + lengths.copy_(torch.tensor([13, 10], device="cuda", dtype=torch.int32)) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + query, _ = indexer.project(hidden, qcos, qsin) + indexer.pool(raw, table, lengths, cos, sin, cache, first, 2) + indexer.select(query, cache, requests, lengths, table, query_chunk_size=1) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + query, _ = indexer.project(hidden, qcos, qsin) + indexer.pool(raw, table, lengths, cos, sin, cache, first, 2) + selected = indexer.select(query, cache, requests, lengths, table, query_chunk_size=1) + for counts, starts in [([13, 10], [1, 1]), ([18, 15], [3, 2])]: + hidden.normal_() + lengths.copy_(torch.tensor(counts, device="cuda", dtype=torch.int32)) + first.copy_(torch.tensor(starts, device="cuda", dtype=torch.int32)) + graph.replay() + expected_cache = reference_cache(indexer, raw, table, counts, cos, sin) + tolerance = 2 * torch.finfo(raw.dtype).eps + torch.testing.assert_close(cache, expected_cache, atol=tolerance, rtol=tolerance, equal_nan=True) + assert_sets(selected, reference_selection(indexer, query, cache, requests, lengths, table)) + torch.testing.assert_close(cache[:, :1], original, atol=0, rtol=0) + + +def test_zero_score_ties_tail_only_empty_and_checkpoint_binding(): + indexer = make_indexer() + prefix = "model.language_model.layers.3.self_attn.indexer" + weights = { + f"{prefix}.index_qk_proj.weight": indexer.projection, + f"{prefix}.q_layernorm.weight": indexer.query_norm, + f"{prefix}.k_layernorm.weight": indexer.key_norm, + } + bound = QSAIndexer.from_checkpoint(weights, prefix, token_budget=8) + assert bound.projection is indexer.projection + query = torch.zeros(3, 4, 32, device="cuda", dtype=torch.bfloat16) + table = torch.arange(29, device="cuda", dtype=torch.int32).view(1, -1) + cache = torch.randn(1, 7, 32, device="cuda", dtype=query.dtype) + requests = torch.zeros(3, device="cuda", dtype=torch.int32) + lengths = torch.tensor([0, 3, 21], device="cuda", dtype=torch.int32) + selected = bound.select(query, cache, requests, lengths, table) + assert torch.all(selected[0] == -1) + assert set(selected[1][selected[1] >= 0].tolist()) == {0, 1, 2} + chosen = selected[2][selected[2] >= 0].tolist() + assert len(chosen) == len(set(chosen)) == 9 + assert chosen[-1] == 20 + for offset in [0, 4]: + block = chosen[offset : offset + 4] + assert block == list(range(block[0], block[0] + 4)) and block[0] % 4 == 0 + empty_query, raw = bound.project(query.new_empty((0, 32)), query.new_empty((0, 0)), query.new_empty((0, 0))) + assert raw.shape == (0, 32) + assert bound.select(empty_query, cache, requests[:0], lengths[:0], table).shape == (0, 11) + + short = bound.select(query[:1], cache[:, :0].contiguous(), requests[:1], lengths[1:2], table[:, :3].contiguous()) + assert set(short[0][short[0] >= 0].tolist()) == {0, 1, 2} + + +def test_invalid_layouts(): + indexer = make_indexer() + with pytest.raises(ValueError, match="budget"): + QSAIndexer(indexer.projection, indexer.query_norm, indexer.key_norm, token_budget=3) + with pytest.raises(ValueError, match="rotary"): + indexer.project( + indexer.projection.new_empty((2, 32)), + indexer.projection.new_empty((2, 3)), + indexer.projection.new_empty((2, 3)), + ) + with pytest.raises(ValueError, match="int32"): + indexer.select( + indexer.projection.new_empty((2, 4, 32)), + indexer.empty_cache(1, 4), + torch.zeros(2, device="cuda", dtype=torch.int64), + torch.ones(2, device="cuda", dtype=torch.int32), + torch.zeros(1, 4, device="cuda", dtype=torch.int32), + ) + + +def test_rejected_suffix_repool_overwrites_candidate_blocks(): + torch.manual_seed(337) + indexer = make_indexer() + table = torch.arange(20, device="cuda", dtype=torch.int32).view(1, -1) + raw = torch.randn(20, 32, device="cuda", dtype=torch.bfloat16) + cos, sin = rotary(20, 16, raw.dtype) + cache = indexer.empty_cache(1, 20).fill_(float("nan")) + indexer.pool( + raw, + table, + torch.tensor([13], device="cuda", dtype=torch.int32), + cos, + sin, + cache, + torch.tensor([0], device="cuda", dtype=torch.int32), + 3, + ) + prefix = cache[:, :1].clone() + + raw[7:12].normal_() + lengths = torch.tensor([12], device="cuda", dtype=torch.int32) + indexer.pool(raw, table, lengths, cos, sin, cache, torch.tensor([7 // 4], device="cuda", dtype=torch.int32), 2) + expected = reference_cache(indexer, raw, table, [12], cos, sin) + torch.testing.assert_close(cache, expected, atol=0.016, rtol=0.016, equal_nan=True) + torch.testing.assert_close(cache[:, :1], prefix, atol=0, rtol=0) + query = torch.randn(1, 4, 32, device="cuda", dtype=raw.dtype) + requests = torch.zeros(1, device="cuda", dtype=torch.int32) + assert_sets( + indexer.select(query, cache, requests, lengths, table), + reference_selection(indexer, query, cache, requests, lengths, table), + ) + + +def test_million_token_context_with_bounded_score_chunks(): + torch.manual_seed(337) + context = 1048579 + indexer = make_indexer(dim=128, budget=2048) + cache = torch.randn(1, context // 4, 128, device="cuda", dtype=torch.bfloat16) + + table = torch.arange(context - 1, -1, -1, device="cuda", dtype=torch.int32).view(1, -1) + query = torch.randn(3, 4, 128, device="cuda", dtype=torch.bfloat16) + requests = torch.zeros(3, device="cuda", dtype=torch.int32) + lengths = torch.tensor([4099, 262147, context], device="cuda", dtype=torch.int32) + actual = indexer.select(query, cache, requests, lengths, table, query_chunk_size=1) + expected = reference_selection(indexer, query, cache, requests, lengths, table) + assert_sets(actual, expected) + + +def test_active_rows_pool_into_shared_request_table(): + torch.manual_seed(337) + indexer = make_indexer() + table = torch.randperm(80, device="cuda").view(4, 20).int() + raw = torch.randn(80, 32, device="cuda", dtype=torch.bfloat16) + cos, sin = rotary(80, 16, raw.dtype) + cache = indexer.empty_cache(4, 20).fill_(float("nan")) + rows = torch.tensor([3, -1, 1], device="cuda", dtype=torch.int32) + lengths = torch.tensor([13, 0, 9], device="cuda", dtype=torch.int32) + indexer.pool(raw, table, lengths, cos, sin, cache, torch.zeros_like(lengths), 4, request_rows=rows) + expected = reference_cache(indexer, raw, table, [0, 9, 0, 13], cos, sin) + torch.testing.assert_close(cache, expected, atol=0.016, rtol=0.016, equal_nan=True) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("queries_per_request", [1, 4]) +@pytest.mark.parametrize("query_chunk_size", [3, 64]) +def test_physical_cache_partial_prefix_forks_and_graph_replay(dtype, queries_per_request, query_chunk_size): + torch.manual_seed(740) + indexer = make_indexer(dtype) + raw = torch.randn(32, 32, device="cuda", dtype=dtype) + cos, sin = rotary(32, 16, dtype) + + table = torch.tensor( + [[0, 1, 2, 3, 4, 5, 6, 7, 9, 10, 11, 12], [0, 1, 2, 3, 4, 5, 6, 8, 13, 14, 15, 16]], + device="cuda", + dtype=torch.int32, + ) + cache = indexer.empty_physical_cache(raw.shape[0]).fill_(float("nan")) + lengths = torch.tensor([7, 7], device="cuda", dtype=torch.int32) + first = torch.zeros_like(lengths) + rows = torch.tensor([1, -1, 0], device="cuda", dtype=torch.int32) + active_lengths = torch.tensor([7, 0, 7], device="cuda", dtype=torch.int32) + active_first = torch.zeros_like(active_lengths) + query = torch.randn(2 * queries_per_request, indexer.query_heads, indexer.head_dim, device="cuda", dtype=dtype) + requests = torch.arange(2, device="cuda", dtype=torch.int32).repeat_interleave(queries_per_request) + + def run(): + indexer.pool(raw, table, active_lengths, cos, sin, cache, active_first, 3, request_rows=rows) + return indexer.select( + query, + cache, + requests, + lengths.repeat_interleave(queries_per_request), + table, + query_chunk_size=query_chunk_size, + ) + + run() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + selected = run() + + for length in (8, 12, 5, 12): + lengths.fill_(length) + active_lengths.copy_(torch.tensor([length, 0, length], device="cuda", dtype=torch.int32)) + graph.replay() + dense = indexer.empty_cache(2, table.shape[1]).fill_(float("nan")) + indexer.pool(raw, table, lengths, cos, sin, dense, first, 3) + for request in range(2): + count = length // indexer.compress_ratio + end_slots = table[ + request, indexer.compress_ratio - 1 : count * indexer.compress_ratio : indexer.compress_ratio + ] + torch.testing.assert_close(cache[end_slots.long()], dense[request, :count], atol=0, rtol=0) + assert_sets( + selected, indexer.select(query, dense, requests, lengths.repeat_interleave(queries_per_request), table) + ) + assert not torch.equal(cache[7], cache[8]) + + untouched = cache[31].clone() + raw[7].add_(3) + run() + dense = indexer.empty_cache(2, table.shape[1]) + indexer.pool(raw, table, lengths, cos, sin, dense, first, 3) + torch.testing.assert_close(cache[7], dense[0, 1], atol=0, rtol=0) + torch.testing.assert_close(cache[31], untouched, equal_nan=True) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("capacity", [4099, 262159]) +def test_sparse_visibility_in_large_cache_and_dynamic_graph(dtype, capacity): + torch.manual_seed(38413) + indexer = make_indexer(dtype, dim=128, heads=4, budget=8) + query = torch.randn(8, 4, 128, device="cuda", dtype=dtype) + table = torch.stack([torch.randperm(capacity, device="cuda") for _ in range(3)]).int() + requests = torch.tensor([0, 0, 0, 0, 1, 1, 2, 2], device="cuda", dtype=torch.int32) + original = torch.tensor([0, 3, 4, 257, 255, 256, capacity - 1, capacity], device="cuda", dtype=torch.int32) + lengths = original.clone() + cache = torch.randn(3, capacity // 4, 128, device="cuda", dtype=dtype) + cache[0, 257 // 4 :] = float("nan") + cache[1, 256 // 4 :] = float("nan") + actual = indexer.select(query, cache, requests, lengths, table) + assert_sets(actual, reference_selection(indexer, query, cache, requests, lengths, table)) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = indexer.select(query, cache, requests, lengths, table) + for pattern in (torch.zeros_like(original), original): + lengths.copy_(pattern) + graph.replay() + assert_sets(output, reference_selection(indexer, query, cache, requests, lengths, table)) + graph.reset() diff --git a/unit_tests/common/qsa/test_topk_compatibility.py b/unit_tests/common/qsa/test_topk_compatibility.py new file mode 100644 index 0000000000..05f2e4f36b --- /dev/null +++ b/unit_tests/common/qsa/test_topk_compatibility.py @@ -0,0 +1,142 @@ +from types import SimpleNamespace + +import pytest +import torch + + +from lightllm.common.qsa import topk +from lightllm.common.qsa.indexer import QSAIndexer + + +@pytest.fixture(autouse=True) +def no_cuda_initialization(): + initially_initialized = torch.cuda.is_initialized() + yield + assert torch.cuda.is_initialized() == initially_initialized + + +@pytest.mark.parametrize("shape,k", [((3, 1027), 512), ((0, 1027), 512), ((2, 11), 0), ((2, 11), 11), ((2, 3, 11), 4)]) +def test_cpu_fallback_keeps_actual_torch_indices(shape, k, monkeypatch): + def unexpected_query(device): + raise AssertionError("CPU fallback queried CUDA") + + monkeypatch.setattr(topk, "_supports_device", unexpected_query) + scores = torch.arange(torch.tensor(shape).prod().item()).reshape(shape).float() + assert torch.equal(topk.cuda_topk_indices(scores, k), torch.topk(scores, k, dim=-1, sorted=False).indices) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64, torch.int64]) +def test_noncontiguous_cpu_fallback(dtype): + scores = torch.arange(22).reshape(11, 2).T.to(dtype) + assert not scores.is_contiguous() + assert torch.equal(topk.cuda_topk_indices(scores, 4), torch.topk(scores, 4, sorted=False).indices) + + +def test_ragged_topk_cpu_and_other_budgets_do_not_query_cuda(monkeypatch): + def unexpected_query(device): + pytest.fail("Unqualified ragged TopK must not initialize CUDA") + + monkeypatch.setattr(torch.cuda, "get_device_capability", unexpected_query) + topk.flashinfer_topk_supported.cache_clear() + assert not topk.flashinfer_topk_supported(torch.device("cpu"), 65542, 512) + assert not topk.flashinfer_topk_supported(torch.device("cuda", 0), 65542, 256) + topk.flashinfer_topk_supported.cache_clear() + + +@pytest.mark.parametrize("k", [-1, 12]) +def test_invalid_budget_keeps_torch_exception(k): + with pytest.raises(RuntimeError): + topk.cuda_topk_indices(torch.ones((2, 11)), k) + + +@pytest.mark.parametrize( + "change", + [ + dict(runtime=False), + dict(dtype=torch.bfloat16), + dict(shape=(3, 1024)), + dict(shape=(3, 8, 1027)), + dict(k=256), + dict(contiguous=False), + dict(device_supported=False), + ], +) +def test_unqualified_cuda_metadata_delegates_without_launch(change, monkeypatch): + config = dict(runtime=True, dtype=torch.float32, shape=(3, 1027), k=512, contiguous=True, device_supported=True) + config.update(change) + scores = SimpleNamespace( + is_cuda=True, + dtype=config["dtype"], + ndim=len(config["shape"]), + shape=config["shape"], + device=torch.device("cuda", 0), + is_contiguous=lambda: config["contiguous"], + ) + sentinel = object() + calls = [] + monkeypatch.setattr(topk, "_CUDA_TOPK_RUNTIME", config["runtime"]) + monkeypatch.setattr(topk, "_supports_device", lambda _: config["device_supported"]) + + def original(value, k, *, dim, sorted): + calls.append((value, k, dim, sorted)) + return SimpleNamespace(indices=sentinel) + + monkeypatch.setattr(topk.torch, "topk", original) + assert topk.cuda_topk_indices(scores, config["k"]) is sentinel + assert calls == [(scores, config["k"], -1, False)] + + +def test_shared_default_and_select_api_remain_compatible(): + import inspect + + assert QSAIndexer.__dataclass_fields__["use_cuda_topk"].default is False + signature = inspect.signature(QSAIndexer.select) + assert list(signature.parameters) == [ + "self", + "query", + "cache", + "requests", + "visible_lengths", + "indices", + "query_chunk_size", + ] + assert signature.parameters["query_chunk_size"].default == 64 + + +@pytest.mark.parametrize( + "change", + [ + dict(dtype=torch.int64), + dict(device=torch.device("cpu")), + dict(ndim=2), + dict(contiguous=False), + dict(first=-1), + dict(first=2), + dict(ratio=0), + ], +) +def test_invalid_visible_metadata_fails_before_allocation(change, monkeypatch): + device = torch.device("cuda", 0) + config = dict(dtype=torch.int32, device=device, ndim=1, contiguous=True, first=0, ratio=4) + config.update(change) + scores = SimpleNamespace( + is_cuda=True, dtype=torch.float32, ndim=2, shape=(3, 1027), device=device, is_contiguous=lambda: True + ) + visible = SimpleNamespace( + dtype=config["dtype"], + device=config["device"], + ndim=config["ndim"], + is_contiguous=lambda: config["contiguous"], + numel=lambda: 3, + ) + monkeypatch.setattr(topk, "_CUDA_TOPK_RUNTIME", True) + monkeypatch.setattr(topk, "_supports_device", lambda _: True) + + def unexpected_allocation(*args, **kwargs): + raise AssertionError("Invalid visible metadata reached GPU allocation") + + monkeypatch.setattr(topk.torch, "empty", unexpected_allocation) + with pytest.raises(ValueError, match="visible lengths"): + topk.cuda_topk_indices( + scores, 512, visible_lengths=visible, first=config["first"], compress_ratio=config["ratio"] + ) diff --git a/unit_tests/common/qsa/test_topk_cuda.py b/unit_tests/common/qsa/test_topk_cuda.py new file mode 100644 index 0000000000..e36249ce3a --- /dev/null +++ b/unit_tests/common/qsa/test_topk_cuda.py @@ -0,0 +1,62 @@ +from dataclasses import replace + +import pytest +import torch + +from lightllm.common.qsa.indexer import QSAIndexer +from lightllm.common.qsa.topk import cuda_topk_supported + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +def components(rows): + if not cuda_topk_supported(torch.device("cuda", 0), 1027, 512): + pytest.skip("The optional selector requires its qualified H100 runtime") + weights = torch.randn((640, 16), device="cuda", dtype=torch.bfloat16) + norm = torch.ones(128, device="cuda", dtype=torch.bfloat16) + original = QSAIndexer(weights, norm, norm) + optimized = replace(original, use_cuda_topk=True) + query = torch.randn((rows, 4, 128), device="cuda", dtype=torch.bfloat16) + cache = torch.randn((16384, 128), device="cuda", dtype=torch.bfloat16) + requests = torch.arange(rows, device="cuda", dtype=torch.int32) % 9 + lengths = torch.zeros(rows, device="cuda", dtype=torch.int32) + table = torch.randint(0, 16384, (9, 4110), device="cuda", dtype=torch.int32) + return original, optimized, (query, cache, requests, lengths, table) + + +@pytest.mark.parametrize("rows", [3, 24]) +@pytest.mark.parametrize("visible", [0, 3, 2047, 2048, 2051, 4096]) +def test_optional_selector_preserves_unsorted_slot_order(rows, visible): + original, optimized, data = components(rows) + data[3].fill_(visible) + assert torch.equal(original.select(*data), optimized.select(*data)) + + +@pytest.mark.parametrize("rows", [3, 24]) +def test_short_long_graph_transitions_ignore_poisoned_short_scores(rows, monkeypatch): + original, optimized, data = components(rows) + allocate = torch.empty + + def poisoned_empty(*args, **kwargs): + value = allocate(*args, **kwargs) + if value.is_cuda and value.dtype == torch.float32 and value.ndim == 2 and value.shape[1] == 1027: + value.fill_(float("nan")) + return value + + monkeypatch.setattr(torch, "empty", poisoned_empty) + data[3].fill_(4096) + original.select(*data) + optimized.select(*data) + torch.cuda.synchronize() + before, after = torch.cuda.CUDAGraph(), torch.cuda.CUDAGraph() + with torch.cuda.graph(before): + expected = original.select(*data) + with torch.cuda.graph(after): + actual = optimized.select(*data) + limits = torch.tensor([0, 3, 2047, 2048, 2051, 4096], device="cuda", dtype=torch.int32) + for step in range(24): + data[0].normal_() + data[3].copy_(limits[(torch.arange(rows, device="cuda") + step) % limits.numel()]) + before.replay() + after.replay() + assert torch.equal(actual, expected) diff --git a/unit_tests/common/quantization/test_aligned_projection.py b/unit_tests/common/quantization/test_aligned_projection.py new file mode 100644 index 0000000000..e0c955c58d --- /dev/null +++ b/unit_tests/common/quantization/test_aligned_projection.py @@ -0,0 +1,97 @@ +import pytest +import torch + +from lightllm.common.basemodel.layer_infer.cache_tensor_manager import g_cache_manager +from lightllm.common.quantization.aligned_projection import OutputAlignedProjection + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +@pytest.fixture(autouse=True) +def current_device(monkeypatch): + monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", str(torch.cuda.current_device())) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test_unaligned_projection_preserves_loading_layout_outputs_and_graph_updates(dtype): + torch.manual_seed(3826) + dimensions = [256, 256, 768, 768, 6, 6] + method = OutputAlignedProjection() + pack, parts = method.create_weight(dimensions, 2560, dtype, 0) + weights = [torch.randn(width, 2560, device="cuda", dtype=dtype) * 0.02 for width in dimensions] + for source, part in zip(weights, parts): + method.load_weight(source, part) + assert pack.weight.shape == (2060, 2560) + assert torch.equal(pack.weight, torch.cat(weights)) + assert method._padded_weight.shape == (2064, 2560) + assert not torch.count_nonzero(method._padded_weight[2060:]) + + hidden = torch.randn(24, 2560, device="cuda", dtype=dtype) + + def project(value, out=None): + return method.apply(value, pack, out=out, use_custom_tensor_mananger=False) + + singles = torch.cat([project(row[None]) for row in hidden]) + for count in (0, 1, 2, 3, 8, 9, 24): + actual = project(hidden[:count]) + assert actual.shape == (count, 2060) and actual.is_contiguous() + torch.testing.assert_close(actual, singles[:count], atol=0, rtol=0) + out = torch.empty(count, 4120, device="cuda", dtype=dtype)[:, ::2] + assert project(hidden[:count], out) is out + torch.testing.assert_close(out, actual, atol=0, rtol=0) + + inputs = hidden[:3].clone() + for _ in range(3): + project(inputs) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = project(inputs) + inputs.add_(0.125) + graph.replay() + torch.testing.assert_close(captured, project(inputs), atol=0, rtol=0) + method.load_weight(weights[-1] + 0.25, parts[-1]) + graph.replay() + torch.testing.assert_close(captured, project(inputs), atol=0, rtol=0) + assert not torch.count_nonzero(method._padded_weight[2060:]) + graph.reset() + + +@pytest.mark.parametrize("dtype,width", [(torch.bfloat16, 2064), (torch.float32, 2060)]) +def test_aligned_and_fp32_paths_keep_native_projection(dtype, width): + method = OutputAlignedProjection() + pack, parts = method.create_weight([width], 128, dtype, 0) + source = torch.randn(width, 128, device="cuda", dtype=dtype) + method.load_weight(source, parts[0]) + assert method._padded_weight is None + hidden = torch.randn(3, 128, device="cuda", dtype=dtype) + result = method.apply(hidden, pack, use_custom_tensor_mananger=False) + torch.testing.assert_close(result, hidden @ source.T, atol=0, rtol=0) + + +def test_cached_outputs_and_graph_replay_keep_live_buffers_independent(): + method = OutputAlignedProjection() + pack, parts = method.create_weight([2060], 2560, torch.bfloat16, 0) + source = torch.randn(2060, 2560, device="cuda", dtype=torch.bfloat16) * 0.02 + method.load_weight(source, parts[0]) + inputs = torch.randn(3, 2560, device="cuda", dtype=torch.bfloat16) + expected = method.apply(inputs, pack, use_custom_tensor_mananger=False) + graph = torch.cuda.CUDAGraph() + g_cache_manager.cache_env_in() + try: + live = method.apply(inputs, pack) + for _ in range(4): + other = method.apply(inputs + 0.25, pack) + assert other.data_ptr() != live.data_ptr() + torch.testing.assert_close(live, expected, atol=0, rtol=0) + del other + with torch.cuda.graph(graph): + captured = method.apply(inputs, pack) + inputs.add_(0.125) + graph.replay() + torch.testing.assert_close( + captured, method.apply(inputs, pack, use_custom_tensor_mananger=False), atol=0, rtol=0 + ) + torch.testing.assert_close(live, expected, atol=0, rtol=0) + finally: + graph.reset() + g_cache_manager.cache_env_out() 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..3f937a11ee --- /dev/null +++ b/unit_tests/common/test_hybrid_req_manager_pd_sharing.py @@ -0,0 +1,144 @@ +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) + if hasattr(state, "get_accepted_state"): + actual = state.get_accepted_state(2) + else: + width = state.config.conv_width + actual = (state.conv_state[:, 2, :, :width], state.ssm_state[:, 2 * (state.mtp_step + 1)]) + 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 + if hasattr(state, "get_accepted_state"): + expected = [tensor.cpu().tolist() for tensor in state.get_accepted_state(1)] + else: + width = state.config.conv_width + 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..a3a11d0b7b --- /dev/null +++ b/unit_tests/common/test_linear_state.py @@ -0,0 +1,268 @@ +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_state_index is not None: + 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: + expected_conv = state.conv_state[:, src, ..., offset : offset + 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 + expected_ssm = (expected_ssm.float() * torch.exp(torch.tensor(-0.5))).to(dtype) + 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), + ): + if state.mtp_state_index is not None: + state.mtp_state_index[dst] = step + if mode == "replay": + state.cursors[dst] = 3 + restore() + if mode == "replay": + actual = state.get_accepted_state(dst) + actual_conv, actual_ssm = actual.conv, actual.ssm + else: + actual_conv = state.conv_state[:, dst, ..., :3] + 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_state_index is not None: + assert state.mtp_state_index[dst].item() == 0 + if mode == "replay": + assert state.cursors[dst].item() == 0 + 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, monkeypatch, 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] + if mode == "replay": + expected = [(item.conv.clone(), item.ssm.clone()) for item in map(state.get_accepted_state, requests)] + else: + expected = [ + ( + state.conv_state[:, req, :, : state.config.conv_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 + + def no_snapshot(*args): + raise AssertionError("batch export must write directly to the destination") + + if mode == "replay": + monkeypatch.setattr(state, "get_accepted_state", no_snapshot) + if history is not state: + monkeypatch.setattr(history, "get_accepted_state", no_snapshot) + 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_state_index is not None: + target.mtp_state_index.fill_(2) + target.copy_runtime_state_from_dp(requests, [0, 1, 2], state) + for req, (conv, ssm) in enumerate(expected): + if mode == "replay": + actual = target.get_accepted_state(req) + actual_conv, actual_ssm = actual.conv, actual.ssm + else: + actual_conv = target.conv_state[:, req, :, : target.config.conv_width] + 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) 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..f720e01f09 --- /dev/null +++ b/unit_tests/common/test_linear_state_contract.py @@ -0,0 +1,201 @@ +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.ssm_autotune import resolve_run_config +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 + + +def test_replay_backend_publishes_detached_run_config(): + ssm_backend = object.__new__(ReplayGDNState) + tuning = {"BV": 32, "num_warps": 2} + ssm_backend.run_config = tuning + ssm_backend._config_is_fixed = True + ssm_backend.compute_mode = "inline" + config = resolve_run_config(ssm_backend, "replay", *([None] * 7), None) + tuning["BV"] = 128 + assert config == {"BV": 32, "num_warps": 2, "num_stages": 3} + + +@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 + + 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, 1, heads, linear_config.key_dim) + v = torch.zeros(1, 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] + 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), + None, + run_config={"BV": bv, "num_warps": 1, "num_stages": 1}, + ) + assert captured.args[33] is precompute + assert captured.args[35] 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) + owner.commit_bounded_att_state( + torch.tensor([1], device="cuda"), torch.tensor([3], dtype=torch.int32, device="cuda") + ) + 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)) + monkeypatch.setattr(history, "accept_updates", lambda *a: pytest.fail("unfused acceptance")) + owner.commit_bounded_att_state(torch.tensor([2, 0, 1]), torch.tensor([1, 2, 3])) + assert len(calls) == 1 + assert calls[0][3] is history.cursors + assert calls[0][0] is owner.mtp_state_index + + +@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.fill_(8) + 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] + owner.commit_accepted(owner.mtp_state_index, reqs, lengths) + assert owner.mtp_state_index.tolist() == [2, 2, 0, 1] + assert history.cursors.tolist() == [11, 8, 9, 8] diff --git a/unit_tests/conftest.py b/unit_tests/conftest.py index 1da9c1e286..622c8087ad 100644 --- a/unit_tests/conftest.py +++ b/unit_tests/conftest.py @@ -22,6 +22,19 @@ def single_rank_environment(monkeypatch: pytest.MonkeyPatch): monkeypatch.setenv(key, value) +@pytest.fixture +def runtime_start_args(monkeypatch: pytest.MonkeyPatch): + import json + from dataclasses import asdict + from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.utils.envs_utils import get_env_start_args + + monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(asdict(StartArgs()))) + get_env_start_args.cache_clear() + yield + get_env_start_args.cache_clear() + + @pytest.fixture def kv_allocator_environment(monkeypatch: pytest.MonkeyPatch): """Provide a service identity and reclaim the allocator's real shared counters.""" diff --git a/unit_tests/distributed/test_ordered_tp_sum.py b/unit_tests/distributed/test_ordered_tp_sum.py new file mode 100644 index 0000000000..bcf4885da4 --- /dev/null +++ b/unit_tests/distributed/test_ordered_tp_sum.py @@ -0,0 +1,110 @@ +import datetime +import json +import os +from pathlib import Path + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + + +def _worker(rank, world, rendezvous): + torch.cuda.set_device(rank) + from lightllm.server.core.objs.start_args_type import StartArgs + + settings = StartArgs() + settings.hardware_platform, settings.data_type = "cuda", "bf16" + os.environ["LIGHTLLM_START_ARGS"] = json.dumps(vars(settings)) + for name, value in { + "LIGHTLLM_DP_WORLD_SIZE": world, + "LIGHTLLM_GLOBAL_WORLD_SIZE": world, + "LIGHTLLM_NODE_WORLD_SIZE": world, + "LIGHTLLM_CURRENT_DEVICE_ID": rank, + "LIGHTLLM_CURRENT_RANK_IN_DP": rank, + "LIGHTLLM_CURRENT_RANK_IN_NODE": rank, + "LIGHTLLM_GLOBAL_RANK": rank, + "LIGHTLLM_DP_SIZE": 1, + "LIGHTLLM_GLOBAL_DP_RANK": 0, + "LIGHTLLM_DP_RANK_IN_NODE": 0, + }.items(): + os.environ[name] = str(value) + dist.init_process_group( + "nccl", + init_method="file://" + rendezvous, + rank=rank, + world_size=world, + timeout=datetime.timedelta(seconds=180), + ) + from lightllm.distributed.communication_op import CustomProcessGroup + from lightllm.distributed.symm_mem_all_reduce import SymmMemAllreduce + from lightllm.common.basemodel.tp_reduce import ordered_tp_sum + + try: + symmetric = SymmMemAllreduce(dist.group.WORLD, rank) + available = torch.tensor(int(not symmetric.disabled), device=rank) + dist.all_reduce(available, op=dist.ReduceOp.MIN) + if not available.item(): + if rank == 0: + Path(rendezvous + ".unavailable").touch() + return + group = CustomProcessGroup.__new__(CustomProcessGroup) + group.device_group, group.dp_world_size = dist.group.WORLD, world + group.symm_mem_reduce, group.flashinfer_reduce = symmetric, None + for rows in (1, 8, 24, 512): + source = torch.empty((rows, 2560), dtype=torch.bfloat16, device=rank) + reference = torch.empty((world * rows, 2560), dtype=source.dtype, device=rank) + scratch = torch.empty_like(source) + + def expected(): + dist.all_gather_into_tensor(reference, source) + return reference.view(world, rows, 2560).double().sum(0).to(source.dtype) + + for iteration in range(4): + if world >= 4 and iteration == 0: + source.fill_((2 ** 40, 1, -(2 ** 40), 1)[rank] if rank < 4 else 0) + else: + source.copy_((torch.arange(source.numel(), device=rank).reshape_as(source) % 257 + rank) / 64) + source.add_(iteration / 8) + gold = expected() + gathered = symmetric.ordered_all_gather(source) + assert gathered is not None + scratch.fill_(rank + 2) + symmetric.all_reduce(scratch) + torch.testing.assert_close(gathered.view_as(reference), reference, rtol=0, atol=0) + actual = ordered_tp_sum(source.clone(), group, world) + torch.testing.assert_close(actual, gold, rtol=0, atol=0) + if rows % world == 0: + actual = ordered_tp_sum(source, group, world, scatter_rank=rank) + torch.testing.assert_close(actual, gold.chunk(world)[rank], rtol=0, atol=0) + + output = torch.empty_like(source) + graph = torch.cuda.CUDAGraph() + torch.cuda.synchronize() + with torch.cuda.graph(graph): + output.copy_(source) + ordered_tp_sum(output, group, world) + for iteration in range(8): + source.fill_(rank + iteration / 16) + gold = expected() + graph.replay() + torch.testing.assert_close(output, gold, rtol=0, atol=0) + graph.reset() + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize("world", [2, 8]) +def test_ordered_sum_preserves_shards_across_collectives_and_graph_replays(world, tmp_path): + if torch.cuda.device_count() < world: + pytest.skip(f"Requires {world} CUDA GPUs") + if any(torch.cuda.get_device_capability(rank)[0] < 9 for rank in range(world)): + pytest.skip("Requires Hopper or newer GPUs with multicast support") + from lightllm.utils.device_utils import has_nvlink + + if not has_nvlink(): + pytest.skip("Requires an NVLink-connected node") + rendezvous = str(tmp_path / "rendezvous") + mp.spawn(_worker, args=(world, rendezvous), nprocs=world, join=True) + if Path(rendezvous + ".unavailable").exists(): + pytest.skip("Symmetric multicast memory is unavailable") diff --git a/unit_tests/models/glm5_next/test_cache.py b/unit_tests/models/glm5_next/test_cache.py index 530b95436a..f583bcf7da 100644 --- a/unit_tests/models/glm5_next/test_cache.py +++ b/unit_tests/models/glm5_next/test_cache.py @@ -48,11 +48,13 @@ def test_nope_cache_config_uses_native_mla_width(): @pytest.mark.parametrize("small_page", [False, True]) @pytest.mark.parametrize("tp_world_size", [1, 4]) @pytest.mark.parametrize("mtp_step", [0, 2]) +@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") @@ -106,6 +108,7 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy( 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, @@ -143,6 +146,16 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy( linear.ssm_state[:, 0].normal_() conv = linear.conv_state[:, 0, ..., :3].clone() ssm = linear.ssm_state[:, 0].clone() + if mode == "replay": + linear.raw_keys.zero_() + linear.raw_values.zero_() + linear.gates.fill_(-0.5) + linear.betas.zero_() + linear.cursors[0] = 1 + ssm = linear.accepted_ssm(0).clone() + 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], @@ -162,12 +175,20 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy( ) torch.cuda.synchronize() assert torch.equal(linear.conv_state[:, 2, ..., :3], conv) - assert torch.equal(linear.ssm_state[:, 2 * (mtp_step + 1)], ssm) + 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() - assert not linear.ssm_state[:, 2 * (mtp_step + 1) : 3 * (mtp_step + 1)].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]]) diff --git a/unit_tests/models/glm5_next/test_mtp.py b/unit_tests/models/glm5_next/test_mtp.py index 8a1d70e3ac..bc6c1428cb 100644 --- a/unit_tests/models/glm5_next/test_mtp.py +++ b/unit_tests/models/glm5_next/test_mtp.py @@ -5,7 +5,9 @@ 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.fla.ops.kda_decode import ( + fused_recurrent_kda, +) from lightllm.models.glm5_next.triton_kernel.kpool import compress_pools @@ -23,7 +25,11 @@ def test_kda_mtp_matches_single_token_decode_and_restores_accepted_state(mtp_ste 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) + 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) @@ -100,7 +106,12 @@ def prefill(values, destination, tail): 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)]: + 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( @@ -117,7 +128,11 @@ def prefill(values, destination, tail): 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)) + 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]) @@ -125,17 +140,26 @@ def prefill(values, destination, tail): @pytest.mark.parametrize("dynamic", [False, True]) @pytest.mark.parametrize("role", ["main", "draft0"]) -def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str): +@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.triton_kernel.linear_att.causal_conv1d import causal_conv1d_update + 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, 3 + heads, dim, width = 2, 128, mtp_step + 1 hidden = heads * dim - lengths = [2, 1, 3] if dynamic else [3, 3, 3] + 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)] @@ -143,10 +167,10 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str): request_rows += [3, 3] offsets += [0, 0] tokens = len(request_rows) - conv = torch.randn(4, 3 * hidden, 5, device="cuda", dtype=torch.bfloat16) + 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], dtype=torch.int32, device="cuda") + 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) @@ -184,7 +208,7 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str): expected_ssm[req * width + token] = state[0] row += 1 backend = SimpleNamespace( - mtp_step=2, + mtp_step=mtp_step, tp_num_heads=heads, head_dim=dim, tp_hidden_size=hidden, @@ -203,11 +227,20 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str): conv_dtype=conv.dtype, ssm_dtype=ssm.dtype, ) + selected = None for model_role in ("main", "draft0"): - linear = KDAState(config, request_capacity=4, mtp_step=2, device="meta") - linear.mtp_state_index = accepted_offsets if model_role == role else torch.zeros_like(accepted_offsets) + 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, linear.ssm_state = conv.unsqueeze(0), ssm.unsqueeze(0) + 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) @@ -223,7 +256,8 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str): linear_A_log=SimpleNamespace(weight=a), linear_dt_bias=SimpleNamespace(weight=bias), ) - state = KDADecodeAttState(backend=backend, _infer_state_weak_ref=weakref.ref(infer)) + 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, @@ -232,11 +266,40 @@ def test_kda_backend_routes_accepted_conv_and_ssm_states(dynamic, role: str): 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 + 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) - torch.testing.assert_close(ssm, expected_ssm, atol=2e-3, rtol=1e-2) - if dynamic: - assert torch.equal(conv[3], old_conv[3]) + 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: + ids = torch.tensor(requests, dtype=torch.int32, device="cuda") + if mtp_step: + accepted = torch.zeros(4, dtype=torch.int32, device="cuda") + accepted[ids] = torch.tensor(lengths, dtype=torch.int32, device="cuda") - 1 + selected.accept_updates(ids, accepted) + for req, length in zip(requests, lengths): + torch.testing.assert_close( + selected.accepted_ssm(req)[0], + 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_pd_cache.py b/unit_tests/models/glm5_next/test_pd_cache.py index e72e4d6bee..7dbe235831 100644 --- a/unit_tests/models/glm5_next/test_pd_cache.py +++ b/unit_tests/models/glm5_next/test_pd_cache.py @@ -46,7 +46,7 @@ def make_mems(monkeypatch, request): get_env_start_args.cache_clear() set_env_start_args(dataclasses.asdict(args)) - def create(tp, target_mtp_step=None): + 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, @@ -87,6 +87,7 @@ def create(tp, target_mtp_step=None): 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), @@ -112,16 +113,21 @@ def create(tp, target_mtp_step=None): @pytest.mark.parametrize("prefill_tp,decode_tp", [(1, 1), (4, 4), (1, 4), (4, 1)]) @pytest.mark.parametrize("remainder", range(4)) @pytest.mark.parametrize("decode_mtp_step", [0, 1, 5]) +@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, dest = make_mems(prefill_tp), make_mems(decode_tp, target_mtp_step=decode_mtp_step) + 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 - dst_mtp_size = decode_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 @@ -159,7 +165,16 @@ def test_pd_kv_and_runtime_state_roundtrip( state.storage.view(torch.uint8)[:, src_indexes] = packed[offset:end] offset = end mem.linear.conv_state[:, src_req, ..., :3].copy_(conv[:, :, heads].reshape(6, -1, 3)) - mem.linear.ssm_state[:, src_req * mtp_size].copy_(ssm[:, heads]) + 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 + ssm[:, heads].copy_(mem.linear.accepted_ssm(src_req)) offset = 0 for state in mem.tails: end = offset + len(state.config.logical_layer_ids) @@ -173,6 +188,10 @@ def test_pd_kv_and_runtime_state_roundtrip( 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_state_index is not None: + 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()) @@ -181,7 +200,7 @@ def test_pd_kv_and_runtime_state_roundtrip( x[:, torch.tensor([0, 1, 3], device="cuda") * stride].clone() for x, stride in ( (mem.linear.conv_state, 1), - (mem.linear.ssm_state, dst_mtp_size), + (mem.linear.ssm_state, dst_ssm_stride), *((state.buffer, 1) for state in mem.tails), ) ] @@ -235,18 +254,20 @@ def test_pd_kv_and_runtime_state_roundtrip( source_field.tensor[:, src_indexes].contiguous().view(torch.uint8), ) assert torch.equal(mem.linear.conv_state[:, dst_req, ..., :3], conv[:, :, heads].reshape(6, -1, 3)) - assert torch.equal(mem.linear.ssm_state[:, dst_req * dst_mtp_size], ssm[:, heads]) + 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_state_index is not None: 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_mtp_size), + (mem.linear.ssm_state, dst_ssm_stride), *((state.buffer, 1) for state in mem.tails), ), untouched[rank], diff --git a/unit_tests/models/qwen2_vl/test_rotary_pos_emb.py b/unit_tests/models/qwen2_vl/test_rotary_pos_emb.py index 34e41578a6..98f770cd77 100644 --- a/unit_tests/models/qwen2_vl/test_rotary_pos_emb.py +++ b/unit_tests/models/qwen2_vl/test_rotary_pos_emb.py @@ -1,4 +1,3 @@ -import math import torch import pytest @@ -51,5 +50,21 @@ def test_triton_matches_reference(shape): assert torch.allclose(out, ref, rtol=1e-2, atol=1e-2), "Triton 与参考实现不一致" +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_strided_qkv_view_uses_output_strides(dtype): + torch.manual_seed(1907) + qkv = torch.randn(37, 3, 3, 72, device="cuda", dtype=dtype) + tensor = qkv[:, 1] + assert not tensor.is_contiguous() + original = qkv.clone() + angles = torch.randn(37, 36, device="cuda", dtype=torch.float32) + cos, sin = angles.cos(), angles.sin() + expected = apply_rotary_pos_emb_vision(tensor.unsqueeze(0), cos, sin).squeeze(0) + actual = apply_rotary_pos_emb_triton(tensor, cos, sin) + tolerance = 0.016 if dtype == torch.bfloat16 else 2e-5 + torch.testing.assert_close(actual, expected, atol=tolerance, rtol=tolerance) + torch.testing.assert_close(qkv, original, atol=0, rtol=0) + + if __name__ == "__main__": pytest.main() diff --git a/unit_tests/models/qwen3next/test_transformer_layer_infer_strided_gate.py b/unit_tests/models/qwen3next/test_transformer_layer_infer_strided_gate.py new file mode 100644 index 0000000000..5bfae4efa7 --- /dev/null +++ b/unit_tests/models/qwen3next/test_transformer_layer_infer_strided_gate.py @@ -0,0 +1,33 @@ +from types import SimpleNamespace + +import torch + +from lightllm.models.qwen3next.layer_infer.transformer_layer_infer import Qwen3NextTransformerLayerInfer + + +def test_linear_post_passes_strided_gate_without_materializing(): + tokens, heads, head_dim = 3, 4, 8 + packed_dim = heads * head_dim + 16 + packed = torch.randn((tokens, packed_dim)) + z = packed[:, 8 : 8 + heads * head_dim].view(tokens, heads, head_dim) + core_attn_out = torch.randn((1, tokens, heads, head_dim)) + calls = [] + + def linear_norm(input, gate_value, eps): + calls.append((input, gate_value, eps)) + return input + + layer_weight = SimpleNamespace( + linear_norm=linear_norm, + linear_out_proj=SimpleNamespace(mm=lambda input: input), + ) + layer_infer = object.__new__(Qwen3NextTransformerLayerInfer) + layer_infer.eps_ = 1e-6 + + output = layer_infer._linear_post(core_attn_out, z, layer_weight) + + assert len(calls) == 1 + assert calls[0][0].shape == (tokens * heads, head_dim) + assert calls[0][1] is z + assert not calls[0][1].is_contiguous() + assert output.shape == (tokens, heads * head_dim) diff --git a/unit_tests/models/qwen4_exp/__init__.py b/unit_tests/models/qwen4_exp/__init__.py new file mode 100644 index 0000000000..2ae28399f5 --- /dev/null +++ b/unit_tests/models/qwen4_exp/__init__.py @@ -0,0 +1 @@ +pass diff --git a/unit_tests/models/qwen4_exp/conftest.py b/unit_tests/models/qwen4_exp/conftest.py new file mode 100644 index 0000000000..074a35b5a1 --- /dev/null +++ b/unit_tests/models/qwen4_exp/conftest.py @@ -0,0 +1,84 @@ +import json + +import pytest + +from .runtime import runtime # noqa: F401 -- shared pytest fixture + +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.utils import envs_utils + + +@pytest.fixture(autouse=True) +def native_runtime(monkeypatch): + args = StartArgs() + args.hardware_platform = "cuda" + args.mtp_step = 0 + args.enable_ep_moe = False + args.auto_update_redundancy_expert = False + args.enable_tpsp_mix_mode = False + monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(vars(args))) + for name, value in { + "LIGHTLLM_DP_WORLD_SIZE": 1, + "LIGHTLLM_CURRENT_RANK_IN_DP": 0, + "LIGHTLLM_CURRENT_DEVICE_ID": 0, + "LIGHTLLM_GLOBAL_RANK": 0, + "LIGHTLLM_GLOBAL_WORLD_SIZE": 1, + }.items(): + monkeypatch.setenv(name, str(value)) + envs_utils.get_env_start_args.cache_clear() + envs_utils.get_added_mtp_kv_layer_num.cache_clear() + envs_utils.get_redundancy_expert_num.cache_clear() + envs_utils.get_redundancy_expert_ids.cache_clear() + yield + envs_utils.get_env_start_args.cache_clear() + envs_utils.get_added_mtp_kv_layer_num.cache_clear() + envs_utils.get_redundancy_expert_num.cache_clear() + envs_utils.get_redundancy_expert_ids.cache_clear() + + +@pytest.fixture +def config(): + return dict( + hidden_size=32, + n_embed=32, + num_hidden_layers=4, + n_layer=4, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=8, + full_attention_interval=4, + layer_types=["linear_attention", "linear_attention", "linear_attention", "full_attention"], + hc_count=2, + hc_lowrank=8, + rms_norm_eps=1e-6, + hidden_act="silu", + output_gate_type="sigmoid", + linear_num_key_heads=1, + linear_num_value_heads=2, + linear_key_head_dim=128, + linear_value_head_dim=128, + linear_conv_kernel_dim=4, + num_experts=2, + num_experts_per_tok=1, + norm_topk_prob=True, + moe_intermediate_size=32, + shared_expert_intermediate_size=32, + indexer_n_heads=2, + indexer_kv_heads=1, + indexer_head_dim=16, + indexer_budget=8, + indexer_compress_ratio=4, + ple_layer_ids=[2], + ple_embed_dim=32, + ple_conv_kernel_size=4, + ngram_size=3, + heads_per_ngram=2, + ngram_vocab_size_base=5, + make_ngram_vocab_size_divisible_by=4, + split_ngram_parts=4, + eos_token_id=63, + vocab_size=64, + partial_rotary_factor=0.5, + rope_scaling={"mrope_section": [1, 1, 0]}, + tie_word_embeddings=False, + ) diff --git a/unit_tests/models/qwen4_exp/helpers.py b/unit_tests/models/qwen4_exp/helpers.py new file mode 100644 index 0000000000..462a5863ea --- /dev/null +++ b/unit_tests/models/qwen4_exp/helpers.py @@ -0,0 +1,133 @@ +import torch + + +def rand(*shape): + return torch.randn(*shape, dtype=torch.bfloat16) * 0.1 + + +def add_hyper(weights, prefix, config, combine=True): + hidden, branches, rank = config["hidden_size"], config["hc_count"], config["hc_lowrank"] + width = hidden * branches + weights[prefix + ".hc_norm.weight"] = rand(width) + weights[prefix + ".input_mix_weight_down.weight"] = rand(rank, width) + weights[prefix + ".input_mix_weight_up.weight"] = rand(width, rank) + if combine: + weights[prefix + ".block_inject_weight.weight"] = rand(branches, width) + + +def checkpoint(config, layer): + hidden, inter, experts = config["hidden_size"], config["moe_intermediate_size"], config["num_experts"] + prefix = f"model.language_model.layers.{layer}" + weights = {} + for block in ["attn_hyper_connection", "mlp_hyper_connection"]: + add_hyper(weights, prefix + "." + block, config) + weights[prefix + ".mlp.gate.weight"] = rand(experts, hidden) + weights[prefix + ".mlp.experts.gate_up_proj"] = rand(experts, 2 * inter, hidden) + weights[prefix + ".mlp.experts.down_proj"] = rand(experts, hidden, inter) + weights[prefix + ".mlp.shared_expert_gate.weight"] = rand(1, hidden) + for proj, shape in [("gate", (inter, hidden)), ("up", (inter, hidden)), ("down", (hidden, inter))]: + weights[prefix + f".mlp.shared_expert.{proj}_proj.weight"] = rand(*shape) + if config["layer_types"][layer] == "linear_attention": + keydim = config["linear_num_key_heads"] * config["linear_key_head_dim"] + vdim = config["linear_num_value_heads"] * config["linear_value_head_dim"] + channels = 2 * keydim + vdim + weights[prefix + ".linear_attn.conv1d.weight"] = rand(channels, 1, 4) + for proj, rows in [ + ("qkv", channels), + ("z", vdim), + ("b", config["linear_num_value_heads"]), + ("a", config["linear_num_value_heads"]), + ]: + weights[prefix + f".linear_attn.in_proj_{proj}.weight"] = rand(rows, hidden) + weights[prefix + ".linear_attn.out_proj.weight"] = rand(hidden, vdim) + weights[prefix + ".linear_attn.A_log"] = torch.zeros(config["linear_num_value_heads"], dtype=torch.float32) + weights[prefix + ".linear_attn.dt_bias"] = torch.ones(config["linear_num_value_heads"], dtype=torch.float32) + weights[prefix + ".linear_attn.norm.weight"] = 1 + rand(config["linear_value_head_dim"]) + else: + dim, qheads, kvheads = config["head_dim"], config["num_attention_heads"], config["num_key_value_heads"] + for name, shape in [ + ("q_proj", (qheads * dim * 2, hidden)), + ("k_proj", (kvheads * dim, hidden)), + ("v_proj", (kvheads * dim, hidden)), + ("o_proj", (hidden, qheads * dim)), + ]: + weights[prefix + f".self_attn.{name}.weight"] = rand(*shape) + for name in ["q_norm", "k_norm"]: + weights[prefix + f".self_attn.{name}.weight"] = rand(dim) + weights[prefix + ".self_attn.indexer.index_qk_proj.weight"] = rand( + (config["indexer_n_heads"] + 1) * config["indexer_head_dim"], hidden + ) + for name in ["q_layernorm", "k_layernorm"]: + weights[prefix + f".self_attn.indexer.{name}.weight"] = rand(config["indexer_head_dim"]) + if layer + 1 in config["ple_layer_ids"]: + width, embed = config["hc_count"] * hidden, config["ple_embed_dim"] + weights[prefix + ".ple.key_proj.weight"] = rand(width, embed) + weights[prefix + ".ple.value_proj.weight"] = rand(hidden, embed) + weights[prefix + ".ple.conv1d.weight"] = rand(width, 1, 4) + for name in ["norm_key", "norm_query", "norm_conv"]: + weights[prefix + f".ple.{name}.weight"] = rand(width) + weights[prefix + ".ple.ple_embedding.layer_multipliers"] = torch.tensor( + [23703573157769, 20109073645365, 8052911324071], dtype=torch.int64 + ) + weights[prefix + ".ple.ple_embedding.ngram_heads_vocab_sizes"] = torch.tensor([5, 7, 11, 13], dtype=torch.int64) + weights[prefix + ".ple.ple_embedding.ngram_heads_offsets"] = torch.tensor([0, 5, 12, 23], dtype=torch.int64) + table = rand(36, 8) + for shard, values in enumerate(table.chunk(4)): + weights[prefix + f".ple.ple_embedding.ngram_embedding.shard_{shard}.weight"] = values + return weights + + +def prepost_checkpoint(config): + weights = { + "model.language_model.embed_tokens.weight": rand(config["vocab_size"], config["hidden_size"]), + "lm_head.weight": rand(config["vocab_size"], config["hidden_size"]), + } + add_hyper(weights, "model.language_model.hyper_connection_mixer", config, False) + return weights + + +def reference_hyper(x, weights, prefix, config, combine=True): + grouped = x.float().unflatten(-1, (config["hc_count"], config["hidden_size"])) + normalized = grouped * torch.rsqrt(grouped.square().mean(-1, keepdim=True) + config["rms_norm_eps"]) + normalized = (normalized.flatten(-2) * (1 + weights[prefix + ".hc_norm.weight"].cuda().float())).to(x.dtype) + low = torch.nn.functional.silu( + torch.nn.functional.linear(normalized, weights[prefix + ".input_mix_weight_down.weight"].cuda()) + / config["hc_count"] + ) + gates = torch.nn.functional.linear(low, weights[prefix + ".input_mix_weight_up.weight"].cuda()).sigmoid() + mixed = (normalized * gates).unflatten(-1, (config["hc_count"], config["hidden_size"])).mean(-2) + if not combine: + return mixed + injection = 2 * torch.sigmoid( + torch.nn.functional.linear(normalized, weights[prefix + ".block_inject_weight.weight"].cuda()) + / config["hc_count"] + ) + return mixed, injection + + +def reference_moe(x, weights, prefix, config): + logits = torch.nn.functional.linear(x, weights[prefix + ".mlp.gate.weight"].cuda()) + prob = logits.float().softmax(-1) + values, indices = prob.topk(config["num_experts_per_tok"], dim=-1) + if config["norm_topk_prob"]: + values /= values.sum(-1, keepdim=True) + values = values.to(x.dtype) + out = torch.zeros_like(x) + inter = config["moe_intermediate_size"] + + for expert in indices.unique(sorted=True).tolist(): + top_position, token = torch.where((indices == expert).T) + gu = torch.nn.functional.linear(x[token], weights[prefix + ".mlp.experts.gate_up_proj"][expert].cuda()) + output = torch.nn.functional.linear( + torch.nn.functional.silu(gu[:, :inter]) * gu[:, inter:], + weights[prefix + ".mlp.experts.down_proj"][expert].cuda(), + ) + weighted = output * values[token, top_position, None] + out.index_add_(0, token, weighted.to(out.dtype)) + gate = torch.nn.functional.linear(x, weights[prefix + ".mlp.shared_expert_gate.weight"].cuda()).sigmoid() + up = torch.nn.functional.linear(x, weights[prefix + ".mlp.shared_expert.up_proj.weight"].cuda()) + g = torch.nn.functional.linear(x, weights[prefix + ".mlp.shared_expert.gate_proj.weight"].cuda()) + shared = torch.nn.functional.linear( + torch.nn.functional.silu(g) * up, weights[prefix + ".mlp.shared_expert.down_proj.weight"].cuda() + ) + return out + gate * shared diff --git a/unit_tests/models/qwen4_exp/lexical_helpers.py b/unit_tests/models/qwen4_exp/lexical_helpers.py new file mode 100644 index 0000000000..ebd26a3389 --- /dev/null +++ b/unit_tests/models/qwen4_exp/lexical_helpers.py @@ -0,0 +1,32 @@ +import torch + +from lightllm.models.qwen4_exp.layer_weights.transformer_layer_weight import Qwen4ExpTransformerLayerWeight +from lightllm.models.qwen4_exp.batch import PackedBatchMetadata +from lightllm.models.qwen4_exp.lexical_cache import LexicalRequestCache +from lightllm.common.quantization import Quantcfg +from .helpers import checkpoint +from .reference import reference_ple + + +def make_cache(config): + source = checkpoint(config, 1) + weights = Qwen4ExpTransformerLayerWeight(1, torch.bfloat16, config, Quantcfg(config)) + weights.load_hf_weights(dict(source)) + weights.verify_load() + return LexicalRequestCache(weights, 8), source + + +def metadata(requests, lengths, queries): + tensors = [torch.tensor(values, device="cuda", dtype=torch.int32) for values in (requests, lengths, queries)] + return PackedBatchMetadata.prepare(*tensors, sum(queries), 7, 4) + + +def reference_segment(residual, tokens, source, config, ngram, convolution): + updates, states = [], [(ngram.clone(), convolution.clone())] + for row in range(tokens.numel()): + update, ngram, convolution = reference_ple( + residual[row : row + 1], tokens[row : row + 1], source, config, ngram, convolution + ) + updates.append(update) + states.append((ngram.clone(), convolution.clone())) + return torch.cat(updates), states diff --git a/unit_tests/models/qwen4_exp/mtp_reference.py b/unit_tests/models/qwen4_exp/mtp_reference.py new file mode 100644 index 0000000000..e9d4332c42 --- /dev/null +++ b/unit_tests/models/qwen4_exp/mtp_reference.py @@ -0,0 +1,27 @@ +import torch.nn.functional as F + +from .reference import norm, reference_layer +from .helpers import reference_hyper + + +def reference_fusion(embedding, hidden, weights, config): + embedding = norm(embedding, weights["mtp.pre_fc_norm_embedding.weight"], config["rms_norm_eps"]) + hidden = norm(hidden, weights["mtp.pre_fc_norm_hidden.weight"], config["rms_norm_eps"]) + projected_embedding = F.linear(embedding, weights["mtp.fc_embedding.weight"].cuda()) + projected_hidden = F.linear( + hidden.view(-1, config["hc_count"], config["hidden_size"]), weights["mtp.fc_hidden.weight"].cuda() + ) + return (projected_hidden + projected_embedding[:, None]).flatten(1) + + +def reference_draft(tokens, hidden, weights, config, cos, sin, lengths): + embeddings = F.embedding(tokens, weights["model.language_model.embed_tokens.weight"].cuda()) + residual = reference_fusion(embeddings, hidden, weights, config) + mapped = { + name.replace("mtp.layers.0.", "model.language_model.layers.3.", 1): value + for name, value in weights.items() + if name.startswith("mtp.layers.0.") + } + residual = reference_layer(residual, mapped, config, cos, sin, lengths) + mixed = reference_hyper(residual, weights, "mtp.hyper_connection_mixer", config, combine=False) + return residual, F.linear(mixed, weights["lm_head.weight"].cuda()).float() diff --git a/unit_tests/models/qwen4_exp/reference.py b/unit_tests/models/qwen4_exp/reference.py new file mode 100644 index 0000000000..2557c479f5 --- /dev/null +++ b/unit_tests/models/qwen4_exp/reference.py @@ -0,0 +1,235 @@ +import torch +import torch.nn.functional as F + +from .helpers import reference_hyper, reference_moe + + +def grouped_norm(x, weight, config): + values = x.float().unflatten(-1, (config["hc_count"], config["hidden_size"])) + normalized = values * torch.rsqrt(values.square().mean(-1, keepdim=True) + config["rms_norm_eps"]) + return (normalized.flatten(-2) * (1 + weight.cuda().float())).to(x.dtype) + + +def reference_ple(residual, tokens, weights, config, ngram_history, conv_history): + prefix = "model.language_model.layers.1.ple" + history = torch.cat([ngram_history, tokens[:, None]], dim=1) + eos = config["eos_token_id"] + eos = eos[0] if isinstance(eos, list) else eos + context = config["ngram_size"] - 1 + shifted = [tokens] + for distance in range(1, context + 1): + preceding = history[:, context - distance : context] + shifted.append(torch.where((preceding == eos).any(-1), eos, history[:, context - distance])) + multipliers = weights[prefix + ".ple_embedding.layer_multipliers"].cuda() + sizes = weights[prefix + ".ple_embedding.ngram_heads_vocab_sizes"].cuda() + offsets = weights[prefix + ".ple_embedding.ngram_heads_offsets"].cuda() + indices = [] + for head in range(sizes.numel()): + order = head // config["heads_per_ngram"] + 2 + mixed = shifted[0] * multipliers[0] + for position in range(1, order): + mixed = mixed ^ (shifted[position] * multipliers[position]) + indices.append(mixed.remainder(sizes[head]) + offsets[head]) + embedding = reference_ngram_lookup(torch.stack(indices, -1), weights, prefix).flatten(-2) + branches, hidden = config["hc_count"], config["hidden_size"] + key = grouped_norm( + F.linear(embedding, weights[prefix + ".key_proj.weight"].cuda()), weights[prefix + ".norm_key.weight"], config + ).view(-1, branches, hidden) + query = grouped_norm(residual, weights[prefix + ".norm_query.weight"], config).view(-1, branches, hidden) + gate = (key * query).sum(-1, keepdim=True) / (hidden ** 0.5) + gate = gate.abs().clamp_min(1e-6).sqrt() * gate.sign() + value = F.linear(embedding, weights[prefix + ".value_proj.weight"].cuda()) + gated = (gate.sigmoid() * value[:, None]).flatten(-2) + normalized = grouped_norm(gated, weights[prefix + ".norm_conv.weight"], config) + full = torch.cat([conv_history, normalized[:, None]], dim=1) + dilation = config["ngram_size"] + conv = F.conv1d( + full.transpose(1, 2), weights[prefix + ".conv1d.weight"].cuda(), groups=branches * hidden, dilation=dilation + ).squeeze(-1) + conv_context = (config["ple_conv_kernel_size"] - 1) * dilation + return ( + gated + F.silu(conv), + history[:, -context:].contiguous(), + full[:, -conv_context:].contiguous() if conv_context else full[:, :0].contiguous(), + ) + + +def reference_ngram_lookup(indices, weights, prefix): + name = prefix + ".ple_embedding.ngram_embedding" + first = weights[name + ".shard_0.weight"] + shard_rows, dim = first.shape + flat = indices.cpu().flatten() + shards = flat // shard_rows + result = torch.empty(flat.numel(), dim, dtype=first.dtype) + for shard in shards.unique().tolist(): + positions = torch.where(shards == shard)[0] + source = weights[name + f".shard_{shard}.weight"] + rows = (flat[positions] % shard_rows).to(source.device) + result[positions] = source[rows].cpu() + return result.to(indices.device).view(*indices.shape, dim) + + +def reference_gdn(mixed, weights, config, conv_history, recurrent): + prefix = "model.language_model.layers.1.linear_attn" + key_heads, value_heads = config["linear_num_key_heads"], config["linear_num_value_heads"] + key_dim, value_dim = config["linear_key_head_dim"], config["linear_value_head_dim"] + kd, vd = key_heads * key_dim, value_heads * value_dim + qkv = F.linear(mixed, weights[prefix + ".in_proj_qkv.weight"].cuda()) + full = torch.cat([conv_history, qkv[:, :, None]], dim=-1) + conv = (full.float() * weights[prefix + ".conv1d.weight"].squeeze(1).cuda().float()[None]).sum(-1).to(mixed.dtype) + conv = F.silu(conv) + q, k, v = torch.split(conv, [kd, kd, vd], dim=-1) + q, k = q.view(-1, key_heads, key_dim).float(), k.view(-1, key_heads, key_dim).float() + q = q * torch.rsqrt(q.square().sum(-1, keepdim=True) + 1e-6) + k = k * torch.rsqrt(k.square().sum(-1, keepdim=True) + 1e-6) + q = q.repeat_interleave(value_heads // key_heads, dim=1) / (key_dim ** 0.5) + k = k.repeat_interleave(value_heads // key_heads, dim=1) + v = v.view(-1, value_heads, value_dim).float() + a = F.linear(mixed, weights[prefix + ".in_proj_a.weight"].cuda()) + b = F.linear(mixed, weights[prefix + ".in_proj_b.weight"].cuda()) + decay = -weights[prefix + ".A_log"].cuda().float().exp() * F.softplus( + a.float() + weights[prefix + ".dt_bias"].cuda() + ) + beta = b.sigmoid().float() + next_state = recurrent * decay.exp()[:, :, None, None] + memory = (next_state * k[:, :, :, None]).sum(-2) + delta = (v - memory) * beta[:, :, None] + next_state = next_state + k[:, :, :, None] * delta[:, :, None, :] + core = (next_state * q[:, :, :, None]).sum(-2).to(mixed.dtype) + normed = core.float() * torch.rsqrt(core.float().square().mean(-1, keepdim=True) + 1e-6) + weighted = weights[prefix + ".norm.weight"].cuda() * normed.to(mixed.dtype) + z = F.linear(mixed, weights[prefix + ".in_proj_z.weight"].cuda()).view(-1, value_heads, value_dim) + gated = (weighted * z.float().sigmoid()).to(mixed.dtype) + output = F.linear(gated.flatten(1), weights[prefix + ".out_proj.weight"].cuda()) + return output, full[:, :, -3:].contiguous(), next_state + + +def reference_step(residual, tokens, source, config, ngram, ple, conv, recurrent): + lexical, ngram, ple = reference_ple(residual, tokens, source, config, ngram, ple) + expected = residual + lexical + mixed, gates = reference_hyper(expected, source, "model.language_model.layers.1.attn_hyper_connection", config) + output, conv, recurrent = reference_gdn(mixed, source, config, conv, recurrent) + expected = expected + (output[:, None] * gates[:, :, None]).flatten(1) + mixed, gates = reference_hyper(expected, source, "model.language_model.layers.1.mlp_hyper_connection", config) + output = reference_moe(mixed, source, "model.language_model.layers.1", config) + expected = expected + (output[:, None] * gates[:, :, None]).flatten(1) + return expected, ngram, ple, conv, recurrent + + +def norm(x, weight, eps=1e-6): + value = x.float() + return (value * torch.rsqrt(value.square().mean(-1, keepdim=True) + eps) * (1 + weight.cuda().float())).to(x.dtype) + + +def rotate(x, cos, sin): + width = cos.shape[-1] + a, b = x[..., :width].chunk(2, -1) + return torch.cat([x[..., :width] * cos + torch.cat([-b, a], -1) * sin, x[..., width:]], -1) + + +def reference_layer(residual, weights, config, cos, sin, lengths): + prefix = "model.language_model.layers.3" + mixed, gates = reference_hyper(residual, weights, prefix + ".attn_hyper_connection", config) + dim, heads, kvheads = config["head_dim"], config["num_attention_heads"], config["num_key_value_heads"] + qg = F.linear(mixed, weights[prefix + ".self_attn.q_proj.weight"].cuda()).view(-1, heads, dim * 2) + query, gate = qg.chunk(2, -1) + key = F.linear(mixed, weights[prefix + ".self_attn.k_proj.weight"].cuda()).view(-1, kvheads, dim) + value = F.linear(mixed, weights[prefix + ".self_attn.v_proj.weight"].cuda()).view(-1, kvheads, dim) + query = rotate(norm(query, weights[prefix + ".self_attn.q_norm.weight"]), cos[:, None], sin[:, None]) + key = rotate(norm(key, weights[prefix + ".self_attn.k_norm.weight"]), cos[:, None], sin[:, None]) + projected = F.linear(mixed, weights[prefix + ".self_attn.indexer.index_qk_proj.weight"].cuda()) + ih, idim = config["indexer_n_heads"], config["indexer_head_dim"] + iq = projected[:, : ih * idim].view(-1, ih, idim) + iq = rotate(norm(iq, weights[prefix + ".self_attn.indexer.q_layernorm.weight"]), cos[:, None], sin[:, None]) + raw = projected[:, ih * idim :] + ratio, budget = config["indexer_compress_ratio"], config["indexer_budget"] // config["indexer_compress_ratio"] + result = [] + start = 0 + for count in lengths: + for position in range(count): + complete = (position + 1) // ratio + selected = [] + if complete: + groups = raw[start : start + complete * ratio].view(complete, ratio, idim).float().mean(1).to(raw.dtype) + pooled = norm(groups, weights[prefix + ".self_attn.indexer.k_layernorm.weight"]) + pooled = rotate( + pooled, cos[start : start + complete * ratio : ratio], sin[start : start + complete * ratio : ratio] + ) + scores = (iq[start + position].float() @ pooled.float().T).relu().sum(0) / (idim ** 0.5) + blocks = scores.topk(min(complete, budget)).indices + selected = (blocks[:, None] * ratio + torch.arange(ratio, device=residual.device)).flatten().tolist() + selected += list(range(complete * ratio, position + 1)) + ids = torch.tensor(selected, device=residual.device, dtype=torch.int64) + start + out = F.scaled_dot_product_attention( + query[start + position].unsqueeze(0).unsqueeze(2), + key[ids].transpose(0, 1).unsqueeze(0), + value[ids].transpose(0, 1).unsqueeze(0), + enable_gqa=True, + ) + result.append(out[0, :, 0]) + start += count + attention = torch.stack(result) + attention = attention * gate.sigmoid() + output = F.linear(attention.flatten(1), weights[prefix + ".self_attn.o_proj.weight"].cuda()) + residual = residual + (output[:, None, :] * gates[:, :, None]).flatten(1) + mixed, gates = reference_hyper(residual, weights, prefix + ".mlp_hyper_connection", config) + update = reference_moe(mixed, weights, prefix, config) + return residual + (update[:, None, :] * gates[:, :, None]).flatten(1) + + +def reference_model(tokens, cos, sin, lengths, weights, config, input_embeddings=None): + if input_embeddings is None: + input_embeddings = F.embedding(tokens, weights["model.language_model.embed_tokens.weight"].cuda()) + residual = input_embeddings.repeat(1, config["hc_count"]) + for layer in range(config["n_layer"]): + prefix = f"model.language_model.layers.{layer}" + if config["layer_types"][layer] == "full_attention": + residual = reference_layer(residual, weights, config, cos, sin, lengths) + continue + if layer + 1 in config["ple_layer_ids"]: + updates, start = [], 0 + for length in lengths: + ngram = torch.full((1, 2), config["eos_token_id"], device="cuda", dtype=torch.int64) + convolution = torch.zeros(1, 9, 64, device="cuda", dtype=torch.bfloat16) + for row in range(start, start + length): + update, ngram, convolution = reference_ple( + residual[row : row + 1], tokens[row : row + 1], weights, config, ngram, convolution + ) + updates.append(update) + start += length + residual = residual + torch.cat(updates) + mixed, gates = reference_hyper(residual, weights, prefix + ".attn_hyper_connection", config) + + local = { + name.replace(prefix, "model.language_model.layers.1", 1): value + for name, value in weights.items() + if name.startswith(prefix + ".") + } + updates, start = [], 0 + for length in lengths: + channels = ( + 2 * config["linear_num_key_heads"] * config["linear_key_head_dim"] + + config["linear_num_value_heads"] * config["linear_value_head_dim"] + ) + convolution = torch.zeros( + 1, channels, config["linear_conv_kernel_dim"] - 1, device="cuda", dtype=torch.bfloat16 + ) + recurrent = torch.zeros( + 1, + config["linear_num_value_heads"], + config["linear_key_head_dim"], + config["linear_value_head_dim"], + device="cuda", + dtype=torch.float32, + ) + for row in range(start, start + length): + update, convolution, recurrent = reference_gdn( + mixed[row : row + 1], local, config, convolution, recurrent + ) + updates.append(update) + start += length + residual = residual + (torch.cat(updates)[:, None] * gates[:, :, None]).flatten(1) + mixed, gates = reference_hyper(residual, weights, prefix + ".mlp_hyper_connection", config) + residual = residual + (reference_moe(mixed, weights, prefix, config)[:, None] * gates[:, :, None]).flatten(1) + hidden = reference_hyper(residual, weights, "model.language_model.hyper_connection_mixer", config, combine=False) + return F.linear(hidden, weights["lm_head.weight"].cuda()).float() diff --git a/unit_tests/models/qwen4_exp/resource_helpers.py b/unit_tests/models/qwen4_exp/resource_helpers.py new file mode 100644 index 0000000000..0813aed3b8 --- /dev/null +++ b/unit_tests/models/qwen4_exp/resource_helpers.py @@ -0,0 +1,73 @@ +from types import SimpleNamespace + +import torch + +from lightllm.models.qwen4_exp.attention_state import Qwen4ExpAttentionState +from lightllm.models.qwen4_exp.lexical_state import Qwen4ExpLexicalState + + +class ResourceView: + def __init__(self, manager): + self.manager = manager + self.allocator = manager.full_att_allocator + self.size = 32 + (self.state,) = manager.get_full_attention_states("main") + self.qsa_states = { + layer: state.qsa_view(layer) + for state in manager.get_full_attention_states() + for group in state.config.logical_layer_groups + for layer in group + } + self.qsa_raw_keys = self.state.indexer.k_storage[:, :, 0] + self.qsa_pooled_keys = self.state.indexer.v_storage[:, :, 0] + self.qsa_key_cos = self.state.positions.k_storage[0, :, 0] + self.qsa_key_sin = self.state.positions.v_storage[0, :, 0] + + def alloc(self, count): + return self.allocator.alloc(count) + + def _slot_buffers(self): + return { + f"{state.state_name}.{name}": (tensor, 1) + for state in self.manager.get_full_attention_states() + for name, tensor in ( + ("k", state.k_storage), + ("v", state.v_storage), + ("raw", state.indexer.k_storage), + ("pooled", state.indexer.v_storage), + ("cos", state.positions.k_storage), + ("sin", state.positions.v_storage), + ) + } + + def get_index_kv_buffer(self, indices): + indices = torch.as_tensor(indices, dtype=torch.int64, device=self.state.device) + return {name: tensor.index_select(axis, indices) for name, (tensor, axis) in self._slot_buffers().items()} + + +def bind_states(model): + manager = model.att_state_manager + manager._get_states.cache_clear() + manager.finish_state_registration() + if manager.full_att_allocator is None: + manager.allocate_storage(32, 1, 8, 8, device="cuda") + else: + for state in manager.get_full_attention_states(): + if state.k_storage is None: + state.device = torch.device("cuda") + state.allocate_storage(32) + req = model.req_manager + req.bind_hold_page(manager.full_att_allocator.hold_token_indices) + lexical = next( + state for state in manager.get_bounded_attention_states("main") if isinstance(state, Qwen4ExpLexicalState) + ) + req.lexical_caches = lexical.lexical_caches + req.init_linear_att_state = lambda request: manager.reset_bounded_att_state([request.req_idx]) + req.snapshot_lexical = lambda request: { + layer: cache.snapshot(request) for layer, cache in lexical.lexical_caches.items() + } + req.restore_lexical = lambda request, snapshots: [ + cache.restore(request, snapshots[layer]) for layer, cache in lexical.lexical_caches.items() + ] + req.get_mamba_cache = lambda layer: manager.get_state("main", "gdn", layer).get_layer_view(layer) + return ResourceView(manager) diff --git a/unit_tests/models/qwen4_exp/runtime.py b/unit_tests/models/qwen4_exp/runtime.py new file mode 100644 index 0000000000..0a02d45d4e --- /dev/null +++ b/unit_tests/models/qwen4_exp/runtime.py @@ -0,0 +1,133 @@ +import json +from types import SimpleNamespace +from contextlib import nullcontext +from uuid import uuid4 + +import pytest +import torch + +from lightllm.models.qwen4_exp.layer_weights.transformer_layer_weight import Qwen4ExpTransformerLayerWeight +from lightllm.models.qwen4_exp.infer_struct import Qwen4ExpInferStateInfo +from lightllm.models.qwen4_exp.model import Qwen4ExpTpPartModel +from lightllm.common.quantization import Quantcfg +from lightllm.utils import envs_utils +from .helpers import checkpoint, prepost_checkpoint + + +@pytest.fixture +def runtime(config, tmp_path, monkeypatch, request): + (tmp_path / "config.json").write_text(json.dumps(config)) + args = envs_utils.get_env_start_args() + args.model_dir = str(tmp_path) + args.data_type = "bf16" + for name, value in getattr(request, "param", {}).items(): + setattr(args, name, value) + args.cpu_cache_token_page_size = args.linear_att_hash_page_size * args.linear_att_page_block_num + args.linear_att_ssm_data_type = "float32" + args.llm_prefill_att_backend = ["triton", "triton"] + args.llm_decode_att_backend = ["triton", "triton"] + monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(vars(args))) + monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_NODE", "0") + monkeypatch.setenv("LIGHTLLM_UNIQUE_SERVICE_NAME_ID", "qwen-runtime-test-" + uuid4().hex) + envs_utils.get_env_start_args.cache_clear() + envs_utils.get_unique_server_name.cache_clear() + weight = Qwen4ExpTransformerLayerWeight(1, torch.bfloat16, config, Quantcfg(config)) + model = Qwen4ExpTpPartModel.__new__(Qwen4ExpTpPartModel) + model.__dict__.update( + config=config, + args=args, + tp_world_size_=1, + tp_rank_=0, + torch_memory_saver=SimpleNamespace(region=lambda **kwargs: nullcontext()), + data_type=torch.bfloat16, + max_req_num=4, + max_seq_length=64, + batch_max_tokens=8, + max_total_token_num=32, + num_kv_heads=2, + mem_fraction=0.9, + trans_layers_weight=[weight], + ) + from lightllm.common.req_manager import ReqManager + from .resource_helpers import bind_states + + model.layers_infer = [SimpleNamespace(layer_num_=i) for i in range(config["n_layer"])] + req = ReqManager(4, 64) + model.bind_req_manager(req) + model.prepare_attention_states() + mem = bind_states(model) + assert not weight.ngram_hasher.verify_load() + torch.manual_seed(1307) + source = checkpoint(config, 1) + weight.load_hf_weights(dict(source)) + weight.verify_load() + angle = torch.randn(64, 2, device="cuda") + + model._cos_cached = angle.cos().bfloat16() + model._sin_cached = angle.sin().bfloat16() + model.mtp_manager = SimpleNamespace(get_decode_draft_step=lambda is_draft: 0) + model._init_att_backend() + model._init_att_backend1() + yield SimpleNamespace(mem=mem, req=req, weight=weight, source=source, model=model) + envs_utils.get_unique_server_name.cache_clear() + + +def state_for(runtime, requests, lengths, tokens, ready=None, mtp_indices=None): + state = Qwen4ExpInferStateInfo() + state.req_manager = runtime.req + state.att_state_manager = runtime.model.att_state_manager + state.att_state_role = runtime.model.att_state_role + state.is_prefill = ready is not None + state.batch_size = len(requests) + state.b_req_idx = torch.tensor(requests, device="cuda", dtype=torch.int32) + state.b_seq_len = torch.tensor(lengths, device="cuda", dtype=torch.int32) + if mtp_indices is not None: + state.b_mtp_index = torch.tensor(mtp_indices, device="cuda", dtype=torch.int32) + state.input_ids = torch.tensor(tokens, device="cuda", dtype=torch.int64) + state.mem_index = runtime.mem.alloc(len(tokens)).cuda(non_blocking=True) + if ready is not None: + state.b_ready_cache_len = torch.tensor(ready, device="cuda", dtype=torch.int32) + query_lengths = [length - prefix for length, prefix in zip(lengths, ready)] + else: + state.b_position_delta = torch.zeros(len(requests), device="cuda", dtype=torch.int64) + query_lengths = [1] * len(requests) + offset = 0 + for request, length, query in zip(requests, lengths, query_lengths): + runtime.req.req_to_token_indexs[request, length - query : length] = state.mem_index[offset : offset + query] + offset += query + state.multimodal_params = [] + state.init_some_extra_state(runtime.model) + if state.is_prefill: + state.prefill_att_state = runtime.model.prefill_att_backend.create_att_prefill_state(state) + if runtime.model.prefill_att_backend1 is not None: + state.prefill_att_state1 = runtime.model.prefill_att_backend1.create_att_prefill_state(state) + else: + state.decode_att_state = runtime.model.decode_att_backend.create_att_decode_state(state) + if runtime.model.decode_att_backend1 is not None: + state.decode_att_state1 = runtime.model.decode_att_backend1.create_att_decode_state(state) + state.init_att_state() + return state + + +def assemble_model(runtime, config): + model, source = runtime.model, prepost_checkpoint(config) + weights = [] + for layer in range(config["n_layer"]): + if layer == 1: + weight = runtime.weight + source.update(runtime.source) + else: + local = checkpoint(config, layer) + source.update(local) + weight = model.transformer_weight_class(layer, torch.bfloat16, config, Quantcfg(config)) + weight.load_hf_weights(dict(local)) + weight.verify_load() + weights.append(weight) + model.trans_layers_weight = weights + model.pre_post_weight = model.pre_and_post_weight_class(torch.bfloat16, config) + model.pre_post_weight.load_hf_weights(dict(source)) + model.pre_post_weight.verify_load() + model._init_infer_layer() + model._init_some_value() + model.prefill_graph = None + return model, source diff --git a/unit_tests/models/qwen4_exp/test_attention_state.py b/unit_tests/models/qwen4_exp/test_attention_state.py new file mode 100644 index 0000000000..62cce52c92 --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_attention_state.py @@ -0,0 +1,205 @@ +from dataclasses import replace +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.attention_state.manager import AttentionStateManager +from lightllm.common.attention_state.types import StateTransferLayout +from lightllm.models import get_model_class +from lightllm.models.draft_registry import get_draft_model_class +from lightllm.models.qwen4_exp.model import Qwen4ExpTpPartModel +from lightllm.models.qwen4_exp.mtp.model import Qwen4ExpMTPModel +from lightllm.models.qwen4_exp.attention_state import Qwen4ExpAttentionState +from lightllm.models.qwen4_exp.lexical_state import Qwen4ExpLexicalState, LexicalStateConfig +from lightllm.utils.envs_utils import get_env_start_args +from lightllm.common.attention_state.bounded_att.gdn.gdn_state import GDNState +from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState + + +@pytest.fixture +def lexical(monkeypatch): + import lightllm.models.qwen4_exp.lexical_state as module + + def copy(*, copies): + for source, destination in copies: + destination.copy_(source) + + monkeypatch.setattr(module, "copy_bounded_att_tensors", copy) + state = Qwen4ExpLexicalState(LexicalStateConfig((1, 5), 2, 6, 8, 63, torch.bfloat16), 4, 2, "cpu") + state.state_name = "main_qwen4_lexical_state" + state.initialize_page_pools(4, 2, pin_memory=False) + return state + + +def test_lazy_registration_and_resource_factories(config): + args = get_env_start_args() + args.data_type, args.linear_att_ssm_data_type = "bf16", "float32" + args.cpu_cache_token_page_size = 8 + args.pd_kv_page_size = 4 + for model_type in ("qwen4_exp", "qwen4_exp_text"): + assert get_model_class({"model_type": model_type}) is Qwen4ExpTpPartModel + assert get_draft_model_class({"model_type": model_type}, "eagle_with_att") is Qwen4ExpMTPModel + manager = AttentionStateManager(request_capacity=4) + for cls, role, layers in ((Qwen4ExpTpPartModel, "main", (0, 1, 2, 3)), (Qwen4ExpMTPModel, "draft0", (4,))): + full = cls.create_full_att_states(config, args, layers) + bounded = cls.create_bounded_att_states(config, args, layers, 4) + for state in (*full, *bounded): + manager.register_state(role, state) + manager.finish_state_registration() + state = manager.get_state("main", "", 3, Qwen4ExpAttentionState) + assert isinstance(state, Qwen4ExpAttentionState) + assert manager.get_state("draft0", "qwen4_gqa", 4) is not state + assert len(manager.get_bounded_attention_states("main")) == 2 + assert not manager.get_bounded_attention_states("draft0") + for state in manager.get_full_attention_states(): + regions = state.get_cpu_transfer_regions() + assert len(regions) == 6 + assert all(left.byte_end <= right.byte_start for left, right in zip(regions, regions[1:])) + assert {region.field for region in regions} == { + "k", + "v", + "indexer.k", + "indexer.v", + "positions.k", + "positions.v", + } + assert state.get_token_bytes() > 0 + + +@pytest.mark.parametrize("mode,state_type", [("native", GDNState), ("replay", ReplayGDNState)]) +def test_mtp_state_modes_keep_qwen4_resources(config, mode, state_type): + args = get_env_start_args() + args.data_type = "bfloat16" + args.mtp_step = 3 + args.ssm_state_mode = mode + args.ssm_replay_config = {"cache_len": 16, "compute_mode": "precompute"} if mode == "replay" else {} + args.linear_att_ssm_data_type = "float32" + layers = (0, 1, 2, 3) + bounded = Qwen4ExpTpPartModel.create_bounded_att_states(config, args, layers, 4, device="cpu") + manager = AttentionStateManager(request_capacity=4) + for state in (*Qwen4ExpTpPartModel.create_full_att_states(config, args, layers), *bounded): + manager.register_state("main", state) + manager.finish_state_registration() + state = manager.get_state("main", "gdn", 1, state_type) + assert state.config.logical_layer_ids == (0, 1, 2) + assert state.get_layer_view(1)[1].dtype == torch.float32 + assert isinstance(manager.get_state("main", "qwen4_gqa", 3), Qwen4ExpAttentionState) + assert any(isinstance(state, Qwen4ExpLexicalState) for state in bounded) + assert not Qwen4ExpMTPModel.create_bounded_att_states(config, args, (4,), 4, device="cpu") + + +@pytest.mark.parametrize("tp", [1, 2, 4, 8]) +def test_replicated_indexer_and_positions_layout(config, tp): + args = get_env_start_args() + args.data_type = "bf16" + for rank in range(tp): + (state,) = Qwen4ExpTpPartModel.create_full_att_states(config, args, (0, 1, 2, 3), rank, tp) + regions = state._get_transfer_regions(8) + for region in regions[2:]: + assert region.write_enabled == (rank == 0) + assert region.shape[2] == 1 + assert sum(region.byte_end - region.byte_start for region in regions) == regions[-1].byte_end + + +def test_state_storage_accounts_for_all_fields_and_restores_positions(config): + args = get_env_start_args() + args.data_type = "bf16" + (state,) = Qwen4ExpTpPartModel.create_full_att_states(config, args, (0, 1, 2, 3)) + state.device = torch.device("cpu") + state.allocate_storage(8) + tensors = ( + state.k_storage, + state.v_storage, + state.indexer.k_storage, + state.indexer.v_storage, + state.positions.k_storage, + state.positions.v_storage, + ) + assert sum(t.numel() * t.element_size() for t in tensors) == (8 + state.page_size) * state.get_token_bytes() + indices = torch.tensor([1, 4]) + cos = torch.arange(8).view(2, 4).bfloat16() + state.store_positions(indices, cos, -cos) + view = state.qsa_view(3) + assert torch.equal(view.key_cos[indices], cos) + assert torch.equal(view.key_sin[indices], -cos) + assert view.raw_keys.data_ptr() == state.indexer.k_storage.data_ptr() + + +@pytest.mark.parametrize("page_kind", ["big", "small"]) +def test_lexical_prefix_round_trip_is_independent_of_runtime(lexical, page_kind): + lexical.ngram[:, 1].fill_(17) + lexical.convolution[:, 1].fill_(3) + lexical.save_bounded_att_runtime_to_pages([1], [128], page_kind, [0]) + lexical.reset_bounded_att_state([1]) + assert bool((lexical.ngram[:, 1] == 63).all()) + lexical.restore_bounded_att_runtime_from_pages([2], [128], page_kind, [0]) + assert bool((lexical.ngram[:, 2] == 17).all()) + assert bool((lexical.convolution[:, 2] == 3).all()) + lexical.reset_bounded_att_state([2]) + assert bool((lexical._pool(page_kind).ngram[0] == 17).all()) + + +def layout_for(state): + regions = state.get_cpu_transfer_regions() + return StateTransferLayout( + state.state_name, + None, + 16, + 16 + regions[-1].byte_end, + tuple(replace(region, byte_start=region.byte_start + 16, byte_end=region.byte_end + 16) for region in regions), + ) + + +def test_cpu_and_pd_lexical_transfers_preserve_unrelated_bytes(lexical): + layout = layout_for(lexical) + pages = torch.full((3, layout.byte_end + 16), 0x5A, dtype=torch.uint8) + lexical.ngram[:, 1].fill_(7) + lexical.convolution[:, 1].fill_(11) + lexical.save_bounded_att_runtime_to_pages([1], [7], "big", [0]) + lexical.write_state_pages_to_cpu("big", [0, 0], pages, [1, 2], layout, [False, True]) + assert bool((pages[2] == 0x5A).all()) + assert bool((pages[:, :16] == 0x5A).all()) + assert bool((pages[:, layout.byte_end :] == 0x5A).all()) + lexical.read_state_pages_from_cpu(pages, [1], [2], layout) + lexical.restore_bounded_att_runtime_from_pages([2], [7], "big", [2]) + assert torch.equal(lexical.ngram[:, 2], lexical.ngram[:, 1]) + assert torch.equal(lexical.convolution[:, 2], lexical.convolution[:, 1]) + lexical.write_to_pd_page(1, pages, 0, 7, layout) + lexical.reset_bounded_att_state([2]) + lexical.read_from_pd_page(pages, 2, 0, 7, layout) + assert torch.equal(lexical.ngram[:, 2], lexical.ngram[:, 1]) + assert torch.equal(lexical.convolution[:, 2], lexical.convolution[:, 1]) + + +def test_dp_copy_includes_runtime_and_prefix(lexical): + other = Qwen4ExpLexicalState(lexical.config, 4, 2, "cpu") + other.initialize_page_pools(4, pin_memory=False) + other.ngram[:, 0].fill_(29) + other.convolution[:, 0].fill_(13) + lexical.copy_runtime_state_from_dp([0], [2], other) + assert bool((lexical.ngram[:, 2] == 29).all()) + other.save_bounded_att_runtime_to_pages([0], [8], "big", [1]) + lexical.copy_big_pages_from_dp([1], [3], other) + lexical.restore_bounded_att_runtime_from_pages([1], [8], "big", [3]) + assert bool((lexical.convolution[:, 1] == 13).all()) + + +def test_mtp_commit_maps_request_rows_and_keeps_hold(lexical): + observed = [] + lexical.speculative_lexical_caches = {1: SimpleNamespace(commit=lambda *args: observed.append(args))} + requests, accepted = torch.tensor([2, 0, 3]), torch.tensor([1, 3, 1]) + lexical.commit_bounded_att_state(requests, accepted) + req, starts, lengths, hold = observed[0] + assert req is requests and lengths is accepted + assert starts.tolist() == [0, 1, 2] and hold == 3 + + +def test_ipc_state_excludes_weights_and_speculative_workspaces(lexical): + lexical.lexical_caches = {1: SimpleNamespace(layer_weight=object())} + lexical.speculative_lexical_caches = {1: object()} + exported = lexical.__getstate__() + assert not exported["lexical_caches"] and not exported["speculative_lexical_caches"] + assert exported["ngram"] is lexical.ngram + assert exported["convolution"] is lexical.convolution + assert lexical.lexical_caches and lexical.speculative_lexical_caches diff --git a/unit_tests/models/qwen4_exp/test_batch.py b/unit_tests/models/qwen4_exp/test_batch.py new file mode 100644 index 0000000000..e4895c294d --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_batch.py @@ -0,0 +1,97 @@ +import pytest +import torch + +from lightllm.models.qwen4_exp.batch import PackedBatchMetadata +from lightllm.models.qwen4_exp.state import QSABatch +from lightllm.common.basemodel.triton_kernel.gen_nsa_ks_ke import gen_nsa_ks_ke + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +def prepare(requests, lengths, query_lengths, hold=7): + tensors = [torch.tensor(value, device="cuda", dtype=torch.int32) for value in (requests, lengths, query_lengths)] + return tensors, PackedBatchMetadata.prepare(*tensors, sum(query_lengths), hold, 4) + + +def check(metadata, requests, lengths, query_lengths, hold=7): + counts = list(query_lengths) + for row, request in enumerate(requests): + if request != hold and row + 1 < len(requests) and requests[row + 1] == request: + counts[row + 1] += counts[row] + counts[row] = 0 + starts, active, key_lengths, first, rows, visible, mask = [0], [], [], [], [], [], [] + for request, length, count in zip(requests, lengths, counts): + real = request != hold and count > 0 + starts.append(starts[-1] + count) + active.append(request if real else -1) + key_lengths.append(length if real else 0) + first.append((length - count) // 4 if real else 0) + rows.extend([request] * count) + visible.extend(range(length - count + 1, length + 1) if real else [0] * count) + mask.extend([real] * count) + expected = [starts, active, key_lengths, first, rows, visible, mask] + for name, values in zip(metadata.__dataclass_fields__, expected): + actual = getattr(metadata, name) + torch.testing.assert_close(actual, torch.tensor(values, device="cuda", dtype=actual.dtype)) + + +@pytest.mark.parametrize( + "requests,lengths,query_lengths", + [ + ([5, 2, 7], [13, 21, 4], [13, 5, 4]), + ([5, 2, 7, 7], [14, 22, 2, 2], [1, 1, 1, 1]), + ([5, 5, 5, 2, 2, 7, 7], [14, 15, 16, 22, 23, 2, 2], [1] * 7), + ([5, 2, 7], [13, 21, 0], [0, 0, 0]), + ([5, 2], [1500, 1900], [1200, 1700]), + ([], [], []), + ], +) +def test_native_packed_batch_metadata(requests, lengths, query_lengths): + _, metadata = prepare(requests, lengths, query_lengths) + check(metadata, requests, lengths, query_lengths) + tokens = sum(query_lengths) + slots = torch.zeros(8, 2000, device="cuda", dtype=torch.int32) + memory = torch.arange(tokens, device="cuda", dtype=torch.int32) + rotary = torch.ones(tokens, 4, device="cuda", dtype=torch.bfloat16) + batch = QSABatch.from_metadata(metadata, slots, memory, rotary, rotary, 4) + assert batch.slot_table.data_ptr() == slots.data_ptr() + assert batch.active_rows.data_ptr() == metadata.active_rows.data_ptr() + assert batch.mem_indices.dtype == torch.int64 + assert batch.max_new_blocks == (tokens + 3) // 4 + + +def test_metadata_visibility_matches_existing_native_speculative_layout(): + requests, lengths, query_lengths = [5, 5, 5, 2, 2, 7, 7], [14, 15, 16, 22, 23, 2, 2], [1] * 7 + tensors, metadata = prepare(requests, lengths, query_lengths) + slots = torch.arange(8 * 32, device="cuda", dtype=torch.int32).reshape(8, 32) + ragged = torch.empty(sum([16, 23, 2, 2]), device="cuda", dtype=torch.int32) + _, _, native_lengths = gen_nsa_ks_ke(tensors[1], tensors[2], tensors[0], slots, 7, ragged, hold_req_idx=7) + torch.testing.assert_close(metadata.visible_lengths[metadata.token_mask], native_lengths[metadata.token_mask]) + + +def test_metadata_graph_replay_and_copy_preserve_buffer_addresses(): + requests, lengths, query_lengths = [5, 5, 2, 7, 7], [14, 15, 22, 2, 2], [1] * 5 + tensors, metadata = prepare(requests, lengths, query_lengths) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = PackedBatchMetadata.prepare(*tensors, 5, 7, 4) + pointers = {name: getattr(captured, name).data_ptr() for name in captured.__dataclass_fields__} + for ids, sizes in [([1, 3, 5, 7, 7], [25, 41, 13, 2, 2]), ([2, 2, 2, 5, 7], [12, 13, 14, 5, 2])]: + tensors[0].copy_(torch.tensor(ids, device="cuda", dtype=torch.int32)) + tensors[1].copy_(torch.tensor(sizes, device="cuda", dtype=torch.int32)) + graph.replay() + check(captured, ids, sizes, query_lengths) + metadata.copy_from(captured) + check(metadata, ids, sizes, query_lengths) + assert pointers == {name: getattr(captured, name).data_ptr() for name in captured.__dataclass_fields__} + _, incompatible = prepare([1], [1], [1]) + with pytest.raises(ValueError, match="Incompatible graph batch metadata"): + metadata.copy_from(incompatible) + + +def test_metadata_rejects_incompatible_input_shapes(): + values = torch.ones(2, device="cuda", dtype=torch.int32) + with pytest.raises(ValueError, match="matching contiguous"): + PackedBatchMetadata.prepare(values, values[:1], values, 2, 7, 4) + with pytest.raises(ValueError, match="compression ratio"): + PackedBatchMetadata.prepare(values, values, values, 2, 7, 0) diff --git a/unit_tests/models/qwen4_exp/test_decoder.py b/unit_tests/models/qwen4_exp/test_decoder.py new file mode 100644 index 0000000000..fc466ec680 --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_decoder.py @@ -0,0 +1,214 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.models.qwen4_exp.layer_weights.transformer_layer_weight import Qwen4ExpTransformerLayerWeight +from lightllm.models.qwen4_exp.layer_infer.transformer_layer_infer import Qwen4ExpTransformerLayerInfer +from lightllm.models.qwen4_exp.state import QSABatch, QSAExecutionState +from lightllm.models.qwen4_exp.batch import PackedBatchMetadata +from lightllm.models.qwen4_exp.position import combine_mrope +from lightllm.common.quantization import Quantcfg +from lightllm.models.qwen4_exp.attention_state import Qwen4ExpAttentionState +from lightllm.common.attention_state.full_att.gqa.gqa_state import GQAAttentionConfig +from .helpers import checkpoint +from .reference import reference_layer + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +def setup(config, lengths, physical_cache=False): + source = checkpoint(config, 3) + weight = Qwen4ExpTransformerLayerWeight(3, torch.bfloat16, config, Quantcfg(config)) + weight.load_hf_weights(dict(source)) + weight.verify_load() + infer = Qwen4ExpTransformerLayerInfer(3, config) + + from lightllm.utils.dist_utils import get_dp_world_size + + manager = Qwen4ExpAttentionState( + GQAAttentionConfig( + logical_layer_groups=((3,),), + global_head_num=config["num_key_value_heads"], + head_dim=config["head_dim"], + dtype=torch.bfloat16, + tp_world_size=get_dp_world_size(), + ), + 16, + 4, + device="cuda", + ) + manager.allocate_storage(47) + rows = [5, 2] + table = torch.full((8, 20), -1, device="cuda", dtype=torch.int32) + slots = torch.randperm(48, device="cuda")[: sum(lengths)] + offset = 0 + for row, count in zip(rows, lengths): + table[row, :count] = slots[offset : offset + count].int() + offset += count + dtype = torch.bfloat16 + angle = torch.randn(3, sum(lengths), 2, device="cuda") + pc, ps = angle.cos().to(dtype), angle.sin().to(dtype) + cos, sin = combine_mrope(pc, ps, [1, 1, 0]) + key_cos, key_sin = torch.zeros(48, 4, device="cuda", dtype=dtype), torch.zeros(48, 4, device="cuda", dtype=dtype) + key_cos.index_copy_(0, slots, cos) + key_sin.index_copy_(0, slots, sin) + metadata = PackedBatchMetadata.prepare( + torch.tensor([5, 7, 2], device="cuda", dtype=torch.int32), + torch.tensor([lengths[0], 0, lengths[1]], device="cuda", dtype=torch.int32), + torch.tensor([lengths[0], 0, lengths[1]], device="cuda", dtype=torch.int32), + sum(lengths), + 7, + 4, + ) + batch = QSABatch.from_metadata(metadata, table, slots, cos, sin, 4) + view = QSAExecutionState( + torch.zeros(48, 16, device="cuda", dtype=dtype), + torch.full((48, 16) if physical_cache else (8, 5, 16), float("nan"), device="cuda", dtype=dtype), + key_cos, + key_sin, + ) + state = SimpleNamespace( + is_prefill=True, + need_dp_prefill_balance=False, + dist_group=None, + position_cos=pc, + position_sin=ps, + mem_index=slots.int(), + att_state_role="main", + att_state_manager=SimpleNamespace(get_state=lambda *args, **kwargs: manager), + qwen4_exp_qsa={3: view}, + qwen4_exp_batch=batch, + ) + return infer, weight, source, state + + +@pytest.mark.parametrize("physical_cache", [False, True]) +def test_native_qsa_decoder_matches_reference(config, physical_cache): + torch.manual_seed(937) + lengths = [13, 5] + infer, weights, source, state = setup(config, lengths, physical_cache) + residual = torch.randn(sum(lengths), 64, device="cuda", dtype=torch.bfloat16) + before = residual.clone() + actual = infer.context_forward(residual, state, weights) + expected = reference_layer(residual, source, config, state.qwen4_exp_batch.cos, state.qwen4_exp_batch.sin, lengths) + torch.testing.assert_close(actual, expected, atol=0.016, rtol=0.016) + torch.testing.assert_close(residual, before, atol=0, rtol=0) + key, value = state.att_state_manager.get_state().get_layer_view(3) + assert key.shape == value.shape + + old_batch = state.qwen4_exp_batch + + used = set(old_batch.mem_indices.tolist()) + slots = torch.tensor([slot for slot in range(48) if slot not in used][:2], device="cuda", dtype=torch.int64) + angle = torch.randn(3, 2, 2, device="cuda") + pc, ps = angle.cos().to(residual.dtype), angle.sin().to(residual.dtype) + cos, sin = combine_mrope(pc, ps, [1, 1, 0]) + view = state.qwen4_exp_qsa[3] + view.key_cos.index_copy_(0, slots, cos) + view.key_sin.index_copy_(0, slots, sin) + for row, count, slot in zip([5, 2], lengths, slots): + old_batch.slot_table[row, count] = slot.int() + metadata = PackedBatchMetadata.prepare( + torch.tensor([5, 2], device="cuda", dtype=torch.int32), + torch.tensor([14, 6], device="cuda", dtype=torch.int32), + torch.ones(2, device="cuda", dtype=torch.int32), + 2, + 7, + 4, + ) + state.qwen4_exp_batch = QSABatch.from_metadata(metadata, old_batch.slot_table, slots, cos, sin, 4) + state.position_cos, state.position_sin, state.mem_index, state.is_prefill = pc, ps, slots.int(), False + new_residual = torch.randn(2, 64, device="cuda", dtype=residual.dtype) + actual = infer.token_forward(new_residual, state, weights) + all_residual = torch.cat([residual[:13], new_residual[:1], residual[13:], new_residual[1:]]) + all_cos = torch.cat([old_batch.cos[:13], cos[:1], old_batch.cos[13:], cos[1:]]) + all_sin = torch.cat([old_batch.sin[:13], sin[:1], old_batch.sin[13:], sin[1:]]) + full_reference = reference_layer(all_residual, source, config, all_cos, all_sin, [14, 6]) + torch.testing.assert_close(actual, full_reference[[13, 19]], atol=0.016, rtol=0.016) + + +@pytest.mark.parametrize("physical_cache", [False, True]) +def test_native_qsa_decoder_cuda_graph_replay(config, physical_cache): + torch.manual_seed(937) + lengths = [13, 5] + infer, weights, source, state = setup(config, lengths, physical_cache) + residual = torch.randn(sum(lengths), 64, device="cuda", dtype=torch.bfloat16) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + infer.context_forward(residual, state, weights) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = infer.context_forward(residual, state, weights) + batch = state.qwen4_exp_batch + view = state.qwen4_exp_qsa[3] + for _ in range(2): + residual.normal_() + angle = torch.randn_like(state.position_cos.float()) + state.position_cos.copy_(angle.cos()) + state.position_sin.copy_(angle.sin()) + cos, sin = combine_mrope(state.position_cos, state.position_sin, [1, 1, 0]) + batch.cos.copy_(cos) + batch.sin.copy_(sin) + view.key_cos.index_copy_(0, batch.mem_indices, cos) + view.key_sin.index_copy_(0, batch.mem_indices, sin) + graph.replay() + expected = reference_layer(residual, source, config, cos, sin, lengths) + torch.testing.assert_close(actual, expected, atol=0.016, rtol=0.016) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("prefill", [False, True]) +def test_long_decode_attention_matches_single_queries(dtype, prefill): + from lightllm.common.qsa.indexer import QSAIndexer + + torch.manual_seed(701) + slots, rows, dim = 4096, 3, 128 + + def rand(*shape): + return torch.randn(*shape, device="cuda", dtype=dtype) + + indexer = QSAIndexer(rand(5 * dim, dim), rand(dim), rand(dim)) + raw = rand(slots, dim) + pooled = rand(1, slots // 4, dim) + angles = torch.empty(slots, 0, device="cuda", dtype=dtype) + view = QSAExecutionState(raw, pooled, angles, angles) + query, key, value = rand(rows, 6, 256), rand(slots, 1, 256), rand(slots, 1, 256) + hidden = rand(rows, dim) + table = torch.arange(slots, device="cuda", dtype=torch.int32)[None, :] + + def execute(start, end, decode=True, attention_splits=None): + count = end - start + batch = QSABatch( + table, + torch.tensor([0], device="cuda", dtype=torch.int32), + torch.tensor([slots], device="cuda", dtype=torch.int32), + torch.tensor([slots // 4], device="cuda", dtype=torch.int32), + torch.zeros(count, device="cuda", dtype=torch.int32), + torch.arange(2050 + start, 2050 + end, device="cuda", dtype=torch.int32), + torch.arange(slots - rows + start, slots - rows + end, device="cuda"), + angles[:count], + angles[:count], + 0, + ) + return view.attend( + indexer, + hidden[start:end], + query[start:end], + key, + value, + batch, + decode=decode, + attention_splits=attention_splits, + ) + + options = {"decode": not prefill, "attention_splits": None} + expected = torch.cat([execute(row, row + 1, **options) for row in range(rows)]) + if prefill: + torch.testing.assert_close(execute(0, rows, **options), expected, atol=1e-5, rtol=1e-5) + else: + torch.testing.assert_close(execute(0, rows, **options), expected, atol=0, rtol=0) + torch.testing.assert_close(execute(0, 1, decode=False), expected[:1], atol=0, rtol=0) diff --git a/unit_tests/models/qwen4_exp/test_embedding_residency.py b/unit_tests/models/qwen4_exp/test_embedding_residency.py new file mode 100644 index 0000000000..9ca54c0478 --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_embedding_residency.py @@ -0,0 +1,45 @@ +import argparse + +import pytest +import torch + +from lightllm.models.qwen4_exp.layer_weights import embedding_weight as module +from lightllm.server.api_cli import add_cli_args +from lightllm.server.core.objs.start_args_type import StartArgs + + +def test_residency_is_opt_in_and_cli_matches_start_args(): + parser = add_cli_args(argparse.ArgumentParser()) + assert not parser.parse_args([]).offload_ngram_embedding + assert not StartArgs().offload_ngram_embedding + assert parser.parse_args(["--offload_ngram_embedding"]).offload_ngram_embedding + + +@pytest.mark.parametrize("args", [{}, {"offload_ngram_embedding": False}]) +def test_gpu_factory_preserves_existing_constructor(args, monkeypatch): + monkeypatch.setattr(module, "get_env_start_args", lambda: args) + monkeypatch.setattr(module, "RowShardedEmbeddingWeight", lambda *values: values) + result = module.ngram_embedding_weight(160, 1024, "table.weight", torch.bfloat16, 4) + assert result == (160, 1024, "table.weight", torch.bfloat16, 4) + + +@pytest.mark.parametrize("torch_fallback", [False, True]) +def test_cpu_factory_injects_model_owned_allocation_policy(monkeypatch, torch_fallback): + monkeypatch.setattr( + module, + "get_env_start_args", + lambda: {"offload_ngram_embedding": True, "hardware_platform": "cuda", "enable_torch_fallback": torch_fallback}, + ) + monkeypatch.setattr(module, "PinnedRowShardedEmbeddingWeight", lambda *values, **kwargs: (values, kwargs)) + from lightllm.models.qwen4_exp.pinned_memory import pinned_zeros + + values, kwargs = module.ngram_embedding_weight(160, 1024, "table.weight", torch.bfloat16, 4) + assert values == (160, 1024, "table.weight", torch.bfloat16, 4) + assert kwargs == {"cpu_allocator": pinned_zeros} + + +@pytest.mark.parametrize("extra", [{"hardware_platform": "musa"}, {"hardware_platform": "cpu"}]) +def test_unsupported_offload_configuration_fails_before_allocation(extra, monkeypatch): + monkeypatch.setattr(module, "get_env_start_args", lambda: dict(offload_ngram_embedding=True, **extra)) + with pytest.raises(ValueError, match="CUDA Triton"): + module.ngram_embedding_weight(160, 1024, "table.weight", torch.bfloat16, 4) diff --git a/unit_tests/models/qwen4_exp/test_eos.py b/unit_tests/models/qwen4_exp/test_eos.py new file mode 100644 index 0000000000..31e15627d9 --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_eos.py @@ -0,0 +1,20 @@ +import json +from types import SimpleNamespace + +import pytest +from transformers import AutoTokenizer + +from lightllm.utils.config_utils import get_eos_token_ids + + +@pytest.mark.parametrize("model_type", ["qwen4_exp", "qwen4_exp_text", "qwen3_5", "llama"]) +@pytest.mark.parametrize("generation_eos", [[248046, 248044], 248046, None]) +def test_chat_end_and_end_of_text_are_preserved(tmp_path, monkeypatch, model_type, generation_eos): + (tmp_path / "config.json").write_text( + json.dumps({"model_type": model_type, "text_config": {"eos_token_id": 248044}}) + ) + if generation_eos is not None: + (tmp_path / "generation_config.json").write_text(json.dumps({"eos_token_id": generation_eos})) + monkeypatch.setattr(AutoTokenizer, "from_pretrained", lambda *args, **kwargs: SimpleNamespace(eos_token_id=248046)) + expected = {248044} if model_type == "llama" else {248044, 248046} + assert set(get_eos_token_ids(str(tmp_path))) == expected diff --git a/unit_tests/models/qwen4_exp/test_lexical_cache.py b/unit_tests/models/qwen4_exp/test_lexical_cache.py new file mode 100644 index 0000000000..fef7202721 --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_lexical_cache.py @@ -0,0 +1,100 @@ +import pytest +import torch + +from lightllm.models.qwen4_exp.lexical_cache import LexicalSnapshot +from .lexical_helpers import make_cache, metadata, reference_segment + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +@pytest.mark.parametrize("capture", [False, True]) +def test_request_histories_acceptance_padding_and_graphs(config, capture): + torch.manual_seed(1207) + cache, source = make_cache(config) + + cache.ngram[5].copy_(torch.tensor([7, 11], device="cuda")) + cache.convolution[5].normal_() + prefix = cache.snapshot(5) + cache.restore(2, prefix) + requests = [5, 5, 5, 2, 7, 7] + batch_meta = metadata(requests, [14, 15, 16, 14, 2, 2], [1] * 6) + accepted = torch.tensor([0, 0, 2, 1, 0, 0], device="cuda", dtype=torch.int32) + tokens = torch.zeros(6, device="cuda", dtype=torch.int64) + residual = torch.randn(6, 64, device="cuda", dtype=torch.bfloat16) + untouched_ngram, untouched_conv = cache.ngram[[0, 7]].clone(), cache.convolution[[0, 7]].clone() + + def run(): + candidate = cache.prepare(batch_meta) + update = candidate.apply(residual, tokens, cache.layer_weight, 1, None) + cache.commit(batch_meta, candidate, accepted) + return update + + if capture: + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + run() + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + result = run() + cache.restore(5, prefix) + cache.restore(2, prefix) + ref_ngram = {request: prefix.ngram[None].clone() for request in (5, 2)} + ref_conv = {request: prefix.convolution[None].clone() for request in (5, 2)} + for ids in ([19, 63, 23, 29, 31, 37], [41, 43, 47, 63, 5, 7], [11, 13, 17, 19, 23, 29]): + tokens.copy_(torch.tensor(ids, device="cuda", dtype=torch.int64)) + residual.normal_() + if capture: + graph.replay() + actual = result + else: + before_ngram, before_conv = cache.ngram.clone(), cache.convolution.clone() + candidate = cache.prepare(batch_meta) + actual = candidate.apply(residual, tokens, cache.layer_weight, 1, None) + + torch.testing.assert_close(cache.ngram, before_ngram, atol=0, rtol=0) + torch.testing.assert_close(cache.convolution, before_conv, atol=0, rtol=0) + cache.commit(batch_meta, candidate, accepted) + for request, start, end, count in ((5, 0, 3, 2), (2, 3, 4, 1)): + expected, states = reference_segment( + residual[start:end], tokens[start:end], source, config, ref_ngram[request], ref_conv[request] + ) + torch.testing.assert_close(actual[start:end], expected, atol=0.016, rtol=0.016) + ref_ngram[request], ref_conv[request] = states[count] + torch.testing.assert_close(cache.ngram[request], ref_ngram[request][0], atol=0, rtol=0) + torch.testing.assert_close(cache.convolution[request], ref_conv[request][0], atol=0.016, rtol=0.016) + torch.testing.assert_close(cache.ngram[[0, 7]], untouched_ngram, atol=0, rtol=0) + torch.testing.assert_close(cache.convolution[[0, 7]], untouched_conv, atol=0, rtol=0) + torch.testing.assert_close(prefix.ngram, torch.tensor([7, 11], device="cuda"), atol=0, rtol=0) + + +def test_lexical_snapshot_reset_cpu_restore_and_rejection(config): + torch.manual_seed(1207) + cache, _ = make_cache(config) + cache.ngram[3].copy_(torch.tensor([7, 11], device="cuda")) + cache.convolution[3].normal_() + snapshot = cache.snapshot(3) + cache.reset(3) + assert torch.all(cache.ngram[3] == config["eos_token_id"]) + assert torch.count_nonzero(cache.convolution[3]) == 0 + cache.restore(5, LexicalSnapshot(snapshot.ngram.cpu(), snapshot.convolution.cpu())) + torch.testing.assert_close(cache.ngram[5], snapshot.ngram) + torch.testing.assert_close(cache.convolution[5], snapshot.convolution) + before_ngram, before_conv = cache.ngram.clone(), cache.convolution.clone() + batch_meta = metadata([5], [12], [3]) + candidate = cache.prepare(batch_meta) + candidate.apply( + torch.randn(3, 64, device="cuda", dtype=torch.bfloat16), + torch.tensor([17, 63, 19], device="cuda"), + cache.layer_weight, + 1, + None, + ) + cache.commit(batch_meta, candidate, torch.zeros(1, device="cuda", dtype=torch.int32)) + torch.testing.assert_close(cache.ngram, before_ngram, atol=0, rtol=0) + torch.testing.assert_close(cache.convolution, before_conv, atol=0, rtol=0) + with pytest.raises(ValueError, match="incompatible"): + cache.restore(3, LexicalSnapshot(snapshot.ngram[:1], snapshot.convolution)) + assert torch.all(cache.ngram[3] == config["eos_token_id"]) diff --git a/unit_tests/models/qwen4_exp/test_linear_decoder.py b/unit_tests/models/qwen4_exp/test_linear_decoder.py new file mode 100644 index 0000000000..952c268a3a --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_linear_decoder.py @@ -0,0 +1,183 @@ +from argparse import Namespace +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.models.qwen4_exp.layer_weights.transformer_layer_weight import Qwen4ExpTransformerLayerWeight +from lightllm.models.qwen4_exp.layer_infer.transformer_layer_infer import Qwen4ExpTransformerLayerInfer +from lightllm.models.qwen4_exp.state import PLEBatch +from lightllm.models.qwen4_exp.batch import PackedBatchMetadata +from lightllm.models.qwen4_exp.lexical_cache import LexicalRequestCache +from lightllm.common.quantization import Quantcfg +from lightllm.common.basemodel.attention.linear.triton import TritonLinearAttBackend +from .helpers import checkpoint +from .reference import reference_step + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +@pytest.mark.parametrize("capture", [False, True]) +def test_native_ple_gdn_and_moe_cached_steps(config, capture): + torch.manual_seed(1037) + source = checkpoint(config, 1) + weights = Qwen4ExpTransformerLayerWeight(1, torch.bfloat16, config, Quantcfg(config)) + weights.load_hf_weights(dict(source)) + weights.verify_load() + infer = Qwen4ExpTransformerLayerInfer(1, config) + channels = 512 + conv = torch.zeros(2, channels, 3, device="cuda", dtype=torch.bfloat16) + recurrent = torch.zeros(2, 2, 128, 128, device="cuda", dtype=torch.float32) + state = Namespace( + is_prefill=False, + need_dp_prefill_balance=False, + dist_group=None, + batch_size=2, + b_req_idx=torch.tensor([0, 1], device="cuda", dtype=torch.int32), + att_state_role="main", + att_state_manager=SimpleNamespace( + get_state=lambda **kwargs: SimpleNamespace(get_layer_view=lambda **kwargs: (conv, recurrent)) + ), + req_manager=SimpleNamespace(get_mamba_cache=lambda layer: (conv, recurrent)), + ) + + class Model: + pass + + model = Model() + model.config, model.tp_world_size_ = config, 1 + model.is_mtp_draft_model = False + model.mtp_manager = SimpleNamespace(get_decode_draft_step=lambda is_draft: 0) + backend = TritonLinearAttBackend(model=model) + state.decode_att_state1 = backend.create_att_decode_state(state) + state.decode_att_state1.init_state() + cache = LexicalRequestCache(weights, 2) + ref_ngram, ref_ple = cache.ngram.clone(), cache.convolution.clone() + ref_conv, ref_recurrent = conv.clone(), recurrent.clone() + ones = torch.ones(2, device="cuda", dtype=torch.int32) + metadata = PackedBatchMetadata.prepare(state.b_req_idx, ones, ones, 2, 7, 4) + + def run(residual, tokens): + batch = cache.prepare(metadata) + state.input_ids, state.qwen4_exp_ple = tokens, {1: batch} + output = infer.token_forward(residual, state, weights) + cache.commit(metadata, batch) + return output + + if capture: + residual = torch.randn(2, 64, device="cuda", dtype=torch.bfloat16) + tokens = torch.zeros(2, device="cuda", dtype=torch.int64) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + run(residual, tokens) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured_output = run(residual, tokens) + conv.zero_() + recurrent.zero_() + for request in range(2): + cache.reset(request) + for ids in [[7, 11], [63, 15], [19, 23]]: + if capture: + residual.normal_() + tokens.copy_(torch.tensor(ids, device="cuda", dtype=torch.int64)) + graph.replay() + actual = captured_output + else: + residual = torch.randn(2, 64, device="cuda", dtype=torch.bfloat16) + tokens = torch.tensor(ids, device="cuda", dtype=torch.int64) + actual = run(residual, tokens) + expected, ref_ngram, ref_ple, ref_conv, ref_recurrent = reference_step( + residual, tokens, source, config, ref_ngram, ref_ple, ref_conv, ref_recurrent + ) + torch.testing.assert_close(actual, expected, atol=0.016, rtol=0.016) + torch.testing.assert_close(cache.ngram, ref_ngram, atol=0, rtol=0) + torch.testing.assert_close(cache.convolution, ref_ple, atol=0.016, rtol=0.016) + torch.testing.assert_close(conv, ref_conv, atol=0.016, rtol=0.016) + + +def test_native_packed_gdn_ple_prefill_matches_recurrent_reference(config): + from lightllm.models.qwen3next.model import Qwen3NextTpPartModel + + torch.manual_seed(1037) + + Qwen3NextTpPartModel._init_triton(None) + source = checkpoint(config, 1) + weights = Qwen4ExpTransformerLayerWeight(1, torch.bfloat16, config, Quantcfg(config)) + weights.load_hf_weights(dict(source)) + weights.verify_load() + infer = Qwen4ExpTransformerLayerInfer(1, config) + conv = torch.zeros(2, 512, 3, device="cuda", dtype=torch.bfloat16) + recurrent = torch.zeros(2, 2, 128, 128, device="cuda", dtype=torch.float32) + starts = torch.tensor([0, 3, 5], device="cuda", dtype=torch.int32) + tokens = torch.tensor([7, 63, 19, 11, 23], device="cuda", dtype=torch.int64) + residual = torch.randn(5, 64, device="cuda", dtype=torch.bfloat16) + batch = PLEBatch(starts, weights.ngram_hasher.component.empty_history(2), weights.ple.component.empty_history(2)) + state = Namespace( + is_prefill=True, + need_dp_prefill_balance=False, + dist_group=None, + batch_size=2, + b_req_idx=torch.tensor([0, 1], device="cuda", dtype=torch.int32), + b1_cu_q_seq_len=starts, + b_ready_cache_len=torch.zeros(2, device="cuda", dtype=torch.int32), + att_state_role="main", + att_state_manager=SimpleNamespace( + get_state=lambda **kwargs: SimpleNamespace(get_layer_view=lambda **kwargs: (conv, recurrent)) + ), + req_manager=SimpleNamespace(get_mamba_cache=lambda layer: (conv, recurrent)), + input_ids=tokens, + qwen4_exp_ple={1: batch}, + ) + + class Model: + pass + + model = Model() + model.config, model.tp_world_size_ = config, 1 + model.is_mtp_draft_model = False + model.mtp_manager = SimpleNamespace(get_decode_draft_step=lambda is_draft: 0) + backend = TritonLinearAttBackend(model=model) + state.prefill_att_state1 = backend.create_att_prefill_state(state) + state.prefill_att_state1.init_state() + actual = infer.context_forward(residual, state, weights) + next_ngram, next_ple = batch.histories_after(weights) + expected = [] + for request, (start, end) in enumerate([(0, 3), (3, 5)]): + ng = weights.ngram_hasher.component.empty_history(1) + ple = weights.ple.component.empty_history(1) + cv = torch.zeros(1, 512, 3, device="cuda", dtype=torch.bfloat16) + rc = torch.zeros(1, 2, 128, 128, device="cuda", dtype=torch.float32) + for token in range(start, end): + out, ng, ple, cv, rc = reference_step( + residual[token : token + 1], tokens[token : token + 1], source, config, ng, ple, cv, rc + ) + expected.append(out) + torch.testing.assert_close(next_ngram[request : request + 1], ng, atol=0, rtol=0) + torch.testing.assert_close(next_ple[request : request + 1], ple, atol=0.016, rtol=0.016) + torch.testing.assert_close(actual, torch.cat(expected), atol=0.016, rtol=0.016) + + +def test_native_ple_empty_packed_batch_preserves_histories(config): + source = checkpoint(config, 1) + weights = Qwen4ExpTransformerLayerWeight(1, torch.bfloat16, config, Quantcfg(config)) + weights.load_hf_weights(dict(source)) + weights.verify_load() + starts = torch.zeros(3, device="cuda", dtype=torch.int32) + ngram = weights.ngram_hasher.component.empty_history(2) + conv = weights.ple.component.empty_history(2) + batch = PLEBatch(starts, ngram, conv) + update = batch.apply( + torch.empty(0, 64, device="cuda", dtype=torch.bfloat16), + torch.empty(0, device="cuda", dtype=torch.int64), + weights, + 1, + None, + ) + assert update.shape == (0, 64) + next_ngram, next_conv = batch.histories_after(weights) + torch.testing.assert_close(next_ngram, ngram, atol=0, rtol=0) + torch.testing.assert_close(next_conv, conv, atol=0, rtol=0) diff --git a/unit_tests/models/qwen4_exp/test_lm_head.py b/unit_tests/models/qwen4_exp/test_lm_head.py new file mode 100644 index 0000000000..27300687aa --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_lm_head.py @@ -0,0 +1,45 @@ +import pytest +import torch + +from lightllm.models.qwen4_exp.layer_weights.lm_head_weight import Qwen4ExpLMHeadWeight + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("rows", [0, 1, 3, 8, 9]) +def test_head_shape_precision_and_explicit_output(dtype, rows): + torch.manual_seed(3107) + head = Qwen4ExpLMHeadWeight(128, 257, "head.weight", dtype) + head.weight.normal_(std=0.02) + hidden = torch.randn(rows, 128, device="cuda", dtype=dtype) + result = head(hidden.T) + assert result.shape == (257, rows) and result.dtype == dtype and result.is_contiguous() + expected = (head.weight.float() @ hidden.float().T).to(dtype) + tolerance = max(1e-6, torch.finfo(dtype).eps) + torch.testing.assert_close(result, expected, atol=tolerance, rtol=tolerance) + out = torch.empty(257, rows * 2, device="cuda", dtype=dtype)[:, ::2] + assert head(hidden.T, out=out) is out + torch.testing.assert_close(out, result, atol=0, rtol=0) + + +def test_trained_geometry_batch_invariance_and_graph_replay(): + torch.manual_seed(3107) + head = Qwen4ExpLMHeadWeight(2560, 62080, "head.weight", torch.bfloat16) + head.weight.normal_(std=0.02) + hidden = torch.randn(24, 2560, device="cuda", dtype=torch.bfloat16) + singles = torch.cat([head(row[:, None]) for row in hidden], dim=1) + for chunk in (2, 3, 5, 7, 8, 9, 15, 24): + result = torch.cat([head(hidden[start : start + chunk].T) for start in range(0, len(hidden), chunk)], dim=1) + torch.testing.assert_close(result, singles, atol=0, rtol=0) + + inputs = hidden[:3].clone() + for _ in range(3): + head(inputs.T) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = head(inputs.T) + inputs.add_(0.125) + graph.replay() + torch.testing.assert_close(captured, head(inputs.T), atol=0, rtol=0) + graph.reset() diff --git a/unit_tests/models/qwen4_exp/test_mtp.py b/unit_tests/models/qwen4_exp/test_mtp.py new file mode 100644 index 0000000000..48c627a3dc --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_mtp.py @@ -0,0 +1,148 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.basemodel.hidden_collector import FinalHiddenCollector +from lightllm.models.draft_registry import get_draft_model_class +from lightllm.models.qwen3next.model import Qwen3NextTpPartModel +from lightllm.models.qwen4_exp.mtp.model import Qwen4ExpMTPModel +from lightllm.models.qwen4_exp.mtp.infer import Qwen4ExpMTPPreLayerInfer +from .helpers import checkpoint, prepost_checkpoint, add_hyper, rand +from .runtime import state_for +from .mtp_reference import reference_fusion, reference_draft + + +def test_public_draft_registry(): + for model_type in ("qwen4_exp", "qwen4_exp_text"): + assert get_draft_model_class({"model_type": model_type}, "eagle_with_att") is Qwen4ExpMTPModel + + +def build_draft(runtime, config): + main = runtime.model + source = prepost_checkpoint(config) + main.pre_post_weight = main.pre_and_post_weight_class(torch.bfloat16, config) + main.pre_post_weight.load_hf_weights(dict(source)) + main.pre_post_weight.verify_load() + main.layers_infer = [None] * config["n_layer"] + draft = Qwen4ExpMTPModel.__new__(Qwen4ExpMTPModel) + draft.__dict__.update( + main_model=main, + mtp_previous_draft_models=[], + args=main.args, + weight_dir_=main.args.model_dir, + finetune_config=None, + tp_world_size_=1, + tp_rank_=0, + torch_memory_saver=main.torch_memory_saver, + data_type=torch.bfloat16, + quant_cfg=runtime.weight.quant_cfg, + max_req_num=main.max_req_num, + prefill_graph=None, + ) + draft._init_config() + draft._init_weights() + from .resource_helpers import bind_states + + draft.bind_req_manager(main.req_manager) + draft._init_infer_layer() + main.att_state_manager._state_registration_finished = False + draft.prepare_attention_states() + runtime.mem = bind_states(main) + main.att_state_manager.get_state.cache_clear() + draft._init_custom() + draft._init_infer_layer() + draft._init_some_value() + draft._init_att_backend() + draft._init_att_backend1() + draft.mtp_manager = SimpleNamespace(get_decode_draft_step=lambda is_draft: 0) + for name, value in checkpoint(config, 3).items(): + source[name.replace("model.language_model.layers.3.", "mtp.layers.0.", 1)] = value + add_hyper(source, "mtp.hyper_connection_mixer", config, False) + source.update( + { + "mtp.pre_fc_norm_embedding.weight": rand(config["hidden_size"]), + "mtp.pre_fc_norm_hidden.weight": rand(config["hidden_size"] * config["hc_count"]), + "mtp.fc_embedding.weight": rand(config["hidden_size"], config["hidden_size"]), + "mtp.fc_hidden.weight": rand(config["hidden_size"], config["hidden_size"]), + } + ) + + for name in reversed(list(source)): + draft.pre_post_weight.load_hf_weights({name: source[name]}) + draft.trans_layers_weight[0].load_hf_weights({name: source[name]}) + draft.pre_post_weight.verify_load() + draft.trans_layers_weight[0].verify_load() + return draft, source + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@pytest.mark.parametrize("runtime", [{"mtp_step": 2, "mtp_mode": "eagle_with_att"}], indirect=True) +def test_native_complete_draft_prefill_cached_decode_and_full_hidden(runtime, config): + torch.manual_seed(2307) + Qwen3NextTpPartModel._init_triton(None) + draft, source = build_draft(runtime, config) + assert draft.req_manager is runtime.req and draft.att_state_manager is runtime.model.att_state_manager + assert draft.pre_post_weight.wte_weight_ is runtime.model.pre_post_weight.wte_weight_ + assert draft.pre_post_weight.lm_head_weight_ is runtime.model.pre_post_weight.lm_head_weight_ + assert draft.layers_infer[0].layer_num_ == 4 + assert not hasattr(draft.trans_layers_weight[0], "ple") + for request in (1, 3): + runtime.req.init_linear_att_state(SimpleNamespace(req_idx=request)) + owner = SimpleNamespace(model=draft, req=runtime.req, mem=runtime.mem) + tokens = torch.tensor([7, 11, 13, 17, 19], device="cuda", dtype=torch.int64) + hidden = torch.randn(5, 64, device="cuda", dtype=torch.bfloat16) + hidden_before = hidden.clone() + state = state_for(owner, [1, 3], [3, 2], tokens.tolist(), ready=[0, 0]) + state.mtp_draft_input_hiddens = hidden + state.hidden_collector = FinalHiddenCollector() + assert not state.qwen4_exp_ple + actual = draft._context_forward(state) + positions = torch.tensor([0, 1, 2, 0, 1], device="cuda") + cos, sin = draft._cos_cached[positions].repeat(1, 2), draft._sin_cached[positions].repeat(1, 2) + multi, logits = reference_draft(tokens, hidden, source, config, cos, sin, [3, 2]) + torch.testing.assert_close(actual.logits, logits[[2, 4]], atol=0.025, rtol=0.025) + torch.testing.assert_close(actual.mtp_collector.spec_hidden, multi, atol=0.016, rtol=0.016) + torch.testing.assert_close(hidden, hidden_before, atol=0, rtol=0) + fusion = draft.pre_post_weight.fusion.component + embeddings = source["model.language_model.embed_tokens.weight"].cuda()[tokens] + torch.testing.assert_close( + fusion(embeddings, hidden), reference_fusion(embeddings, hidden, source, config), atol=0.016, rtol=0.016 + ) + with pytest.raises(ValueError, match="full pre-mixer"): + fusion(embeddings, hidden[:, :32].contiguous()) + + image_pre = Qwen4ExpMTPPreLayerInfer(dict(config, image_token_id=60)) + virtual = torch.tensor([7, 1000000, 1000001], device="cuda", dtype=torch.int64) + image_state = SimpleNamespace( + multimodal_params=[{"images": [{"token_id": 1000000, "token_num": 2}]}], + mtp_draft_input_hiddens=hidden[:3], + dist_group=None, + ) + canonical = torch.tensor([7, 60, 60], device="cuda") + image_input = image_pre.context_forward(virtual, image_state, draft.pre_post_weight) + image_embedding = source["model.language_model.embed_tokens.weight"].cuda()[canonical] + torch.testing.assert_close( + image_input, reference_fusion(image_embedding, hidden[:3], source, config), atol=0.016, rtol=0.016 + ) + assert virtual.tolist() == [7, 1000000, 1000001] + + state = state_for(owner, [3, 1], [3, 4], [23, 29]) + new_hidden = torch.randn(2, 64, device="cuda", dtype=torch.bfloat16) + state.mtp_draft_input_hiddens = new_hidden + state.hidden_collector = FinalHiddenCollector() + actual = draft._token_forward(state) + full_tokens = torch.tensor([7, 11, 13, 29, 17, 19, 23], device="cuda") + full_hidden = torch.cat([hidden[:3], new_hidden[1:2], hidden[3:], new_hidden[:1]]) + positions = torch.tensor([0, 1, 2, 3, 0, 1, 2], device="cuda") + multi, logits = reference_draft( + full_tokens, + full_hidden, + source, + config, + draft._cos_cached[positions].repeat(1, 2), + draft._sin_cached[positions].repeat(1, 2), + [4, 3], + ) + torch.testing.assert_close(actual.logits, logits[[6, 3]], atol=0.025, rtol=0.025) + torch.testing.assert_close(actual.mtp_collector.spec_hidden, multi[[6, 3]], atol=0.016, rtol=0.016) diff --git a/unit_tests/models/qwen4_exp/test_multimodal.py b/unit_tests/models/qwen4_exp/test_multimodal.py new file mode 100644 index 0000000000..6cb7bc2b10 --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_multimodal.py @@ -0,0 +1,213 @@ +from copy import deepcopy +import json +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F + +from lightllm.models.qwen4_exp.configuration import normalize_model_config +from lightllm.models.qwen4_exp.multimodal import canonical_image_ids +from lightllm.server.router.model_infer.infer_batch import g_infer_context +from lightllm.common.basemodel.hidden_collector import NoopHiddenCollector +from .runtime import assemble_model, state_for +from .reference import reference_model + + +def test_model_types_resolve_through_public_registry(): + from lightllm.models.registry import ModelRegistry + from lightllm.models.qwen4_exp.model import Qwen4ExpTpPartModel + + for model_type in ("qwen4_exp", "qwen4_exp_text"): + assert ModelRegistry.get_model_class({"model_type": model_type}) is Qwen4ExpTpPartModel + + +@pytest.mark.parametrize("model_type", ["qwen4_exp", "qwen4_exp_text"]) +def test_startup_detects_hybrid_cache_and_optional_vision(model_type, tmp_path): + from lightllm.utils.config_utils import is_linear_att_mixed_model, has_vision_module, auto_set_fused_shared_experts + + config = {"model_type": model_type, "architectures": ["Qwen4ExpForConditionalGeneration"]} + if model_type == "qwen4_exp": + config["vision_config"] = {"hidden_size": 32} + (tmp_path / "config.json").write_text(json.dumps(config)) + assert is_linear_att_mixed_model(str(tmp_path)) + assert has_vision_module(str(tmp_path)) == (model_type == "qwen4_exp") + args = SimpleNamespace(model_dir=str(tmp_path), enable_ep_moe=False, enable_fused_shared_experts=False) + auto_set_fused_shared_experts(args) + assert args.enable_fused_shared_experts + + +def test_tokenizer_uses_native_processor_and_virtual_image_ranges(monkeypatch): + from lightllm.server import tokenizer as tokenizer_module + from lightllm.models.qwen4_exp.image_processing import Qwen4ExpImageProcessor + from lightllm.models.qwen4_exp.tokenizer import Qwen4ExpTokenizer + + config = dict( + model_type="qwen4_exp", + architectures=["Qwen4ExpForConditionalGeneration"], + vision_config={}, + vision_start_token_id=61, + vision_end_token_id=62, + image_token_id=60, + ) + monkeypatch.setattr(tokenizer_module.AutoTokenizer, "from_pretrained", lambda *args, **kwargs: object()) + monkeypatch.setattr(tokenizer_module.PretrainedConfig, "get_config_dict", lambda *args: (config, {})) + processor = SimpleNamespace(min_pixels=16, max_pixels=64, patch_size=2, merge_size=2) + monkeypatch.setattr(Qwen4ExpImageProcessor, "from_pretrained", lambda *args: processor) + tokenizer = tokenizer_module.get_tokenizer("checkpoint") + assert type(tokenizer) is Qwen4ExpTokenizer + image = SimpleNamespace(token_id=1000000, token_num=2, start_idx=None) + actual = tokenizer.encode([7, 61, 60, 62, 11], SimpleNamespace(images=[image])) + assert actual == [7, 61, 1000000, 1000001, 62, 11] + assert image.start_idx == 2 + assert tokenizer.image_processor is processor + + +@pytest.mark.parametrize("image_count", [0, 1, 2]) +@pytest.mark.parametrize("qwen4", [False, True]) +def test_token_count_hook_matches_qwen4_expansion_and_preserves_legacy(image_count, qwen4): + from lightllm.models.qwen2_vl.model import QWen2VLTokenizer + from lightllm.models.qwen4_exp.tokenizer import Qwen4ExpTokenizer + from lightllm.server.httpserver.manager import HttpServerManager + + cls = Qwen4ExpTokenizer if qwen4 else QWen2VLTokenizer + processor = SimpleNamespace(min_pixels=16, max_pixels=64, patch_size=2, merge_size=2) + tokenizer = cls( + tokenizer=object(), + image_processor=processor, + model_cfg=dict(vision_start_token_id=61, vision_end_token_id=62, image_token_id=60), + ) + images = [] + for index in range(image_count): + image = SimpleNamespace(image_w=4 * (index + 1), image_h=4, token_id=1000 + index * 100) + image.token_num = tokenizer.get_image_token_length(image) + images.append(image) + multimodal = SimpleNamespace(images=images, audios=[]) + prompt = [7] + [61, 60, 62] * image_count + [11] + expanded = tokenizer.encode(prompt, multimodal) + manager = HttpServerManager.__new__(HttpServerManager) + manager.tokenizer = tokenizer + checked = [] + manager._assert_image_token_count = checked.append + actual = manager.tokens(prompt, multimodal, SimpleNamespace()) + assert checked == [image.token_num for image in images] + assert actual == len(expanded) + + +def test_plain_text_counting_fallback_is_unchanged(): + from lightllm.server.httpserver.manager import HttpServerManager + + manager = HttpServerManager.__new__(HttpServerManager) + manager.tokenizer = SimpleNamespace(encode=lambda prompt, *_args, **_kwargs: prompt) + assert manager.tokens([7, 11], SimpleNamespace(images=[], audios=[]), SimpleNamespace()) == 2 + + +def test_token_count_hook_cannot_bypass_image_limits(): + from lightllm.server.httpserver.manager import HttpServerManager + from lightllm.models.qwen4_exp.tokenizer import Qwen4ExpTokenizer + + tokenizer = Qwen4ExpTokenizer( + tokenizer=object(), + image_processor=SimpleNamespace(min_pixels=16, max_pixels=64, patch_size=2, merge_size=2), + model_cfg=dict(vision_start_token_id=61, vision_end_token_id=62, image_token_id=60), + ) + manager = HttpServerManager.__new__(HttpServerManager) + manager.tokenizer = tokenizer + + def reject(_count): + raise ValueError("image token limit") + + manager._assert_image_token_count = reject + with pytest.raises(ValueError, match="image token limit"): + manager.tokens( + [61, 60, 62], + SimpleNamespace(images=[SimpleNamespace(image_w=4, image_h=4)], audios=[]), + SimpleNamespace(), + ) + + +def test_image_processor_loads_real_config_without_runtime_enum_serialization(tmp_path): + from lightllm.models.qwen4_exp.image_processing import Qwen4ExpImageProcessor + + (tmp_path / "preprocessor_config.json").write_text( + json.dumps( + dict( + size={"shortest_edge": 16, "longest_edge": 64}, + patch_size=2, + merge_size=2, + image_mean=[0.5] * 3, + image_std=[0.5] * 3, + image_processor_type="Qwen2VLImageProcessorFast", + ) + ) + ) + processor = Qwen4ExpImageProcessor.from_pretrained(str(tmp_path)) + assert processor.min_pixels == 16 and processor.max_pixels == 64 + assert processor.patch_size == processor.merge_size == 2 + assert processor.image_mean == processor.image_std == [0.5] * 3 + + +def test_normalize_outer_and_text_config_preserves_checkpoint(config): + text = dict(config, rope_parameters={"rope_theta": 1000, "partial_rotary_factor": 0.5}) + text.pop("n_layer") + outer = {"text_config": text, "vision_config": {"hidden_size": 32}, "image_token_id": 60} + original = deepcopy(outer) + actual, vision = normalize_model_config(outer) + standalone, no_vision = normalize_model_config(text) + assert actual["n_layer"] == standalone["n_layer"] == 4 + assert actual["image_token_id"] == 60 + assert actual["rope_theta"] == standalone["rope_theta"] == 1000 + assert vision == outer["vision_config"] and no_vision is None + assert outer == original + customized, _ = normalize_model_config(text, SimpleNamespace(vocab_size=128)) + assert customized["vocab_size"] == 128 and text["vocab_size"] == 64 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@pytest.mark.parametrize("dtype", [torch.int32, torch.int64]) +def test_image_ranges_preserve_text_and_physical_ids(dtype): + second = 2 ** 34 if dtype == torch.int64 else 2000000 + tokens = torch.tensor( + [7, 60, 999999, 1000000, 1000001, 1000002, second + 4, second + 5, 63], device="cuda", dtype=dtype + ) + starts = torch.tensor([1000000, second], device="cuda", dtype=torch.int64) + lengths = torch.tensor([2, 5], device="cuda", dtype=torch.int64) + before = tokens.clone() + actual = canonical_image_ids(tokens, starts, lengths, 60) + expected = torch.tensor([7, 60, 999999, 60, 60, 1000002, 60, second + 5, 63], device="cuda", dtype=dtype) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + torch.testing.assert_close(tokens, before, atol=0, rtol=0) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = canonical_image_ids(tokens, starts, lengths, 60) + tokens[3:5] = 11 + graph.replay() + expected[3:5] = 11 + torch.testing.assert_close(captured, expected, atol=0, rtol=0) + graph.reset() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_image_cache_ids_are_canonicalized_before_ple_commit(runtime, config, monkeypatch): + config["image_token_id"] = 60 + model, source = assemble_model(runtime, config) + runtime.req.init_linear_att_state(SimpleNamespace(req_idx=1)) + image_cache = torch.randn(4, 1, 32, dtype=torch.bfloat16).pin_memory() + monkeypatch.setattr(g_infer_context, "cpu_embed_cache_client", SimpleNamespace(cpu_embed_cache_tensor=image_cache)) + virtual = [7, 1000000, 1000001, 11] + canonical = torch.tensor([7, 60, 60, 11], device="cuda", dtype=torch.int64) + state = state_for(runtime, [1], [4], virtual, ready=[0]) + state.hidden_collector = NoopHiddenCollector() + state.multimodal_params = [ + {"images": [{"token_id": 1000000, "token_num": 2, "start_index_in_embed_cache": 1}], "audios": []} + ] + actual = model._context_forward(state) + torch.testing.assert_close(state.input_ids, torch.tensor(virtual, device="cuda"), atol=0, rtol=0) + torch.testing.assert_close(state.qwen4_exp_ple[1].input_tokens, canonical, atol=0, rtol=0) + torch.testing.assert_close(runtime.req.lexical_caches[1].ngram[1], canonical[-2:], atol=0, rtol=0) + embeddings = F.embedding(canonical.cpu(), source["model.language_model.embed_tokens.weight"]) + embeddings[1:3] = image_cache[1:3, 0] + positions = torch.arange(4, device="cuda") + cos, sin = model._cos_cached[positions].repeat(1, 2), model._sin_cached[positions].repeat(1, 2) + expected = reference_model(canonical, cos, sin, [4], source, config, input_embeddings=embeddings.cuda())[-1:] + torch.testing.assert_close(actual.logits, expected, atol=0.025, rtol=0.025) diff --git a/unit_tests/models/qwen4_exp/test_norm.py b/unit_tests/models/qwen4_exp/test_norm.py new file mode 100644 index 0000000000..50ccc871c0 --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_norm.py @@ -0,0 +1,48 @@ +import pytest +import torch + +from lightllm.models.qwen4_exp.layer_weights.norm_weight import Qwen4ExpGatedNormWeight + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("activation", ["sigmoid", "silu"]) +def test_native_gdn_norm_rounding(dtype, activation): + torch.manual_seed(837) + native = Qwen4ExpGatedNormWeight(128, "linear.norm.weight", dtype, activation) + weight = 1 + torch.randn(128, dtype=dtype) * 0.1 + native.load_hf_weights({"linear.norm.weight": weight}) + values = torch.randn(19, 128, device="cuda", dtype=dtype) + gate = torch.randn_like(values) * 3 + normalized = values.float() * torch.rsqrt(values.float().square().mean(-1, keepdim=True) + 1e-6) + weighted = native.weight * normalized.to(dtype) + gate_activation = gate.float().sigmoid() + if activation == "silu": + gate_activation *= gate.float() + expected = (weighted * gate_activation).to(dtype) + out = torch.empty_like(values) + actual = native(values, gate, 1e-6, out=out) + assert actual is out + tolerance = 2e-6 if dtype == torch.float32 else torch.finfo(dtype).eps * 2 + torch.testing.assert_close(actual, expected, atol=tolerance, rtol=tolerance) + + +@pytest.mark.parametrize("activation", ["sigmoid", "silu"]) +def test_gdn_norm_accepts_strided_gate(activation): + from lightllm.models.qwen4_exp.triton_kernel.gated_norm import gated_norm + + torch.manual_seed(123) + tokens, heads, dim = 5, 4, 128 + values = torch.randn(tokens * heads, dim, device="cuda", dtype=torch.bfloat16) + packed = torch.randn(tokens, heads * dim + 256, device="cuda", dtype=torch.bfloat16) + gate = packed[:, 128 : 128 + heads * dim].view(tokens, heads, dim) + weight = torch.randn(dim, device="cuda", dtype=torch.bfloat16) + + actual = gated_norm(values, gate, weight, 1e-6, activation) + normalized = values.float() * torch.rsqrt(values.float().square().mean(-1, keepdim=True) + 1e-6) + expected = normalized * weight.float() + gate = gate.reshape(-1, dim).float() + gate = gate.sigmoid() if activation == "sigmoid" else torch.nn.functional.silu(gate) + expected = (expected * gate).to(values.dtype) + torch.testing.assert_close(actual, expected, atol=0.008, rtol=0.008) diff --git a/unit_tests/models/qwen4_exp/test_pinned_memory.py b/unit_tests/models/qwen4_exp/test_pinned_memory.py new file mode 100644 index 0000000000..cffc72cf36 --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_pinned_memory.py @@ -0,0 +1,112 @@ +import gc +import weakref + +import pytest +import torch +import triton +import triton.language as tl + +from lightllm.models.qwen4_exp import pinned_memory + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +@triton.jit +def read_ends(Source, Output, SIZE: tl.constexpr): + row = tl.arange(0, 2) + position = tl.where(row == 0, 0, SIZE - 1) + tl.store(Output + row, tl.load(Source + position)) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +def test_registered_snapshot_views_graph_and_lifetime(dtype, monkeypatch): + monkeypatch.setattr(pinned_memory, "_LARGE_BYTES", 1) + monkeypatch.setattr(pinned_memory, "_CHUNK_BYTES", 4096) + value = pinned_memory.pinned_zeros((6144,), dtype=dtype) + assert value.is_pinned() and value.is_contiguous() + assert torch.count_nonzero(value) == 0 + value[0], value[-1] = 3, 7 + output = torch.empty(2, device="cuda", dtype=torch.float32) + read_ends[(1,)](value, output, value.numel()) + torch.testing.assert_close(output, torch.tensor([3, 7], device="cuda", dtype=torch.float32), atol=0, rtol=0) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + read_ends[(1,)](value, output, value.numel()) + value[-1] = -7 + graph.replay() + torch.testing.assert_close(output, torch.tensor([3, -7], device="cuda", dtype=torch.float32), atol=0, rtol=0) + graph.reset() + owner = weakref.ref(value._qwen4_pinned_allocation) + view = value[4096:] + del value + gc.collect() + assert owner() is not None and view.is_pinned() + read_ends[(1,)](view, output, view.numel()) + torch.testing.assert_close(output, torch.tensor([0, -7], device="cuda", dtype=torch.float32), atol=0, rtol=0) + del view + gc.collect() + assert owner() is None + + +def test_partial_registration_failure_releases_registered_ranges(monkeypatch): + monkeypatch.setattr(pinned_memory, "_LARGE_BYTES", 1) + monkeypatch.setattr(pinned_memory, "_CHUNK_BYTES", 4096) + runtime = torch.cuda.cudart() + registered, released = [], [] + + class FailingRuntime: + def cudaHostRegister(self, address, size, flags): + if registered: + return 1 + status = runtime.cudaHostRegister(address, size, flags) + assert int(status) == 0 + registered.append(address) + return status + + def cudaHostUnregister(self, address): + released.append(address) + return runtime.cudaHostUnregister(address) + + def cudaGetErrorString(self, status): + return "intentional registration failure" + + monkeypatch.setattr(torch.cuda, "cudart", lambda: FailingRuntime()) + + def allocate(): + try: + pinned_memory.pinned_zeros((6144,), dtype=torch.float32) + except RuntimeError as error: + assert "intentional registration failure" in str(error) + else: + pytest.fail("Expected registration failure") + + allocate() + gc.collect() + assert registered and released == registered + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +def test_snapshot_dma_restore_never_crosses_registration(dtype, monkeypatch): + monkeypatch.setattr(pinned_memory, "_LARGE_BYTES", 1) + monkeypatch.setattr(pinned_memory, "_CHUNK_BYTES", 16384) + value = pinned_memory.pinned_zeros((7, 3, 512), dtype=dtype) + cell_bytes = value[0].numel() * value.element_size() + starts = value._qwen4_pinned_allocation.addresses + assert len(starts) > 1 + assert all((address - value.data_ptr()) % cell_bytes == 0 for address in starts) + ends = starts[1:] + [value.data_ptr() + value.numel() * value.element_size()] + assert all(end - start <= 16384 for start, end in zip(starts, ends)) + destination = torch.empty((3, 5, 512), dtype=dtype, device="cuda") + for index in range(value.shape[0]): + value[index].fill_(index + 1) + + destination[:, 1, ...] = value[index] + torch.cuda.synchronize() + torch.testing.assert_close(destination[:, 1].cpu(), value[index], atol=0, rtol=0) + + +def test_snapshot_larger_than_registration_bound_is_rejected(monkeypatch): + monkeypatch.setattr(pinned_memory, "_LARGE_BYTES", 1) + monkeypatch.setattr(pinned_memory, "_CHUNK_BYTES", 16384) + with pytest.raises(ValueError, match="one bounded host registration"): + pinned_memory.pinned_zeros((2, 8192), dtype=torch.float32) diff --git a/unit_tests/models/qwen4_exp/test_reduction.py b/unit_tests/models/qwen4_exp/test_reduction.py new file mode 100644 index 0000000000..2eaf564f06 --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_reduction.py @@ -0,0 +1,104 @@ +import datetime +import json +from types import SimpleNamespace + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + + +def _check_reductions(rank, rendezvous, cuda=False): + import os + + from lightllm.models.qwen4_exp.layer_infer.transformer_layer_infer import Qwen4ExpTransformerLayerInfer + from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.utils.envs_utils import get_env_start_args + + args = StartArgs() + args.enable_tpsp_mix_mode = False + os.environ["LIGHTLLM_START_ARGS"] = json.dumps(vars(args)) + get_env_start_args.cache_clear() + device = torch.device("cuda", rank) if cuda else torch.device("cpu") + if cuda: + torch.cuda.set_device(rank) + dist.init_process_group( + "nccl" if cuda else "gloo", + init_method="file://" + rendezvous, + rank=rank, + world_size=4, + timeout=datetime.timedelta(seconds=60), + ) + try: + layer = Qwen4ExpTransformerLayerInfer.__new__(Qwen4ExpTransformerLayerInfer) + layer.tp_world_size_ = 4 + layer.tp_rank_ = rank + state = SimpleNamespace(dist_group=dist.group.WORLD) + + contribution = (256.0, 1.0, -256.0, 0.25)[rank] + for dtype in (torch.bfloat16, torch.float16, torch.float32): + for rows in (1, 3, 8): + value = torch.full((rows, 16), contribution, dtype=dtype, device=device) + value[:, 1::2].neg_() + expected = torch.full_like(value, 1.25) + expected[:, 1::2].neg_() + for reduce in (layer._reduce_linear_output, layer._tpsp_reduce): + actual = reduce(value.clone(), state) + assert actual.dtype == dtype + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + midpoint_parts = (1.0, 2 ** -8, 2 ** -24, 2 ** -24) + for rows in (1, 3, 256, 512): + value = torch.full((rows, 16), midpoint_parts[rank], dtype=torch.bfloat16, device=device) + expected = torch.full_like(value, 1.0078125) + for reduce in (layer._reduce_linear_output, layer._tpsp_reduce): + actual = reduce(value.clone(), state) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + value = torch.full((16, 3), contribution, dtype=torch.bfloat16, device=device).t() + assert layer._reduce_linear_output(value, state) is value + torch.testing.assert_close(value, torch.full_like(value, 1.25), atol=0, rtol=0) + + args.enable_tpsp_mix_mode = True + os.environ["LIGHTLLM_START_ARGS"] = json.dumps(vars(args)) + get_env_start_args.cache_clear() + for rows in (4, 12): + common = torch.arange(rows * 16, device=device, dtype=torch.float32).reshape(rows, 16) / 16 + value = (common + rank / 4).to(torch.bfloat16) + expected = (common * 4 + 1.5).to(value.dtype).chunk(4)[rank] + actual = layer._tpsp_reduce(value, state) + assert actual.shape == (rows // 4, 16) and actual.dtype == value.dtype + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + if cuda: + for _ in range(3): + layer._tpsp_reduce(value, state) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + result = layer._tpsp_reduce(value, state) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(result, expected, atol=0, rtol=0) + graph.reset() + finally: + dist.destroy_process_group() + + +@pytest.mark.skipif(not dist.is_available() or not dist.is_gloo_available(), reason="Gloo is required") +def test_tp_and_sequence_parallel_accumulation(tmp_path): + mp.spawn(_check_reductions, args=(str(tmp_path / "rendezvous"),), nprocs=4, join=True) + + +@pytest.mark.skipif(torch.cuda.device_count() < 4 or not dist.is_nccl_available(), reason="Four CUDA GPUs are required") +def test_cuda_tp_and_sequence_parallel_accumulation(tmp_path): + mp.spawn(_check_reductions, args=(str(tmp_path / "nccl-rendezvous"), True), nprocs=4, join=True) + + +def test_single_rank_preserves_tensor(): + from lightllm.models.qwen4_exp.layer_infer.transformer_layer_infer import Qwen4ExpTransformerLayerInfer + + layer = Qwen4ExpTransformerLayerInfer.__new__(Qwen4ExpTransformerLayerInfer) + layer.tp_world_size_ = 1 + value = torch.tensor([[1.25]], dtype=torch.bfloat16) + assert layer._reduce_linear_output(value, None) is value + assert layer._tpsp_reduce(value, None) is value diff --git a/unit_tests/models/qwen4_exp/test_runtime.py b/unit_tests/models/qwen4_exp/test_runtime.py new file mode 100644 index 0000000000..25fe2c0f79 --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_runtime.py @@ -0,0 +1,228 @@ +import json +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.models.qwen4_exp.layer_weights.transformer_layer_weight import Qwen4ExpTransformerLayerWeight +from lightllm.models.qwen4_exp.model import Qwen4ExpTpPartModel +from lightllm.models.qwen4_exp.layer_infer.transformer_layer_infer import Qwen4ExpTransformerLayerInfer +from lightllm.common.quantization import Quantcfg +from lightllm.models.qwen3next.model import Qwen3NextTpPartModel +from .helpers import checkpoint +from .reference import reference_step, reference_layer, reference_model +from .runtime import state_for, assemble_model +from lightllm.common.basemodel.hidden_collector import NoopHiddenCollector + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +@pytest.mark.parametrize("runtime", [{}, {"ssm_state_mode": "replay"}], ids=["native", "replay"], indirect=True) +def test_inference_state_prefill_decode_and_automatic_commit(runtime, config): + Qwen3NextTpPartModel._init_triton(None) + req, weight = runtime.req, runtime.weight + for request in (1, 3): + req.init_linear_att_state(SimpleNamespace(req_idx=request)) + state = state_for(runtime, [1, 3], [3, 2], [7, 63, 19, 11, 23], ready=[0, 0]) + assert state.qwen4_exp_metadata.starts.tolist() == [0, 3, 5, 5, 5, 5] + assert state.qwen4_exp_batch.slot_table.data_ptr() == req.req_to_token_indexs.data_ptr() + torch.testing.assert_close( + runtime.mem.qsa_key_cos[state.mem_index.long()], state.qwen4_exp_batch.cos, atol=0, rtol=0 + ) + infer = Qwen4ExpTransformerLayerInfer(1, config) + residual = torch.randn(5, 64, device="cuda", dtype=torch.bfloat16) + reference = {} + expected = [] + for request, start, end in ((1, 0, 3), (3, 3, 5)): + ngram = weight.ngram_hasher.component.empty_history(1) + ple = weight.ple.component.empty_history(1) + conv = torch.zeros(1, 512, 3, device="cuda", dtype=torch.bfloat16) + ssm = torch.zeros(1, 2, 128, 128, device="cuda", dtype=torch.float32) + for row in range(start, end): + output, ngram, ple, conv, ssm = reference_step( + residual[row : row + 1], state.input_ids[row : row + 1], runtime.source, config, ngram, ple, conv, ssm + ) + expected.append(output) + reference[request] = (ngram, ple, conv, ssm) + actual = infer.context_forward(residual, state, weight) + torch.testing.assert_close(actual, torch.cat(expected), atol=0.016, rtol=0.016) + + state = state_for(runtime, [1, 3], [4, 3], [29, 31]) + residual = torch.randn(2, 64, device="cuda", dtype=torch.bfloat16) + actual = infer.token_forward(residual, state, weight) + for row, request in enumerate((1, 3)): + expected, ngram, ple, _, _ = reference_step( + residual[row : row + 1], state.input_ids[row : row + 1], runtime.source, config, *reference[request] + ) + torch.testing.assert_close(actual[row : row + 1], expected, atol=0.016, rtol=0.016) + torch.testing.assert_close(req.lexical_caches[1].ngram[request], ngram[0], atol=0, rtol=0) + torch.testing.assert_close(req.lexical_caches[1].convolution[request], ple[0], atol=0.016, rtol=0.016) + + +def test_native_prefill_graph_state_copy_changes_request_count(runtime, config): + source = checkpoint(config, 3) + weight = Qwen4ExpTransformerLayerWeight(3, torch.bfloat16, config, Quantcfg(config)) + weight.load_hf_weights(dict(source)) + weight.verify_load() + infer = Qwen4ExpTransformerLayerInfer(3, config) + old = state_for(runtime, [1], [4], [7, 11, 13, 17], ready=[0]) + residual = torch.randn(4, 64, device="cuda", dtype=torch.bfloat16) + for _ in range(3): + infer.context_forward(residual, old, weight) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = infer.context_forward(residual, old, weight) + pointers = [ + getattr(old.qwen4_exp_metadata, name).data_ptr() for name in old.qwen4_exp_metadata.__dataclass_fields__ + ] + before = runtime.mem.get_index_kv_buffer(old.mem_index) + new = state_for(runtime, [2, 3], [2, 2], [19, 23, 29, 31], ready=[0, 0]) + old.copy_for_prefill_cuda_graph(new) + residual.normal_() + graph.replay() + expected = reference_layer(residual, source, config, new.qwen4_exp_batch.cos, new.qwen4_exp_batch.sin, [2, 2]) + torch.testing.assert_close(output, expected, atol=0.016, rtol=0.016) + assert pointers == [ + getattr(old.qwen4_exp_metadata, name).data_ptr() for name in old.qwen4_exp_metadata.__dataclass_fields__ + ] + for name, values in runtime.mem.get_index_kv_buffer(torch.arange(4, device="cuda")).items(): + torch.testing.assert_close(values, before[name], atol=0, rtol=0, equal_nan=True) + + +def test_native_config_normalization_and_periodic_layout(runtime, config, tmp_path): + outer = dict(text_config=dict(config), vision_config={"hidden_size": 16}) + outer["text_config"]["rope_parameters"] = dict( + rope_type="default", rope_theta=10000, partial_rotary_factor=0.5, mrope_section=[1, 1, 0] + ) + (tmp_path / "config.json").write_text(json.dumps(outer)) + model = Qwen4ExpTpPartModel.__new__(Qwen4ExpTpPartModel) + model.weight_dir_, model.finetune_config, model.tp_world_size_ = str(tmp_path), None, 1 + model._init_config() + assert model.config["n_layer"] == 4 and model.config["n_embed"] == 32 + assert model.config["rope_theta"] == 10000 + assert model.num_kv_heads == 2 + invalid = dict(config, layer_types=["full_attention"] * 4) + with pytest.raises(ValueError, match="periodic"): + Qwen4ExpTpPartModel.normalize_model_config(invalid) + + +def test_native_decode_graph_replay_updates_positions_slots_and_request_order(runtime, config): + source = checkpoint(config, 3) + weight = Qwen4ExpTransformerLayerWeight(3, torch.bfloat16, config, Quantcfg(config)) + weight.load_hf_weights(dict(source)) + weight.verify_load() + infer = Qwen4ExpTransformerLayerInfer(3, config) + prefix = state_for(runtime, [1, 3], [4, 4], [7, 11, 13, 17, 19, 23, 29, 31], ready=[0, 0]) + prefix_residual = torch.randn(8, 64, device="cuda", dtype=torch.bfloat16) + infer.context_forward(prefix_residual, prefix, weight) + old = state_for(runtime, [1, 3], [5, 5], [37, 41]) + residual = torch.randn(2, 64, device="cuda", dtype=torch.bfloat16) + for _ in range(3): + infer.token_forward(residual, old, weight) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = infer.token_forward(residual, old, weight) + previous_residual = residual.clone() + previous_cos, previous_sin = old.qwen4_exp_batch.cos.clone(), old.qwen4_exp_batch.sin.clone() + new = state_for(runtime, [3, 1], [6, 6], [43, 47]) + old.copy_for_cuda_graph(new) + residual.normal_() + graph.replay() + all_residual, all_cos, all_sin = [], [], [] + for row, (start, end) in enumerate(((0, 4), (4, 8))): + incoming = 1 - row + all_residual.extend( + [prefix_residual[start:end], previous_residual[row : row + 1], residual[incoming : incoming + 1]] + ) + all_cos.extend( + [ + prefix.qwen4_exp_batch.cos[start:end], + previous_cos[row : row + 1], + new.qwen4_exp_batch.cos[incoming : incoming + 1], + ] + ) + all_sin.extend( + [ + prefix.qwen4_exp_batch.sin[start:end], + previous_sin[row : row + 1], + new.qwen4_exp_batch.sin[incoming : incoming + 1], + ] + ) + expected = reference_layer(torch.cat(all_residual), source, config, torch.cat(all_cos), torch.cat(all_sin), [6, 6]) + torch.testing.assert_close(output, expected[[11, 5]], atol=0.016, rtol=0.016) + + +@pytest.mark.parametrize("runtime", [{}, {"ssm_state_mode": "replay"}], ids=["native", "replay"], indirect=True) +def test_complete_small_model_prefill_and_cached_logits(runtime, config): + Qwen3NextTpPartModel._init_triton(None) + model, source = assemble_model(runtime, config) + for request in (1, 3): + runtime.req.init_linear_att_state(SimpleNamespace(req_idx=request)) + state = state_for(runtime, [1, 3], [3, 2], [7, 63, 19, 11, 23], ready=[0, 0]) + state.hidden_collector = NoopHiddenCollector() + output = model._context_forward(state) + expected = reference_model( + state.input_ids, state.qwen4_exp_batch.cos, state.qwen4_exp_batch.sin, [3, 2], source, config + ) + torch.testing.assert_close(output.logits, expected[[2, 4]], atol=0.025, rtol=0.025) + decode = state_for(runtime, [1, 3], [4, 3], [29, 31]) + decode.hidden_collector = NoopHiddenCollector() + output = model._token_forward(decode) + all_tokens = torch.tensor([7, 63, 19, 29, 11, 23, 31], device="cuda", dtype=torch.int64) + all_cos = torch.cat( + [ + state.qwen4_exp_batch.cos[:3], + decode.qwen4_exp_batch.cos[:1], + state.qwen4_exp_batch.cos[3:], + decode.qwen4_exp_batch.cos[1:], + ] + ) + all_sin = torch.cat( + [ + state.qwen4_exp_batch.sin[:3], + decode.qwen4_exp_batch.sin[:1], + state.qwen4_exp_batch.sin[3:], + decode.qwen4_exp_batch.sin[1:], + ] + ) + expected = reference_model(all_tokens, all_cos, all_sin, [4, 3], source, config) + torch.testing.assert_close(output.logits, expected[[3, 6]], atol=0.025, rtol=0.025) + + third = state_for(runtime, [1, 3], [5, 4], [37, 41]) + third.hidden_collector = NoopHiddenCollector() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = model._token_forward(third) + graph.replay() + third_cos, third_sin = third.qwen4_exp_batch.cos.clone(), third.qwen4_exp_batch.sin.clone() + all_tokens = torch.tensor([7, 63, 19, 29, 37, 11, 23, 31, 41], device="cuda", dtype=torch.int64) + all_cos = torch.cat([all_cos[:4], third_cos[:1], all_cos[4:], third_cos[1:]]) + all_sin = torch.cat([all_sin[:4], third_sin[:1], all_sin[4:], third_sin[1:]]) + expected = reference_model(all_tokens, all_cos, all_sin, [5, 4], source, config) + torch.testing.assert_close(captured.logits, expected[[4, 8]], atol=0.025, rtol=0.025) + fourth = state_for(runtime, [3, 1], [5, 6], [47, 43]) + fourth.hidden_collector = NoopHiddenCollector() + third.copy_for_cuda_graph(fourth) + graph.replay() + all_tokens = torch.tensor([7, 63, 19, 29, 37, 43, 11, 23, 31, 41, 47], device="cuda", dtype=torch.int64) + all_cos = torch.cat([all_cos[:5], fourth.qwen4_exp_batch.cos[1:], all_cos[5:], fourth.qwen4_exp_batch.cos[:1]]) + all_sin = torch.cat([all_sin[:5], fourth.qwen4_exp_batch.sin[1:], all_sin[5:], fourth.qwen4_exp_batch.sin[:1]]) + expected = reference_model(all_tokens, all_cos, all_sin, [6, 5], source, config) + torch.testing.assert_close(captured.logits, expected[[10, 5]], atol=0.025, rtol=0.025) + + +def test_fresh_prefill_ignores_stale_lexical_slot_without_mutating_accepted_state(runtime): + cache = runtime.req.lexical_caches[1] + cache.ngram[1].fill_(17) + cache.convolution[1].normal_() + accepted = cache.snapshot(1) + fresh = state_for(runtime, [1], [2], [7, 11], ready=[0]) + batch = fresh.qwen4_exp_ple[1] + assert torch.all(batch.ngram_history == 63) + assert torch.count_nonzero(batch.conv_history) == 0 + torch.testing.assert_close(cache.ngram[1], accepted.ngram, atol=0, rtol=0) + torch.testing.assert_close(cache.convolution[1], accepted.convolution, atol=0, rtol=0) + resumed = state_for(runtime, [1], [4], [13, 19], ready=[2]) + batch = resumed.qwen4_exp_ple[1] + torch.testing.assert_close(batch.ngram_history[0], accepted.ngram, atol=0, rtol=0) + torch.testing.assert_close(batch.conv_history[0], accepted.convolution, atol=0, rtol=0) diff --git a/unit_tests/models/qwen4_exp/test_sequence_parallel.py b/unit_tests/models/qwen4_exp/test_sequence_parallel.py new file mode 100644 index 0000000000..08a4aa8e67 --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_sequence_parallel.py @@ -0,0 +1,87 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.models.qwen3_5.layer_infer.transformer_layer_infer import Qwen35TransformerLayerInfer +from lightllm.models.qwen4_exp.layer_infer import transformer_layer_infer as decoder + + +def partitioned_layer(world, rank, enabled, full): + layer = decoder.Qwen4ExpTransformerLayerInfer.__new__(decoder.Qwen4ExpTransformerLayerInfer) + layer.tp_world_size_, layer.tp_rank_ = world, rank + layer.layer_num_, layer.is_linear_attention_layer = 1, True + layer._tpsp_allgather = lambda value, state: full if enabled else value + layer._tpsp_sp_split = lambda value, state: value.chunk(world)[rank].contiguous() if enabled else value + return layer + + +@pytest.mark.parametrize("world,rank", [(2, 0), (2, 1), (4, 0), (4, 3)]) +@pytest.mark.parametrize("enabled", [False, True]) +@pytest.mark.parametrize("method", ["context_attention_forward", "token_attention_forward"]) +def test_gdn_uses_global_rows_and_returns_local_partition(monkeypatch, world, rank, enabled, method): + full = torch.arange(64, dtype=torch.float32).reshape(8, 8) + layer = partitioned_layer(world, rank, enabled, full) + monkeypatch.setattr(decoder, "get_env_start_args", lambda: SimpleNamespace(enable_tpsp_mix_mode=enabled)) + + def ordered(value, group, size, scatter_rank=None): + total = value * size + return total if scatter_rank is None else total.chunk(size)[scatter_rank].contiguous() + + def gdn(self, value, state, weight): + assert torch.equal(value, full) + return self._reduce_linear_output(value, state) + + monkeypatch.setattr(decoder, "ordered_tp_sum", ordered) + monkeypatch.setattr(Qwen35TransformerLayerInfer, method, gdn) + value = full.chunk(world)[rank].contiguous() if enabled else full + actual = getattr(layer, method)(value, SimpleNamespace(dist_group=None), None) + expected = (full * world).chunk(world)[rank] if enabled else full * world + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + +@pytest.mark.parametrize("world,rank", [(2, 0), (2, 1), (4, 0), (4, 3)]) +@pytest.mark.parametrize("enabled", [False, True]) +@pytest.mark.parametrize("prefill", [False, True]) +@pytest.mark.parametrize("commit", [None, False, True]) +def test_ple_keeps_global_metadata_and_local_residual(world, rank, enabled, prefill, commit): + full = torch.arange(64, dtype=torch.float32).reshape(8, 8) + lexical = full / 16 + layer = partitioned_layer(world, rank, enabled, full) + layer.context_attention_forward = layer.token_attention_forward = lambda value, *args: torch.zeros_like(value) + layer._ffn = lambda value, *args: torch.zeros_like(value) + tokens, metadata, batch = torch.arange(8), object(), SimpleNamespace() + calls = [] + + def apply(value, ids, weight, size, group): + assert torch.equal(value, full) and ids is tokens and size == world + return lexical + + def record(kind, observed_metadata, observed_batch): + assert observed_metadata is metadata and observed_batch is batch + calls.append(kind) + + batch.apply = apply + cache = SimpleNamespace(commit=lambda *args: record("commit", *args)) + state = SimpleNamespace( + qwen4_exp_ple={1: batch}, + input_ids=tokens, + qwen4_exp_metadata=metadata, + qwen4_exp_auto_commit=commit, + dist_group=None, + qwen4_exp_lexical_state=SimpleNamespace( + lexical_caches={1: cache}, + stage_speculative_lexical=lambda number, *args: record("stage", *args), + ), + ) + residual = SimpleNamespace(read=lambda value: (value, None), write=lambda value, update, gates: value + update) + weight = SimpleNamespace( + ple=True, + attn_hyper_connection=SimpleNamespace(component=residual), + mlp_hyper_connection=SimpleNamespace(component=residual), + ) + value = full.chunk(world)[rank].contiguous() if enabled else full + actual = layer._forward(value, state, weight, prefill) + expected = (full + lexical).chunk(world)[rank] if enabled else full + lexical + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + assert calls == ([] if commit is None else ["commit" if commit else "stage"]) diff --git a/unit_tests/models/qwen4_exp/test_speculative_cache.py b/unit_tests/models/qwen4_exp/test_speculative_cache.py new file mode 100644 index 0000000000..7b95cc4c84 --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_speculative_cache.py @@ -0,0 +1,159 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.models.qwen4_exp.lexical_state import Qwen4ExpLexicalState +from lightllm.models.qwen4_exp.speculative_cache import LexicalSpeculativeCache +from lightllm.server.router.model_infer.mtp_speculative.utils import verify_mtp_tokens +from .lexical_helpers import make_cache, metadata, reference_segment +from .reference import reference_step +from lightllm.models.qwen3next.model import Qwen3NextTpPartModel +from lightllm.models.qwen4_exp.layer_infer.transformer_layer_infer import Qwen4ExpTransformerLayerInfer + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +@pytest.mark.parametrize("capture", [False, True]) +def test_actual_verifier_commits_only_accepted_lexical_candidates(config, capture): + torch.manual_seed(1907) + cache, source = make_cache(config) + cache.ngram[2].copy_(torch.tensor([7, 11], device="cuda")) + cache.ngram[5].copy_(torch.tensor([13, 17], device="cuda")) + cache.convolution.normal_() + untouched = cache.snapshot(7) + speculative = LexicalSpeculativeCache(cache, 4) + req = SimpleNamespace() + req.speculative_lexical_caches = {1: speculative} + req.req_sampling_params_manager = SimpleNamespace( + req_to_next_token_ids=torch.full((8, 4), -1, device="cuda", dtype=torch.int64) + ) + req.req_to_mtp_state_index = torch.zeros(8, device="cuda", dtype=torch.int32) + req.HOLD_REQUEST_ID = 7 + lexical = Qwen4ExpLexicalState.__new__(Qwen4ExpLexicalState) + lexical.request_capacity = 8 + lexical.speculative_lexical_caches = {1: speculative} + + def commit(requests, accepted, role): + assert role == "main" + lexical.commit_bounded_att_state(requests, accepted) + req.req_to_mtp_state_index[requests] = accepted - 1 + + backend = SimpleNamespace( + model=SimpleNamespace(req_manager=req, att_state_manager=SimpleNamespace(commit_bounded_att_state=commit)), + is_linear_att_mixed_model=True, + max_draft_step=3, + ) + meta = metadata([5, 5, 5, 2, 2, 7], [11, 12, 13, 11, 12, 1], [1] * 6) + tokens = torch.empty(6, device="cuda", dtype=torch.int64) + residual = torch.empty(6, 64, device="cuda", dtype=torch.bfloat16) + + def run(): + batch = cache.prepare(meta) + result = batch.apply(residual, tokens, cache.layer_weight, 1, None) + speculative.stage(meta, batch) + return result + + tokens.fill_(19) + residual.normal_() + if capture: + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + run() + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + result = run() + + cases = ( + ([5, 5, 5, 2, 2, 7], [2, 1], [19, 63, 23, 29, 31, 37]), + ([2, 2, 5, 5, 5, 7], [2, 3], [41, 43, 47, 63, 5, 7]), + ([5, 5, 5, 5, 2, 7], [1, 1], [11, 13, 17, 19, 23, 29]), + ) + for requests, accepts, ids in cases: + starts = [0] + [i for i in range(1, len(requests)) if requests[i] != requests[i - 1]] + counts = [end - start for start, end in zip(starts, starts[1:] + [len(requests)])] + lengths = [] + for count in counts: + lengths.extend(range(11, 11 + count)) + meta.copy_from(metadata(requests, lengths, [1] * 6)) + tokens.copy_(torch.tensor(ids, device="cuda")) + residual.normal_() + before_ngram, before_conv = cache.ngram.clone(), cache.convolution.clone() + if capture: + graph.replay() + actual = result + else: + actual = run() + torch.testing.assert_close(cache.ngram, before_ngram, atol=0, rtol=0) + torch.testing.assert_close(cache.convolution, before_conv, atol=0, rtol=0) + + next_ids = torch.arange(101, 107, device="cuda", dtype=torch.int64) + proposals = req.req_sampling_params_manager.req_to_next_token_ids + proposals.fill_(-1) + expected_mask = [] + for start, count, accepted in zip(starts, counts, accepts + [1]): + request = requests[start] + proposals[request, 1:accepted] = next_ids[start : start + accepted - 1] + expected_mask.extend([1] * accepted + [0] * (count - accepted)) + accept_lengths, accepted_index = verify_mtp_tokens( + backend, + next_ids, + torch.tensor(requests, device="cuda", dtype=torch.int32), + torch.tensor(starts, device="cuda", dtype=torch.int32), + torch.tensor(sum((list(range(count)) for count in counts), []), device="cuda", dtype=torch.int32), + ) + assert accept_lengths.tolist() == accepts + [1] + assert accepted_index.tolist() == expected_mask + for start, count, accepted in zip(starts[:2], counts[:2], accepts): + request = requests[start] + expected, states = reference_segment( + residual[start : start + count], + tokens[start : start + count], + source, + config, + before_ngram[request : request + 1], + before_conv[request : request + 1], + ) + torch.testing.assert_close(actual[start : start + count], expected, atol=0.016, rtol=0.016) + ngram, convolution = states[accepted] + torch.testing.assert_close(cache.ngram[request], ngram[0], atol=0, rtol=0) + torch.testing.assert_close(cache.convolution[request], convolution[0], atol=0.016, rtol=0.016) + assert req.req_to_mtp_state_index[request] == accepted - 1 + torch.testing.assert_close(cache.ngram[7], untouched.ngram, atol=0, rtol=0) + torch.testing.assert_close(cache.convolution[7], untouched.convolution, atol=0, rtol=0) + if capture: + graph.reset() + + +def test_microbatch_candidates_survive_reverse_verification(config): + cache, source = make_cache(config) + speculative = LexicalSpeculativeCache(cache, 3) + batches = [] + for request, ids in ((5, [19, 63, 23]), (2, [29, 31, 37])): + meta = metadata([request] * 3, [11, 12, 13], [1] * 3) + tokens = torch.tensor(ids, device="cuda") + residual = torch.randn(3, 64, device="cuda", dtype=torch.bfloat16) + batch = cache.prepare(meta) + batch.apply(residual, tokens, cache.layer_weight, 1, None) + speculative.stage(meta, batch) + batches.append((request, residual, tokens)) + for request, residual, tokens in reversed(batches): + _, states = reference_segment( + residual, + tokens, + source, + config, + cache.ngram[request : request + 1].clone(), + cache.convolution[request : request + 1].clone(), + ) + speculative.commit( + torch.tensor([request] * 3, device="cuda", dtype=torch.int32), + torch.tensor([0], device="cuda", dtype=torch.int32), + torch.tensor([2], device="cuda", dtype=torch.int32), + 7, + ) + torch.testing.assert_close(cache.ngram[request], states[2][0][0], atol=0, rtol=0) + torch.testing.assert_close(cache.convolution[request], states[2][1][0], atol=0.016, rtol=0.016) diff --git a/unit_tests/models/qwen4_exp/test_vision.py b/unit_tests/models/qwen4_exp/test_vision.py new file mode 100644 index 0000000000..ccf34a0521 --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_vision.py @@ -0,0 +1,49 @@ +import pytest +import torch + +from lightllm.models.qwen4_exp.vision import Qwen4ExpVisionTransformer +from lightllm.server import visualserver +from lightllm.common.basemodel.attention_vit.create_utils import get_vit_att_backend_class +from .vision_reference import interpolate_positions, vision_reference + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +def test_vision_encoder_matches_independent_packed_frames(dtype, monkeypatch): + monkeypatch.setattr(visualserver, "VIT_ATTN_BACKEND", get_vit_att_backend_class("sdpa")) + config = dict( + depth=2, + hidden_size=32, + out_hidden_size=16, + intermediate_size=48, + hidden_act="gelu_pytorch_tanh", + num_heads=4, + in_channels=3, + patch_size=2, + temporal_patch_size=2, + spatial_merge_size=2, + num_position_embeddings=16, + deepstack_visual_indexes=[], + ) + torch.manual_seed(1907) + model = Qwen4ExpVisionTransformer({"data_type": dtype}, **config).eval().to(device="cuda", dtype=dtype) + grid = torch.tensor([[1, 4, 6], [2, 2, 4]], device="cuda") + pixels = torch.randn(40, 24, device="cuda", dtype=dtype) + source = {name: value.detach().clone() for name, value in model.state_dict().items()} + tolerance = 0.016 if dtype == torch.bfloat16 else 2e-5 + with torch.no_grad(): + actual_positions = model.fast_pos_embed_interpolate(grid) + expected_positions = interpolate_positions(source["pos_embed.weight"], grid, 2) + torch.testing.assert_close(actual_positions, expected_positions, atol=tolerance, rtol=tolerance) + actual, deepstack = model(pixels, grid) + expected = vision_reference(pixels, grid, source, config) + torch.testing.assert_close(actual_positions, expected_positions, atol=tolerance, rtol=tolerance) + torch.testing.assert_close(actual, expected, atol=tolerance, rtol=tolerance) + assert not deepstack + assert actual.shape == (10, 16) + + +def test_vision_rejects_unsupported_deepstack(): + with pytest.raises(ValueError, match="deepstack"): + Qwen4ExpVisionTransformer({"data_type": "bf16"}, deepstack_visual_indexes=[1]) diff --git a/unit_tests/models/qwen4_exp/test_weights.py b/unit_tests/models/qwen4_exp/test_weights.py new file mode 100644 index 0000000000..51bebb72ee --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_weights.py @@ -0,0 +1,99 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.models.qwen4_exp.checkpoint import normalize_text_weight_names, ngram_table_shape +from lightllm.models.qwen4_exp.layer_weights.transformer_layer_weight import Qwen4ExpTransformerLayerWeight +from lightllm.models.qwen4_exp.layer_weights.pre_and_post_layer_weight import Qwen4ExpPreAndPostLayerWeight +from lightllm.models.qwen4_exp.layer_infer.pre_layer_infer import Qwen4ExpPreLayerInfer +from lightllm.models.qwen4_exp.layer_infer.post_layer_infer import Qwen4ExpPostLayerInfer +from lightllm.common.quantization import Quantcfg +from .helpers import checkpoint, prepost_checkpoint, reference_hyper + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +def test_scalar_gate_dispatch_is_model_owned(config): + from lightllm.common.quantization.no_quant import NoQuantization + from lightllm.common.quantization.scalar_projection import StableScalarProjection + from lightllm.models.qwen3_5_moe.layer_weights.transformer_layer_weight import Qwen35MOETransformerLayerWeight + + native = Qwen4ExpTransformerLayerWeight(0, torch.bfloat16, config, Quantcfg(config)) + legacy = Qwen35MOETransformerLayerWeight(0, torch.bfloat16, config, Quantcfg(config)) + assert isinstance(native.shared_expert_gate.quant_method, StableScalarProjection) + assert type(legacy.shared_expert_gate.quant_method) is NoQuantization + + +@pytest.mark.parametrize("layer", [0, 1, 3]) +def test_incremental_native_checkpoint_loading(config, layer): + torch.manual_seed(737) + weights = checkpoint(config, layer) + native = Qwen4ExpTransformerLayerWeight(layer, torch.bfloat16, config, Quantcfg(config)) + assert not native.attn_hyper_connection.verify_load() + with pytest.raises(RuntimeError, match="not fully loaded"): + _ = native.attn_hyper_connection.component + + for key in reversed(list(weights)): + native.load_hf_weights({key: weights[key]}) + native.verify_load() + assert not hasattr(native, "att_norm_weight_") + assert not hasattr(native, "ffn_norm_weight_") + if layer == 3: + assert native.qk_norm_weight_.q_weight.dtype == torch.float32 + expected = 1 + weights[f"model.language_model.layers.{layer}.self_attn.q_norm.weight"].float().cuda() + torch.testing.assert_close(native.qk_norm_weight_.q_weight, expected, atol=0, rtol=0) + assert native.indexer.component.head_dim == 16 + else: + assert native.linear_norm.activation == "sigmoid" + hidden = torch.randn(3, 32, device="cuda", dtype=torch.bfloat16) + actual = native.linear_in_proj.mm(hidden, use_custom_tensor_mananger=False) + p = f"model.language_model.layers.{layer}.linear_attn" + expected = torch.cat( + [ + torch.nn.functional.linear(hidden, weights[p + f".in_proj_{name}.weight"].cuda()) + for name in ["qkv", "z", "b", "a"] + ], + -1, + ) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + if layer == 1: + assert native.ngram_embedding.weight.shape == (36, 8) + assert native.ngram_hasher.component.heads == 4 + assert native.ple.component.context_len == 9 + + +def test_embedding_widening_and_final_native_logits(config): + torch.manual_seed(737) + source = prepost_checkpoint(config) + native = Qwen4ExpPreAndPostLayerWeight(torch.bfloat16, config) + native.load_hf_weights(dict(source)) + native.verify_load() + assert not hasattr(native, "final_norm_weight_") + pre, post = Qwen4ExpPreLayerInfer(config), Qwen4ExpPostLayerInfer(config) + state = SimpleNamespace(is_prefill=False, is_draft_model=False, batch_size=2, prompt_logics=None, dist_group=None) + ids = torch.tensor([2, 17], device="cuda", dtype=torch.int64) + widened = pre.token_forward(ids, state, native) + torch.testing.assert_close( + widened, source["model.language_model.embed_tokens.weight"].cuda()[ids].repeat(1, 2), atol=0, rtol=0 + ) + widened = widened + torch.randn_like(widened) * 0.1 + actual = post.token_forward(widened, state, native).logits + mixed = reference_hyper(widened, source, "model.language_model.hyper_connection_mixer", config, False) + expected = torch.nn.functional.linear(mixed, source["lm_head.weight"].cuda()).float() + torch.testing.assert_close(actual, expected, atol=0.016, rtol=0.016) + + +def test_namespace_and_allocation_metadata(): + value = torch.zeros(1) + weights = { + "model.language_model.x": value, + "model.visual.x": value, + "mtp.x": value, + "adapter.model.language_model.x": value, + } + normalize_text_weight_names(weights) + assert set(weights) == {"model.x", "model.visual.x", "mtp.x", "adapter.model.language_model.x"} + with pytest.raises(ValueError, match="Conflicting"): + normalize_text_weight_names({"model.x": value, "model.language_model.x": value.clone()}) + assert ngram_table_shape(20000000, 16, 0, 2560, 128) == (320001536, 160) diff --git a/unit_tests/models/qwen4_exp/vision_reference.py b/unit_tests/models/qwen4_exp/vision_reference.py new file mode 100644 index 0000000000..bd0f2b68b1 --- /dev/null +++ b/unit_tests/models/qwen4_exp/vision_reference.py @@ -0,0 +1,93 @@ +import torch +import torch.nn.functional as F + + +def patch_coordinates(grid, merge, device): + coordinates = [] + for frames, height, width in grid.cpu().tolist(): + coordinates.extend( + (y + dy, x + dx) + for _ in range(frames) + for y in range(0, height, merge) + for x in range(0, width, merge) + for dy in range(merge) + for dx in range(merge) + ) + return torch.tensor(coordinates, dtype=torch.int64, device=device) + + +def interpolate_positions(table, grid, merge): + side = int(table.shape[0] ** 0.5) + result = [] + for frames, height, width in grid.cpu().tolist(): + coords = patch_coordinates(torch.tensor([[frames, height, width]]), merge, table.device) + ys = torch.linspace(0, side - 1, height, device=table.device)[coords[:, 0]] + xs = torch.linspace(0, side - 1, width, device=table.device)[coords[:, 1]] + y0, x0 = ys.floor().long(), xs.floor().long() + y1, x1 = (y0 + 1).clamp_max(side - 1), (x0 + 1).clamp_max(side - 1) + dy, dx = ys - y0, xs - x0 + indices = torch.stack([y0 * side + x0, y0 * side + x1, y1 * side + x0, y1 * side + x1], -1) + coefficients = torch.stack([(1 - dy) * (1 - dx), (1 - dy) * dx, dy * (1 - dx), dy * dx], -1) + result.append((table[indices] * coefficients[:, :, None]).sum(1).to(table.dtype)) + return torch.cat(result) + + +def vision_reference(pixels, grid, weights, config): + hidden, heads = config["hidden_size"], config["num_heads"] + patch, temporal, merge = config["patch_size"], config["temporal_patch_size"], config["spatial_merge_size"] + values = F.conv3d( + pixels.view(-1, config["in_channels"], temporal, patch, patch), + weights["patch_embed.proj.weight"], + weights["patch_embed.proj.bias"], + stride=(temporal, patch, patch), + ).view(-1, hidden) + values = values + interpolate_positions(weights["pos_embed.weight"], grid, merge) + dim = hidden // heads + theta = config.get("rope_parameters", {}).get("rope_theta", 10000.0) + invfreq = theta ** (-torch.arange(0, dim // 2, 2, device=pixels.device).float() / (dim // 2)) + coordinates = patch_coordinates(grid, merge, pixels.device) + angles = (coordinates[:, :, None] * invfreq).flatten(1).repeat(1, 2) + cos, sin = angles.cos()[:, None], angles.sin()[:, None] + + def norm(value, prefix): + return F.layer_norm(value, (value.shape[-1],), weights[prefix + ".weight"], weights[prefix + ".bias"], 1e-6) + + def rotate(value): + value = value.float() + a, b = value.chunk(2, -1) + return (value * cos + torch.cat([-b, a], -1) * sin).to(pixels.dtype) + + lengths = [height * width for frames, height, width in grid.cpu().tolist() for _ in range(frames)] + for layer in range(config["depth"]): + prefix = f"blocks.{layer}" + qkv = F.linear( + norm(values, prefix + ".norm1"), weights[prefix + ".attn.qkv.weight"], weights[prefix + ".attn.qkv.bias"] + ) + q, k, v = qkv.view(-1, 3, heads, dim).unbind(1) + q, k = rotate(q), rotate(k) + outputs, offset = [], 0 + for count in lengths: + outputs.append( + F.scaled_dot_product_attention( + q[offset : offset + count].transpose(0, 1)[None], + k[offset : offset + count].transpose(0, 1)[None], + v[offset : offset + count].transpose(0, 1)[None], + )[0].transpose(0, 1) + ) + offset += count + attention = torch.cat(outputs).reshape(-1, hidden) + values = values + F.linear( + attention, weights[prefix + ".attn.proj.weight"], weights[prefix + ".attn.proj.bias"] + ) + mlp = F.linear( + norm(values, prefix + ".norm2"), + weights[prefix + ".mlp.linear_fc1.weight"], + weights[prefix + ".mlp.linear_fc1.bias"], + ) + mlp = F.gelu(mlp, approximate="tanh" if config["hidden_act"] == "gelu_pytorch_tanh" else "none") + values = values + F.linear( + mlp, weights[prefix + ".mlp.linear_fc2.weight"], weights[prefix + ".mlp.linear_fc2.bias"] + ) + values = norm(values, "merger.norm").view(-1, hidden * merge * merge) + values = F.gelu(F.linear(values, weights["merger.linear_fc1.weight"], weights["merger.linear_fc1.bias"])) + return F.linear(values, weights["merger.linear_fc2.weight"], weights["merger.linear_fc2.bias"]) 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..a31c0acaf8 --- /dev/null +++ b/unit_tests/models/sensenova_loop/test_gdn_state.py @@ -0,0 +1,150 @@ +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() + assert resource.get_storage_layer_index(logical_layer) == storage_layer + 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.triton_kernel.linear_att.fla.ops.fused_recurrent_gated_delta_rule", recurrent + ) + monkeypatch.setattr( + "lightllm.common.basemodel.triton_kernel.linear_att.mtp_fused_recurrent." + "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/server/httpserver/test_slot_availability.py b/unit_tests/server/httpserver/test_slot_availability.py new file mode 100644 index 0000000000..74b687a72e --- /dev/null +++ b/unit_tests/server/httpserver/test_slot_availability.py @@ -0,0 +1,274 @@ +import asyncio +from types import SimpleNamespace +import pytest +import torch +from lightllm.server.httpserver.slot_availability import RequestSlotAvailability +from lightllm.server.core.objs.shm_req_manager import ShmReqManager + +from lightllm.server.httpserver.manager import HttpServerManager +import lightllm.server.httpserver.manager as module + + +class Pool(ShmReqManager): + def __init__(self, indices=()): + self.free = list(indices) + self.claimed = set() + self.empty = asyncio.Event() + self.fail_release = False + + def alloc_req_index(self): + if not self.free: + self.empty.set() + return None + index = self.free.pop(0) + assert index not in self.claimed + self.claimed.add(index) + return index + + def alloc_req_indexes(self, count): + if len(self.free) < count: + self.empty.set() + return None + return [self.alloc_req_index() for _ in range(count)] + + def release_req_index(self, index): + if self.fail_release: + raise RuntimeError("test release failed") + self.claimed.discard(index) + self.free.append(index) + + async def async_put_back_req_obj(self, req): + pass + + def __del__(self): + pass + + +def manager(pool, availability=True): + result = HttpServerManager.__new__(HttpServerManager) + result.shm_req_manager = pool + result._request_slot_availability = RequestSlotAvailability(availability) + result.is_multinode_tp_slave = False + result.args = SimpleNamespace(running_max_req_size=8, run_mode="normal") + return result + + +@pytest.fixture(autouse=True) +def no_cuda(): + assert not torch.cuda.is_initialized() + yield + assert not torch.cuda.is_initialized() + + +def test_original_async_allocator_keeps_immediate_none_contract(): + async def scenario(): + assert await Pool().async_alloc_req_index() is None + + asyncio.run(scenario()) + + +def test_generation_change_before_wait_is_observed(): + async def scenario(): + notifier = RequestSlotAvailability() + generation = notifier.generation + notifier.notify_release() + assert await notifier.wait(generation, 10) + + asyncio.run(scenario()) + + +def test_waiting_task_is_notified_after_release(): + async def scenario(): + notifier = RequestSlotAvailability() + task = asyncio.create_task(notifier.wait(0, 10)) + await asyncio.sleep(0) + notifier.notify_release() + assert await task + + asyncio.run(scenario()) + + +def test_unobserved_external_release_uses_timeout_fallback(): + async def scenario(): + notifier = RequestSlotAvailability() + assert not await notifier.wait(notifier.generation, 0.001) + + asyncio.run(scenario()) + + +def test_actual_http_retry_observes_release_between_check_and_wait(): + async def scenario(): + pool = Pool() + http = manager(pool) + availability = http._request_slot_availability + original = availability.wait + + async def release_before_wait(generation, timeout): + await pool.async_release_req_index(7) + http._notify_req_slot_release() + return await original(generation, timeout) + + availability.wait = release_before_wait + assert await http._alloc_shm_req_indexes(1) == [7] + assert pool.claimed == {7} + + asyncio.run(scenario()) + + +def test_recycler_notifies_only_after_actual_release(): + async def scenario(): + pool = Pool() + pool.claimed = {7} + http = manager(pool) + http.recycle_event = asyncio.Event() + req = SimpleNamespace(request_id=8, index_in_shm_mem=7) + group = SimpleNamespace(group_req_id=8, shm_req_objs=[req], multimodal_params=None) + http.req_id_to_out_inf = {8: SimpleNamespace(can_release=lambda: True, group_req_objs=group)} + + async def release_media(params): + pass + + http._release_multimodal_resources = release_media + waiter = asyncio.create_task(http._alloc_shm_req_indexes(1)) + await pool.empty.wait() + recycler = asyncio.create_task(http.recycle_resource_loop()) + http.recycle_event.set() + try: + assert await waiter == [7] + assert not http.req_id_to_out_inf + assert http._request_slot_availability.generation == 1 + finally: + recycler.cancel() + with pytest.raises(asyncio.CancelledError): + await recycler + + asyncio.run(scenario()) + + +def test_http_cancellation_keeps_unallocated_group_free(): + async def scenario(): + pool = Pool([3]) + http = manager(pool) + task = asyncio.create_task(http._alloc_shm_req_indexes(2)) + await pool.empty.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert pool.free == [3] and not pool.claimed + assert http._request_slot_availability.generation == 0 + + asyncio.run(scenario()) + + +def test_waiter_cancellation_does_not_consume_notification(): + async def scenario(): + notifier = RequestSlotAvailability() + task = asyncio.create_task(notifier.wait(0, 10)) + await asyncio.sleep(0) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + notifier.notify_release() + assert await notifier.wait(0, 10) + + asyncio.run(scenario()) + + +def test_multiple_http_waiters_do_not_duplicate_allocation(): + async def scenario(): + pool = Pool() + http = manager(pool) + tasks = [asyncio.create_task(http._alloc_shm_req_indexes(1)) for _ in range(2)] + await pool.empty.wait() + await pool.async_release_req_index(4) + http._notify_req_slot_release() + await pool.async_release_req_index(5) + http._notify_req_slot_release() + assert {result[0] for result in await asyncio.gather(*tasks)} == {4, 5} + assert pool.claimed == {4, 5} + + asyncio.run(scenario()) + + +def test_failed_release_does_not_increment_generation(): + async def scenario(): + pool = Pool([3]) + http = manager(pool) + pool.fail_release = True + + task = asyncio.create_task(http._alloc_shm_req_indexes(2)) + await pool.empty.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert http._request_slot_availability.generation == 0 + + asyncio.run(scenario()) + + +@pytest.mark.parametrize( + "change", + [ + dict(run_mode="prefill"), + dict(run_mode="decode"), + dict(dp=2), + dict(nnodes=2), + dict(enable_rl=True), + dict(httpserver_workers=2), + ], +) +def test_unverified_modes_keep_legacy_backoff(change): + args = SimpleNamespace(run_mode="normal", dp=1, nnodes=1, enable_rl=False, httpserver_workers=1) + for key, value in change.items(): + setattr(args, key, value) + notifier = RequestSlotAvailability.from_args(args) + assert not notifier.enabled + notifier.notify_release() + assert notifier.generation == 0 + + +@pytest.mark.parametrize("priority,expected", [(False, 0.1), (True, 0.02)]) +def test_original_priority_backoff_interval_is_preserved(priority, expected): + async def scenario(): + pool = Pool() + http = manager(pool, False) + observed = [] + + async def wait_and_release(generation, timeout): + observed.append(timeout) + await pool.async_release_req_index(7) + return False + + http._request_slot_availability.wait = wait_and_release + assert await http._alloc_shm_req_indexes(1, high_priority_request=priority) == [7] + assert observed == pytest.approx([expected]) + + asyncio.run(scenario()) + + +def test_existing_timeout_preserves_unallocated_group(): + async def scenario(): + pool = Pool([3]) + http = manager(pool, False) + http.pd_node_request_limit_enabled = True + http.pd_node_shm_req_alloc_timeout_seconds = 0 + http.args = SimpleNamespace(run_mode="prefill", running_max_req_size=8) + with pytest.raises(module.ServerBusyError): + await http._alloc_shm_req_indexes(2, pd_node_resource_wait_timeout_seconds=0) + assert pool.free == [3] and not pool.claimed + + asyncio.run(scenario()) + + +@pytest.mark.parametrize("mode,workers,enabled", [("normal", 1, True), ("decode", 1, False), ("normal", 2, False)]) +def test_real_constructor_initializes_notifier_before_external_io(monkeypatch, mode, workers, enabled): + args = SimpleNamespace(run_mode=mode, dp=1, nnodes=1, enable_rl=False, httpserver_workers=workers) + + def stop_before_external_io(): + raise RuntimeError("stop_before_external_io") + + monkeypatch.setattr(module, "get_shm_port_args", stop_before_external_io) + http = HttpServerManager.__new__(HttpServerManager) + with pytest.raises(RuntimeError, match="stop_before_external_io"): + http.__init__(args) + assert http._request_slot_availability.enabled is enabled 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 index 2e488ebb30..9a64edf769 100644 --- a/unit_tests/server/router/model_infer/mtp_speculative/test_planner.py +++ b/unit_tests/server/router/model_infer/mtp_speculative/test_planner.py @@ -962,3 +962,86 @@ def test_topk_prefix_sums_only_computes_requested_counts(): assert set(result) == {0, 2, 4} np.testing.assert_allclose([result[0], result[2], result[4]], [0.0, 1.6, 2.1]) + + +def build_stalled_zero_draft_planner(req_num, spec_mode="eagle_no_att"): + planner = build_lightspec_planner(max_draft_step=4, spec_mode=spec_mode) + for batch_size in range(1, 41): + planner.target_infer_costs.update(batch_size, 1.0 + 0.1 * batch_size) + planner.draft_infer_costs.update(batch_size, 0.2 + 0.01 * batch_size) + planner._update_verified_batch( + accept_lengths=[1] * req_num, + req_num=req_num, + dynamic_batch_size=5 * req_num, + verified_draft_step=4, + ) + assert planner._select_draft_step(req_num) == planner.draft_steps[0] + planner.pre_draft_step = planner.draft_steps[0] + return planner + + +@pytest.mark.parametrize("req_num", [1, 2, 8]) +@pytest.mark.parametrize("all_new", [False, True]) +def test_lightspec_new_request_recovers_from_zero_draft_feedback(req_num, all_new): + planner = build_stalled_zero_draft_planner(req_num) + target_costs = dict(planner.target_infer_costs.infer_cost_ms_table) + draft_costs = dict(planner.draft_infer_costs.infer_cost_ms_table) + new_requests = req_num if all_new else 1 + + plan = planner.plan( + decode_reqs=build_decode_reqs(req_num, req_num_with_proposals=req_num - new_requests), + origin_batch_size=5 * req_num, + ) + + assert plan.pre_draft_step == 0 + assert plan.dynamic_batch_size == req_num + assert plan.draft_step == 4 + assert not plan.all_reqs_have_proposals + assert dict(planner.target_infer_costs.infer_cost_ms_table) == target_costs + assert dict(planner.draft_infer_costs.infer_cost_ms_table) == draft_costs + + planner.pre_draft_step = plan.draft_step + next_plan = planner.plan(decode_reqs=build_decode_reqs(req_num), origin_batch_size=5 * req_num) + assert next_plan.dynamic_batch_size == 5 * req_num + planner._update_verified_batch( + accept_lengths=[5] * req_num, + req_num=req_num, + dynamic_batch_size=next_plan.dynamic_batch_size, + verified_draft_step=4, + ) + assert planner._select_draft_step(req_num) == 4 + + +@pytest.mark.parametrize("empty", [False, True]) +def test_lightspec_existing_or_empty_batch_retains_zero_draft_evidence(empty): + planner = build_stalled_zero_draft_planner(2) + progress = dict(planner.progress_ema_by_config) + survival = list(planner.prefix_survival_by_depth) + + plan = planner.plan(decode_reqs=[] if empty else build_decode_reqs(2), origin_batch_size=0 if empty else 10) + + assert plan.draft_step == 0 + assert planner.progress_ema_by_config == progress + assert planner.prefix_survival_by_depth == survival + + +@pytest.mark.parametrize("req_num", [1, 2, 8]) +@pytest.mark.parametrize("spec_mode", ["eagle_with_att", "eagle3", "vanilla_with_att"]) +def test_lightspec_new_workload_escapes_target_only_verification(req_num, spec_mode): + planner = build_stalled_zero_draft_planner(req_num, spec_mode=spec_mode) + previous_depth = planner.pre_draft_step + assert previous_depth > 0 + target_costs = dict(planner.target_infer_costs.infer_cost_ms_table) + draft_costs = dict(planner.draft_infer_costs.infer_cost_ms_table) + + first = planner.plan(decode_reqs=build_decode_reqs(req_num, 0), origin_batch_size=5 * req_num) + + assert first.pre_draft_step == previous_depth + assert first.dynamic_batch_size == req_num + assert first.draft_step == 4 + assert not first.all_reqs_have_proposals + assert dict(planner.target_infer_costs.infer_cost_ms_table) == target_costs + assert dict(planner.draft_infer_costs.infer_cost_ms_table) == draft_costs + planner.pre_draft_step = first.draft_step + following = planner.plan(decode_reqs=build_decode_reqs(req_num), origin_batch_size=5 * req_num) + assert following.dynamic_batch_size == 5 * req_num diff --git a/unit_tests/server/router/req_queue/test_cached_prefix_queue.py b/unit_tests/server/router/req_queue/test_cached_prefix_queue.py new file mode 100644 index 0000000000..b941f91475 --- /dev/null +++ b/unit_tests/server/router/req_queue/test_cached_prefix_queue.py @@ -0,0 +1,224 @@ +import random +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +import lightllm.server.router.req_queue as module +from lightllm.server.router.req_queue.chunked_prefill.cached_prefix_queue import ( + CachedPrefixChunkedPrefillQueue, + supports_cached_prefix_reservations, +) +from lightllm.server.router.req_queue.chunked_prefill.impl import ChunkedPrefillQueue +from lightllm.server.router.req_queue.cached_prefix import CachedPrefixReservations + + +class Load: + def set_estimated_peak_token_count(self, value, index): + self.peak = value + + def set_dynamic_max_load(self, value, index): + self.ratio = value + + +class Request: + def __init__(self, identity, base=1800, remaining=350, cached=512): + self.request_id, self.input_len, self.shm_cur_kv_len = identity, base, base + self.prompt_cache_len = cached + self.is_paused = self.is_aborted = False + self.sample_params = SimpleNamespace(suggested_dp_index=0, disable_prompt_cache=False) + self.multimodal_params = SimpleNamespace(images=[], audios=[]) + self.shm_prompt_ids = SimpleNamespace(arr=np.arange(base, dtype=np.int64)) + self.linear_att_token_hash_list = SimpleNamespace(get_all=lambda: [17] * (base // 512)) + self.remaining = remaining + + def get_tuple_tokens(self, is_busy, ema): + return self.input_len, self.remaining + + def get_first_router_need_tokens(self): + return 512 + + +def arguments(**changes): + result = SimpleNamespace( + dp=1, + nnodes=1, + run_mode="normal", + enable_rl=False, + disable_dynamic_prompt_cache=False, + enable_cpu_cache=False, + linear_att_page_block_num=1, + linear_att_hash_page_size=512, + short_prefill_token_threshold=None, + diverse_mode=False, + output_constraint_mode="none", + first_token_constraint_mode=False, + disable_chunked_prefill=False, + model_dir="unused", + ) + for name, value in changes.items(): + setattr(result, name, value) + return result + + +def queue(requests, cls=CachedPrefixChunkedPrefillQueue): + result = cls.__new__(cls) + result.args = arguments() + result.dp_index = 0 + result.max_total_tokens, result.running_max_req_size, result.batch_max_tokens = 16384, 8, 1024 + result.router = SimpleNamespace(router_statics=SimpleNamespace(ema_req_out_len=2048), shared_token_load=Load()) + result.is_busy = lambda: True + if cls is CachedPrefixChunkedPrefillQueue: + result.prefix_reservations = CachedPrefixReservations(result.args, 0) + result._init_cache_list(SimpleNamespace(reqs=requests), True) + return result + + +@pytest.fixture(autouse=True) +def namespace(monkeypatch): + assert not torch.cuda.is_initialized() + monkeypatch.setenv("LIGHTLLM_UNIQUE_SERVICE_NAME_ID", "owned-canonical-test") + yield + assert not torch.cuda.is_initialized() + + +def test_known_cached_prefix_changes_actual_admission_at_same_capacity(): + legacy = queue([Request(i) for i in range(7)], ChunkedPrefillQueue) + candidate = queue([Request(i) for i in range(7)]) + assert not legacy._can_add_new_req(Request(7), True) + assert candidate._can_add_new_req(Request(7), True) is True + assert candidate.router.shared_token_load.peak == 13616 + assert candidate.max_total_tokens == 16384 and candidate.running_max_req_size == 8 + + +def test_refactored_legacy_peak_matches_original_formula_for_ties_and_empty_batches(): + rng = random.Random(8245) + for _ in range(1500): + entries = [(rng.randrange(1, 3000), rng.randrange(0, 100)) for _ in range(rng.randrange(0, 17))] + legacy = queue([], ChunkedPrefillQueue) + legacy.cache_len_list = entries.copy() + expected = max( + ( + sum(a for a, b in sorted(entries, key=lambda e: -e[1])[:i]) + i * pair[1] + for i, pair in enumerate(sorted(entries, key=lambda e: -e[1]), 1) + ), + default=0, + ) + assert legacy._estimate_peak_tokens() == expected + + +@pytest.mark.parametrize( + "change", + [ + "images", + "audios", + "paused", + "aborted", + "disabled_cache", + "unresident", + "missing_hashes", + "different_hashes", + "different_tokens", + ], +) +def test_uncertain_or_distinct_prefixes_preserve_legacy_rejection(change): + requests = [Request(i) for i in range(8)] + for request in requests: + if change in ("images", "audios"): + setattr(request.multimodal_params, change, [object()]) + elif change in ("paused", "aborted"): + setattr(request, "is_" + change, True) + elif change == "disabled_cache": + request.sample_params.disable_prompt_cache = True + elif change == "unresident": + request.shm_cur_kv_len = 128 + elif change == "missing_hashes": + request.linear_att_token_hash_list = SimpleNamespace(get_all=lambda: []) + elif change == "different_hashes": + request.linear_att_token_hash_list = SimpleNamespace(get_all=lambda i=request.request_id: [i]) + elif change == "different_tokens": + request.shm_prompt_ids.arr += request.request_id * 1800 + candidate = queue(requests[:7]) + assert not candidate._can_add_new_req(requests[7], True) + + +def test_reset_and_load_preserve_tuple_identity_and_request_lifetime(): + requests = [Request(i, remaining=350 - i * 3) for i in range(7)] + candidate = queue(requests) + assert candidate._can_add_new_req(Request(7, remaining=420), True) + peak, ratio = candidate._calcu_batch_token_load_batch_not_none(SimpleNamespace(reqs=requests)) + assert peak < 16384 and ratio == peak / 16384 + candidate._init_cache_list(SimpleNamespace(reqs=[Request(50)]), True) + assert set(candidate.prefix_reservations.memo) == {50} + assert len(candidate.prefix_reservations.by_tuple) == 1 + candidate._init_cache_list(None, True) + assert not candidate.prefix_reservations.memo and candidate._estimate_peak_tokens() == 0 + + +@pytest.mark.parametrize( + "change", + [ + dict(dp=2), + dict(nnodes=2), + dict(run_mode="decode"), + dict(enable_rl=True), + dict(disable_dynamic_prompt_cache=True), + dict(enable_cpu_cache=True), + dict(linear_att_page_block_num=4), + dict(linear_att_hash_page_size=0), + ], +) +def test_unverified_cache_geometry_and_modes_are_not_enabled(change): + assert not supports_cached_prefix_reservations(arguments(**change)) + + +@pytest.mark.parametrize( + "model_type,expected", + [("qwen4_exp", True), ("qwen4_exp_text", True), ("qwen3_5", False), ("llama", False), (None, False)], +) +def test_factory_selects_only_verified_models(monkeypatch, model_type, expected): + monkeypatch.setattr("lightllm.utils.config_utils.get_model_type", lambda path: model_type) + assert (module._get_req_queue_class(arguments(), None, 1) is CachedPrefixChunkedPrefillQueue) == expected + + +@pytest.mark.parametrize( + "change", + [ + dict(linear_att_page_block_num=4), + dict(enable_cpu_cache=True), + dict(dp=2), + dict(output_constraint_mode="x"), + dict(first_token_constraint_mode=True), + ], +) +def test_factory_retains_legacy_queue_for_unsupported_configurations(monkeypatch, change): + monkeypatch.setattr("lightllm.utils.config_utils.get_model_type", lambda path: "qwen4_exp") + assert module._get_req_queue_class(arguments(**change), None, 1) is ChunkedPrefillQueue + + +def test_request_limit_is_unchanged(): + candidate = queue([Request(i) for i in range(8)]) + assert not candidate._can_add_new_req(Request(8), True) + + +def test_unknown_namespace_gets_no_credit(monkeypatch): + monkeypatch.delenv("LIGHTLLM_UNIQUE_SERVICE_NAME_ID") + candidate = queue([Request(i) for i in range(7)]) + assert not candidate._can_add_new_req(Request(7), True) + + +def test_missing_prompt_handle_attaches_via_existing_request_api(): + requests = [Request(i) for i in range(8)] + linked = [] + for request in requests: + del request.shm_prompt_ids + + def attach(request=request): + linked.append(request.request_id) + request.shm_prompt_ids = SimpleNamespace(arr=np.arange(request.input_len, dtype=np.int64)) + + request.link_prompt_ids_shm_array = attach + candidate = queue(requests[:7]) + assert candidate._can_add_new_req(requests[7], True) + assert linked == list(range(8)) diff --git a/unit_tests/server/router/req_queue/test_peak_reservation.py b/unit_tests/server/router/req_queue/test_peak_reservation.py new file mode 100644 index 0000000000..713b30a44e --- /dev/null +++ b/unit_tests/server/router/req_queue/test_peak_reservation.py @@ -0,0 +1,96 @@ +import random + +import pytest + +from lightllm.server.router.req_queue.peak_reservation import Reservation, estimate_peak + + +def simulated_peak(entries, page_tokens): + totals = [] + for progress in range(max((entry.remaining_tokens for entry in entries), default=0) + 1): + pages, private = set(), 0 + for entry in entries: + if entry.remaining_tokens < progress: + continue + private += entry.base_tokens + progress + if entry.cache_domain is not None: + private -= len(entry.prefix_pages) * page_tokens + for depth in range(1, len(entry.prefix_pages) + 1): + pages.add((entry.cache_domain, entry.prefix_pages[:depth])) + totals.append(private + len(pages) * page_tokens) + return max(totals) + + +def test_shared_prefix_reservation_changes_admission_without_increasing_capacity(): + entries = [Reservation(1800, 350, (b"same exact page",), b"one cache") for _ in range(8)] + peak = estimate_peak(entries, 512) + assert peak.legacy_tokens == 17200 + assert peak.shared_tokens == simulated_peak(entries, 512) == 13616 + assert peak.shared_tokens < 16384 < peak.legacy_tokens + + +def test_prefix_ancestry_and_cache_domain_prevent_false_sharing(): + entries = [ + Reservation(20, 3, (b"a", b"same"), b"cache"), + Reservation(20, 3, (b"b", b"same"), b"cache"), + Reservation(20, 3, (b"a", b"same"), b"other cache"), + ] + peak = estimate_peak(entries, 4) + assert peak.legacy_tokens == peak.shared_tokens == simulated_peak(entries, 4) + + +def test_unknown_domains_keep_legacy_reservation(): + entries = [Reservation(20, remaining, (b"same",)) for remaining in (1, 3, 7)] + peak = estimate_peak(entries, 4) + assert peak.legacy_tokens == peak.shared_tokens == simulated_peak(entries, 4) + + +def test_credits_follow_each_live_subset_when_requests_complete(): + entries = [ + Reservation(25, 2, (b"a", b"b", b"c"), b"cache"), + Reservation(25, 8, (b"a",), b"cache"), + Reservation(25, 5, (b"a", b"b"), b"cache"), + ] + assert estimate_peak(entries, 4).shared_tokens == simulated_peak(entries, 4) + + +def test_reservation_matches_independent_simulated_lifetimes(): + rng = random.Random(381020) + for _ in range(1500): + page_tokens = rng.choice((1, 2, 4, 512)) + entries = [] + for _ in range(rng.randrange(1, 9)): + pages = tuple(rng.choice((b"a", b"b", b"c")) for _ in range(rng.randrange(5))) + entries.append( + Reservation( + len(pages) * page_tokens + rng.randrange(8), + rng.randrange(13), + pages, + rng.choice((None, b"x", b"y")), + ) + ) + peak = estimate_peak(entries, page_tokens) + assert peak.shared_tokens == simulated_peak(entries, page_tokens) + assert 0 <= peak.shared_tokens <= peak.legacy_tokens + + +@pytest.mark.parametrize( + "entry,page_tokens", + [ + (Reservation(-1, 1), 1), + (Reservation(1, -1), 1), + (Reservation(True, 1), 1), + (Reservation(1, 1), 0), + (Reservation(1, 1), True), + (Reservation(1, 1, (b"a", b"b"), b"cache"), 1), + (Reservation(1, 1, ("text is not an exact byte identity",), b"cache"), 1), + (Reservation(1, 1, (b"a",), "mutable or ambiguous domain"), 1), + ], +) +def test_rejects_invalid_accounting_metadata(entry, page_tokens): + with pytest.raises(ValueError): + estimate_peak([entry], page_tokens) + + +def test_empty_batch_has_no_reservation(): + assert estimate_peak([], 512).shared_tokens == 0 diff --git a/unit_tests/server/test_hybrid_start_args.py b/unit_tests/server/test_hybrid_start_args.py index 230e23b0e5..1c391c09cd 100644 --- a/unit_tests/server/test_hybrid_start_args.py +++ b/unit_tests/server/test_hybrid_start_args.py @@ -102,3 +102,37 @@ def test_whole_image_blocks_fit_batch_at_startup( 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_native_stream_details.py b/unit_tests/server/test_native_stream_details.py new file mode 100644 index 0000000000..0d983258a0 --- /dev/null +++ b/unit_tests/server/test_native_stream_details.py @@ -0,0 +1,69 @@ +import asyncio +import copy +import json +from types import SimpleNamespace + +import pytest + +from lightllm.server import api_lightllm + + +@pytest.mark.parametrize("details", [None, False, True]) +@pytest.mark.parametrize("reported", [False, True]) +def test_stream_details_are_opt_in_and_preserve_metadata(monkeypatch, details, reported): + class Sampling: + best_of = 1 + + def init(self, **kwargs): + assert "return_details" not in kwargs + + async def verify_cpu_async(self, compiler): + pass + + monkeypatch.setattr(api_lightllm, "SamplingParams", Sampling) + parameters = {"do_sample": False} + if details is not None: + parameters["return_details"] = details + + class Request: + async def json(self): + return {"inputs": [11, 12], "parameters": parameters} + + metadata = dict( + id=7, + prompt_tokens=2, + count_output_tokens=1, + logprob=-0.2, + logprobs={}, + prompt_cache_len=1, + mtp_accepted_token_num=2, + mtp_verify_token_num=3, + mtp_verify_step_num=1, + input_usage=dict(input_text_tokens=2), + ) + fields = ("prompt_cache_len", "mtp_accepted_token_num", "mtp_verify_token_num", "mtp_verify_step_num") + if not reported: + for name in fields: + metadata.pop(name) + before = copy.deepcopy(metadata) + + async def generate(*args, **kwargs): + yield 0, "token", metadata, SimpleNamespace(is_finished=lambda: True, get_finish_reason=lambda: "length") + + async def collect(): + response = await api_lightllm.lightllm_generate_stream( + Request(), SimpleNamespace(tokenizer=None, generate=generate, output_grammar_compiler=None) + ) + return [json.loads(chunk.decode().removeprefix("data:")) async for chunk in response.body_iterator] + + bodies = asyncio.run(collect()) + assert len(bodies) == 1 + body = bodies[0] + for name in fields: + if details and reported: + assert body["token"][name] == metadata[name] + else: + assert name not in body["token"] + assert body["input_usage"] == metadata["input_usage"] and body["details"] is None + assert body["token"]["id"] == 7 and body["finish_reason"] == "length" + assert metadata == before 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..da786bddef --- /dev/null +++ b/unit_tests/server/test_ssm_state_mode.py @@ -0,0 +1,167 @@ +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", 3, "qwen4_exp", 16, None), + ("replay", "bfloat16", 3, "qwen4_exp_text", 8, 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", "qwen4_exp", "qwen4_exp_text"]) +@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", "qwen4_exp", "qwen4_exp_text"]) +@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]) From 08588bb2174e7e2a2132c1f4f93d03c624e7cf30 Mon Sep 17 00:00:00 2001 From: sufubao Date: Thu, 12 Nov 2026 05:08:48 +0800 Subject: [PATCH 2/6] feat: add Triton causal convolution fallback --- .../common/basemodel/attention/linear/gdn.py | 1 + .../triton_kernel/linear_att/causal_conv1d.py | 112 +++-- .../linear_att/causal_conv1d_triton.py | 435 ++++++++++++++++++ .../linear_att/test_causal_conv1d.py | 378 +++++++++++++++ 4 files changed, 875 insertions(+), 51 deletions(-) create mode 100644 lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d_triton.py create mode 100644 unit_tests/common/basemodel/triton_kernel/linear_att/test_causal_conv1d.py diff --git a/lightllm/common/basemodel/attention/linear/gdn.py b/lightllm/common/basemodel/attention/linear/gdn.py index b715b77128..48e3eb64d4 100644 --- a/lightllm/common/basemodel/attention/linear/gdn.py +++ b/lightllm/common/basemodel/attention/linear/gdn.py @@ -178,6 +178,7 @@ def _gdn_prefill_kernel( has_initial_state=infer_state.b_ready_cache_len > 0, conv_states=conv_states, activation=backend.activation, + max_seqlen=infer_state.max_q_seq_len, ) mixed_qkv = out_tensor.transpose(0, 1) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d.py b/lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d.py index 2bf325340f..31435fa744 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d.py @@ -17,38 +17,41 @@ def causal_conv1d_fn( pad_slot_id: int = -1, **kwargs, ): - """ - x: (batch, dim, seqlen) or (dim,cu_seq_len) for varlen - sequences are concatenated from left to right for varlen - weight: (dim, width) - bias: (dim,) - query_start_loc: (batch + 1) int32 - The cumulative sequence lengths of the sequences in - the batch, used to index into sequence. prepended by 0. - for example: query_start_loc = torch.Tensor([0,10,16,17]), - x.shape=(dim,17) - cache_indices: (batch) int32 - indicates the corresponding state index, - like so: conv_state = conv_states[cache_indices[batch_id]] - has_initial_state: (batch) bool - indicates whether should the kernel take the current state as initial - state for the calculations - conv_states: (...,dim,width - 1) itype - updated inplace if provided - activation: either None or "silu" or "swish" - pad_slot_id: int - if cache_indices is passed, lets the kernel identify padded - entries that will not be processed, - for example: cache_indices = [pad_slot_id, 1, 20, pad_slot_id] - in this case, the kernel will not process entries at - indices 0 and 3 - - - out: (batch, dim, seqlen) + """Run causal depthwise convolution with SGL Kernel or a local Triton fallback. + + x is [batch, channels, tokens] or packed [channels, total_tokens]. The Triton + path requires packed inputs and all four request/state tensors below. + weight is [channels, width]; optional bias is [channels]. + + query_start_loc: cumulative token offsets [batch + 1], starting at zero. + cache_indices: maps each request to a slot in conv_states; pad_slot_id + entries keep their input and state unchanged. + has_initial_state: one bool per request; false means a zero-filled prefix. + conv_states: [slots, channels, width - 1], updated in place with input tails. + + activation is None, "silu" or "swish". The optional max_seqlen keyword bounds + the longest packed query for Triton without reading GPU lengths on the CPU. + Returns the same shape as x. SGL may overwrite x; Triton allocates an output. """ if activation not in [None, "silu", "swish"]: raise NotImplementedError("activation must be None, silu, or swish") - from sgl_kernel import causal_conv1d_fwd + try: + from sgl_kernel import causal_conv1d_fwd + except (ImportError, OSError): + from .causal_conv1d_triton import causal_conv1d_fn as triton_causal_conv1d_fn + + return triton_causal_conv1d_fn( + x, + weight, + bias, + query_start_loc=query_start_loc, + cache_indices=cache_indices, + has_initial_state=has_initial_state, + conv_states=conv_states, + activation=activation, + pad_slot_id=pad_slot_id, + **kwargs, + ) if x.stride(-1) != 1: x = x.contiguous() @@ -78,31 +81,38 @@ def causal_conv1d_update( conv_state_indices: Optional[torch.Tensor] = None, pad_slot_id: int = -1, ): - """ - x: (batch, dim) or (batch, dim, seqlen) - conv_state: (batch, dim, state_len), where state_len >= width - 1 - weight: (dim, width) - bias: (dim,) - cache_seqlens: (batch,), dtype int32. - If not None, the conv_state is treated as a circular buffer. - The conv_state will be updated by copying x to the conv_state - starting at the index - @cache_seqlens % state_len. - conv_state_indices: (batch,), dtype int32 - If not None, the conv_state is a larger tensor along the batch dim, - and we are selecting the batch coords specified by conv_state_indices. - Useful for a continuous batching scenario. - pad_slot_id: int - if cache_indices is passed, lets the kernel identify padded - entries that will not be processed, - for example: cache_indices = [pad_slot_id, 1 ,20 ,pad_slot_id] - in this case, the kernel will not process entries at - indices 0 and 3 - out: (batch, dim) or (batch, dim, seqlen) + """Decode with SGL Kernel when available, otherwise the local Triton kernel. + + x is [batch, channels] or [batch, channels, tokens]. conv_state is + [slots, channels, state_len], where state_len >= width - 1, and is updated + in place. weight is [channels, width]; optional bias is [channels]. + + cache_seqlens supplies one cursor per request for circular state buffers; + new tokens are written starting at cursor % state_len. Without cursors, + states shift left and append the new tokens. + conv_state_indices optionally maps requests to cache slots. Entries equal + to pad_slot_id keep their input and state unchanged. + + activation is None, "silu" or "swish". Returns the same shape as x. + SGL updates x in place; Triton allocates an output. """ if activation not in [None, "silu", "swish"]: raise NotImplementedError(f"activation must be None, silu, or swish, actual: {activation}") - from sgl_kernel import causal_conv1d_update as causal_conv1d_update_kernel + try: + from sgl_kernel import causal_conv1d_update as causal_conv1d_update_kernel + except (ImportError, OSError): + from .causal_conv1d_triton import causal_conv1d_update as triton_causal_conv1d_update + + return triton_causal_conv1d_update( + x, + conv_state, + weight, + bias, + activation=activation, + cache_seqlens=cache_seqlens, + conv_state_indices=conv_state_indices, + pad_slot_id=pad_slot_id, + ) activation_val = activation in ["silu", "swish"] unsqueeze = x.dim() == 2 diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d_triton.py b/lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d_triton.py new file mode 100644 index 0000000000..162cb36c60 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d_triton.py @@ -0,0 +1,435 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Causal depthwise convolution without external CUDA extensions. + +Packed prefill reads strided [channels, total_tokens] projection views. +Outputs are allocated separately; convolution states are updated in place. +""" + +from typing import Optional + +import torch +import triton +import triton.language as tl + +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, autotune + + +def _prefill_configs(): + return [ + {"BLOCK_TOKENS": block_tokens, "BLOCK_CHANNELS": block_channels, "num_warps": num_warps} + for block_tokens in (8, 16, 32, 64) + for block_channels in (64, 128, 256) + for num_warps in (2, 4, 8) + if block_tokens * block_channels <= num_warps * 32 * 64 + ] + + +def _update_configs(): + return [ + {"BLOCK_CHANNELS": block_channels, "num_warps": num_warps} + for block_channels in (64, 128, 256) + for num_warps in (1, 2, 4) + ] + + +def _conv_static_key(x, weight, bias, activation, conv_state): + # Group by memory order so nearby token counts can reuse a configuration. + if x.stride(-2) == 1: + layout = "channel" + elif x.stride(-1) == 1: + layout = "token" + else: + layout = "strided" + return { + "channels": weight.shape[0], + "width": weight.shape[1], + "dtype": str(x.dtype), + "layout": layout, + "state_stride": f"{conv_state.stride(-2)}_{conv_state.stride(-1)}", + "bias": bias is not None, + "silu": activation is not None, + } + + +def _prefill_static_key(x, weight, bias, conv_states, activation): + return _conv_static_key(x, weight, bias, activation, conv_states) + + +def _prefill_run_key(x, query_start_loc, max_seqlen): + batch_size = query_start_loc.numel() - 1 + max_query_len = x.shape[-1] if max_seqlen is None else max_seqlen + return f"{batch_size}_{max_query_len}" + + +def _prefill_key_distance(run_key, config_key): + batch_size, max_query_len = map(int, run_key.split("_")) + config_batch_size, config_max_query_len = map(int, config_key.split("_")) + return ( + abs(batch_size * max_query_len - config_batch_size * config_max_query_len), + abs(batch_size - config_batch_size), + ) + + +def _update_static_key(x, conv_state, weight, bias, activation, cache_seqlens): + input_3d = x.unsqueeze(-1) if x.ndim == 2 else x + key = _conv_static_key(input_3d, weight, bias, activation, conv_state) + key.update( + tokens=input_3d.shape[-1], + state_len=conv_state.shape[-1], + circular=cache_seqlens is not None, + ) + return key + + +@triton.jit +def _causal_conv1d_fwd_kernel( + x_ptr, + weight_ptr, + bias_ptr, + conv_states_ptr, + query_start_loc_ptr, + cache_indices_ptr, + has_initial_state_ptr, + out_ptr, + num_channels, + stride_x_channel, + stride_x_token, + stride_weight_channel, + stride_weight_tap, + stride_state_cache, + stride_state_channel, + stride_state_token, + stride_out_channel, + stride_out_token, + CONV_WIDTH: tl.constexpr, + HAS_BIAS: tl.constexpr, + APPLY_SILU: tl.constexpr, + PAD_SLOT_ID: tl.constexpr, + BLOCK_TOKENS: tl.constexpr, + BLOCK_CHANNELS: tl.constexpr, +): + HISTORY_LEN: tl.constexpr = CONV_WIDTH - 1 + tl.static_assert(CONV_WIDTH >= 2 and CONV_WIDTH <= 4, "Supported convolution widths are 2, 3 and 4") + tl.static_assert(BLOCK_TOKENS >= HISTORY_LEN, "Only the first token block may read initial state") + + sequence_index = tl.program_id(0) + token_block_index = tl.program_id(1) + channel_indices = tl.program_id(2) * BLOCK_CHANNELS + tl.arange(0, BLOCK_CHANNELS) + valid_channels = channel_indices < num_channels + sequence_start = tl.load(query_start_loc_ptr + sequence_index) + sequence_length = tl.load(query_start_loc_ptr + sequence_index + 1) - sequence_start + if token_block_index * BLOCK_TOKENS >= sequence_length: + return + + token_indices = token_block_index * BLOCK_TOKENS + tl.arange(0, BLOCK_TOKENS) + valid_output = (token_indices[:, None] < sequence_length) & valid_channels[None, :] + input_base = x_ptr + sequence_start * stride_x_token + channel_indices * stride_x_channel + output_base = out_ptr + sequence_start * stride_out_token + channel_indices * stride_out_channel + output_ptrs = output_base[None, :] + token_indices[:, None] * stride_out_token + cache_index = tl.load(cache_indices_ptr + sequence_index) + if cache_index == PAD_SLOT_ID: + # Padding requests retain their input, as in the CUDA implementation. + padded_input = tl.load(input_base[None, :] + token_indices[:, None] * stride_x_token, valid_output, other=0) + tl.store(output_ptrs, padded_input, valid_output) + return + + state_base = conv_states_ptr + cache_index * stride_state_cache + channel_indices * stride_state_channel + weight_base = weight_ptr + channel_indices * stride_weight_channel + use_initial_state = tl.load(has_initial_state_ptr + sequence_index) + accumulator = tl.full((BLOCK_TOKENS, BLOCK_CHANNELS), 0, tl.float32) + if HAS_BIAS: + bias_values = tl.load(bias_ptr + channel_indices, valid_channels, other=0).to(tl.float32) + accumulator += bias_values[None, :] + + # Negative token positions come from cached history, or zero for a new request. + for tap in tl.static_range(CONV_WIDTH): + source_tokens = token_indices[:, None] + tap - HISTORY_LEN + input_values = tl.load( + input_base[None, :] + source_tokens * stride_x_token, valid_output & (source_tokens >= 0), other=0 + ).to(tl.float32) + history_values = tl.load( + state_base[None, :] + (source_tokens + HISTORY_LEN) * stride_state_token, + valid_output & (source_tokens < 0) & use_initial_state, + other=0, + ).to(tl.float32) + window_values = tl.where(source_tokens < 0, history_values, input_values) + conv_weight = tl.load(weight_base + tap * stride_weight_tap, valid_channels, other=0).to(tl.float32) + accumulator = tl.fma(window_values, conv_weight[None, :], accumulator) + if APPLY_SILU: + accumulator = accumulator / (1.0 + tl.exp(-accumulator)) + tl.store(output_ptrs, accumulator, valid_output) + + if token_block_index == 0: + # Only the first token block reads/writes initial state. Later blocks + # obtain their history from input, avoiding cross-CTA state races. + state_offsets = tl.arange(0, triton.next_power_of_2(HISTORY_LEN))[:, None] + tail_tokens = sequence_length - HISTORY_LEN + state_offsets + valid_state = (state_offsets < HISTORY_LEN) & valid_channels[None, :] + tail_inputs = tl.load( + input_base[None, :] + tail_tokens * stride_x_token, valid_state & (tail_tokens >= 0), other=0 + ) + tail_history = tl.load( + state_base[None, :] + (tail_tokens + HISTORY_LEN) * stride_state_token, + valid_state & (tail_tokens < 0) & use_initial_state, + other=0, + ) + updated_state = tl.where(tail_tokens >= 0, tail_inputs, tail_history) + # Finish all prefix reads before overwriting their state slots, + # including requests shorter than the history window. + tl.debug_barrier() + tl.store(state_base[None, :] + state_offsets * stride_state_token, updated_state, valid_state) + + +@autotune( + kernel_name="causal_conv1d_prefill:v1", + configs_gen_func=_prefill_configs, + static_key_func=_prefill_static_key, + run_key_func=_prefill_run_key, + run_key_distance_func=_prefill_key_distance, + mutates_args=["conv_states"], +) +def causal_conv1d_fn( + x: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor] = None, + *, + query_start_loc: torch.Tensor, + cache_indices: torch.Tensor, + has_initial_state: torch.Tensor, + conv_states: torch.Tensor, + activation: Optional[str] = "silu", + pad_slot_id: int = -1, + max_seqlen: Optional[int] = None, + run_config: Optional[dict] = None, + **kwargs, +): + """Convolve packed [channels, total_tokens] inputs and update mapped states. + + Request boundaries, cache indices and initial-state flags are required, + matching the GDN prefill caller. States have shape [slots, channels, width - 1] + and may be strided MTP slices. Metadata tensors are contiguous. + max_seqlen bounds the longest query; total tokens are a safe default that + avoids reading lengths back from the GPU. Empty/padding requests keep state. + """ + if activation not in (None, "silu", "swish"): + raise NotImplementedError(f"Unsupported activation: {activation}") + num_channels, total_tokens = x.shape + batch_size = query_start_loc.numel() - 1 + out = torch.empty_like(x) + if batch_size == 0 or x.numel() == 0: + return out + + bias = bias.contiguous() if bias is not None else None + max_query_len = total_tokens if max_seqlen is None else max_seqlen + if run_config is None: + run_config = {"BLOCK_TOKENS": 16, "BLOCK_CHANNELS": 128, "num_warps": 4} + block_tokens = run_config["BLOCK_TOKENS"] + block_channels = run_config["BLOCK_CHANNELS"] + grid = (batch_size, triton.cdiv(max_query_len, block_tokens), triton.cdiv(num_channels, block_channels)) + _causal_conv1d_fwd_kernel[grid]( + x, + weight, + bias, + conv_states, + query_start_loc, + cache_indices, + has_initial_state, + out, + num_channels, + *x.stride(), + *weight.stride(), + *conv_states.stride(), + *out.stride(), + CONV_WIDTH=weight.shape[1], + HAS_BIAS=bias is not None, + APPLY_SILU=activation is not None, + PAD_SLOT_ID=pad_slot_id, + BLOCK_TOKENS=block_tokens, + BLOCK_CHANNELS=block_channels, + num_warps=run_config["num_warps"], + ) + return out + + +@triton.jit +def _causal_conv1d_update_kernel( + x_ptr, + weight_ptr, + bias_ptr, + conv_state_ptr, + cache_seqlens_ptr, + conv_state_indices_ptr, + out_ptr, + num_channels, + stride_x_batch, + stride_x_channel, + stride_x_token, + stride_weight_channel, + stride_weight_tap, + stride_state_cache, + stride_state_channel, + stride_state_token, + stride_out_batch, + stride_out_channel, + stride_out_token, + CONV_WIDTH: tl.constexpr, + NUM_TOKENS: tl.constexpr, + STATE_LEN: tl.constexpr, + HAS_INDICES: tl.constexpr, + HAS_BIAS: tl.constexpr, + CIRCULAR_BUFFER: tl.constexpr, + APPLY_SILU: tl.constexpr, + PAD_SLOT_ID: tl.constexpr, + BLOCK_CHANNELS: tl.constexpr, +): + tl.static_assert(CONV_WIDTH >= 2 and CONV_WIDTH <= 4, "Supported convolution widths are 2, 3 and 4") + tl.static_assert(STATE_LEN >= CONV_WIDTH - 1, "State must hold the full convolution history") + + batch_index = tl.program_id(0) + channel_indices = tl.program_id(1) * BLOCK_CHANNELS + tl.arange(0, BLOCK_CHANNELS) + valid_channels = channel_indices < num_channels + input_base = x_ptr + batch_index * stride_x_batch + channel_indices * stride_x_channel + output_base = out_ptr + batch_index * stride_out_batch + channel_indices * stride_out_channel + cache_index = tl.load(conv_state_indices_ptr + batch_index) if HAS_INDICES else batch_index + if cache_index == PAD_SLOT_ID: + for token_index in range(NUM_TOKENS): + padded_input = tl.load(input_base + token_index * stride_x_token, valid_channels, other=0) + tl.store(output_base + token_index * stride_out_token, padded_input, valid_channels) + return + state_base = conv_state_ptr + cache_index * stride_state_cache + channel_indices * stride_state_channel + weight_base = weight_ptr + channel_indices * stride_weight_channel + state_cursor = tl.load(cache_seqlens_ptr + batch_index) % STATE_LEN if CIRCULAR_BUFFER else STATE_LEN + history_start = state_cursor + STATE_LEN - (CONV_WIDTH - 1) + + # Keep the short convolution window and weights in registers across tokens. + history_0 = tl.load( + state_base + (history_start % STATE_LEN) * stride_state_token, + valid_channels, + other=0, + ).to(tl.float32) + weight_0 = tl.load(weight_base, valid_channels, other=0).to(tl.float32) + weight_1 = tl.load(weight_base + stride_weight_tap, valid_channels, other=0).to(tl.float32) + if CONV_WIDTH >= 3: + history_1 = tl.load( + state_base + ((history_start + 1) % STATE_LEN) * stride_state_token, + valid_channels, + other=0, + ).to(tl.float32) + weight_2 = tl.load(weight_base + 2 * stride_weight_tap, valid_channels, other=0).to(tl.float32) + if CONV_WIDTH == 4: + history_2 = tl.load( + state_base + ((history_start + 2) % STATE_LEN) * stride_state_token, valid_channels, other=0 + ).to(tl.float32) + weight_3 = tl.load(weight_base + 3 * stride_weight_tap, valid_channels, other=0).to(tl.float32) + bias_values = tl.full((BLOCK_CHANNELS,), 0, tl.float32) + if HAS_BIAS: + bias_values = tl.load(bias_ptr + channel_indices, valid_channels, other=0).to(tl.float32) + + for token_index in range(NUM_TOKENS): + input_values = tl.load(input_base + token_index * stride_x_token, valid_channels, other=0).to(tl.float32) + accumulator = tl.fma(weight_0, history_0, bias_values) + if CONV_WIDTH == 2: + accumulator = tl.fma(weight_1, input_values, accumulator) + history_0 = input_values + elif CONV_WIDTH == 3: + accumulator = tl.fma(weight_1, history_1, accumulator) + accumulator = tl.fma(weight_2, input_values, accumulator) + history_0, history_1 = history_1, input_values + else: + accumulator = tl.fma(weight_1, history_1, accumulator) + accumulator = tl.fma(weight_2, history_2, accumulator) + accumulator = tl.fma(weight_3, input_values, accumulator) + history_0, history_1, history_2 = history_1, history_2, input_values + if APPLY_SILU: + accumulator = accumulator / (1.0 + tl.exp(-accumulator)) + tl.store(output_base + token_index * stride_out_token, accumulator, valid_channels) + if CIRCULAR_BUFFER: + tl.store( + state_base + ((state_cursor + token_index) % STATE_LEN) * stride_state_token, + input_values, + valid_channels, + ) + if not CIRCULAR_BUFFER: + # Keep the last STATE_LEN tokens of [old state, new input]. + state_offsets = tl.arange(0, triton.next_power_of_2(STATE_LEN))[:, None] + valid_state = (state_offsets < STATE_LEN) & valid_channels[None, :] + source_offsets = state_offsets + NUM_TOKENS + old_history = tl.load( + state_base[None, :] + source_offsets * stride_state_token, + valid_state & (source_offsets < STATE_LEN), + other=0, + ) + new_inputs = tl.load( + input_base[None, :] + (source_offsets - STATE_LEN) * stride_x_token, + valid_state & (source_offsets >= STATE_LEN), + other=0, + ) + updated_state = tl.where(source_offsets < STATE_LEN, old_history, new_inputs) + tl.debug_barrier() + tl.store(state_base[None, :] + state_offsets * stride_state_token, updated_state, valid_state) + + +@autotune( + kernel_name="causal_conv1d_update:v1", + kernel_type=AutotuneKernelType.DECODE_ATTENTION, + configs_gen_func=_update_configs, + static_key_func=_update_static_key, + run_key_func=lambda x: x.shape[0], + mutates_args=["conv_state"], +) +def causal_conv1d_update( + x: torch.Tensor, + conv_state: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor] = None, + activation: Optional[str] = None, + cache_seqlens: Optional[torch.Tensor] = None, + conv_state_indices: Optional[torch.Tensor] = None, + pad_slot_id: int = -1, + run_config: Optional[dict] = None, +): + """Decode [batch, channels] or [batch, channels, tokens], updating linear/ring states. + + State shape is [slots, channels, state_len], with state_len >= width - 1. + Cache indices and sequence lengths, when supplied, are contiguous batch vectors. + """ + if activation not in (None, "silu", "swish"): + raise NotImplementedError(f"Unsupported activation: {activation}") + out = torch.empty_like(x) + if x.numel() == 0: + return out + + input_3d = x.unsqueeze(-1) if x.ndim == 2 else x + output_3d = out.unsqueeze(-1) if x.ndim == 2 else out + batch_size, num_channels, num_tokens = input_3d.shape + bias = bias.contiguous() if bias is not None else None + if run_config is None: + run_config = {"BLOCK_CHANNELS": 128, "num_warps": 4} + block_channels = run_config["BLOCK_CHANNELS"] + grid = (batch_size, triton.cdiv(num_channels, block_channels)) + _causal_conv1d_update_kernel[grid]( + input_3d, + weight, + bias, + conv_state, + cache_seqlens, + conv_state_indices, + output_3d, + num_channels, + *input_3d.stride(), + *weight.stride(), + *conv_state.stride(), + *output_3d.stride(), + CONV_WIDTH=weight.shape[1], + NUM_TOKENS=num_tokens, + STATE_LEN=conv_state.shape[2], + HAS_INDICES=conv_state_indices is not None, + HAS_BIAS=bias is not None, + CIRCULAR_BUFFER=cache_seqlens is not None, + APPLY_SILU=activation is not None, + PAD_SLOT_ID=pad_slot_id, + BLOCK_CHANNELS=block_channels, + num_warps=run_config["num_warps"], + ) + return out diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_causal_conv1d.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_causal_conv1d.py new file mode 100644 index 0000000000..3166cc1692 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_causal_conv1d.py @@ -0,0 +1,378 @@ +import builtins +from collections import defaultdict + +import pytest +import torch +import torch.nn.functional as F + +from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import ( + causal_conv1d_fn as dispatch_prefill, + causal_conv1d_update as dispatch_update, +) +from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_triton import ( + _prefill_configs, + _update_configs, + causal_conv1d_fn, + causal_conv1d_update, +) +from lightllm.common.triton_utils import autotuner + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + + +def _prefill_reference(x, weight, bias, states, lengths, indices, initial, activation): + """Independent grouped-convolution reference, including cached prefixes.""" + width = weight.shape[-1] + output = x.clone() + start = 0 + for length, slot, use_initial in zip(lengths, indices, initial): + if length and slot != -1: + prefix = states[slot] if use_initial else torch.zeros_like(states[slot]) + extended = torch.cat((prefix, x[:, start : start + length]), dim=-1) + with torch.backends.cudnn.flags(allow_tf32=False): + convolved = F.conv1d( + extended.float().unsqueeze(0), + weight.float().unsqueeze(1), + bias.float() if bias is not None else None, + groups=x.shape[0], + ).squeeze(0) + output[:, start : start + length] = F.silu(convolved) if activation else convolved + states[slot] = extended[:, -(width - 1) :] + start += length + return output + + +def _assert_output(actual, expected): + if actual.dtype == torch.float32: + torch.testing.assert_close(actual, expected, rtol=2e-5, atol=2e-6) + else: + torch.testing.assert_close(actual, expected, rtol=1e-2, atol=2e-3) + + +def _check_prefill_ragged_prefix(dtype, width, layout, run_config=None): + torch.manual_seed(42) + dim = 131 # Exercise masked channels and unaligned strides. + lengths = [1, 0, 3, 19, 67] + tokens = sum(lengths) + if layout == "projection": + x = torch.randn(tokens, dim + 32, device="cuda", dtype=dtype)[:, :dim].T + else: + x = torch.randn(dim, tokens * 2, device="cuda", dtype=dtype)[:, ::2] + weight = torch.randn(dim, width * 2, device="cuda", dtype=dtype)[:, ::2] * 0.2 + bias = torch.randn(dim * 2, device="cuda", dtype=dtype)[::2] if layout == "projection" else None + activation = "silu" if bias is not None else None + state_storage = torch.randn(6, dim, width + 2, device="cuda", dtype=dtype) + original_storage = state_storage.clone() + states = state_storage[..., : width - 1] + ref_states = states.clone() + indices = [3, 1, -1, 4, 0] + initial = [True, False, False, True, False] + starts = torch.tensor([0, 1, 1, 4, 23, 90], device="cuda", dtype=torch.int32) + expected = _prefill_reference(x, weight, bias, ref_states, lengths, indices, initial, activation) + actual = causal_conv1d_fn( + x, + weight, + bias, + query_start_loc=starts, + cache_indices=torch.tensor(indices, device="cuda"), + has_initial_state=torch.tensor(initial, device="cuda"), + conv_states=states, + activation=activation, + max_seqlen=max(lengths), + run_config=run_config, + ) + _assert_output(actual, expected) + torch.testing.assert_close(states, ref_states, rtol=0, atol=0) + torch.testing.assert_close(state_storage[..., width - 1 :], original_storage[..., width - 1 :], rtol=0, atol=0) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("width", [2, 3, 4]) +@pytest.mark.parametrize("layout", ["projection", "sequence"]) +def test_prefill_ragged_prefix_and_strided_state(dtype, width, layout): + _check_prefill_ragged_prefix(dtype, width, layout) + + +@pytest.mark.parametrize("run_config", _prefill_configs()) +def test_prefill_autotune_configurations(run_config): + _check_prefill_ragged_prefix(torch.bfloat16, 4, "projection", run_config) + + +def test_prefill_new_requests_ignore_stale_state(): + x = torch.randn(2, 65, 33, device="cuda") + weight = torch.randn(65, 4, device="cuda") + states = torch.randn(2, 65, 3, device="cuda") + with torch.backends.cudnn.flags(allow_tf32=False): + expected = F.silu(F.conv1d(F.pad(x, (3, 0)), weight[:, None], groups=65)) + actual = causal_conv1d_fn( + x.transpose(0, 1).reshape(65, -1), + weight, + query_start_loc=torch.tensor([0, 33, 66], device="cuda", dtype=torch.int32), + cache_indices=torch.arange(2, device="cuda", dtype=torch.int32), + has_initial_state=torch.zeros(2, device="cuda", dtype=torch.bool), + conv_states=states, + activation="swish", + ) + _assert_output(actual, expected.transpose(0, 1).reshape(65, -1)) + torch.testing.assert_close(states, x[..., -3:], rtol=0, atol=0) + + +def _update_reference(x, states, weight, bias, indices, cache_lengths, activation): + output = x.clone() + input_3d = x.unsqueeze(-1) if x.ndim == 2 else x + output_3d = output.unsqueeze(-1) if output.ndim == 2 else output + width, state_len = weight.shape[-1], states.shape[-1] + for batch, slot in enumerate(indices): + if slot == -1: + continue + for token in range(input_3d.shape[-1]): + if cache_lengths is None: + history = states[slot, :, -(width - 1) :].clone() + states[slot] = torch.cat((states[slot, :, 1:], input_3d[batch, :, token : token + 1]), dim=-1) + else: + cursor = cache_lengths[batch] + token + history = torch.stack( + [states[slot, :, (cursor - width + 1 + j) % state_len] for j in range(width - 1)], dim=-1 + ) + states[slot, :, cursor % state_len] = input_3d[batch, :, token] + values = torch.cat((history, input_3d[batch, :, token : token + 1]), dim=-1).double() + # Round after each FP32 FMA, independently of Triton's implementation. + accumulator = bias.float().clone() if bias is not None else torch.zeros(input_3d.shape[1], device=x.device) + for tap in range(width): + accumulator = (accumulator.double() + values[:, tap] * weight[:, tap].double()).float() + output_3d[batch, :, token] = F.silu(accumulator) if activation else accumulator + return output + + +def _check_update_buffers(dtype, width, circular, tokens, run_config=None): + torch.manual_seed(17) + dim = 131 + x = torch.randn(3, dim + 16, tokens * 2, device="cuda", dtype=dtype)[:, :dim, ::2] + if tokens == 1: + x = x.squeeze(-1) + weight = torch.randn(dim, width, device="cuda", dtype=dtype) * 0.2 + bias = torch.randn(dim * 2, device="cuda", dtype=dtype)[::2] + state_len = width + 1 + storage = torch.randn(5, dim, state_len + 3, device="cuda", dtype=dtype) + original = storage.clone() + states = storage[..., :state_len] + ref_states = states.clone() + indices = [2, -1, 0] + cache_lengths = [0, 3, 13] if circular else None + activation = "silu" if tokens == 1 else None + expected = _update_reference(x, ref_states, weight, bias, indices, cache_lengths, activation) + actual = causal_conv1d_update( + x, + states, + weight, + bias, + activation, + torch.tensor(cache_lengths, device="cuda", dtype=torch.int32) if circular else None, + torch.tensor(indices, device="cuda", dtype=torch.int32), + run_config=run_config, + ) + _assert_output(actual, expected) + torch.testing.assert_close(states, ref_states, rtol=0, atol=0) + torch.testing.assert_close(storage[..., state_len:], original[..., state_len:], rtol=0, atol=0) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("width", [2, 3, 4]) +@pytest.mark.parametrize("circular,tokens", [(False, 1), (False, 7), (True, 1), (True, 7)]) +def test_update_linear_and_ring_buffers(dtype, width, circular, tokens): + _check_update_buffers(dtype, width, circular, tokens) + + +@pytest.mark.parametrize("run_config", _update_configs()) +@pytest.mark.parametrize("circular", [False, True]) +def test_update_autotune_configurations(run_config, circular): + _check_update_buffers(torch.bfloat16, 4, circular, 7, run_config) + + +@pytest.mark.parametrize("mode", ["prefill", "linear", "circular"]) +def test_autotune_and_cached_config_update_state_once(monkeypatch, tmp_path, mode): + torch.manual_seed(13) + tuner = causal_conv1d_fn if mode == "prefill" else causal_conv1d_update + monkeypatch.setattr(autotuner, "get_triton_autotune_level", lambda: autotuner.AutotuneLevel.ADAPTIVE_AUTOTUNE) + monkeypatch.setattr(tuner, "_cache_dir", str(tmp_path), raising=False) + monkeypatch.setattr(tuner, "cached_configs", {}) + monkeypatch.setattr(tuner, "fast_match_configs", defaultdict(dict)) + monkeypatch.setattr(tuner, "warmuped_configs_set", set()) + configs = tuner.configs_gen_func() + monkeypatch.setattr(tuner, "configs_gen_func", lambda: [configs[0], configs[-1]]) + + dim, state_len = 65, 3 if mode == "prefill" else 5 + storage = torch.randn(4, dim, state_len + 2, device="cuda", dtype=torch.bfloat16) + original = storage.clone() + states = storage[..., :state_len] + reference_states = states.clone() + weight = torch.randn(dim, 4, device="cuda", dtype=states.dtype) * 0.2 + indices = torch.tensor([3, 1], device="cuda", dtype=torch.int32) + if mode == "prefill": + x = torch.randn(3, dim + 16, device="cuda", dtype=states.dtype)[:, :dim].T + starts = torch.tensor([0, 1, 3], device="cuda", dtype=torch.int32) + initial = torch.tensor([True, True], device="cuda") + + def run(): + return tuner( + x, + weight, + query_start_loc=starts, + cache_indices=indices, + has_initial_state=initial, + conv_states=states, + max_seqlen=2, + ) + + def reference(): + return _prefill_reference(x, weight, None, reference_states, [1, 2], [3, 1], [True, True], "silu") + + else: + x = torch.randn(2, dim, device="cuda", dtype=states.dtype) + lengths = [0, 13] if mode == "circular" else None + cache_lengths = torch.tensor(lengths, device="cuda", dtype=torch.int32) if lengths is not None else None + + def run(): + return tuner(x, states, weight, activation="silu", cache_seqlens=cache_lengths, conv_state_indices=indices) + + def reference(): + return _update_reference(x, reference_states, weight, None, [3, 1], lengths, "silu") + + clone_mutated_args = tuner._mutate_args_clone + scratch_states = [] + + def check_clones(args, kwargs): + new_args, new_kwargs, original_states, cloned_states = clone_mutated_args(args, kwargs) + assert len(cloned_states) == 1 and original_states[0] is states + assert cloned_states[0].data_ptr() != states.data_ptr() + scratch_states.append(cloned_states[0]) + return new_args, new_kwargs, original_states, cloned_states + + monkeypatch.setattr(tuner, "_mutate_args_clone", check_clones) + expected = reference() + with autotuner.Autotuner.autotune_warmup(tuner.kernel_type): + actual = run() + assert len(scratch_states) == 2 # One protected copy per candidate. + assert len(list(tmp_path.glob("*.json"))) == 1 + assert all(config is not None for cached in tuner.cached_configs.values() for config in cached.values()) + _assert_output(actual, expected) + torch.testing.assert_close(states, reference_states, rtol=0, atol=0) + + # Historical-config warmup must also use protected state, including when + # the first cache load is captured in a CUDA graph. + monkeypatch.setattr(autotuner, "get_triton_autotune_level", lambda: autotuner.AutotuneLevel.USE_AUTOTUNE_HIS_CONFIG) + monkeypatch.setattr(tuner, "cached_configs", {}) + monkeypatch.setattr(tuner, "fast_match_configs", defaultdict(dict)) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = run() + expected = reference() + graph.replay() + _assert_output(output, expected) + torch.testing.assert_close(states, reference_states, rtol=0, atol=0) + torch.testing.assert_close(storage[..., state_len:], original[..., state_len:], rtol=0, atol=0) + assert len(scratch_states) > 2 + + +def test_prefill_graph_replays_with_changed_lengths_and_prefixes(): + torch.manual_seed(1) + x = torch.randn(25, 96, device="cuda", dtype=torch.bfloat16).T + weight = torch.randn(96, 4, device="cuda", dtype=x.dtype) * 0.2 + original = torch.randn(2, 96, 3, device="cuda", dtype=x.dtype) + states = original.clone() + starts = torch.tensor([0, 1, 25], dtype=torch.int32, device="cuda") + initial = torch.tensor([True, False], device="cuda") + indices = torch.arange(2, device="cuda", dtype=torch.int32) + + def run(): + return causal_conv1d_fn( + x, + weight, + query_start_loc=starts, + cache_indices=indices, + has_initial_state=initial, + conv_states=states, + max_seqlen=25, + ) + + run() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = run() + for lengths, flags in [([1, 24], [True, False]), ([25, 0], [False, True]), ([2, 23], [True, True])]: + starts.copy_(torch.tensor([0, lengths[0], 25], dtype=starts.dtype, device="cuda")) + initial.copy_(torch.tensor(flags, device="cuda")) + states.copy_(original) + reference_states = original.clone() + expected = _prefill_reference(x, weight, None, reference_states, lengths, [0, 1], flags, "silu") + graph.replay() + _assert_output(output, expected) + torch.testing.assert_close(states, reference_states, rtol=0, atol=0) + + +def test_dispatch_without_sgl_kernel_or_vllm(monkeypatch): + import_fn = builtins.__import__ + + def without_optional_kernels(name, *args, **kwargs): + if name.split(".")[0] in ("sgl_kernel", "vllm"): + raise ModuleNotFoundError(name) + return import_fn(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", without_optional_kernels) + x = torch.randn(13, 65, dtype=torch.bfloat16, device="cuda").T + weight = torch.randn(65, 4, dtype=x.dtype, device="cuda") * 0.2 + states = torch.zeros(1, 65, 3, dtype=x.dtype, device="cuda") + ref_states = states.clone() + starts = torch.tensor([0, 13], dtype=torch.int32, device="cuda") + expected = _prefill_reference(x, weight, None, ref_states, [13], [0], [False], "silu") + _assert_output( + dispatch_prefill( + x, + weight, + query_start_loc=starts, + conv_states=states, + cache_indices=torch.zeros(1, device="cuda", dtype=torch.int32), + has_initial_state=torch.zeros(1, device="cuda", dtype=torch.bool), + ), + expected, + ) + token = torch.randn(1, 65, dtype=x.dtype, device="cuda") + expected = _update_reference(token, ref_states, weight, None, [0], None, "silu") + _assert_output(dispatch_update(token, states, weight, activation="silu"), expected) + torch.testing.assert_close(states, ref_states, rtol=0, atol=0) + + +def test_bfloat16_products_accumulate_in_float32(): + # BF16 operands whose product cannot be represented in BF16. With a unit + # residual the premature product rounding also changes the BF16 output. + x = torch.tensor([[1.0078125, 1.0]], device="cuda", dtype=torch.bfloat16) + weight = torch.tensor([[1.0078125, -1.0]], device="cuda", dtype=torch.bfloat16) + expected = torch.tensor([[-1.0078125, 0.01568603515625]], device="cuda", dtype=x.dtype) + actual = causal_conv1d_fn( + x, + weight, + activation=None, + query_start_loc=torch.tensor([0, 2], device="cuda", dtype=torch.int32), + cache_indices=torch.zeros(1, device="cuda", dtype=torch.int32), + has_initial_state=torch.zeros(1, device="cuda", dtype=torch.bool), + conv_states=torch.ones(1, 1, 1, device="cuda", dtype=x.dtype), + ) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +def test_empty_inputs_do_not_touch_state(): + state = torch.randn(2, 17, 3, device="cuda") + original = state.clone() + weight = torch.randn(17, 4, device="cuda") + output = causal_conv1d_fn( + torch.empty(17, 0, device="cuda"), + weight, + conv_states=state, + query_start_loc=torch.zeros(3, device="cuda", dtype=torch.int32), + cache_indices=torch.arange(2, device="cuda", dtype=torch.int32), + has_initial_state=torch.ones(2, device="cuda", dtype=torch.bool), + ) + assert output.numel() == 0 + assert causal_conv1d_update(torch.empty(2, 17, 0, device="cuda"), state, weight).numel() == 0 + torch.testing.assert_close(state, original, rtol=0, atol=0) From 962adfd0f304db09707cad64568ea52c2aeec823 Mon Sep 17 00:00:00 2001 From: sufubao Date: Thu, 12 Nov 2026 07:51:19 +0800 Subject: [PATCH 3/6] perf: use vLLM causal convolution kernels when available --- .../triton_kernel/linear_att/causal_conv1d.py | 134 ++++++++++++++---- 1 file changed, 106 insertions(+), 28 deletions(-) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d.py b/lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d.py index 31435fa744..e8582016b9 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d.py @@ -5,6 +5,45 @@ import torch +def _load_causal_conv1d_fwd(): + """Return the first available optimized prefill implementation. + + The official vLLM image ships its Triton causal-convolution kernel even + when ``sgl_kernel`` is not installed. Keep both optional imports lazy so + importing LightLLM does not make either package a hard dependency. + """ + try: + from sgl_kernel import causal_conv1d_fwd + + return causal_conv1d_fwd + except (ImportError, OSError): + try: + from vllm.model_executor.layers.mamba.ops.causal_conv1d import ( + causal_conv1d_fn as vllm_causal_conv1d_fn, + ) + + return vllm_causal_conv1d_fn + except (ImportError, OSError): + return None + + +def _load_causal_conv1d_update(): + """Return an optimized single-step convolution helper when available.""" + try: + from sgl_kernel import causal_conv1d_update + + return causal_conv1d_update + except (ImportError, OSError): + try: + from vllm.model_executor.layers.mamba.ops.causal_conv1d import ( + causal_conv1d_update, + ) + + return causal_conv1d_update + except (ImportError, OSError): + return None + + def causal_conv1d_fn( x: torch.Tensor, weight: torch.Tensor, @@ -35,9 +74,8 @@ def causal_conv1d_fn( """ if activation not in [None, "silu", "swish"]: raise NotImplementedError("activation must be None, silu, or swish") - try: - from sgl_kernel import causal_conv1d_fwd - except (ImportError, OSError): + optimized_causal_conv1d = _load_causal_conv1d_fwd() + if optimized_causal_conv1d is None: from .causal_conv1d_triton import causal_conv1d_fn as triton_causal_conv1d_fn return triton_causal_conv1d_fn( @@ -57,17 +95,33 @@ def causal_conv1d_fn( x = x.contiguous() bias = bias.contiguous() if bias is not None else None - causal_conv1d_fwd( - x, - weight, - bias, - conv_states, - query_start_loc, - cache_indices, - has_initial_state, - activation in ["silu", "swish"], - pad_slot_id, - ) + if optimized_causal_conv1d.__module__.startswith("sgl_kernel"): + optimized_causal_conv1d( + x, + weight, + bias, + conv_states, + query_start_loc, + cache_indices, + has_initial_state, + activation in ["silu", "swish"], + pad_slot_id, + ) + else: + # vLLM's implementation has the same packed [channels, tokens] and + # state-cache contract, but exposes it as a functional helper. + return optimized_causal_conv1d( + x, + weight, + bias, + conv_states, + query_start_loc, + cache_indices=cache_indices, + has_initial_state=has_initial_state, + activation=activation, + pad_slot_id=pad_slot_id, + null_block_id=None, + ) return x @@ -98,9 +152,25 @@ def causal_conv1d_update( """ if activation not in [None, "silu", "swish"]: raise NotImplementedError(f"activation must be None, silu, or swish, actual: {activation}") - try: - from sgl_kernel import causal_conv1d_update as causal_conv1d_update_kernel - except (ImportError, OSError): + optimized_causal_conv1d_update = _load_causal_conv1d_update() + if optimized_causal_conv1d_update is None: + from .causal_conv1d_triton import causal_conv1d_update as triton_causal_conv1d_update + + return triton_causal_conv1d_update( + x, + conv_state, + weight, + bias, + activation=activation, + cache_seqlens=cache_seqlens, + conv_state_indices=conv_state_indices, + pad_slot_id=pad_slot_id, + ) + + if cache_seqlens is not None: + # The vLLM/SGL interfaces use a linear speculative cursor rather than + # LightLLM's circular cache cursor. Keep the local implementation for + # this less common state layout. from .causal_conv1d_triton import causal_conv1d_update as triton_causal_conv1d_update return triton_causal_conv1d_update( @@ -115,19 +185,27 @@ def causal_conv1d_update( ) activation_val = activation in ["silu", "swish"] - unsqueeze = x.dim() == 2 - if unsqueeze: - x = x.unsqueeze(-1) - causal_conv1d_update_kernel( + if optimized_causal_conv1d_update.__module__.startswith("sgl_kernel"): + unsqueeze = x.dim() == 2 + if unsqueeze: + x = x.unsqueeze(-1) + result = optimized_causal_conv1d_update( + x, + conv_state, + weight, + bias, + activation_val, + None, + conv_state_indices, + pad_slot_id=pad_slot_id, + ) + return result.squeeze(-1) if unsqueeze else result + return optimized_causal_conv1d_update( x, conv_state, weight, bias, - activation_val, - cache_seqlens, - conv_state_indices, - pad_slot_id, + activation=activation, + conv_state_indices=conv_state_indices, + null_block_id=pad_slot_id, ) - if unsqueeze: - x = x.squeeze(-1) - return x From bf6490df9266e93951f457bc82adcdc54d3c0704 Mon Sep 17 00:00:00 2001 From: sufubao Date: Thu, 12 Nov 2026 12:14:51 +0800 Subject: [PATCH 4/6] fix: return in-place causal convolution output --- .../triton_kernel/linear_att/causal_conv1d.py | 16 ++---------- .../linear_att/test_causal_conv1d.py | 26 +++++++++++++++++++ 2 files changed, 28 insertions(+), 14 deletions(-) diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d.py b/lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d.py index e8582016b9..0055ea6d93 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d.py @@ -6,12 +6,6 @@ def _load_causal_conv1d_fwd(): - """Return the first available optimized prefill implementation. - - The official vLLM image ships its Triton causal-convolution kernel even - when ``sgl_kernel`` is not installed. Keep both optional imports lazy so - importing LightLLM does not make either package a hard dependency. - """ try: from sgl_kernel import causal_conv1d_fwd @@ -28,7 +22,6 @@ def _load_causal_conv1d_fwd(): def _load_causal_conv1d_update(): - """Return an optimized single-step convolution helper when available.""" try: from sgl_kernel import causal_conv1d_update @@ -108,8 +101,6 @@ def causal_conv1d_fn( pad_slot_id, ) else: - # vLLM's implementation has the same packed [channels, tokens] and - # state-cache contract, but exposes it as a functional helper. return optimized_causal_conv1d( x, weight, @@ -168,9 +159,6 @@ def causal_conv1d_update( ) if cache_seqlens is not None: - # The vLLM/SGL interfaces use a linear speculative cursor rather than - # LightLLM's circular cache cursor. Keep the local implementation for - # this less common state layout. from .causal_conv1d_triton import causal_conv1d_update as triton_causal_conv1d_update return triton_causal_conv1d_update( @@ -189,7 +177,7 @@ def causal_conv1d_update( unsqueeze = x.dim() == 2 if unsqueeze: x = x.unsqueeze(-1) - result = optimized_causal_conv1d_update( + optimized_causal_conv1d_update( x, conv_state, weight, @@ -199,7 +187,7 @@ def causal_conv1d_update( conv_state_indices, pad_slot_id=pad_slot_id, ) - return result.squeeze(-1) if unsqueeze else result + return x.squeeze(-1) if unsqueeze else x return optimized_causal_conv1d_update( x, conv_state, diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_causal_conv1d.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_causal_conv1d.py index 3166cc1692..1bbc9627b0 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_causal_conv1d.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_causal_conv1d.py @@ -343,6 +343,32 @@ def without_optional_kernels(name, *args, **kwargs): torch.testing.assert_close(states, ref_states, rtol=0, atol=0) +@pytest.mark.parametrize("tokens", [1, 3]) +def test_sgl_update(tokens): + pytest.importorskip("sgl_kernel") + torch.manual_seed(29) + x = torch.randn(3, 64, tokens, dtype=torch.bfloat16, device="cuda") + if tokens == 1: + x = x.squeeze(-1) + states = torch.randn(5, 64, 3, dtype=x.dtype, device="cuda") + ref_states = states.clone() + weight = torch.randn(64, 4, dtype=x.dtype, device="cuda") * 0.2 + bias = torch.randn(64, dtype=x.dtype, device="cuda") + indices = [2, -1, 0] + expected = _update_reference(x, ref_states, weight, bias, indices, None, "silu") + actual = dispatch_update( + x, + states, + weight, + bias, + activation="silu", + conv_state_indices=torch.tensor(indices, dtype=torch.int32, device="cuda"), + ) + assert actual.shape == x.shape + _assert_output(actual, expected) + torch.testing.assert_close(states, ref_states, rtol=0, atol=0) + + def test_bfloat16_products_accumulate_in_float32(): # BF16 operands whose product cannot be represented in BF16. With a unit # residual the premature product rounding also changes the BF16 output. From 66465e199d675d81740c3c55f7c4997385490ccb Mon Sep 17 00:00:00 2001 From: sufubao Date: Thu, 12 Nov 2026 13:58:53 +0800 Subject: [PATCH 5/6] test: make AgentX control replay repeatable --- test/benchmark/qwen38/README.agentx.md | 165 +++++++++++ test/benchmark/qwen38/audit_agentx.py | 158 +++++++++++ test/benchmark/qwen38/prepare_agentx.py | 190 +++++++++++++ test/benchmark/qwen38/run_agentx.py | 257 ++++++++++++++++++ .../benchmark/qwen38/test_agentx_protocol.py | 181 ++++++++++++ 5 files changed, 951 insertions(+) create mode 100644 test/benchmark/qwen38/README.agentx.md create mode 100644 test/benchmark/qwen38/audit_agentx.py create mode 100644 test/benchmark/qwen38/prepare_agentx.py create mode 100644 test/benchmark/qwen38/run_agentx.py create mode 100644 unit_tests/benchmark/qwen38/test_agentx_protocol.py diff --git a/test/benchmark/qwen38/README.agentx.md b/test/benchmark/qwen38/README.agentx.md new file mode 100644 index 0000000000..f880763462 --- /dev/null +++ b/test/benchmark/qwen38/README.agentx.md @@ -0,0 +1,165 @@ +# Repeatable AgentX control replay + +Use this protocol before comparing or tuning engines. A run that completes is not +automatically a valid performance result: all request audits and the three-trial +repeatability gate must pass. + +## What is measured + +The same pinned AgentX installation drives every engine. `prepare_agentx.py` uses +its Weka loader and synthetic prompt generator offline, then freezes full request +bodies and synthetic conversation history into `inputs.json`. Measurement uses +the installation's AIPerf `inputs_json` loader, which sends these payloads verbatim. +Live responses cannot change subsequent inputs. + +The default workload is 32 independent streams, eight sequential turns per stream, +256 requests in total, with concurrency 32. Output budgets come from the source +trace, with `ignore_eos=true`, temperature zero and seed 42. Each stream runs once +and drains completely. There is no duration cutoff, resampling, request clipping, +think delay or DAG join. Prefix reuse within and across the frozen streams remains +enabled. Tail drain is included in throughput, so concurrency falls as streams +finish; this is a fixed-work control, not a sustained-capacity measurement. + +This is **Weka-derived control replay**, not the official AgentX score. The +official scenario retains trace delays, branching and response-dependent history; +even the same seed does not make those future requests byte-identical. Neither +mode establishes real business quality or production capacity. Do not mix their +results in one ranking. + +## Frozen inputs + +- Tool: `agentx-harness==1.0.13`. The manifest records all installed dependency + versions and hashes every Python source file in `aiperf` and `agentx_cli.py`. + Use Python 3.11--3.13 and a dedicated environment. Install + `agentx-harness==1.0.13` before preparation; thereafter reuse the same environment + or restore the exact dependency versions in the manifest. A matching version + number alone is insufficient: the runner also requires the source hashes and + complete installed distribution list to match. +- Dataset: `semianalysisai/cc-traces-weka-062126-256k`, revision + `8fecd2fc56694469f758f0afbbb6335ad3043740`, `traces.jsonl` SHA-256 + `e39cd2ff3eba21d4a3664be51da743ac3d2149a1933898cafc7bfeac8147eeef`. +- Default model name: `Qwen3.8-Flash-Next`. The manifest hashes configuration, + tokenizer, chat template, acquisition metadata and weight index, and records + every weight shard's size. Weight contents are not rehashed by this tool; use + the same verified, read-only checkpoint mount for all engines. +- Every selected request is checked with the checkpoint's chat template against + the 262,144-token prompt-plus-output limit. An overflow fails preparation. +- `sources.json` maps each turn to its original Weka trace and request index. + Frozen histories are synthetic reconstructions, not real customer text. + +Prepare once and reuse that exact output directory for every engine. Never +regenerate a different workload for a slower engine. Preparation is deterministic; +verify a second independent generation has identical `inputs.json` bytes. + +```sh +exp -m "Freeze AgentX control; record source checkpoint and purpose" \ + /path/to/agentx-venv/bin/python -m test.benchmark.qwen38.prepare_agentx \ + --source /path/to/snapshots/8fecd2fc56694469f758f0afbbb6335ad3043740/traces.jsonl \ + --model-directory /path/to/verified/model \ + --output /path/to/new/frozen-workload +``` + +## Server and trial procedure + +Review previous `/data6/exp_logs` records (on this H100 host, +`/home/devsft/experiments/runs`) before each experiment. Launch the engine through +`exp` and record its source revision, image digest, complete command, checkpoint, +GPU assignment, precision, KV format, EP/TP, MTP and cache settings. Engine tuning +must use a separate screening set; freeze the final configuration before a suite. + +The runner is restricted to this task's isolated eight-GPU container, labelled +`qwen38.task=a3844d54-8a2b-4532-bf54-f657c9ca72a0`, and the localhost endpoint +`http://127.0.0.1:28081`. It refuses another running Docker container or a foreign +GPU compute process at phase boundaries. Use a dedicated host; stop image pulls, +image imports, builds, weight copies and other benchmark clients before starting. +Do not change clocks, power limits or engine configuration between trials. + +Enable and verify the engine's cache-reset API before measurement: + +| Engine | Required API | Required confirmation | Runtime status | +| --- | --- | --- | --- | +| vLLM 0.33 | `POST /reset_prefix_cache` | `{"success":true}` | Validated on H100 | +| LightLLM | `POST /flush_cache`, requires `--enable_rl` | `{"success":true}` | Adapter awaits runtime validation | +| SGLang | `POST /flush_cache` | HTTP 200, empty/null/true body | Adapter awaits runtime validation | + +The tested vLLM image requires `VLLM_SERVER_DEV_MODE=1` to expose the reset +endpoint. Bind this benchmark service to localhost. A missing endpoint, unexpected +response or failed reset aborts the suite; it is never silently skipped. + +Each of at least three trials executes: + +1. Confirm the same container and frozen tool/checkpoint/workload. +2. Reset the server cache and require success. +3. Take three idle GPU samples, ten seconds apart: utilization at most 5%, + temperature at most 60 C, all eight GPUs present, no foreign compute PID. +4. Replay and audit the entire fixed workload as warmup. +5. Reset cache again, repeat the idle checks, then replay and audit the same full + workload for measurement. + +Warmup is identical in every trial; no implicit first-run-only pressure warmup is +used. Only the measurement phase enters the reported three-trial statistics. + +```sh +exp -m "Qualify repeatability; record engine/image/flags, hardware, workload and expected question" \ + /path/to/agentx-venv/bin/python -m test.benchmark.qwen38.run_agentx \ + --engine vllm \ + --container owned-engine-container \ + --workload /path/to/frozen-workload \ + --model-directory /path/to/verified/model \ + --output /path/to/new/suite --repeats 3 +``` + +For SGLang or LightLLM, change `--engine` and the container only after independently +validating the reset and OpenAI streaming adapters. Keep the same model name, +tokenizer, full request file, load and measurement rules. Disclose precision and +speculation differences rather than treating them as equivalent accuracy. + +## Acceptance and evidence + +`audit_agentx.py` rejects missing/duplicate/changed payloads, incorrect conversation +turn mappings or order, overlapping turns in one stream, HTTP errors, tool errors, +cancellations, skipped contexts, missing SSE `[DONE]` or usage, non-length finish, +and output counts different from the frozen budget. It independently computes +output tokens / (last response end - first request start), requiring agreement +with the tool's throughput within 2%. Per-request prompt counts must also match +across repeated trials. + +The numerator is server `usage.completion_tokens`, including reasoning tokens. +AgentX `output_token_count` counts visible output separately and must not be used +as the total for reasoning models. Use `usage_completion_tokens` or +`output_sequence_length` when reconciling the summary with raw SSE usage. + +All trials must pass these predeclared limits: + +| Metric | Sample standard deviation / mean (CV) | (max - min) / mean | +| --- | --- | --- | +| Output token throughput | <= 5% | <= 10% | +| Request throughput | <= 5% | <= 10% | +| TTFT p50 | <= 10% | <= 20% | + +The limits are protocol choices, not guarantees provided by AgentX. Do not drop +the slowest trial, select the best trial or loosen limits after viewing results. +Failure exits nonzero and preserves evidence. Investigate it before ranking +engines. Passing establishes repeatability for this workload and configuration; +it does not establish repeatability across service restarts or hosts. Qualify a +fresh-launch suite separately before publishing broader claims. + +`suite.json` includes the manifest, protocol hashes, container image/command, +reset replies, idle samples, per-phase audits and final variance gate. Each trial +retains its config, exact command, PID, log, raw request/SSE exports, tool summary +and server metrics before/after. The suite also keeps `nvidia-smi -q`, five-second +GPU utilization/memory/power/temperature samples and server logs. Review telemetry +and host activity for interference; phase-boundary checks cannot prove that no +short-lived external process ran during measurement. + +Recompute the variance gate without model requests: + +```sh +python -m test.benchmark.qwen38.audit_agentx \ + --suite /path/to/suite/suite.json --workload /path/to/frozen-workload +``` + +For engine comparisons, report all trial values, mean, CV, range, lengths, errors, +TTFT p50/p95 and evidence locations. Alternate engine order in a later independent +validation block to check order effects. Older runs with different warmup, +cache or history conditions must remain excluded. diff --git a/test/benchmark/qwen38/audit_agentx.py b/test/benchmark/qwen38/audit_agentx.py new file mode 100644 index 0000000000..eb60fbd46d --- /dev/null +++ b/test/benchmark/qwen38/audit_agentx.py @@ -0,0 +1,158 @@ +import argparse +from collections import Counter +import json +import math +from pathlib import Path +import statistics + +from test.benchmark.qwen38.prepare_agentx import canonical, digest + + +def audit(directory, inputs): + expected = Counter(digest(canonical(p)) for s in inputs["data"] for p in s["payloads"]) + turns = {(s["session_id"], i): p for s in inputs["data"] for i, p in enumerate(s["payloads"])} + observed_turns, prompt_counts, intervals = set(), {}, {} + seen, prompt_tokens, output_tokens, starts, ends = Counter(), 0, 0, [], [] + merged = directory / "profile_export_raw.jsonl" + files = [merged] if merged.is_file() else sorted((directory / "raw_records").glob("*.jsonl")) + if not files: + raise ValueError("Missing raw request/SSE export shards") + for row in raw_rows(files): + metadata = row["metadata"] + if metadata["benchmark_phase"] != "profiling": + raise ValueError("Unexpected warmup records in a standalone measurement") + if ( + row["status"] != 200 + or row.get("error") + or metadata.get("was_cancelled") + or metadata.get("context_overflow_skip") + ): + raise ValueError("Failed, cancelled or skipped request") + payload = row["payload"] + key = (metadata["conversation_id"], metadata["turn_index"]) + if key not in turns or key in observed_turns or canonical(payload) != canonical(turns[key]): + raise ValueError("Missing, duplicate or changed session turn") + observed_turns.add(key) + seen[digest(canonical(payload))] += 1 + usage, done, reason = None, 0, None + for response in row["responses"]: + for packet in response.get("packets", []): + if packet.get("name") != "data": + continue + if done: + raise ValueError("SSE data after DONE") + if packet["value"] == "[DONE]": + done += 1 + continue + body = json.loads(packet["value"]) + usage = body.get("usage") or usage + for choice in body.get("choices", []): + reason = choice.get("finish_reason") or reason + if done != 1 or reason != "length" or not usage: + raise ValueError("Missing usage, DONE or length termination") + if type(usage.get("completion_tokens")) is not int or usage["completion_tokens"] != payload["max_tokens"]: + raise ValueError("Server output count differs from fixed work") + if type(usage.get("prompt_tokens")) is not int or usage["prompt_tokens"] <= 0: + raise ValueError("Missing actual prompt count") + prompt_tokens += usage["prompt_tokens"] + output_tokens += usage["completion_tokens"] + prompt_counts[key] = usage["prompt_tokens"] + if metadata["request_start_ns"] >= metadata["request_end_ns"]: + raise ValueError("Invalid request interval") + intervals[key] = (metadata["request_start_ns"], metadata["request_end_ns"]) + starts.append(metadata["request_start_ns"]) + ends.append(metadata["request_end_ns"]) + if seen != expected: + raise ValueError("Missing, duplicate or changed payloads") + for (session, turn), (start, end) in intervals.items(): + if turn and start < intervals[(session, turn - 1)][1]: + raise ValueError("Session turns overlapped or ran out of order") + summary = json.loads((directory / "profile_export_aiperf.json").read_text()) + if summary.get("error_summary") or summary.get("was_cancelled"): + raise ValueError("AgentX reported errors or cancellation") + if summary["request_count"]["avg"] != sum(expected.values()): + raise ValueError("AgentX request count mismatch") + if summary["usage_completion_tokens"]["sum"] != output_tokens: + raise ValueError("AgentX output token total differs from SSE usage") + if summary["usage_prompt_tokens"]["sum"] != prompt_tokens: + raise ValueError("AgentX prompt token total differs from SSE usage") + seconds = (max(ends) - min(starts)) / 1e9 + if not math.isfinite(seconds) or seconds <= 0: + raise ValueError("Invalid request wall time") + throughput = output_tokens / seconds + reported = summary["output_token_throughput"]["avg"] + if not math.isfinite(reported) or abs(reported / throughput - 1) > 0.02: + raise ValueError("AgentX throughput differs from independent count / request wall time by >2%") + request_tps = summary["request_throughput"]["avg"] + if not math.isfinite(request_tps) or abs(request_tps / (len(observed_turns) / seconds) - 1) > 0.02: + raise ValueError("AgentX request throughput differs from independent count / request wall time by >2%") + return dict( + payloads_sha256=digest(canonical(sorted(expected.items()))), + prompt_counts_sha256=digest(canonical(sorted(prompt_counts.items()))), + requests=sum(expected.values()), + prompt_tokens=prompt_tokens, + output_tokens=output_tokens, + request_wall_s=seconds, + output_tps=reported, + independent_output_tps=throughput, + request_tps=request_tps, + ttft_p50_ms=summary["time_to_first_token"]["p50"], + ttft_p95_ms=summary["time_to_first_token"]["p95"], + ) + + +def raw_rows(files): + for path in files: + with path.open() as handle: + for line in handle: + yield json.loads(line) + + +def repeatability(runs): + if len(runs) < 3: + raise ValueError("Require at least three audited runs") + for key in ("payloads_sha256", "prompt_counts_sha256", "requests", "prompt_tokens", "output_tokens"): + if len({r[key] for r in runs}) != 1: + raise ValueError(f"Run workload mismatch: {key}") + metrics = {} + for key, limit in (("output_tps", 0.05), ("request_tps", 0.05), ("ttft_p50_ms", 0.10)): + values = [r[key] for r in runs] + if any(not math.isfinite(v) or v <= 0 for v in values): + raise ValueError(f"Invalid metric: {key}") + average = statistics.mean(values) + metrics[key] = dict( + values=values, + mean=average, + cv=statistics.stdev(values) / average, + relative_range=(max(values) - min(values)) / average, + cv_limit=limit, + ) + passed = all(m["cv"] <= m["cv_limit"] and m["relative_range"] <= 2 * m["cv_limit"] for m in metrics.values()) + return dict(status="passed" if passed else "nonrepeatable", runs=len(runs), metrics=metrics) + + +def main(): + parser = argparse.ArgumentParser(description="Reaudit fixed AgentX exports and trial variance") + parser.add_argument("--suite", type=Path, required=True) + parser.add_argument("--workload", type=Path, required=True) + args = parser.parse_args() + suite = json.loads(args.suite.read_text()) + raw = (args.workload / "inputs.json").read_bytes() + if digest(raw) != suite["manifest"]["inputs_sha256"]: + raise ValueError("Workload differs from recorded suite") + inputs = json.loads(raw) + runs = [] + for trial in suite["runs"]: + for phase in ("warmup", "measure"): + current = audit(args.suite.parent / f"trial-{trial['trial']:02}" / phase, inputs) + if current != trial[phase + "_audit"]: + raise ValueError("Retained raw audit differs from the suite") + if phase == "measure": + runs.append(current) + result = repeatability(runs) + print(json.dumps(result, indent=2)) + raise SystemExit(0 if result["status"] == "passed" else 1) + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/qwen38/prepare_agentx.py b/test/benchmark/qwen38/prepare_agentx.py new file mode 100644 index 0000000000..e1c0d0eb0a --- /dev/null +++ b/test/benchmark/qwen38/prepare_agentx.py @@ -0,0 +1,190 @@ +import argparse +import hashlib +import importlib.metadata +import json +import os +from pathlib import Path + + +VERSION = "1.0.13" +DATASET_REVISION = "8fecd2fc56694469f758f0afbbb6335ad3043740" +DATASET_SHA256 = "e39cd2ff3eba21d4a3664be51da743ac3d2149a1933898cafc7bfeac8147eeef" + + +def digest(value): + return hashlib.sha256(value).hexdigest() + + +def canonical(value): + return json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":")).encode() + + +def prompt_length(tokenizer, messages): + encoded = tokenizer.apply_chat_template(messages, add_generation_prompt=True) + return len(encoded["input_ids"] if hasattr(encoded, "keys") else encoded) + + +def tool_identity(): + import aiperf + import agentx_cli + + version = importlib.metadata.version("agentx-harness") + if version != VERSION: + raise ValueError(f"Require agentx-harness {VERSION}, got {version}") + package = Path(aiperf.__file__).parent + files = {str(p.relative_to(package)): digest(p.read_bytes()) for p in sorted(package.rglob("*.py"))} + files["../agentx_cli.py"] = digest(Path(agentx_cli.__file__).read_bytes()) + dependencies = sorted((d.metadata["Name"], d.version) for d in importlib.metadata.distributions()) + return dict(version=version, source_sha256=digest(canonical(files)), files=files, dependencies=dependencies) + + +def frozen_sessions(conversations, model, streams, turns): + sessions, sources = [], [] + for conversation in sorted(conversations, key=lambda c: c.session_id): + if len(conversation.turns) < turns: + continue + messages, payloads = [], [] + for turn in conversation.turns[:turns]: + if turn.reset_context: + messages = [] + messages.extend(turn.raw_messages) + payloads.append( + dict( + model=model, + messages=json.loads(json.dumps(messages)), + max_tokens=turn.max_tokens, + temperature=0, + seed=42, + ignore_eos=True, + stream=True, + stream_options=dict(include_usage=True), + ) + ) + sessions.append(dict(session_id=conversation.session_id, payloads=payloads)) + sources.append( + dict( + session_id=conversation.session_id, + turns=[ + dict( + source_trace_id=t.source_trace_id, + source_outer_idx=t.source_outer_idx, + source_kind=t.source_kind, + output_tokens=t.max_tokens, + payload_sha256=digest(canonical(p)), + ) + for t, p in zip(conversation.turns[:turns], payloads) + ], + ) + ) + if len(sessions) == streams: + return dict(data=sessions), sources + raise ValueError(f"Only {len(sessions)} eligible streams, require {streams}") + + +def main(): + parser = argparse.ArgumentParser(description="Freeze AgentX and a fixed Weka control workload") + parser.add_argument("--source", type=Path, required=True, help="Pinned HF snapshot traces.jsonl") + parser.add_argument("--model-directory", type=Path, required=True) + parser.add_argument("--model", default="Qwen3.8-Flash-Next") + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--streams", type=int, default=32) + parser.add_argument("--turns", type=int, default=8) + args = parser.parse_args() + if not os.environ.get("EXP_ACTIVE"): + parser.error('Run through exp -m "purpose"') + if not 1 <= args.streams <= 256 or not 1 <= args.turns <= 100: + parser.error("Invalid stream or turn count") + if DATASET_REVISION not in args.source.parts: + parser.error("Use the pinned HF snapshot, not refs/main") + + from aiperf.common import random_generator as rng + from aiperf.common.tokenizer import Tokenizer + from aiperf.config import BenchmarkConfig, BenchmarkRun + from aiperf.dataset.generator.prompt import PromptGenerator + from aiperf.dataset.loader.weka_trace import WekaTraceLoader, _trace_peak_context_length + from aiperf.dataset.loader.weka_trace_models import WekaTrace + + identity = tool_identity() + raw = args.source.read_bytes() + if digest(raw) != DATASET_SHA256: + raise ValueError("Dataset bytes differ from the pinned snapshot") + traces = [WekaTrace.model_validate(json.loads(line)) for line in raw.splitlines() if line.strip()] + eligible = [t for t in traces if _trace_peak_context_length(t, max_osl=None) <= 262144] + cfg = BenchmarkConfig.model_validate( + dict( + models=[args.model], + endpoint=dict(urls=["http://127.0.0.1:28081"], type="chat", streaming=True), + datasets=[dict(name="main", type="file", path=str(args.source), format="weka_trace")], + tokenizer=dict(name=str(args.model_directory)), + phases=[dict(name="profiling", type="concurrency", concurrency=32, requests=1)], + ) + ) + rng.init(42) + tokenizer = Tokenizer.from_pretrained(str(args.model_directory), resolve_alias=False) + generator = PromptGenerator(prompts=None, prefix_prompts=None, tokenizer=tokenizer) + run = BenchmarkRun(benchmark_id="qwen38-fixed-work-v1", cfg=cfg, random_seed=42, artifact_dir=args.output) + loader = WekaTraceLoader(run=run, prompt_generator=generator) + conversations = [] + for trace in eligible: + conversations.extend(loader.convert_to_conversations({trace.id: [trace]})) + if sum(len(c.turns) >= args.turns for c in conversations) >= args.streams: + break + inputs, sources = frozen_sessions(conversations, args.model, args.streams, args.turns) + from transformers import AutoTokenizer + + chat_tokenizer = AutoTokenizer.from_pretrained(str(args.model_directory), local_files_only=True) + for session, source in zip(inputs["data"], sources): + for payload, turn in zip(session["payloads"], source["turns"]): + count = prompt_length(chat_tokenizer, payload["messages"]) + if count + payload["max_tokens"] > 262144: + raise ValueError("Templated request exceeds context limit; never clip input or output") + turn["templated_prompt_tokens"] = count + args.output.mkdir(parents=True, exist_ok=False) + (args.output / "inputs.json").write_bytes(canonical(inputs) + b"\n") + (args.output / "sources.json").write_bytes(canonical(sources) + b"\n") + model_files = { + p.name: digest(p.read_bytes()) + for p in ( + args.model_directory / name + for name in ( + "config.json", + "model.safetensors.index.json", + "tokenizer.json", + "tokenizer_config.json", + "chat_template.jinja", + "generation_config.json", + "vocab.json", + "merges.txt", + "acquisition.json", + ) + ) + if p.is_file() + } + weights = sorted( + set(json.loads((args.model_directory / "model.safetensors.index.json").read_text())["weight_map"].values()) + ) + manifest = dict( + schema=1, + tool=identity, + dataset_revision=DATASET_REVISION, + source_sha256=digest(raw), + inputs_sha256=digest((args.output / "inputs.json").read_bytes()), + sources_sha256=digest((args.output / "sources.json").read_bytes()), + model_files=model_files, + weight_files={name: (args.model_directory / name).stat().st_size for name in weights}, + streams=args.streams, + turns_per_stream=args.turns, + requests=args.streams * args.turns, + output_tokens=sum(p["max_tokens"] for s in inputs["data"] for p in s["payloads"]), + scope="Fixed Weka-derived multi-turn control; frozen synthetic history; no think delays or DAG joins. " + "This is not the official AgentX scenario or a business-quality test.", + ) + (args.output / "manifest.json").write_text(json.dumps(manifest, indent=2) + "\n") + print( + json.dumps({k: manifest[k] for k in ("inputs_sha256", "sources_sha256", "requests", "output_tokens", "scope")}), + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/qwen38/run_agentx.py b/test/benchmark/qwen38/run_agentx.py new file mode 100644 index 0000000000..b22bdc4c37 --- /dev/null +++ b/test/benchmark/qwen38/run_agentx.py @@ -0,0 +1,257 @@ +import argparse +import datetime +import json +import os +import signal +import time +from pathlib import Path +import subprocess +import sys +import urllib.request + +from test.benchmark.qwen38.audit_agentx import audit, repeatability +from test.benchmark.qwen38.prepare_agentx import canonical, digest, tool_identity + + +TASK = "a3844d54-8a2b-4532-bf54-f657c9ca72a0" + + +def http(base, path, method="GET"): + opener = urllib.request.build_opener(urllib.request.ProxyHandler({})) + with opener.open(urllib.request.Request(base + path, method=method), timeout=120) as response: + return response.read() + + +def container_info(name): + result = subprocess.check_output(["sudo", "-n", "docker", "inspect", name], text=True) + info = json.loads(result)[0] + if not info["State"]["Running"] or info["Config"].get("Labels", {}).get("qwen38.task") != TASK: + raise ValueError("Require a running container owned by this task") + others = subprocess.check_output(["sudo", "-n", "docker", "ps", "-q"], text=True).split() + if others != [info["Id"][:12]]: + raise ValueError("Unexpected concurrent container; run on the isolated benchmark host") + return info + + +def idle_window(container): + samples = [] + for _ in range(3): + owners = subprocess.check_output( + ["sudo", "-n", "docker", "top", container, "-eo", "pid"], text=True + ).splitlines()[1:] + allowed = {p.strip() for p in owners} + processes = subprocess.check_output( + ["nvidia-smi", "--query-compute-apps=pid", "--format=csv,noheader,nounits"], text=True + ).split() + if set(processes) - allowed: + raise ValueError("GPU process outside the benchmark container") + raw = subprocess.check_output( + [ + "nvidia-smi", + "--query-gpu=index,utilization.gpu,temperature.gpu", + "--format=csv,noheader,nounits", + ], + text=True, + ) + rows = [[int(v.strip()) for v in line.split(",")] for line in raw.splitlines()] + if len(rows) != 8 or any(util > 5 or temperature > 60 for _, util, temperature in rows): + raise ValueError("Require eight idle GPUs (utilization <=5%, temperature <=60C)") + samples.append(rows) + time.sleep(10) + return samples + + +def reset_cache(base, engine): + path = "/reset_prefix_cache" if engine == "vllm" else "/flush_cache" + raw = http(base, path, "POST") + if engine in ("vllm", "lightllm"): + if json.loads(raw).get("success") is not True: + raise ValueError("Server did not confirm cache reset") + elif raw.strip().lower() not in (b"true", b"", b"null"): + raise ValueError("Unexpected SGLang cache reset response") + return dict(path=path, response=raw.decode()) + + +def fixed_config(base, model, tokenizer, inputs, count, streams, output): + return { + "schemaVersion": "2.0", + "randomSeed": 42, + "benchmark": dict( + model=model, + endpoint=dict( + url=base, + type="chat", + streaming=True, + use_server_token_count=True, + timeout=1200, + wait_for_model_timeout=0, + connection_reuse="pooled", + ), + tokenizer=dict(name=str(tokenizer)), + dataset=dict(type="file", format="inputs_json", path=str(inputs), sampling="sequential"), + phases=[dict(name="profiling", type="concurrency", concurrency=streams, sessions=count)], + artifacts=dict( + dir=str(output), + summary=["json"], + records=["jsonl"], + raw=True, + export_outputs_json=False, + auto_plot=False, + ), + runtime=dict(ui="none", stats_interval=30), + gpu_telemetry=dict(enabled=False), + ), + } + + +def execute(command, log, environment): + with log.open("w") as handle: + process = subprocess.Popen(command, stdout=handle, stderr=subprocess.STDOUT, env=environment) + log.with_suffix(".pid").write_text(str(process.pid) + "\n") + try: + returncode = process.wait() + except BaseException: + process.terminate() + try: + process.wait(timeout=30) + except subprocess.TimeoutExpired: + process.kill() + process.wait(timeout=10) + raise + if returncode: + raise RuntimeError(f"AgentX exited {returncode}; inspect {log}") + + +def main(): + parser = argparse.ArgumentParser(description="Run three independently warmed fixed AgentX trials") + parser.add_argument("--engine", choices=["vllm", "sglang", "lightllm"], required=True) + parser.add_argument("--container", required=True) + parser.add_argument("--workload", type=Path, required=True) + parser.add_argument("--model-directory", type=Path, required=True) + parser.add_argument("--model", default="Qwen3.8-Flash-Next") + parser.add_argument("--url", default="http://127.0.0.1:28081") + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--repeats", type=int, default=3) + args = parser.parse_args() + if not os.environ.get("EXP_ACTIVE"): + parser.error('Run through exp -m "purpose"') + if args.repeats < 3 or args.repeats > 10: + parser.error("Require three to ten trials") + if args.url != "http://127.0.0.1:28081": + parser.error("This protocol is restricted to the isolated local benchmark endpoint") + args.output = args.output.resolve() + args.workload = args.workload.resolve() + manifest = json.loads((args.workload / "manifest.json").read_text()) + if canonical(tool_identity()) != canonical(manifest["tool"]): + raise ValueError("Installed tool/dependencies changed from the frozen manifest") + inputs_path = args.workload / "inputs.json" + if digest(inputs_path.read_bytes()) != manifest["inputs_sha256"]: + raise ValueError("Frozen workload changed") + if digest((args.workload / "sources.json").read_bytes()) != manifest["sources_sha256"]: + raise ValueError("Frozen workload provenance changed") + for name, expected in manifest["model_files"].items(): + if digest((args.model_directory / name).read_bytes()) != expected: + raise ValueError(f"Checkpoint/tokenizer changed: {name}") + for name, size in manifest["weight_files"].items(): + if (args.model_directory / name).stat().st_size != size: + raise ValueError(f"Checkpoint shard size changed: {name}") + info = container_info(args.container) + args.output.mkdir(parents=True, exist_ok=False) + inputs = json.loads(inputs_path.read_text()) + report = dict( + status="running", + mode="fixed", + engine=args.engine, + manifest=manifest, + container=info, + runs=[], + scope="Fixed Weka-derived control; not the official response-dependent AgentX scenario", + protocol_files={p.name: digest(p.read_bytes()) for p in sorted(Path(__file__).parent.glob("*_agentx.py"))}, + cache_policy="Confirmed reset before every complete warmup; same full fixed-work warmup every trial; " + "reset again before measurement.", + ) + (args.output / "hardware.txt").write_text(subprocess.check_output(["nvidia-smi", "-q"], text=True)) + environment = os.environ.copy() + environment.update( + NO_PROXY="127.0.0.1,localhost", + no_proxy="127.0.0.1,localhost", + HF_HUB_OFFLINE="1", + TRANSFORMERS_OFFLINE="1", + PYTHONHASHSEED="42", + AIPERF_HTTP_KEEPALIVE_TIMEOUT="4", + ) + entry = "from aiperf.cli import app; app()" + + def save(): + (args.output / "suite.json").write_text(json.dumps(report, indent=2) + "\n") + + monitor = None + + def interrupted(signum, frame): + raise KeyboardInterrupt(f"Signal {signum}") + + signal.signal(signal.SIGTERM, interrupted) + try: + save() + gpu_log = (args.output / "gpu-samples.csv").open("w") + monitor = subprocess.Popen( + [ + "nvidia-smi", + "--query-gpu=timestamp,index,utilization.gpu,memory.used,power.draw,temperature.gpu", + "--format=csv", + "--loop=5", + ], + stdout=gpu_log, + ) + for index in range(args.repeats): + if container_info(args.container)["Id"] != info["Id"]: + raise ValueError("Server container changed between trials") + trial = args.output / f"trial-{index + 1:02}" + trial.mkdir() + result = dict(trial=index + 1) + report["runs"].append(result) + save() + result["reset_before_warmup"] = reset_cache(args.url, args.engine) + for phase in ("warmup", "measure"): + if phase == "measure": + result["reset_before_measure"] = reset_cache(args.url, args.engine) + result[phase + "_idle_samples"] = idle_window(args.container) + save() + output = trial / phase + config = fixed_config( + args.url, args.model, args.model_directory, inputs_path, manifest["streams"], 32, output + ) + config_path = trial / (phase + ".json") + config_path.write_text(json.dumps(config, indent=2) + "\n") + command = [sys.executable, "-c", entry, "profile", "--config", str(config_path)] + (trial / (phase + "-command.json")).write_text(json.dumps(command, indent=2) + "\n") + print(datetime.datetime.now().isoformat(), index + 1, phase, output, flush=True) + (trial / (phase + "-metrics-before.txt")).write_bytes(http(args.url, "/metrics")) + execute(command, trial / (phase + ".log"), environment) + if container_info(args.container)["Id"] != info["Id"]: + raise ValueError("Container changed during trial") + (trial / (phase + "-metrics-after.txt")).write_bytes(http(args.url, "/metrics")) + result[phase + "_audit"] = audit(output, inputs) + save() + result["audit"] = result["measure_audit"] + save() + report["repeatability"] = repeatability([r["audit"] for r in report["runs"]]) + report["status"] = report["repeatability"]["status"] + except BaseException as error: + report.update(status="failed", error=repr(error)) + raise + finally: + if monitor is not None: + monitor.terminate() + monitor.wait(timeout=20) + gpu_log.close() + logs = subprocess.run(["sudo", "-n", "docker", "logs", args.container], capture_output=True, text=True) + (args.output / "server.log").write_text(logs.stdout + logs.stderr) + save() + print(json.dumps(report.get("repeatability", {"status": report["status"]}), indent=2), flush=True) + if report["status"] == "nonrepeatable": + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/unit_tests/benchmark/qwen38/test_agentx_protocol.py b/unit_tests/benchmark/qwen38/test_agentx_protocol.py new file mode 100644 index 0000000000..aa16095995 --- /dev/null +++ b/unit_tests/benchmark/qwen38/test_agentx_protocol.py @@ -0,0 +1,181 @@ +import copy +import json + +import pytest + +from test.benchmark.qwen38.audit_agentx import audit, repeatability +from test.benchmark.qwen38.run_agentx import fixed_config +from test.benchmark.qwen38.prepare_agentx import prompt_length +from test.benchmark.qwen38.prepare_agentx import frozen_sessions +from types import SimpleNamespace + + +def fixture(tmp_path, change=None): + payload = dict(model="model", messages=[dict(role="user", content="hello")], max_tokens=2) + inputs = dict(data=[dict(session_id="one", payloads=[payload])]) + row = dict( + metadata=dict( + benchmark_phase="profiling", + conversation_id="one", + turn_index=0, + request_start_ns=1_000_000_000, + request_end_ns=2_000_000_000, + ), + status=200, + payload=copy.deepcopy(payload), + responses=[ + dict( + packets=[ + dict( + name="data", + value=json.dumps( + dict( + choices=[dict(finish_reason="length")], usage=dict(prompt_tokens=3, completion_tokens=2) + ) + ), + ), + dict(name="data", value="[DONE]"), + ] + ) + ], + ) + summary = dict( + request_count=dict(avg=1), + usage_completion_tokens=dict(sum=2), + usage_prompt_tokens=dict(sum=3), + output_token_throughput=dict(avg=2), + request_throughput=dict(avg=1), + time_to_first_token=dict(p50=100, p95=120), + error_summary=[], + ) + if change: + change(row, summary) + (tmp_path / "raw_records").mkdir() + (tmp_path / "raw_records" / "worker.jsonl").write_text(json.dumps(row) + "\n") + (tmp_path / "profile_export_aiperf.json").write_text(json.dumps(summary)) + return inputs + + +def test_complete_fixed_work_is_independently_counted(tmp_path): + result = audit(tmp_path, fixture(tmp_path)) + assert result["requests"] == 1 + assert result["output_tokens"] == 2 + assert result["independent_output_tps"] == 2 + + +@pytest.mark.parametrize( + "failure", ["payload", "usage", "done", "error", "throughput", "request_throughput", "cancelled", "turn", "time"] +) +def test_audit_rejects_changed_work_or_miscounted_stream(tmp_path, failure): + def change(row, summary): + if failure == "payload": + row["payload"]["messages"][0]["content"] = "different" + elif failure == "usage": + packet = row["responses"][0]["packets"][0] + body = json.loads(packet["value"]) + body["usage"]["completion_tokens"] = 1 + packet["value"] = json.dumps(body) + elif failure == "done": + row["responses"][0]["packets"].pop() + elif failure == "error": + summary["error_summary"] = [dict(count=1)] + elif failure == "throughput": + summary["output_token_throughput"]["avg"] = 3 + elif failure == "request_throughput": + summary["request_throughput"]["avg"] = 2 + elif failure == "turn": + row["metadata"]["turn_index"] = 1 + elif failure == "time": + row["metadata"]["request_end_ns"] = 0 + else: + row["metadata"]["was_cancelled"] = True + + inputs = fixture(tmp_path, change) + with pytest.raises(ValueError): + audit(tmp_path, inputs) + + +def test_repeatability_rejects_large_variance_and_workload_mismatch(): + base = dict( + payloads_sha256="same", + prompt_counts_sha256="same-counts", + requests=100, + prompt_tokens=1000, + output_tokens=200, + output_tps=100, + request_tps=50, + ttft_p50_ms=10, + ) + assert repeatability([base, base, base])["status"] == "passed" + varied = dict(base, output_tps=190) + assert repeatability([base, varied, base])["status"] == "nonrepeatable" + with pytest.raises(ValueError, match="workload mismatch"): + repeatability([base, dict(base, prompt_tokens=1001), base]) + + +def test_fixed_work_has_no_time_cutoff_or_hidden_warmup(tmp_path): + config = fixed_config("http://localhost", "model", tmp_path, tmp_path / "inputs.json", 32, 32, tmp_path) + phases = config["benchmark"]["phases"] + assert len(phases) == 1 and phases[0]["sessions"] == 32 + assert "duration" not in phases[0] + assert "multiRun" not in config and "scenario" not in config["benchmark"] + + +@pytest.mark.parametrize("encoded", [[1, 2, 3], dict(input_ids=[1, 2, 3], attention_mask=[1, 1, 1])]) +def test_chat_template_length_handles_list_and_batch_encoding(encoded): + class Tokenizer: + def apply_chat_template(self, messages, add_generation_prompt): + return encoded + + assert prompt_length(Tokenizer(), []) == 3 + + +def test_freeze_reconstructs_deltas_and_resets_without_live_responses(): + def turn(text, reset=False): + return SimpleNamespace( + raw_messages=[dict(role="user", content=text)], + reset_context=reset, + max_tokens=5, + source_trace_id="trace", + source_outer_idx=0, + source_kind="weka_main", + ) + + conversation = SimpleNamespace(session_id="one", turns=[turn("first"), turn("second"), turn("new", True)]) + inputs, sources = frozen_sessions([conversation], "model", 1, 3) + payloads = inputs["data"][0]["payloads"] + assert [p["messages"] for p in payloads] == [ + [dict(role="user", content="first")], + [dict(role="user", content="first"), dict(role="user", content="second")], + [dict(role="user", content="new")], + ] + assert len(sources[0]["turns"]) == 3 + + +def test_audit_uses_merged_export_without_double_counting_worker_shards(tmp_path): + inputs = fixture(tmp_path) + row = (tmp_path / "raw_records" / "worker.jsonl").read_text() + (tmp_path / "profile_export_raw.jsonl").write_text(row) + assert audit(tmp_path, inputs)["requests"] == 1 + + +def test_audit_counts_completion_usage_including_reasoning(tmp_path): + inputs = fixture(tmp_path) + path = tmp_path / "profile_export_aiperf.json" + summary = json.loads(path.read_text()) + summary["output_token_count"] = dict(sum=0) + summary["reasoning_token_count"] = dict(sum=2) + path.write_text(json.dumps(summary)) + assert audit(tmp_path, inputs)["output_tokens"] == 2 + + +def test_audit_rejects_overlapping_turns_with_identical_payloads(tmp_path): + inputs = fixture(tmp_path) + inputs["data"][0]["payloads"] *= 2 + path = tmp_path / "raw_records" / "worker.jsonl" + first = json.loads(path.read_text()) + second = copy.deepcopy(first) + second["metadata"].update(turn_index=1, request_start_ns=1_500_000_000, request_end_ns=2_500_000_000) + path.write_text(json.dumps(first) + "\n" + json.dumps(second) + "\n") + with pytest.raises(ValueError, match="overlapped"): + audit(tmp_path, inputs) From 9faba16b137eef3909b360a712defbe1c3f678bf Mon Sep 17 00:00:00 2001 From: sufubao Date: Thu, 12 Nov 2026 15:50:30 +0800 Subject: [PATCH 6/6] fix: wait for idle GPUs before AgentX replay --- test/benchmark/qwen38/README.agentx.md | 12 +++-- test/benchmark/qwen38/run_agentx.py | 28 ++++++++---- .../benchmark/qwen38/test_agentx_protocol.py | 44 ++++++++++++++++++- 3 files changed, 71 insertions(+), 13 deletions(-) diff --git a/test/benchmark/qwen38/README.agentx.md b/test/benchmark/qwen38/README.agentx.md index f880763462..ff05d234a0 100644 --- a/test/benchmark/qwen38/README.agentx.md +++ b/test/benchmark/qwen38/README.agentx.md @@ -80,7 +80,7 @@ Enable and verify the engine's cache-reset API before measurement: | --- | --- | --- | --- | | vLLM 0.33 | `POST /reset_prefix_cache` | `{"success":true}` | Validated on H100 | | LightLLM | `POST /flush_cache`, requires `--enable_rl` | `{"success":true}` | Adapter awaits runtime validation | -| SGLang | `POST /flush_cache` | HTTP 200, empty/null/true body | Adapter awaits runtime validation | +| SGLang 0.5.23 | `POST /flush_cache` | HTTP 200, `Cache flushed.` text or empty/null/true body | Streaming/reset adapter validated on H100 | The tested vLLM image requires `VLLM_SERVER_DEV_MODE=1` to expose the reset endpoint. Bind this benchmark service to localhost. A missing endpoint, unexpected @@ -92,6 +92,8 @@ Each of at least three trials executes: 2. Reset the server cache and require success. 3. Take three idle GPU samples, ten seconds apart: utilization at most 5%, temperature at most 60 C, all eight GPUs present, no foreign compute PID. + Wait up to five minutes for three consecutive qualifying samples, retaining + every check. Missing GPUs or a foreign compute PID abort immediately. 4. Replay and audit the entire fixed workload as warmup. 5. Reset cache again, repeat the idle checks, then replay and audit the same full workload for measurement. @@ -126,8 +128,12 @@ across repeated trials. The numerator is server `usage.completion_tokens`, including reasoning tokens. AgentX `output_token_count` counts visible output separately and must not be used -as the total for reasoning models. Use `usage_completion_tokens` or -`output_sequence_length` when reconciling the summary with raw SSE usage. +as the total for reasoning models. Reconcile `usage_completion_tokens` with raw +SSE usage and use the independently counted throughput for engine comparisons. +SGLang 0.5.23 can report reasoning tokens above completion tokens; AgentX 1.0.13 +then slightly inflates `output_sequence_length` and its derived throughput. Keep +the tool and server unchanged during a suite, disclose this discrepancy and +retain the 2% independent cross-check. All trials must pass these predeclared limits: diff --git a/test/benchmark/qwen38/run_agentx.py b/test/benchmark/qwen38/run_agentx.py index b22bdc4c37..297b34adea 100644 --- a/test/benchmark/qwen38/run_agentx.py +++ b/test/benchmark/qwen38/run_agentx.py @@ -33,9 +33,10 @@ def container_info(name): return info -def idle_window(container): +def idle_window(container, history, timeout=300): samples = [] - for _ in range(3): + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: owners = subprocess.check_output( ["sudo", "-n", "docker", "top", container, "-eo", "pid"], text=True ).splitlines()[1:] @@ -54,11 +55,17 @@ def idle_window(container): text=True, ) rows = [[int(v.strip()) for v in line.split(",")] for line in raw.splitlines()] - if len(rows) != 8 or any(util > 5 or temperature > 60 for _, util, temperature in rows): - raise ValueError("Require eight idle GPUs (utilization <=5%, temperature <=60C)") - samples.append(rows) + if sorted(index for index, _, _ in rows) != list(range(8)): + raise ValueError("Require all eight benchmark GPUs") + history.append(dict(timestamp=datetime.datetime.now().isoformat(), gpus=rows)) + if any(util > 5 or temperature > 60 for _, util, temperature in rows): + samples.clear() + else: + samples.append(rows) time.sleep(10) - return samples + if len(samples) == 3: + return samples + raise TimeoutError(f"GPUs did not stay idle (utilization <=5%, temperature <=60C) within {timeout} seconds") def reset_cache(base, engine): @@ -67,8 +74,10 @@ def reset_cache(base, engine): if engine in ("vllm", "lightllm"): if json.loads(raw).get("success") is not True: raise ValueError("Server did not confirm cache reset") - elif raw.strip().lower() not in (b"true", b"", b"null"): - raise ValueError("Unexpected SGLang cache reset response") + else: + confirmation = raw.strip().lower() + if confirmation not in (b"true", b"", b"null") and not confirmation.startswith(b"cache flushed.\n"): + raise ValueError("Unexpected SGLang cache reset response") return dict(path=path, response=raw.decode()) @@ -215,7 +224,8 @@ def interrupted(signum, frame): for phase in ("warmup", "measure"): if phase == "measure": result["reset_before_measure"] = reset_cache(args.url, args.engine) - result[phase + "_idle_samples"] = idle_window(args.container) + result[phase + "_idle_checks"] = [] + result[phase + "_idle_samples"] = idle_window(args.container, result[phase + "_idle_checks"]) save() output = trial / phase config = fixed_config( diff --git a/unit_tests/benchmark/qwen38/test_agentx_protocol.py b/unit_tests/benchmark/qwen38/test_agentx_protocol.py index aa16095995..4da19e1f81 100644 --- a/unit_tests/benchmark/qwen38/test_agentx_protocol.py +++ b/unit_tests/benchmark/qwen38/test_agentx_protocol.py @@ -4,7 +4,7 @@ import pytest from test.benchmark.qwen38.audit_agentx import audit, repeatability -from test.benchmark.qwen38.run_agentx import fixed_config +from test.benchmark.qwen38.run_agentx import fixed_config, idle_window, reset_cache from test.benchmark.qwen38.prepare_agentx import prompt_length from test.benchmark.qwen38.prepare_agentx import frozen_sessions from types import SimpleNamespace @@ -179,3 +179,45 @@ def test_audit_rejects_overlapping_turns_with_identical_payloads(tmp_path): path.write_text(json.dumps(first) + "\n" + json.dumps(second) + "\n") with pytest.raises(ValueError, match="overlapped"): audit(tmp_path, inputs) + + +def test_sglang_cache_reset_text(monkeypatch): + target = "test.benchmark.qwen38.run_agentx.http" + monkeypatch.setattr(target, lambda *args: b"Cache flushed.\nPlease check backend logs for more details.\n") + assert reset_cache("http://localhost", "sglang")["path"] == "/flush_cache" + monkeypatch.setattr(target, lambda *args: b"Cache not flushed because there are pending requests.") + with pytest.raises(ValueError, match="Unexpected SGLang"): + reset_cache("http://localhost", "sglang") + + +@pytest.mark.parametrize("failure", ["transient", "busy", "foreign"]) +def test_idle_gate(monkeypatch, failure): + target = "test.benchmark.qwen38.run_agentx" + now = [0] + calls = [0] + + def output(command, **kwargs): + if command[0] == "sudo": + return "PID\n123\n" + if "--query-compute-apps=pid" in command: + return "456\n" if failure == "foreign" else "123\n" + calls[0] += 1 + busy = failure == "busy" or calls[0] in (1, 3) + return "\n".join(f"{i}, {9 if busy else 0}, 54" for i in range(8)) + + monkeypatch.setattr(target + ".subprocess.check_output", output) + monkeypatch.setattr(target + ".time.monotonic", lambda: now[0]) + monkeypatch.setattr(target + ".time.sleep", lambda seconds: now.__setitem__(0, now[0] + seconds)) + history = [] + if failure == "foreign": + with pytest.raises(ValueError, match="outside"): + idle_window("container", history) + assert not history + elif failure == "busy": + with pytest.raises(TimeoutError, match="idle"): + idle_window("container", history, timeout=60) + assert len(history) == 6 + else: + samples = idle_window("container", history) + assert len(history) == 6 and len(samples) == 3 + assert all(util == 0 for sample in samples for _, util, _ in sample)