From 1c6734284aa12cdadc4a0221d5f106cdd37bd5c0 Mon Sep 17 00:00:00 2001 From: sufubao Date: Mon, 2 Nov 2026 15:57:42 +0800 Subject: [PATCH] feat: support Qwen3.8 Flash Next --- docs/qwen38_adaptation.rst | 48 +++ docs/qwen38_h100_performance.rst | 76 ++++ .../common/basemodel/attention/linear/gdn.py | 12 +- .../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 +++ .../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 | 108 ++++++ lightllm/common/qsa/attention_kernels.py | 136 +++++++ lightllm/common/qsa/indexer.py | 264 +++++++++++++ lightllm/common/qsa/indexer_kernels.py | 169 +++++++++ lightllm/common/qsa/topk.py | 99 +++++ .../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 + .../qwen2_vl/triton_kernel/rotary_pos_emb.py | 9 +- .../layer_infer/transformer_layer_infer.py | 12 +- 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 | 92 +++++ .../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 | 68 ++++ lightllm/models/qwen4_exp/position.py | 18 + .../models/qwen4_exp/speculative_cache.py | 108 ++++++ lightllm/models/qwen4_exp/state.py | 128 +++++++ 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 | 60 +++ lightllm/models/qwen4_exp/vision.py | 36 ++ lightllm/server/api_cli.py | 5 + lightllm/server/api_lightllm.py | 13 +- lightllm/server/core/objs/start_args_type.py | 1 + lightllm/server/httpserver/manager.py | 19 +- .../server/httpserver/slot_availability.py | 38 ++ .../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 +- .../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 ++++++++++++++ .../basemodel/attention/linear/test_gdn.py | 86 +++++ .../linear_att/test_causal_conv1d_mtp.py | 25 ++ .../test_mtp_fused_recurrent_equiv.py | 52 +++ .../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_indexer.py | 353 ++++++++++++++++++ .../common/qsa/test_topk_compatibility.py | 131 +++++++ unit_tests/common/qsa/test_topk_cuda.py | 62 +++ .../quantization/test_aligned_projection.py | 97 +++++ unit_tests/distributed/test_ordered_tp_sum.py | 110 ++++++ .../models/qwen2_vl/test_rotary_pos_emb.py | 17 +- 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 | 72 ++++ unit_tests/models/qwen4_exp/runtime.py | 133 +++++++ .../models/qwen4_exp/test_attention_state.py | 181 +++++++++ unit_tests/models/qwen4_exp/test_batch.py | 97 +++++ unit_tests/models/qwen4_exp/test_decoder.py | 211 +++++++++++ .../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 | 182 +++++++++ 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 | 28 ++ .../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 | 226 +++++++++++ .../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 +++++ .../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 +++++ .../server/test_native_stream_details.py | 69 ++++ 176 files changed, 16254 insertions(+), 40 deletions(-) create mode 100644 docs/qwen38_adaptation.rst create mode 100644 docs/qwen38_h100_performance.rst 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/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/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/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_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/distributed/test_ordered_tp_sum.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/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 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/basemodel/attention/linear/gdn.py b/lightllm/common/basemodel/attention/linear/gdn.py index 986214568f..4cb7b812da 100644 --- a/lightllm/common/basemodel/attention/linear/gdn.py +++ b/lightllm/common/basemodel/attention/linear/gdn.py @@ -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 @@ -418,7 +417,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 +425,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/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..e3c7baa8e2 --- /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, enable_fp_fusion=False + ) + 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/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..6f3cd841db --- /dev/null +++ b/lightllm/common/qsa/attention.py @@ -0,0 +1,108 @@ +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, +) -> 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 splits is None: + splits = min(32, max(1, triton.cdiv(64, rows * kv_heads)), triton.cdiv(selected, 64)) + 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 + _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)), + 64, + short_rows_direct, + skip_empty_splits, + ) + 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, + 64, + short_rows_direct, + skip_empty_splits, + ) + return output diff --git a/lightllm/common/qsa/attention_kernels.py b/lightllm/common/qsa/attention_kernels.py new file mode 100644 index 0000000000..94c47c6f40 --- /dev/null +++ b/lightllm/common/qsa/attention_kernels.py @@ -0,0 +1,136 @@ +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, +): + 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="ieee") * 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="ieee") + 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..373cd0b777 --- /dev/null +++ b/lightllm/common/qsa/indexer.py @@ -0,0 +1,264 @@ +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 + + +@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, + ) + return output, projected[:, self.query_heads * self.head_dim :].contiguous() + + 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) + score_buffer = torch.empty((min(tokens, query_chunk_size), blocks), dtype=torch.float32, device=query.device) + for first in range(0, tokens, query_chunk_size): + count = min(query_chunk_size, tokens - first) + scores = score_buffer[:count] + if blocks: + _scores[(count, triton.cdiv(blocks, 64))]( + 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)), + 64, + block_budget if cuda_topk else 0, + ) + selected = ( + cuda_topk_indices( + scores, topk, visible_lengths=visible_lengths, first=first, compress_ratio=self.compress_ratio + ) + if cuda_topk + else 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..6d8eaf69e0 --- /dev/null +++ b/lightllm/common/qsa/indexer_kernels.py @@ -0,0 +1,169 @@ +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, +): + 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) // 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="ieee") + 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..8c2d6a4047 --- /dev/null +++ b/lightllm/common/qsa/topk.py @@ -0,0 +1,99 @@ +from functools import lru_cache + +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) + + +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/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 39f9c42919..758c4d199f 100644 --- a/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py @@ -183,7 +183,6 @@ def _linear_prefill_cuda_graph_wrapper( infer_state: Qwen3NextInferStateInfo, layer_weight: Qwen3NextTransformerLayerWeight, ) -> Tuple[torch.Tensor, torch.Tensor]: - backend = infer_state.prefill_att_state1.backend mixed_qkvzba = mixed_qkvzba.contiguous() @@ -266,9 +265,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 +293,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, 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..4327f0f8f9 --- /dev/null +++ b/lightllm/models/qwen4_exp/layer_infer/transformer_layer_infer.py @@ -0,0 +1,92 @@ +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, + attention_splits=32, + ) + 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..ecff4a2af4 --- /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 + + for state in self.att_state_manager.get_bounded_attention_states(self.att_state_role): + 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..f5a1222fad --- /dev/null +++ b/lightllm/models/qwen4_exp/pinned_memory.py @@ -0,0 +1,68 @@ +import math +import mmap + +import numpy as np +import torch + +_CHUNK_BYTES = 256 * 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..585f21fff7 --- /dev/null +++ b/lightllm/models/qwen4_exp/state.py @@ -0,0 +1,128 @@ +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 + ) + + splits = attention_splits + if splits is None and decode: + splits = min(32, triton.cdiv(selected.shape[1], 64)) + 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, + ) + + +@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..8b56491ddc --- /dev/null +++ b/lightllm/models/qwen4_exp/triton_kernel/gated_norm.py @@ -0,0 +1,60 @@ +import torch +import triton +import triton.language as tl + + +@triton.jit +def _gated_norm(X, Gate, Weight, Out, 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 * D + 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 = tl.load(Gate + row * D + col, col < D, other=0).to(tl.float32) + activation = tl.sigmoid(gate) + if not SIGMOID: + activation = gate * activation + tl.store(Out + row * D + 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.shape != gate.shape + or input.shape[-1] != weight.numel() + or not input.is_contiguous() + or not gate.is_contiguous() + ): + raise ValueError("GDN norm requires contiguous, matching input/gate tensors 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.numel() // input.shape[-1] + if rows: + _gated_norm[(rows,)]( + input, + gate, + weight, + result, + input.shape[-1], + eps, + activation == "sigmoid", + 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 627f2c741a..ada53f7a9a 100644 --- a/lightllm/server/api_cli.py +++ b/lightllm/server/api_cli.py @@ -935,6 +935,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, 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/core/objs/start_args_type.py b/lightllm/server/core/objs/start_args_type.py index baaefae2e3..673fbe3c7c 100644 --- a/lightllm/server/core/objs/start_args_type.py +++ b/lightllm/server/core/objs/start_args_type.py @@ -242,6 +242,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) 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/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 9457e77391..38f1f7221a 100644 --- a/lightllm/utils/config_utils.py +++ b/lightllm/utils/config_utils.py @@ -215,6 +215,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}") @@ -287,7 +289,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 = [] @@ -418,6 +428,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: @@ -463,6 +475,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", ] @@ -558,6 +572,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/qwen38/benchmark_gated_residual.py b/test/benchmark/qwen38/benchmark_gated_residual.py new file mode 100644 index 0000000000..5dfd3a20e0 --- /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..e62f0fe87b --- /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..5c4cf97268 --- /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/basemodel/attention/linear/test_gdn.py b/unit_tests/common/basemodel/attention/linear/test_gdn.py index 8ec77c6f4d..7d45a44123 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) 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_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/test_scalar_linear.py b/unit_tests/common/basemodel/triton_kernel/test_scalar_linear.py new file mode 100644 index 0000000000..9c6cc3078c --- /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..19cc1d677d --- /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..1c8e5d26e5 --- /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_indexer.py b/unit_tests/common/qsa/test_indexer.py new file mode 100644 index 0000000000..d601388974 --- /dev/null +++ b/unit_tests/common/qsa/test_indexer.py @@ -0,0 +1,353 @@ +import math + +import pytest +import torch +import torch.nn.functional as F + +from lightllm.common.qsa.indexer import QSAIndexer + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +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) + 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]) +def test_physical_cache_partial_prefix_forks_and_graph_replay(dtype): + 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, indexer.query_heads, indexer.head_dim, device="cuda", dtype=dtype) + requests = torch.arange(2, device="cuda", dtype=torch.int32) + + def run(): + indexer.pool(raw, table, active_lengths, cos, sin, cache, active_first, 3, request_rows=rows) + return indexer.select(query, cache, requests, lengths, table) + + 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, 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..9a0d9a89af --- /dev/null +++ b/unit_tests/common/qsa/test_topk_compatibility.py @@ -0,0 +1,131 @@ +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) + + +@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/distributed/test_ordered_tp_sum.py b/unit_tests/distributed/test_ordered_tp_sum.py new file mode 100644 index 0000000000..cf6df14ef5 --- /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/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/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..f7963d715c --- /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..4167f43dff --- /dev/null +++ b/unit_tests/models/qwen4_exp/resource_helpers.py @@ -0,0 +1,72 @@ +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.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..6ceeed9998 --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_attention_state.py @@ -0,0 +1,181 @@ +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 + + +@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("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..d2365948e2 --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_decoder.py @@ -0,0 +1,211 @@ +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": 32 if prefill else None} + expected = torch.cat([execute(row, row + 1, **options) for row in range(rows)]) + 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..ef1b32eaf9 --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_linear_decoder.py @@ -0,0 +1,182 @@ +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 = SimpleNamespace( + 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 = SimpleNamespace( + 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..1856f7c8a0 --- /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..766bc6394d --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_norm.py @@ -0,0 +1,28 @@ +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) 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..bce672b9fc --- /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..e30106f3ce --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_runtime.py @@ -0,0 +1,226 @@ +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") + + +def test_native_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) + + +def test_complete_small_native_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..018a96c14f --- /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, 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) + 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/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_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