From 6535891be2e6a6f709297059828dbb1bf81cd43f Mon Sep 17 00:00:00 2001 From: sufubao Date: Fri, 13 Nov 2026 02:23:57 +0800 Subject: [PATCH] feat(model): support Qwen3.8 Flash Next --- docs/qwen38_adaptation.rst | 85 +++ .../attention_state/bounded_att/utils.py | 60 ++- .../attention_state/full_att/gqa/gqa_state.py | 27 +- .../attention/create_linear_utils.py | 35 +- .../basemodel/attention/linear/flashinfer.py | 63 +++ .../common/basemodel/attention/linear/gdn.py | 51 +- .../basemodel/attention/linear/replay_gdn.py | 17 + lightllm/common/basemodel/basemodel.py | 31 +- lightllm/common/basemodel/batch_objs.py | 6 + lightllm/common/basemodel/cuda_graph.py | 24 +- lightllm/common/basemodel/infer_struct.py | 1 + .../fused_moe/fused_moe_weight.py | 4 + .../meta_weights/fused_moe/impl/base_impl.py | 13 + .../fused_moe/impl/deepgemm_impl.py | 2 + .../meta_weights/pinned_embedding_weight.py | 84 +++ .../row_sharded_embedding_weight.py | 64 +++ .../common/basemodel/multimodal_tokenizer.py | 3 + lightllm/common/basemodel/tp_reduce.py | 73 +++ .../triton_kernel/destindex_copy_kv.py | 73 +++ .../fused_moe/grouped_fused_moe_ep/bf16.py | 250 +++++++++ .../prefill_decode_sync.py | 16 + .../triton_kernel/fused_moe/mask_topk.py | 27 + .../common/basemodel/triton_kernel/kv_move.py | 97 ++++ .../triton_kernel/linear_att/causal_conv1d.py | 144 ++++- .../linear_att/causal_conv1d_triton.py | 407 +++++++++++++++ .../linear_att/fla/ops/__init__.py | 3 +- .../linear_att/fla/ops/fused_recurrent.py | 154 ++++++ .../linear_att/mtp_state_params.py | 2 +- .../basemodel/triton_kernel/mtp_utils.py | 6 +- .../basemodel/triton_kernel/scalar_linear.py | 42 ++ lightllm/common/gated_residual/__init__.py | 93 ++++ lightllm/common/gated_residual/kernels.py | 146 ++++++ lightllm/common/gated_residual/projection.py | 208 ++++++++ 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 | 127 +++++ lightllm/common/qsa/attention_kernels.py | 139 +++++ lightllm/common/qsa/indexer.py | 302 +++++++++++ lightllm/common/qsa/indexer_kernels.py | 219 ++++++++ lightllm/common/qsa/topk.py | 129 +++++ .../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/communication_op.py | 19 +- lightllm/distributed/symm_mem_all_reduce.py | 60 +++ lightllm/models/builtin.py | 5 + .../llama/layer_infer/post_layer_infer.py | 3 + .../qwen2_vl/triton_kernel/rotary_pos_emb.py | 9 +- .../layer_infer/transformer_layer_infer.py | 5 +- .../layer_infer/transformer_layer_infer.py | 12 +- lightllm/models/qwen4_exp/__init__.py | 1 + lightllm/models/qwen4_exp/attention_state.py | 198 +++++++ lightllm/models/qwen4_exp/batch.py | 194 +++++++ 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 | 96 ++++ .../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 | 25 + .../layer_infer/replicated_decoder.py | 220 ++++++++ .../layer_infer/transformer_layer_infer.py | 211 ++++++++ .../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 | 141 +++++ lightllm/models/qwen4_exp/lexical_state.py | 257 +++++++++ lightllm/models/qwen4_exp/model.py | 204 ++++++++ lightllm/models/qwen4_exp/mtp/__init__.py | 1 + lightllm/models/qwen4_exp/mtp/infer.py | 31 ++ lightllm/models/qwen4_exp/mtp/model.py | 53 ++ 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 | 73 +++ lightllm/models/qwen4_exp/qsa_compressor.py | 185 +++++++ .../models/qwen4_exp/speculative_cache.py | 108 ++++ lightllm/models/qwen4_exp/state.py | 177 +++++++ lightllm/models/qwen4_exp/tokenizer.py | 6 + .../qwen4_exp/triton_kernel/__init__.py | 1 + .../qwen4_exp/triton_kernel/cache_pages.py | 71 +++ .../qwen4_exp/triton_kernel/gated_norm.py | 92 ++++ .../qwen4_exp/triton_kernel/replicated_moe.py | 184 +++++++ lightllm/models/qwen4_exp/vision.py | 36 ++ lightllm/server/api_cli.py | 11 +- lightllm/server/api_lightllm.py | 13 +- lightllm/server/api_start.py | 15 +- lightllm/server/core/objs/start_args_type.py | 1 + lightllm/server/httpserver/manager.py | 24 +- .../server/httpserver/slot_availability.py | 38 ++ .../model_infer/mode_backend/base_backend.py | 3 + .../server/router/model_infer/model_rpc.py | 3 +- .../mtp_speculative/planner/lightspec.py | 4 + .../proposers/eagle_with_att.py | 21 +- 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 +- lightllm/utils/envs_utils.py | 22 +- lightllm/utils/net_utils.py | 1 + test/__init__.py | 0 .../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 ++++++ .../att_resources/test_page_pool_storage.py | 92 +++- .../attention/linear/test_flashinfer_gdn.py | 105 ++++ .../basemodel/attention/linear/test_gdn.py | 99 +++- .../basemodel/attention/test_lazy_import.py | 34 +- .../linear_att/test_causal_conv1d.py | 398 ++++++++++++++ .../linear_att/test_causal_conv1d_mtp.py | 27 + .../test_fused_recurrent_packed_decode.py | 62 +++ .../test_mtp_fused_recurrent_equiv.py | 52 ++ .../test_gated_rmsnorm_strided_gate.py | 41 ++ .../triton_kernel/test_scalar_linear.py | 70 +++ .../test_append_shared_expert_topk.py | 83 +++ unit_tests/common/fused_moe/test_bf16_ep.py | 173 ++++++ .../fused_moe/test_bf16_ep_activation.py | 102 ++++ .../gated_residual/test_gated_residual.py | 287 ++++++++++ .../common/ngram_embedding/test_hash.py | 206 ++++++++ .../ngram_embedding/test_pinned_embedding.py | 123 +++++ unit_tests/common/ngram_embedding/test_ple.py | 259 +++++++++ .../test_row_sharded_embedding.py | 161 ++++++ unit_tests/common/qsa/test_attention.py | 194 +++++++ unit_tests/common/qsa/test_flashinfer_topk.py | 88 ++++ unit_tests/common/qsa/test_indexer.py | 494 ++++++++++++++++++ .../common/qsa/test_topk_compatibility.py | 146 ++++++ unit_tests/common/qsa/test_topk_cuda.py | 62 +++ .../quantization/test_aligned_projection.py | 97 ++++ unit_tests/distributed/test_ordered_tp_sum.py | 192 +++++++ unit_tests/distributed/test_qsa_selection.py | 98 ++++ .../models/qwen2_vl/test_rotary_pos_emb.py | 19 +- ...st_transformer_layer_infer_strided_gate.py | 33 ++ unit_tests/models/qwen4_exp/__init__.py | 1 + unit_tests/models/qwen4_exp/conftest.py | 84 +++ unit_tests/models/qwen4_exp/helpers.py | 133 +++++ .../models/qwen4_exp/lexical_helpers.py | 32 ++ unit_tests/models/qwen4_exp/mtp_reference.py | 27 + unit_tests/models/qwen4_exp/reference.py | 235 +++++++++ .../models/qwen4_exp/resource_helpers.py | 74 +++ unit_tests/models/qwen4_exp/runtime.py | 156 ++++++ .../models/qwen4_exp/test_attention_state.py | 303 +++++++++++ unit_tests/models/qwen4_exp/test_batch.py | 121 +++++ unit_tests/models/qwen4_exp/test_decoder.py | 218 ++++++++ .../qwen4_exp/test_embedding_residency.py | 45 ++ unit_tests/models/qwen4_exp/test_eos.py | 20 + .../models/qwen4_exp/test_lexical_cache.py | 125 +++++ .../models/qwen4_exp/test_linear_decoder.py | 184 +++++++ unit_tests/models/qwen4_exp/test_lm_head.py | 45 ++ unit_tests/models/qwen4_exp/test_mtp.py | 311 +++++++++++ .../models/qwen4_exp/test_multimodal.py | 213 ++++++++ unit_tests/models/qwen4_exp/test_norm.py | 48 ++ .../models/qwen4_exp/test_pinned_memory.py | 112 ++++ unit_tests/models/qwen4_exp/test_position.py | 28 + .../models/qwen4_exp/test_qsa_compressor.py | 218 ++++++++ unit_tests/models/qwen4_exp/test_reduction.py | 104 ++++ .../models/qwen4_exp/test_replicated_moe.py | 231 ++++++++ unit_tests/models/qwen4_exp/test_runtime.py | 397 ++++++++++++++ .../qwen4_exp/test_sequence_parallel.py | 133 +++++ .../qwen4_exp/test_speculative_cache.py | 159 ++++++ unit_tests/models/qwen4_exp/test_vision.py | 49 ++ unit_tests/models/qwen4_exp/test_weights.py | 99 ++++ .../models/qwen4_exp/vision_reference.py | 93 ++++ .../models/sensenova_loop/test_gdn_state.py | 15 +- .../core/objs/test_stop_sequence_output.py | 24 + .../httpserver/test_slot_availability.py | 274 ++++++++++ .../mtp_speculative/test_eagle_with_att.py | 10 +- .../mtp_speculative/test_planner.py | 83 +++ .../req_queue/test_cached_prefix_queue.py | 224 ++++++++ .../router/req_queue/test_peak_reservation.py | 96 ++++ unit_tests/server/test_api_start_deepep.py | 49 +- .../server/test_native_stream_details.py | 69 +++ unit_tests/server/test_ssm_state_mode.py | 6 +- unit_tests/utils/test_envs_utils.py | 41 ++ unit_tests/utils/test_net_utils.py | 45 ++ .../utils/test_startup_process_lifecycle.py | 46 ++ 191 files changed, 16925 insertions(+), 217 deletions(-) create mode 100644 docs/qwen38_adaptation.rst create mode 100644 lightllm/common/basemodel/attention/linear/flashinfer.py create mode 100644 lightllm/common/basemodel/layer_weights/meta_weights/pinned_embedding_weight.py create mode 100644 lightllm/common/basemodel/layer_weights/meta_weights/row_sharded_embedding_weight.py create mode 100644 lightllm/common/basemodel/tp_reduce.py create mode 100644 lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep/bf16.py create mode 100644 lightllm/common/basemodel/triton_kernel/fused_moe/mask_topk.py create mode 100644 lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d_triton.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/replicated_decoder.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/qsa_compressor.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/triton_kernel/replicated_moe.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/__init__.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 unit_tests/common/basemodel/attention/linear/test_flashinfer_gdn.py create mode 100644 unit_tests/common/basemodel/triton_kernel/linear_att/test_causal_conv1d.py create mode 100644 unit_tests/common/basemodel/triton_kernel/linear_att/test_fused_recurrent_packed_decode.py create mode 100644 unit_tests/common/basemodel/triton_kernel/test_gated_rmsnorm_strided_gate.py create mode 100644 unit_tests/common/basemodel/triton_kernel/test_scalar_linear.py create mode 100644 unit_tests/common/fused_moe/test_bf16_ep.py create mode 100644 unit_tests/common/fused_moe/test_bf16_ep_activation.py create mode 100644 unit_tests/common/gated_residual/test_gated_residual.py create mode 100644 unit_tests/common/ngram_embedding/test_hash.py create mode 100644 unit_tests/common/ngram_embedding/test_pinned_embedding.py create mode 100644 unit_tests/common/ngram_embedding/test_ple.py create mode 100644 unit_tests/common/ngram_embedding/test_row_sharded_embedding.py create mode 100644 unit_tests/common/qsa/test_attention.py create mode 100644 unit_tests/common/qsa/test_flashinfer_topk.py create mode 100644 unit_tests/common/qsa/test_indexer.py create mode 100644 unit_tests/common/qsa/test_topk_compatibility.py create mode 100644 unit_tests/common/qsa/test_topk_cuda.py create mode 100644 unit_tests/common/quantization/test_aligned_projection.py create mode 100644 unit_tests/distributed/test_ordered_tp_sum.py create mode 100644 unit_tests/distributed/test_qsa_selection.py create mode 100644 unit_tests/models/qwen3next/test_transformer_layer_infer_strided_gate.py create mode 100644 unit_tests/models/qwen4_exp/__init__.py create mode 100644 unit_tests/models/qwen4_exp/conftest.py create mode 100644 unit_tests/models/qwen4_exp/helpers.py create mode 100644 unit_tests/models/qwen4_exp/lexical_helpers.py create mode 100644 unit_tests/models/qwen4_exp/mtp_reference.py create mode 100644 unit_tests/models/qwen4_exp/reference.py create mode 100644 unit_tests/models/qwen4_exp/resource_helpers.py create mode 100644 unit_tests/models/qwen4_exp/runtime.py create mode 100644 unit_tests/models/qwen4_exp/test_attention_state.py create mode 100644 unit_tests/models/qwen4_exp/test_batch.py create mode 100644 unit_tests/models/qwen4_exp/test_decoder.py create mode 100644 unit_tests/models/qwen4_exp/test_embedding_residency.py create mode 100644 unit_tests/models/qwen4_exp/test_eos.py create mode 100644 unit_tests/models/qwen4_exp/test_lexical_cache.py create mode 100644 unit_tests/models/qwen4_exp/test_linear_decoder.py create mode 100644 unit_tests/models/qwen4_exp/test_lm_head.py create mode 100644 unit_tests/models/qwen4_exp/test_mtp.py create mode 100644 unit_tests/models/qwen4_exp/test_multimodal.py create mode 100644 unit_tests/models/qwen4_exp/test_norm.py create mode 100644 unit_tests/models/qwen4_exp/test_pinned_memory.py create mode 100644 unit_tests/models/qwen4_exp/test_position.py create mode 100644 unit_tests/models/qwen4_exp/test_qsa_compressor.py create mode 100644 unit_tests/models/qwen4_exp/test_reduction.py create mode 100644 unit_tests/models/qwen4_exp/test_replicated_moe.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 create mode 100644 unit_tests/utils/test_net_utils.py diff --git a/docs/qwen38_adaptation.rst b/docs/qwen38_adaptation.rst new file mode 100644 index 0000000000..1aaec169ab --- /dev/null +++ b/docs/qwen38_adaptation.rst @@ -0,0 +1,85 @@ +Qwen3.8-Flash-Next +================== + +Qwen3.8-Flash-Next checkpoints use the ``qwen4_exp`` model type; +text-only checkpoints use ``qwen4_exp_text``. 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; enable it with +``--offload_ngram_embedding``. Shared-prefix admission is limited to the +model's single-node, DP1 normal serving mode with one-page state blocks. + +QSA supports BF16 KV storage and ``--llm_kv_type fp8kv_spt``. The latter +stores main KV and pooled indexer keys as FP8 e4m3 with unit scales; raw +indexer keys and rotary inputs retain the model dtype. Calibration files +are not supported for this QSA format. With ``--page_size 4``, pooled +indexer keys use one entry per compression group and unfinished groups +use request-local ring storage that includes MTP candidate capacity. + +``--enable_ep_moe`` supports unquantized BF16 experts through DeepEP and +DeepGEMM. It does not require ``--expert_dtype fp8``. The BF16 path supports +ordinary prefill/decode and CUDA Graphs; split dispatch/GEMM interfaces +used by microbatch overlap still require quantized experts. + +GDN supports ``--ssm_state_mode replay`` with FP32 SSM state and a bounded +replay history. It uses the upstream GDN state lifecycle: decode advances +the verification window, and MTP acceptance trims rejected tokens using +each request's verification width. Cache pages store prefill Conv/SSM state; +Qwen lexical state participates in the same acceptance call. + +FlashInfer 0.7.2 or later provides an optional GDN prefill backend through +``--llm_prefill_att_backend auto flashinfer``. It accepts the existing KV +state layout through explicit transposes and preserves zero-length states. +Prefill CUDA Graphs reuse the captured GDN subgraph when the request count +matches; other request counts use the current batch state outside the graph. +QSA prefill replays its captured subgraph when the current visible context +fits the captured bound, including batches with a different request count. +Longer contexts use eager QSA attention with the current batch metadata. +The default linear prefill backend selection is unchanged. + +Prefill graphs require workspace in addition to the KV cache. The evaluated +TP8 profile on eight H100 80GB GPUs uses +``--enable_prefill_cudagraph --prefill_cudagraph_max_handle_token 1024`` +with ``--max_total_token_num 2097152``, 64 request slots and a 16K context +limit. Larger prefill batches retain the eager path. The explicit KV limit +reduces spare prefix-cache capacity; it is a deployment setting, not a +model requirement. + +MTP projects draft logits only for accepted-tail rows while retaining all +hidden rows and KV writes. Its graph cache keys include both input and output +capacity, and all supported capacities are captured during startup before +target hidden states can occupy the shared graph pool. Ordered TP reduction +retains rank-ordered FP64 accumulation; the H100 peer path synchronizes +readers before reusing its shared buffer. + +On H100, TP8/DP1 with BF16 weights, FP8 QSA, FP32 ReplaySSM with a +16-token precomputed history and fixed MTP2, MTP3 or MTP4 uses a replicated +small-batch decode path. Complete verification groups of up to 64 rows +retain the complete residual on each TP rank while experts remain +distributed across eight ranks. Graph capacities and padding follow the +verification width; groups that exceed 64 rows use the existing TP/SP +alignment. Each graph shape owns its expert workspace. Shared-expert gating +is combined with the routed-expert sum while preserving BF16 intermediate +rounding. Larger batches, unsupported expert layouts, quantized weights, +microbatch overlap and routed-expert collection use the existing path. + +Within the MTP3 profile, QSA decode with up to 128 rows packs the attention +and indexer projections into one GEMM; prefill and draft execution retain +separate projections. MTP3 also fixes the QSA split count and ReplaySSM kernel +configuration across decode batch sizes. MTP2 and MTP4 retain their existing +QSA and ReplaySSM kernel settings. These optimizations do not establish +end-to-end deterministic generation. + +For gated residuals with up to 64 rows, width at least 8192 and projection +rank at most 512, the down projection and injection gates share a kernel +launch. Their reductions and intermediate rounding remain separate. +This uses the existing weights and split workspace. Larger batches and +FP32 inputs retain separate gate execution. The projection kernels take +the row count at runtime to avoid recompilation for each prefill batch size. + +Kernel microbenchmarks are in ``test/benchmark/qwen38/``. Numerical and +cache-state regression tests are in ``unit_tests/models/qwen4_exp/`` and +the corresponding ``unit_tests/common/`` directories. diff --git a/lightllm/common/attention_state/bounded_att/utils.py b/lightllm/common/attention_state/bounded_att/utils.py index 93f5b2ebe3..c298c4a565 100644 --- a/lightllm/common/attention_state/bounded_att/utils.py +++ b/lightllm/common/attention_state/bounded_att/utils.py @@ -2,7 +2,8 @@ from __future__ import annotations -from math import prod +from math import lcm, prod +from mmap import PAGESIZE from weakref import finalize import torch @@ -17,6 +18,7 @@ CUDA_HOST_REGISTER_PORTABLE = 1 # Mapped：建立设备可访问映射，供 CUDA/Triton kernel 直接访问主机存储。 CUDA_HOST_REGISTER_MAPPED = 2 +_HOST_REGISTRATION_BYTES = 1024 * 1024 * 1024 # ======================================================================== @@ -43,16 +45,25 @@ def create_page_pool_tensor( # 阶段一：未启用 DP fetch 时沿用 torch 分配；GPU 页池也保持原设备分配方式。 if not shared: - tensor = torch.empty( - size=size, - dtype=dtype, - device=device, - pin_memory=pin_memory, - ) + try: + tensor = torch.empty( + size=size, + dtype=dtype, + device=device, + pin_memory=pin_memory, + ) + except RuntimeError: + if device.type != "cpu" or not pin_memory or not torch.cuda.is_available(): + raise + tensor = torch.empty(size=size, dtype=dtype, device=device, pin_memory=False) + register_page_pool_tensor(tensor=tensor) + allocation = "torch.empty + cudaHostRegister" + else: + allocation = "torch.empty" reason = "GPU device allocation" if device.type != "cpu" else "DP prompt cache fetch disabled" logger.info( f"Allocated bounded attention page pool: shape={size}, dtype={dtype}, bytes={nbytes}, " - f"device={device}, shared=False, pin_memory={pin_memory}, allocator=torch.empty, reason={reason}" + f"device={device}, shared=False, pin_memory={pin_memory}, allocator={allocation}, reason={reason}" ) return tensor @@ -91,24 +102,42 @@ def register_page_pool_tensor(tensor: torch.Tensor) -> None: """ if tensor.device.type != "cpu" or tensor.numel() == 0 or tensor.is_pinned(): return - assert tensor.is_shared(), "page pool registration requires shared CPU storage" storage = tensor.untyped_storage() ptr = storage.data_ptr() nbytes = storage.nbytes() device_id = torch.cuda.current_device() flags = CUDA_HOST_REGISTER_PORTABLE | CUDA_HOST_REGISTER_MAPPED - torch.cuda.check_error(torch.cuda.cudart().cudaHostRegister(ptr, nbytes, flags)) + registration_unit = lcm(prod(tensor.shape[1:]) * tensor.element_size(), PAGESIZE) + chunk_bytes = max(registration_unit, _HOST_REGISTRATION_BYTES // registration_unit * registration_unit) + if not tensor.is_shared(): + chunk_bytes = nbytes + addresses = [] + runtime = torch.cuda.cudart() + try: + for offset in range(0, nbytes, chunk_bytes): + address = ptr + offset + torch.cuda.check_error(runtime.cudaHostRegister(address, min(chunk_bytes, nbytes - offset), flags)) + addresses.append(address) + except BaseException: + logger.error( + f"Failed to register bounded attention page pool: bytes={nbytes}, cuda_device={device_id}, " + f"chunk_bytes={chunk_bytes}, registered_chunks={len(addresses)}" + ) + if addresses: + _unregister_page_pool_storage(tuple(addresses), device_id) + raise # finalizer 绑定 Storage，所有本地 Tensor 及视图释放后才注销，不能绑在单个 Tensor 上。 # 正常释放时等待 GPU 完成使用；解释器退出时交给进程清理，避免 CUDA 已关闭后再调用。 cleanup = finalize( storage, _unregister_page_pool_storage, - ptr, + tuple(addresses), device_id, ) cleanup.atexit = False logger.info( - f"Registered shared bounded attention page pool: bytes={nbytes}, cuda_device={device_id}, " + f"Registered shared bounded attention page pool: bytes={nbytes}, chunks={len(addresses)}, " + f"cuda_device={device_id}, " "flags=Portable|Mapped; registration applies to this process's mapping" ) @@ -118,8 +147,9 @@ def register_page_pool_tensor(tensor: torch.Tensor) -> None: # ======================================================================== -def _unregister_page_pool_storage(ptr: int, device_id: int) -> None: +def _unregister_page_pool_storage(addresses: tuple[int, ...], device_id: int) -> None: """等待当前 GPU 使用结束，再注销本进程的映射；共享内存本身由 Storage 管理。""" torch.cuda.synchronize(device=device_id) - torch.cuda.check_error(torch.cuda.cudart().cudaHostUnregister(ptr)) - logger.debug(f"Unregistered shared bounded attention page pool: cuda_device={device_id}, ptr={ptr}") + for ptr in reversed(addresses): + torch.cuda.check_error(torch.cuda.cudart().cudaHostUnregister(ptr)) + logger.debug(f"Unregistered shared bounded attention page pool: cuda_device={device_id}, chunks={len(addresses)}") diff --git a/lightllm/common/attention_state/full_att/gqa/gqa_state.py b/lightllm/common/attention_state/full_att/gqa/gqa_state.py index fe184b4110..c658533abc 100644 --- a/lightllm/common/attention_state/full_att/gqa/gqa_state.py +++ b/lightllm/common/attention_state/full_att/gqa/gqa_state.py @@ -146,14 +146,13 @@ def write_full_att_kv(self, logical_layer_id: int, token_indices: torch.Tensor, if not len(token_indices): return assert self.k_storage.is_cuda, "KV writes require CUDA storage" - from lightllm.common.basemodel.triton_kernel.destindex_copy_kv import destindex_copy_kv + from lightllm.common.basemodel.triton_kernel.destindex_copy_kv import destindex_copy_kv_fused - # TODO：将同一层 K/V 的两次索引写入融合到一个 kernel，直接写入独立 K/V 存储。 - destindex_copy_kv( - K=kv[:, : self.config.head_num], DestLoc=token_indices, Out=self.k_storage[storage_layer_index] - ) - destindex_copy_kv( - K=kv[:, self.config.head_num :], DestLoc=token_indices, Out=self.v_storage[storage_layer_index] + destindex_copy_kv_fused( + KV=kv, + DestLoc=token_indices, + OutK=self.k_storage[storage_layer_index], + OutV=self.v_storage[storage_layer_index], ) def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Tensor) -> None: @@ -169,22 +168,16 @@ def copy_full_att_tokens(self, src_indices: torch.Tensor, dst_indices: torch.Ten assert self.k_storage.is_cuda, "token copies require CUDA storage" if not src_indices.numel(): return - from lightllm.common.basemodel.triton_kernel.kv_move import copy_kv_buffer_to_kv_buffer + from lightllm.common.basemodel.triton_kernel.kv_move import copy_kv_buffer_pair_to_kv_buffer_pair # 直接使用调用方提供的 CUDA 索引，保留原 dtype 和 stride。 # 按字节直接从源槽位写到目标槽位，覆盖量化 scale 和 packed padding，不创建中间 KV。 # 调用方保证源、目标集合不相交，且目标索引不重复，避免并行读写竞争。 - # TODO：将 K/V 的同索引 token 复制融合为一次 kernel 发射， - # 保持原始字节编码及各字段 stride，仍要求源目标不重叠、目标索引不重复。 - copy_kv_buffer_to_kv_buffer( - src_mem_index=src_indices, - dst_mem_index=dst_indices, - kv_buffer=self.k_storage, - ) - copy_kv_buffer_to_kv_buffer( + copy_kv_buffer_pair_to_kv_buffer_pair( src_mem_index=src_indices, dst_mem_index=dst_indices, - kv_buffer=self.v_storage, + k_buffer=self.k_storage, + v_buffer=self.v_storage, ) # ======================================================================== diff --git a/lightllm/common/basemodel/attention/create_linear_utils.py b/lightllm/common/basemodel/attention/create_linear_utils.py index 08528f98f6..6653072593 100644 --- a/lightllm/common/basemodel/attention/create_linear_utils.py +++ b/lightllm/common/basemodel/attention/create_linear_utils.py @@ -1,11 +1,4 @@ -from lightllm.common.basemodel.attention.linear.flashqla import FlashQlaLinearAttBackend -from lightllm.common.basemodel.attention.linear.kda import KDALinearAttBackend -from lightllm.common.basemodel.attention.linear.replay_gdn import ( - ReplayGDNFlashQlaAttBackend, - ReplayGDNTritonAttBackend, -) -from lightllm.common.basemodel.attention.linear.replay_kda import ReplayKDALinearAttBackend -from lightllm.common.basemodel.attention.linear.triton import TritonLinearAttBackend +from .import_utils import load_backend_class from lightllm.utils.backend_validator import validate from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.log_utils import init_logger @@ -13,29 +6,31 @@ logger = init_logger(__name__) gdn_prefill_backends = { - "flashqla": FlashQlaLinearAttBackend, - "triton": TritonLinearAttBackend, + "flashqla": "linear.flashqla:FlashQlaLinearAttBackend", + "flashinfer": "linear.flashinfer:FlashInferLinearAttBackend", + "triton": "linear.triton:TritonLinearAttBackend", } gdn_replay_prefill_backends = { - "flashqla": ReplayGDNFlashQlaAttBackend, - "triton": ReplayGDNTritonAttBackend, + "flashqla": "linear.replay_gdn:ReplayGDNFlashQlaAttBackend", + "flashinfer": "linear.replay_gdn:ReplayGDNFlashInferAttBackend", + "triton": "linear.replay_gdn:ReplayGDNTritonAttBackend", } gdn_decode_backends = { - "triton": TritonLinearAttBackend, + "triton": "linear.triton:TritonLinearAttBackend", } gdn_replay_decode_backends = { - "triton": ReplayGDNTritonAttBackend, + "triton": "linear.replay_gdn:ReplayGDNTritonAttBackend", } kda_backends = { - "triton": KDALinearAttBackend, + "triton": "linear.kda:KDALinearAttBackend", } kda_replay_backends = { - "triton": ReplayKDALinearAttBackend, + "triton": "linear.replay_kda:ReplayKDALinearAttBackend", } @@ -54,13 +49,13 @@ def get_gdn_linear_decode_att_backend_class(index=1, priority_list=("triton",)): def get_kda_linear_att_backend_class(): args = get_env_start_args() classes = kda_replay_backends if args.ssm_state_mode == "replay" else kda_backends - return classes["triton"] + return load_backend_class(classes["triton"]) def _select_backend(classes, backend_args, index, priority_list, stage): backend_str = _get_backend_str(backend_args, index) if backend_str != "auto": - return classes[backend_str] + return load_backend_class(classes[backend_str]) for backend_name in priority_list: if backend_name == "triton" or validate(backend_name): @@ -68,10 +63,10 @@ def _select_backend(classes, backend_args, index, priority_list, stage): logger.info(f"Linear {stage} attention backend: triton.") else: logger.info(f"Linear {stage} attention backend: {backend_name} (validated).") - return classes[backend_name] + return load_backend_class(classes[backend_name]) logger.warning(f"No linear {stage} attention backend validation succeeded, falling back to Triton.") - return classes["triton"] + return load_backend_class(classes["triton"]) def _get_backend_str(backend_args, index): diff --git a/lightllm/common/basemodel/attention/linear/flashinfer.py b/lightllm/common/basemodel/attention/linear/flashinfer.py new file mode 100644 index 0000000000..d4217144b2 --- /dev/null +++ b/lightllm/common/basemodel/attention/linear/flashinfer.py @@ -0,0 +1,63 @@ +import torch + +from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.l2norm import l2norm_fwd + +from .gdn import LinearAttBackend + + +def chunk_gated_delta_rule( + q, + k, + v, + g, + beta, + initial_state=None, + output_final_state=False, + cu_seqlens=None, + use_qk_l2norm_in_kernel=False, + scale=None, +): + from flashinfer.gdn_prefill import chunk_gated_delta_rule as flashinfer_gdn + + if q.shape[0] != 1 or cu_seqlens is None: + raise ValueError("FlashInfer GDN requires packed variable-length inputs") + q, k, v = (tensor.squeeze(0).contiguous() for tensor in (q, k, v)) + if use_qk_l2norm_in_kernel: + q, k = l2norm_fwd(q), l2norm_fwd(k) + state = None + if initial_state is not None: + state = initial_state.transpose(-1, -2).to(torch.float32).contiguous() + final = None + if output_final_state: + final = ( + state.clone() + if state is not None + else torch.zeros( + (cu_seqlens.numel() - 1, max(q.shape[1], v.shape[1]), v.shape[2], k.shape[2]), + device=q.device, + dtype=torch.float32, + ) + ) + result = flashinfer_gdn( + q=q, + k=k, + v=v, + g=g.squeeze(0).float().exp(), + beta=beta.squeeze(0).float().contiguous(), + scale=scale, + initial_state=state, + output_final_state=output_final_state, + output_state=final, + cu_seqlens=cu_seqlens.to(torch.int64), + use_cp=False if q.shape[0] <= 1024 * (cu_seqlens.numel() - 1) else "auto", + backend="flashinfer", + ) + if output_final_state: + output, state = result + return output.unsqueeze(0), state.transpose(-1, -2).contiguous() + return result.unsqueeze(0), None + + +class FlashInferLinearAttBackend(LinearAttBackend): + def get_prefill_kernel(self): + return chunk_gated_delta_rule diff --git a/lightllm/common/basemodel/attention/linear/gdn.py b/lightllm/common/basemodel/attention/linear/gdn.py index 986214568f..344d58866e 100644 --- a/lightllm/common/basemodel/attention/linear/gdn.py +++ b/lightllm/common/basemodel/attention/linear/gdn.py @@ -8,7 +8,6 @@ from lightllm.utils.envs_utils import get_env_start_args from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import causal_conv1d_fn from lightllm.common.basemodel.triton_kernel.linear_att.fused_gdn_gating import fused_gdn_gating -from lightllm.common.basemodel.triton_kernel.linear_att.gdn_decode_pack import conv_pack_gdn_decode_inputs from lightllm.common.basemodel.triton_kernel.linear_att.mtp_fused_recurrent import ( mtp_fused_recurrent_gated_delta_rule, ) @@ -16,7 +15,7 @@ build_dynamic_mtp_linear_att_state_params, ) from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops import ( - fused_recurrent_gated_delta_rule, + fused_recurrent_gated_delta_rule_packed_decode, ) if TYPE_CHECKING: @@ -37,7 +36,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 +107,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 @@ -180,6 +177,7 @@ def _gdn_prefill_kernel( has_initial_state=infer_state.b_ready_cache_len > 0, conv_states=conv_states, activation=backend.activation, + max_seqlen=infer_state.max_q_seq_len, ) mixed_qkv = out_tensor.transpose(0, 1) @@ -207,11 +205,11 @@ def _gdn_prefill_kernel( @dataclasses.dataclass class LinearAttDecodeAttState(BaseDecodeAttState): - b_conv_buffer_idx: torch.Tensor = None b_ssm_buffer_idx: torch.Tensor = None b1_mtp_cu_q_seq_len: torch.Tensor = None b_num_accepted_tokens: torch.Tensor = None + b1_mtp_cu_q_seq_len_int64: torch.Tensor = None def init_state(self): draft_step = self.backend.model.mtp_manager.get_decode_draft_step(self.backend.model.is_mtp_draft_model) @@ -232,6 +230,7 @@ def init_state(self): self._init_dynamic_mtp_decode_state(draft_step + 1, req_to_mtp_state_index=state.mtp_state_index) else: self._init_fixed_mtp_decode_state(draft_step, req_to_mtp_state_index=state.mtp_state_index) + self.b1_mtp_cu_q_seq_len_int64 = self.b1_mtp_cu_q_seq_len.to(torch.int64) def _init_normal_decode_state(self): self.b_conv_buffer_idx = self.infer_state.b_req_idx @@ -345,37 +344,26 @@ def _gdn_decode_kernel( ): backend: LinearAttBackend = self.backend - # Recurrent processing with fused gating. Decode uses a specialized - # conv+pack kernel to avoid materializing the post-conv qkv tensor - # before immediately splitting it into q/k/v. - query, key, value, z, a, b = conv_pack_gdn_decode_inputs( + from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import causal_conv1d_update + + mixed_qkv = causal_conv1d_update( mixed_qkv, - z, - a, - b, conv_states, layer_weight.linear_conv1d.mm_param.weight, layer_weight.linear_conv1d.bias, - self.b_conv_buffer_idx, - backend.activation, - backend.conv_kernel_dim, - backend.tp_num_k_heads, - backend.head_k_dim, - backend.tp_num_v_heads, - backend.head_v_dim, + activation=backend.activation, + conv_state_indices=self.b_conv_buffer_idx, ) - core_attn_out, _ = fused_recurrent_gated_delta_rule( - q=query, - k=key, - v=value, + core_attn_out, _ = fused_recurrent_gated_delta_rule_packed_decode( + mixed_qkv=mixed_qkv, + a=a, + b=b, + A_log=layer_weight.linear_A_log.weight, + dt_bias=layer_weight.linear_dt_bias.weight, + scale=backend.head_k_dim ** -0.5, initial_state=ssm_states, - inplace_final_state=True, ssm_state_indices=self.b_ssm_buffer_idx, use_qk_l2norm_in_kernel=True, - A_log=layer_weight.linear_A_log.weight, - dt_bias=layer_weight.linear_dt_bias.weight, - a_raw=a, - b_raw=b, ) return core_attn_out, z @@ -418,7 +406,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 +414,10 @@ def _gdn_mtp_kernel( dt_bias=layer_weight.linear_dt_bias.weight, a_raw=a, b_raw=b, + run_config=( + {"BV": 32, "num_warps": 4, "num_stages": 1} + if getattr(backend.model, "linear_att_mtp_match_decode", False) + else None + ), ) return core_attn_out diff --git a/lightllm/common/basemodel/attention/linear/replay_gdn.py b/lightllm/common/basemodel/attention/linear/replay_gdn.py index 8e4a081f73..690f7d9a98 100644 --- a/lightllm/common/basemodel/attention/linear/replay_gdn.py +++ b/lightllm/common/basemodel/attention/linear/replay_gdn.py @@ -76,6 +76,20 @@ def create_att_decode_state(self, infer_state: "InferStateInfo"): return ReplayGDNDecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) +class ReplayGDNFlashInferAttBackend(BaseAttBackend): + def __init__(self, model): + super().__init__(model=model) + from .flashinfer import chunk_gated_delta_rule + + _configure_gdn(self, model, chunk_gated_delta_rule) + + def create_att_prefill_state(self, infer_state: "InferStateInfo"): + return ReplayGDNPrefillAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) + + def create_att_decode_state(self, infer_state: "InferStateInfo"): + return ReplayGDNDecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state)) + + @dataclasses.dataclass class ReplayGDNPrefillAttState(BasePrefillAttState): b_conv_buffer_idx: torch.Tensor = None @@ -252,6 +266,7 @@ def _decode(self, mixed_qkv, z, conv_states, a, b, layer_weight, layer_num, stat layer_weight.linear_dt_bias.weight, self.b_conv_buffer_idx, self.b_cu_seqlens, + run_config=backend.replay_run_config, ) return output, z @@ -288,6 +303,7 @@ def _mtp(self, mixed_qkv, conv_states, a, b, layer_weight, layer_num, draft_step layer_weight.linear_dt_bias.weight, self.b_conv_buffer_idx, self.b_cu_seqlens, + run_config=backend.replay_run_config, ) return output @@ -326,6 +342,7 @@ def _configure_gdn(backend, model, prefill_kernel) -> None: ssm_dtypes = {"bfloat16": torch.bfloat16, "float32": torch.float32} backend.ssm_state_dtype = ssm_dtypes.get(get_env_start_args().linear_att_ssm_data_type, torch.bfloat16) backend.prefill_kernel = prefill_kernel + backend.replay_run_config = None def _split_qkvzba(backend, mixed_qkvzba): diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index 0edbdadb94..08f9023274 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -58,6 +58,7 @@ class TpPartBaseModel: + decode_graph_class = CudaGraph is_mtp_draft_model = False # weight class @@ -368,7 +369,7 @@ def _init_cudagraph(self): self.graph = ( None if self.disable_cudagraph - else CudaGraph( + else self.decode_graph_class( batch_step_size_before_split=cuda_graph_grow_step_size, split_batch_size=self.args.graph_split_batch_size * decode_tokens_per_request, batch_step_size_after_split=self.args.graph_grow_step_size * cuda_graph_grow_step_size, @@ -466,6 +467,7 @@ def _create_inferstate(self, model_input: ModelInput, microbatch_index: int = 0) # 特殊模型，特殊模式的特定变量初始化操作。 infer_state.mtp_draft_input_hiddens = model_input.mtp_draft_input_hiddens + infer_state.logits_indices = model_input.logits_indices if infer_state.is_prefill: infer_state.prefill_att_state = ( @@ -585,14 +587,19 @@ def _create_model_output(self, post_output: PostLayerOutput, infer_state: InferS ) return output - def _create_unpad_decode_model_output(self, model_output: ModelOutput, origin_batch_size: int): + def _create_unpad_decode_model_output( + self, model_output: ModelOutput, origin_batch_size: int, origin_logits_num=None + ): padded_batch_size = model_output.logits.shape[0] - if padded_batch_size == origin_batch_size: + if origin_logits_num is not None: + padded_batch_size = model_output.mtp_collector.spec_hidden.shape[0] + if padded_batch_size == origin_batch_size and origin_logits_num is None: return model_output + logits_num = origin_batch_size if origin_logits_num is None else origin_logits_num new_model_output = copy.copy(model_output) - new_model_output.logits = new_model_output.logits[0:origin_batch_size] + new_model_output.logits = new_model_output.logits[:logits_num] if new_model_output.logits_token_ids is not None: - new_model_output.logits_token_ids = new_model_output.logits_token_ids[:origin_batch_size] + new_model_output.logits_token_ids = new_model_output.logits_token_ids[:logits_num] new_model_output.mtp_collector = model_output.mtp_collector.unpad_decode( padded_batch_size=padded_batch_size, origin_batch_size=origin_batch_size, @@ -686,6 +693,8 @@ def _decode( ) origin_batch_size = model_input.batch_size + origin_logits_num = None if model_input.logits_indices is None else model_input.logits_indices.numel() + padded_logits_num = None if origin_logits_num is None else triton.next_power_of_2(max(1, origin_logits_num)) # 空 DP rank 也需要完整的 dummy verify group；TP/SP 切分不能破坏 MTP 分组。 infer_batch_size = self._align_decode_batch_size(max(1, origin_batch_size)) @@ -699,9 +708,15 @@ def _decode( need_capture = False if use_cuda_graph: infer_batch_size = self.graph.find_closest_graph_batch_size(batch_size=infer_batch_size) - need_capture = self.graph.need_capture(infer_batch_size) + if origin_logits_num is None: + need_capture = self.graph.need_capture(infer_batch_size) + else: + need_capture = self.graph.need_capture(infer_batch_size, padded_logits_num) model_input = self._create_padded_decode_model_input(model_input=model_input, new_batch_size=infer_batch_size) + if origin_logits_num is not None: + model_input = copy.copy(model_input) + model_input.logits_indices = F.pad(model_input.logits_indices, (0, padded_logits_num - origin_logits_num)) infer_state = self._create_inferstate(model_input) # attention backend 会根据执行方式准备 capture、replay 或 eager 状态， # 因此必须在 init_att_state 之前完成赋值。 @@ -718,7 +733,9 @@ def _decode( else: model_output = self._token_forward(infer_state) - return self._create_unpad_decode_model_output(model_output, origin_batch_size=origin_batch_size) + return self._create_unpad_decode_model_output( + model_output, origin_batch_size=origin_batch_size, origin_logits_num=origin_logits_num + ) @final def _context_forward(self, infer_state: InferStateInfo): diff --git a/lightllm/common/basemodel/batch_objs.py b/lightllm/common/basemodel/batch_objs.py index d8f2329b24..feb3909af2 100644 --- a/lightllm/common/basemodel/batch_objs.py +++ b/lightllm/common/basemodel/batch_objs.py @@ -51,9 +51,12 @@ class ModelInput: # mtp_draft_input_hiddens 用于模型 mtp 模式下 # 的 draft 模型的输入 mtp_draft_input_hiddens: Optional[torch.Tensor] = None + logits_indices: Optional[torch.Tensor] = None def to_cuda(self): self.check_input() + if self.logits_indices is not None: + self.logits_indices = self.logits_indices.cuda(non_blocking=True) # Prefill 和 decode 都必须提供的公共张量。 self.b_req_idx = self.b_req_idx.cuda(non_blocking=True) @@ -80,6 +83,9 @@ def __post_init__(self): self.check_input() def check_input(self): + if self.logits_indices is not None: + assert not self.is_prefill + assert self.logits_indices.ndim == 1 and self.logits_indices.dtype == torch.int64 if self.input_ids is not None: assert ( self.input_ids.dtype == torch.int64 diff --git a/lightllm/common/basemodel/cuda_graph.py b/lightllm/common/basemodel/cuda_graph.py index ab0e0c0721..3654349ffd 100644 --- a/lightllm/common/basemodel/cuda_graph.py +++ b/lightllm/common/basemodel/cuda_graph.py @@ -87,10 +87,11 @@ def __init__( def can_run(self, batch_size, max_len_in_batch): return batch_size <= self.max_batch_size and max_len_in_batch <= self.graph_max_len_in_batch - def need_capture(self, batch_size): + def need_capture(self, batch_size, logits_num=None): find_batch_size = self.find_closest_graph_batch_size(batch_size) if find_batch_size is not None: - return find_batch_size not in self.graph + key = find_batch_size if logits_num is None else (find_batch_size, logits_num) + return key not in self.graph else: assert False, "dead code" @@ -131,9 +132,12 @@ def _capture_decode(self, decode_func, infer_state: InferStateInfo): with self.torch_memory_saver.cuda_graph(graph_obj, pool=self.mempool): model_output = decode_func(infer_state) - self.graph[batch_size] = (graph_obj, infer_state, model_output) + indices = getattr(infer_state, "logits_indices", None) + key = batch_size if indices is None else (batch_size, indices.numel()) + self.graph[key] = (graph_obj, infer_state, model_output) graph_obj.replay() - self._measure_replay_cost(graph_obj=graph_obj, batch_size=batch_size) + if indices is None: + self._measure_replay_cost(graph_obj=graph_obj, batch_size=batch_size) return model_output def _capture_decode_overlap( @@ -219,7 +223,9 @@ def capture_decode( def _replay(self, infer_state: InferStateInfo): batch_size = infer_state.input_ids.shape[0] - graph_obj, graph_infer_state, graph_output = self.graph[batch_size] + indices = getattr(infer_state, "logits_indices", None) + key = batch_size if indices is None else (batch_size, indices.numel()) + graph_obj, graph_infer_state, graph_output = self.graph[key] graph_infer_state.copy_for_cuda_graph(infer_state) graph_obj.replay() return graph_output @@ -288,6 +294,14 @@ def warmup(self, model): ) model_output: ModelOutput = model.forward(model_input) del model_output + if getattr(model, "supports_selected_logits", False): + logits_num = 1 + while logits_num <= triton.next_power_of_2(batch_size): + model_input.logits_indices = torch.zeros(logits_num, dtype=torch.int64, device="cuda") + model_output = model.forward(model_input) + del model_output + logits_num *= 2 + model_input.logits_indices = None del input_ids del b_req_idx del b_seq_len diff --git a/lightllm/common/basemodel/infer_struct.py b/lightllm/common/basemodel/infer_struct.py index 83b312f991..12766ba244 100755 --- a/lightllm/common/basemodel/infer_struct.py +++ b/lightllm/common/basemodel/infer_struct.py @@ -96,6 +96,7 @@ def __init__(self): # 在开启 mtp_mode 时，mtp draft model # 的输入会用到，其他模型和场景都不会用到 self.mtp_draft_input_hiddens: Optional[torch.Tensor] = None + self.logits_indices: Optional[torch.Tensor] = None # 在单节点多dp的运行模式下，在进行prefill的阶段，如果出现了dp之间数据不平衡的现象， # 可以将推理的数据，进行重新分配到各个dp，在做 att 之前，重新 all to all 到各自的 diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py index 4c79d34e0e..0756a6da0c 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py @@ -129,6 +129,8 @@ def experts( alpha: Optional[float] = None, limit: Optional[float] = None, clamp_up_add_one: bool = True, + token_mask: Optional[torch.Tensor] = None, + shared_expert_func=None, ) -> torch.Tensor: assert is_prefill is not None, "is_prefill must be explicitly specified for fused MoE execution" # Captures MoE topk expert ids for routed-experts metadata when enabled. @@ -152,6 +154,8 @@ def experts( alpha=alpha, limit=limit, clamp_up_add_one=clamp_up_add_one, + token_mask=token_mask, + shared_expert_func=shared_expert_func, ) def experts_with_topk( diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/base_impl.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/base_impl.py index d40157fdd8..aa10055e98 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/base_impl.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/base_impl.py @@ -34,6 +34,8 @@ class FuseMoeBaseImpl(ABC): 专家路径中则是完成修复后的执行 ID。 """ + execution_index_dtype = None + def __init__( self, n_routed_experts: int, @@ -72,6 +74,8 @@ def __call__( alpha: Optional[float] = None, limit: Optional[float] = None, clamp_up_add_one: bool = True, + token_mask: Optional[torch.Tensor] = None, + shared_expert_func=None, ) -> torch.Tensor: assert is_prefill is not None, "is_prefill must be explicitly specified for fused MoE execution" topk_weights, topk_ids = self._select_experts( @@ -94,6 +98,14 @@ def __call__( shared_expert_gate=shared_expert_gate, is_prefill=is_prefill, ) + if token_mask is not None: + if token_mask.shape != (input_tensor.shape[0],) or token_mask.dtype != torch.bool: + raise ValueError("MoE token mask must contain one boolean per input row") + if token_mask.device != input_tensor.device: + raise ValueError("MoE token mask must be on the input device") + from lightllm.common.basemodel.triton_kernel.fused_moe.mask_topk import mask_topk + + topk_weights, topk_ids = mask_topk(topk_weights, topk_ids, token_mask, self.execution_index_dtype) return self._fused_experts( input_tensor=input_tensor, w13=w13, @@ -105,6 +117,7 @@ def __call__( alpha=alpha, limit=limit, clamp_up_add_one=clamp_up_add_one, + shared_expert_func=shared_expert_func, ) @abstractmethod diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py index f2ac6c6fa9..50dec8df2d 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py @@ -28,6 +28,8 @@ class FuseMoeDeepGEMM(FuseMoeBaseImpl): + execution_index_dtype = torch.int64 + def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) world_size = get_global_world_size() 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..222901da06 --- /dev/null +++ b/lightllm/common/basemodel/layer_weights/meta_weights/pinned_embedding_weight.py @@ -0,0 +1,84 @@ +import torch +import triton + +from .row_sharded_embedding_weight import RowShardedEmbeddingWeight +from lightllm.common.basemodel.triton_kernel.embedding import embedding_kernel +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 + embedding_kernel[(input_ids.numel(),)]( + self.weight, + input_ids, + out, + self.tp_vocab_start_id, + self.tp_vocab_end_id, + self.weight.stride(0), + out.stride(0), + input_ids.numel(), + self.dim, + triton.next_power_of_2(self.dim), + 1, + 1, + num_warps=1, + num_stages=1, + ) + return out + + 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..a0cf9a6811 --- /dev/null +++ b/lightllm/common/basemodel/tp_reduce.py @@ -0,0 +1,73 @@ +import torch +import torch.distributed as dist +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() + shard_size = input.numel() + symmetric = getattr(group, "symm_mem_reduce", None) if input.is_cuda else None + if ( + symmetric is not None + and scatter_rank is not None + and symmetric.ordered_reduce_scatter(source, output, scatter_rank) + ): + return output + gathered = symmetric.ordered_all_gather(source) if symmetric is not None else None + device_group = getattr(group, "device_group", group) + if gathered is None and scatter_rank is not None and input.is_cuda: + if scatter_rank != dist.get_rank(group=device_group): + raise ValueError("Sequence-parallel scatter rank must match the process group rank") + gathered = torch.empty_like(source) + dist.all_to_all_single(gathered, source, group=device_group) + shard_size, start = output.numel(), 0 + elif 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, shard_size, start, output.numel(), world_size, 256 + ) + else: + accumulator = torch.zeros(output.numel(), dtype=torch.float64, device=input.device) + shards = gathered.reshape(world_size, -1) + for rank in range(world_size): + accumulator.add_(shards[rank, start : start + output.numel()]) + output.copy_(accumulator.reshape(output.shape)) + if scatter_rank is None and output is not input: + input.copy_(output) + return input + return output diff --git a/lightllm/common/basemodel/triton_kernel/destindex_copy_kv.py b/lightllm/common/basemodel/triton_kernel/destindex_copy_kv.py index 060a92bf7c..88394c48e9 100644 --- a/lightllm/common/basemodel/triton_kernel/destindex_copy_kv.py +++ b/lightllm/common/basemodel/triton_kernel/destindex_copy_kv.py @@ -34,6 +34,43 @@ def _fwd_kernel_destindex_copy_kv( return +@triton.jit +def _fwd_kernel_destindex_copy_kv_fused( + KV, + Dest_loc, + OutK, + OutV, + stride_kv_bs, + stride_kv_h, + stride_kv_d, + stride_ok_bs, + stride_ok_h, + stride_ok_d, + stride_ov_bs, + stride_ov_h, + stride_ov_d, + head_num, + head_dim, + BLOCK_DMODEL: tl.constexpr, + BLOCK_HEAD: tl.constexpr, +): + cur_index = tl.program_id(0) + offs_h = tl.arange(0, BLOCK_HEAD) + offs_d = (tl.arange(0, BLOCK_DMODEL)) % head_dim + + dest_index = tl.load(Dest_loc + cur_index).to(tl.int64) + mask = offs_h[:, None] < head_num + + kv_base = KV + cur_index * stride_kv_bs + offs_h[:, None] * stride_kv_h + offs_d[None, :] * stride_kv_d + v_base = KV + cur_index * stride_kv_bs + (head_num + offs_h[:, None]) * stride_kv_h + offs_d[None, :] * stride_kv_d + k_out = OutK + dest_index * stride_ok_bs + offs_h[:, None] * stride_ok_h + offs_d[None, :] * stride_ok_d + v_out = OutV + dest_index * stride_ov_bs + offs_h[:, None] * stride_ov_h + offs_d[None, :] * stride_ov_d + + tl.store(k_out, tl.load(kv_base, mask=mask, other=0.0), mask=mask) + tl.store(v_out, tl.load(v_base, mask=mask, other=0.0), mask=mask) + return + + @torch.no_grad() def destindex_copy_kv(K, DestLoc, Out): seq_len = DestLoc.shape[0] @@ -62,3 +99,39 @@ def destindex_copy_kv(K, DestLoc, Out): num_stages=1, ) return + + +@torch.no_grad() +def destindex_copy_kv_fused(KV, DestLoc, OutK, OutV): + seq_len = DestLoc.shape[0] + head_num = OutK.shape[1] + head_dim = OutK.shape[2] + assert KV.ndim == 3 and KV.shape[0] == seq_len + assert OutK.shape == OutV.shape + assert KV.shape[1] == 2 * head_num and KV.shape[2] == head_dim + + BLOCK_HEAD = triton.next_power_of_2(head_num) + grid = (seq_len,) + + _fwd_kernel_destindex_copy_kv_fused[grid]( + KV, + DestLoc, + OutK, + OutV, + KV.stride(0), + KV.stride(1), + KV.stride(2), + OutK.stride(0), + OutK.stride(1), + OutK.stride(2), + OutV.stride(0), + OutV.stride(1), + OutV.stride(2), + head_num, + head_dim, + BLOCK_DMODEL=triton.next_power_of_2(head_dim), + BLOCK_HEAD=BLOCK_HEAD, + num_warps=1, + num_stages=1, + ) + return diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep/bf16.py b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep/bf16.py new file mode 100644 index 0000000000..fca5d61da7 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep/bf16.py @@ -0,0 +1,250 @@ +import torch +import triton +import triton.language as tl + +from lightllm.distributed import dist_group_manager +from lightllm.utils.envs_utils import ( + get_deepep_num_max_dispatch_tokens_per_rank_prefill, + get_deepep_num_max_dispatch_tokens_per_rank_decode, +) +from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import silu_and_mul_fwd +from lightllm.common.basemodel.triton_kernel.fused_moe.deepep_expanded_layout_kernels import ep_reduce_decode_output +from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType +from lightllm.utils.config_utils import ffn_use_tanh_approximate_gelu +from .common import get_ep_num_sms +from .deepgemm import deep_gemm, get_mk_alignment_for_contiguous_layout + + +@triton.jit +def _expert_tile( + Input, + Weight, + Output, + expert, + base, + end, + K: tl.constexpr, + N: tl.constexpr, + BM: tl.constexpr, + BN: tl.constexpr, + BK: tl.constexpr, +): + col = tl.program_id(1) * BN + tl.arange(0, BN) + inner = tl.arange(0, BK) + row = base + tl.arange(0, BM) + acc = tl.full((BM, BN), 0, tl.float32) + for first in range(tl.cdiv(K, BK)): + k = first * BK + inner + x = tl.load( + Input + row[:, None].to(tl.int64) * K + k[None, :], + (row[:, None] < end) & (k[None, :] < K), + other=0, + ) + w = tl.load( + Weight + (expert * N + col[None, :]).to(tl.int64) * K + k[:, None], + (col[None, :] < N) & (k[:, None] < K), + other=0, + ) + acc = tl.dot(x, w, acc) + tl.store( + Output + row[:, None].to(tl.int64) * N + col[None, :], + acc, + (row[:, None] < end) & (col[None, :] < N), + ) + + +@triton.jit +def _grouped_gemm( + Input, + Weight, + Output, + ExpertEnds, + K: tl.constexpr, + N: tl.constexpr, + ALIGNMENT: tl.constexpr, + BM: tl.constexpr, + BN: tl.constexpr, + BK: tl.constexpr, +): + expert = tl.program_id(0) + end = tl.load(ExpertEnds + expert) + previous = tl.load(ExpertEnds + expert - 1, expert > 0, other=0) + start = tl.cdiv(previous, ALIGNMENT) * ALIGNMENT + if end > start: + if end - start <= BM: + _expert_tile(Input, Weight, Output, expert, start, end, K, N, BM, BN, BK) + else: + for base in range(start, end, BM): + _expert_tile(Input, Weight, Output, expert, base, end, K, N, BM, BN, BK) + + +def _expert_gemm(input, weight, output, expert_ends, expected_m, is_prefill): + if not is_prefill and expected_m <= 8: + _grouped_gemm[(weight.shape[0], triton.cdiv(weight.shape[1], 64))]( + input, + weight, + output, + expert_ends, + input.shape[1], + weight.shape[1], + get_mk_alignment_for_contiguous_layout(), + 16, + 64, + 256 if input.shape[1] >= 2048 and weight.shape[1] <= 2048 else 128, + num_warps=4, + num_stages=3, + ) + else: + deep_gemm.m_grouped_bf16_gemm_nt_contiguous( + input, + weight, + output, + expert_ends, + use_psum_layout=True, + expected_m_for_psum_layout=expected_m, + ) + + +@triton.jit +def _activate_valid_rows( + Up, + Activated, + ExpertEnds, + N: tl.constexpr, + ALIGNMENT: tl.constexpr, + ALPHA: tl.constexpr, + LIMIT: tl.constexpr, + ADD_ONE: tl.constexpr, + GELU: tl.constexpr, + BM: tl.constexpr, + BN: tl.constexpr, + SHARDS: tl.constexpr = 1, +): + expert = tl.program_id(0) + end = tl.load(ExpertEnds + expert) + previous = tl.load(ExpertEnds + expert - 1, expert > 0, other=0) + start = tl.cdiv(previous, ALIGNMENT) * ALIGNMENT + col = tl.program_id(1) * BN + tl.arange(0, BN) + for base in range(start + tl.program_id(2) * BM, end, SHARDS * BM): + row = base + tl.arange(0, BM) + mask = (row[:, None] < end) & (col[None, :] < N) + gate = tl.load(Up + row[:, None].to(tl.int64) * (2 * N) + col[None, :], mask, other=0.0).to(tl.float32) + value = tl.load(Up + row[:, None].to(tl.int64) * (2 * N) + N + col[None, :], mask, other=0.0) + if LIMIT is not None: + gate = tl.minimum(gate, LIMIT) + value = tl.minimum(tl.maximum(value, -LIMIT), LIMIT) + gate = gate / (1.0 + tl.exp(-gate * ALPHA)) + if ADD_ONE: + value += 1 + elif GELU: + arg = 0.7978845608028654 * (gate + 0.044715 * gate * gate * gate) + gate = 0.5 * gate * (1.0 + (2.0 / (1.0 + tl.exp(-2.0 * arg)) - 1.0)) + else: + gate = gate / (1.0 + tl.exp(-gate)) + gate = gate.to(Up.dtype.element_ty) + tl.store(Activated + row[:, None].to(tl.int64) * N + col[None, :], value * gate, mask) + + +def fused_experts_bf16( + hidden_states, + w13, + w2, + topk_weights, + topk_idx, + num_experts, + is_prefill, + alpha=None, + limit=None, + clamp_up_add_one=True, + shared_expert_func=None, +): + from .prefill_decode_sync import _combine_with_optional_shared_expert + + if hidden_states.dtype != torch.bfloat16 or w13.dtype != torch.bfloat16 or w2.dtype != torch.bfloat16: + raise ValueError("Unquantized DeepEP experts require BF16 activations and weights") + sync_prefill = is_prefill and not torch.cuda.is_current_stream_capturing() + buffer = dist_group_manager.ep_buffer + capacity = ( + get_deepep_num_max_dispatch_tokens_per_rank_prefill() + if is_prefill + else get_deepep_num_max_dispatch_tokens_per_rank_decode() + ) + recv_x, _, recv_weights, handle, _ = buffer.dispatch( + hidden_states, + topk_idx=topk_idx, + topk_weights=topk_weights, + num_experts=num_experts, + num_max_tokens_per_rank=capacity, + expert_alignment=get_mk_alignment_for_contiguous_layout(), + num_sms=get_ep_num_sms(overlap_with_compute=False), + async_with_compute_stream=False, + allocate_on_comm_stream=False, + do_cpu_sync=sync_prefill, + do_handle_copy=False, + do_expand=True, + do_zero_padding=False, + ) + rows = recv_x.shape[0] + if rows: + expected_m = max(1, triton.cdiv(hidden_states.shape[0] * buffer.num_ranks * topk_idx.shape[1], num_experts)) + up = torch.empty((rows, w13.shape[1]), device=recv_x.device, dtype=recv_x.dtype) + _expert_gemm( + recv_x, + w13, + up, + handle.psum_num_recv_tokens_per_expert, + expected_m, + is_prefill, + ) + activated = torch.empty((rows, w13.shape[1] // 2), device=recv_x.device, dtype=recv_x.dtype) + if sync_prefill: + silu_and_mul_fwd(up, activated, alpha=alpha, limit=limit, clamp_up_add_one=clamp_up_add_one) + else: + activation_shards = 8 if is_prefill and hidden_states.shape[0] >= 128 else 1 + _activate_valid_rows[(w13.shape[0], triton.cdiv(activated.shape[1], 128), activation_shards)]( + up, + activated, + handle.psum_num_recv_tokens_per_expert, + activated.shape[1], + get_mk_alignment_for_contiguous_layout(), + alpha, + limit, + clamp_up_add_one, + ffn_use_tanh_approximate_gelu(), + 16, + 128, + activation_shards, + num_warps=4, + ) + del recv_x, up + expert_output = torch.empty((rows, w2.shape[1]), device=activated.device, dtype=activated.dtype) + _expert_gemm( + activated, + w2, + expert_output, + handle.psum_num_recv_tokens_per_expert, + expected_m, + is_prefill, + ) + del activated + else: + expert_output = torch.empty((0, w2.shape[1]), device=hidden_states.device, dtype=hidden_states.dtype) + if Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.GENERAL): + up = torch.zeros((1, w13.shape[1]), device=hidden_states.device, dtype=hidden_states.dtype) + activated = torch.empty((1, w13.shape[1] // 2), device=hidden_states.device, dtype=hidden_states.dtype) + silu_and_mul_fwd(up, activated, alpha=alpha, limit=limit, clamp_up_add_one=clamp_up_add_one) + dense, metadata = ep_reduce_decode_output( + expert_output=expert_output, + route_weights=recv_weights, + recv_src_metadata=handle.recv_src_metadata, + num_valid_recv_tokens=handle.psum_num_recv_tokens_per_scaleup_rank[-1:], + ) + handle.recv_src_metadata = metadata + return _combine_with_optional_shared_expert( + buffer, + dense, + handle, + None, + False, + shared_expert_func, + ) diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep/prefill_decode_sync.py b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep/prefill_decode_sync.py index 3a4909261d..b9123e96ce 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep/prefill_decode_sync.py +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep/prefill_decode_sync.py @@ -42,6 +42,22 @@ def fused_experts( shared_expert_func: Optional[Callable[[], torch.Tensor]] = None, ): assert (limit is None and alpha is None) or (limit is not None and alpha is not None) + if quant_method.method_name == "none": + from .bf16 import fused_experts_bf16 + + return fused_experts_bf16( + hidden_states, + w13.weight, + w2.weight, + topk_weights, + topk_idx, + num_experts, + is_prefill, + alpha=alpha, + limit=limit, + clamp_up_add_one=clamp_up_add_one, + shared_expert_func=shared_expert_func, + ) check_ep_expert_dtype(quant_method) if use_sm100_mega_moe(quant_method): if limit is not None: diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/mask_topk.py b/lightllm/common/basemodel/triton_kernel/fused_moe/mask_topk.py new file mode 100644 index 0000000000..840e1d716d --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/mask_topk.py @@ -0,0 +1,27 @@ +import torch +import triton +import triton.language as tl + + +@triton.jit +def _mask_topk(Weights, Ids, Mask, OutWeights, OutIds, ROWS: tl.constexpr, K: tl.constexpr, BLOCK: tl.constexpr): + index = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + valid = index < ROWS * K + active = tl.load(Mask + index // K, valid, other=False) + weights = tl.load(Weights + index, valid, other=0) + ids = tl.load(Ids + index, valid, other=-1).to(OutIds.dtype.element_ty) + tl.store(OutWeights + index, tl.where(active, weights, 0), valid) + tl.store(OutIds + index, tl.where(active, ids, -1), valid) + + +def mask_topk(weights, ids, mask, index_dtype=None): + index_dtype = ids.dtype if index_dtype is None else index_dtype + if not weights.is_cuda or not weights.is_contiguous() or not ids.is_contiguous() or not mask.is_contiguous(): + return weights.masked_fill(~mask[:, None], 0), ids.masked_fill(~mask[:, None], -1).to(index_dtype) + output_weights = torch.empty_like(weights) + output_ids = torch.empty_like(ids, dtype=index_dtype) + if ids.numel(): + _mask_topk[(triton.cdiv(ids.numel(), 128),)]( + weights, ids, mask, output_weights, output_ids, ids.shape[0], ids.shape[1], 128, num_warps=4 + ) + return output_weights, output_ids diff --git a/lightllm/common/basemodel/triton_kernel/kv_move.py b/lightllm/common/basemodel/triton_kernel/kv_move.py index ce99a007bd..34b812327b 100644 --- a/lightllm/common/basemodel/triton_kernel/kv_move.py +++ b/lightllm/common/basemodel/triton_kernel/kv_move.py @@ -42,6 +42,64 @@ def _copy_kv_buffer_to_kv_buffer( return +@triton.jit +def _copy_kv_buffer_pair_to_kv_buffer_pair( + mem_num, + src_mem_index, + dst_mem_index, + src_index_stride: tl.constexpr, + dst_index_stride: tl.constexpr, + k_buffer, + v_buffer, + k_buffer_stride_l, + k_buffer_stride_s, + k_buffer_stride_d, + v_buffer_stride_l, + v_buffer_stride_s, + v_buffer_stride_d, + k_buffer_tail_dim, + v_buffer_tail_dim, + BLOCK: tl.constexpr, +): + layer_index = tl.program_id(0).to(tl.int64) + start_index = tl.program_id(1).to(tl.int64) + grid_num = tl.num_programs(1).to(tl.int64) + + k_buffer_stride_l = tl.cast(k_buffer_stride_l, dtype=tl.int64) + k_buffer_stride_s = tl.cast(k_buffer_stride_s, dtype=tl.int64) + k_buffer_stride_d = tl.cast(k_buffer_stride_d, dtype=tl.int64) + v_buffer_stride_l = tl.cast(v_buffer_stride_l, dtype=tl.int64) + v_buffer_stride_s = tl.cast(v_buffer_stride_s, dtype=tl.int64) + v_buffer_stride_d = tl.cast(v_buffer_stride_d, dtype=tl.int64) + + for i in range(start_index, mem_num, grid_num): + src_mem = tl.load(src_mem_index + i * src_index_stride) + dst_mem = tl.load(dst_mem_index + i * dst_index_stride) + for j in range(tl.cdiv(k_buffer_tail_dim, BLOCK)): + offs = j * BLOCK + tl.arange(0, BLOCK) + mask = offs < k_buffer_tail_dim + k_buffer_data = tl.load( + k_buffer + layer_index * k_buffer_stride_l + src_mem * k_buffer_stride_s + offs, mask=mask + ) + tl.store( + k_buffer + layer_index * k_buffer_stride_l + dst_mem * k_buffer_stride_s + offs, + k_buffer_data, + mask=mask, + ) + for j in range(tl.cdiv(v_buffer_tail_dim, BLOCK)): + offs = j * BLOCK + tl.arange(0, BLOCK) + mask = offs < v_buffer_tail_dim + v_buffer_data = tl.load( + v_buffer + layer_index * v_buffer_stride_l + src_mem * v_buffer_stride_s + offs, mask=mask + ) + tl.store( + v_buffer + layer_index * v_buffer_stride_l + dst_mem * v_buffer_stride_s + offs, + v_buffer_data, + mask=mask, + ) + return + + def copy_kv_buffer_to_kv_buffer( src_mem_index: torch.Tensor, dst_mem_index: torch.Tensor, @@ -70,3 +128,42 @@ def copy_kv_buffer_to_kv_buffer( BLOCK=BLOCK, ) return + + +def copy_kv_buffer_pair_to_kv_buffer_pair( + src_mem_index: torch.Tensor, + dst_mem_index: torch.Tensor, + k_buffer: torch.Tensor, + v_buffer: torch.Tensor, +): + assert len(src_mem_index) == len(dst_mem_index) + assert src_mem_index.is_cuda and dst_mem_index.is_cuda and k_buffer.is_cuda and v_buffer.is_cuda + assert k_buffer.shape == v_buffer.shape + k_buffer = k_buffer.view(k_buffer.shape[0], k_buffer.shape[1], -1).view(dtype=torch.uint8) + v_buffer = v_buffer.view(v_buffer.shape[0], v_buffer.shape[1], -1).view(dtype=torch.uint8) + assert k_buffer.shape == v_buffer.shape + BLOCK = 4096 + layer_num = k_buffer.shape[0] + grid = ( + layer_num, + 1024, + ) + _copy_kv_buffer_pair_to_kv_buffer_pair[grid]( + mem_num=len(src_mem_index), + src_mem_index=src_mem_index, + dst_mem_index=dst_mem_index, + src_index_stride=src_mem_index.stride(0), + dst_index_stride=dst_mem_index.stride(0), + k_buffer=k_buffer, + v_buffer=v_buffer, + k_buffer_stride_l=k_buffer.stride(0), + k_buffer_stride_s=k_buffer.stride(1), + k_buffer_stride_d=k_buffer.stride(2), + v_buffer_stride_l=v_buffer.stride(0), + v_buffer_stride_s=v_buffer.stride(1), + v_buffer_stride_d=v_buffer.stride(2), + k_buffer_tail_dim=k_buffer.shape[-1], + v_buffer_tail_dim=v_buffer.shape[-1], + BLOCK=BLOCK, + ) + return diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d.py b/lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d.py index 2bf325340f..44e116391b 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d.py @@ -5,6 +5,38 @@ import torch +def _load_causal_conv1d_fwd(): + try: + from sgl_kernel import causal_conv1d_fwd + + return causal_conv1d_fwd + except (ImportError, OSError): + try: + from vllm.model_executor.layers.mamba.ops.causal_conv1d import ( + causal_conv1d_fn as vllm_causal_conv1d_fn, + ) + + return vllm_causal_conv1d_fn + except (ImportError, OSError): + return None + + +def _load_causal_conv1d_update(): + try: + from sgl_kernel import causal_conv1d_update + + return causal_conv1d_update + except (ImportError, OSError): + try: + from vllm.model_executor.layers.mamba.ops.causal_conv1d import ( + causal_conv1d_update, + ) + + return causal_conv1d_update + except (ImportError, OSError): + return None + + def causal_conv1d_fn( x: torch.Tensor, weight: torch.Tensor, @@ -48,23 +80,52 @@ def causal_conv1d_fn( """ if activation not in [None, "silu", "swish"]: raise NotImplementedError("activation must be None, silu, or swish") - from sgl_kernel import causal_conv1d_fwd + optimized_causal_conv1d = _load_causal_conv1d_fwd() + if optimized_causal_conv1d is None: + from .causal_conv1d_triton import causal_conv1d_fn as triton_causal_conv1d_fn + + return triton_causal_conv1d_fn( + x, + weight, + bias, + query_start_loc=query_start_loc, + cache_indices=cache_indices, + has_initial_state=has_initial_state, + conv_states=conv_states, + activation=activation, + pad_slot_id=pad_slot_id, + **kwargs, + ) if x.stride(-1) != 1: x = x.contiguous() bias = bias.contiguous() if bias is not None else None - causal_conv1d_fwd( - x, - weight, - bias, - conv_states, - query_start_loc, - cache_indices, - has_initial_state, - activation in ["silu", "swish"], - pad_slot_id, - ) + if optimized_causal_conv1d.__module__.startswith("sgl_kernel"): + optimized_causal_conv1d( + x, + weight, + bias, + conv_states, + query_start_loc, + cache_indices, + has_initial_state, + activation in ["silu", "swish"], + pad_slot_id, + ) + else: + return optimized_causal_conv1d( + x, + weight, + bias, + conv_states, + query_start_loc, + cache_indices=cache_indices, + has_initial_state=has_initial_state, + activation=activation, + pad_slot_id=pad_slot_id, + null_block_id=None, + ) return x @@ -102,22 +163,57 @@ def causal_conv1d_update( """ if activation not in [None, "silu", "swish"]: raise NotImplementedError(f"activation must be None, silu, or swish, actual: {activation}") - from sgl_kernel import causal_conv1d_update as causal_conv1d_update_kernel + optimized_causal_conv1d_update = _load_causal_conv1d_update() + if optimized_causal_conv1d_update is None: + from .causal_conv1d_triton import causal_conv1d_update as triton_causal_conv1d_update + + return triton_causal_conv1d_update( + x, + conv_state, + weight, + bias, + activation=activation, + cache_seqlens=cache_seqlens, + conv_state_indices=conv_state_indices, + pad_slot_id=pad_slot_id, + ) + + if cache_seqlens is not None: + from .causal_conv1d_triton import causal_conv1d_update as triton_causal_conv1d_update + + return triton_causal_conv1d_update( + x, + conv_state, + weight, + bias, + activation=activation, + cache_seqlens=cache_seqlens, + conv_state_indices=conv_state_indices, + pad_slot_id=pad_slot_id, + ) activation_val = activation in ["silu", "swish"] - unsqueeze = x.dim() == 2 - if unsqueeze: - x = x.unsqueeze(-1) - causal_conv1d_update_kernel( + if optimized_causal_conv1d_update.__module__.startswith("sgl_kernel"): + unsqueeze = x.dim() == 2 + if unsqueeze: + x = x.unsqueeze(-1) + optimized_causal_conv1d_update( + x, + conv_state, + weight, + bias, + activation_val, + None, + conv_state_indices, + pad_slot_id=pad_slot_id, + ) + return x.squeeze(-1) if unsqueeze else x + return optimized_causal_conv1d_update( x, conv_state, weight, bias, - activation_val, - cache_seqlens, - conv_state_indices, - pad_slot_id, + activation=activation, + conv_state_indices=conv_state_indices, + null_block_id=pad_slot_id, ) - if unsqueeze: - x = x.squeeze(-1) - return x diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d_triton.py b/lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d_triton.py new file mode 100644 index 0000000000..f417ad5c45 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/linear_att/causal_conv1d_triton.py @@ -0,0 +1,407 @@ +# SPDX-License-Identifier: Apache-2.0 + +from typing import Optional + +import torch +import triton +import triton.language as tl + +from lightllm.common.triton_utils.autotuner import AutotuneKernelType, autotune + + +def _prefill_configs(): + return [ + {"BLOCK_TOKENS": block_tokens, "BLOCK_CHANNELS": block_channels, "num_warps": num_warps} + for block_tokens in (8, 16, 32, 64) + for block_channels in (64, 128, 256) + for num_warps in (2, 4, 8) + if block_tokens * block_channels <= num_warps * 32 * 64 + ] + + +def _update_configs(): + return [ + {"BLOCK_CHANNELS": block_channels, "num_warps": num_warps} + for block_channels in (64, 128, 256) + for num_warps in (1, 2, 4) + ] + + +def _conv_static_key(x, weight, bias, activation, conv_state): + if x.stride(-2) == 1: + layout = "channel" + elif x.stride(-1) == 1: + layout = "token" + else: + layout = "strided" + return { + "channels": weight.shape[0], + "width": weight.shape[1], + "dtype": str(x.dtype), + "layout": layout, + "state_stride": f"{conv_state.stride(-2)}_{conv_state.stride(-1)}", + "bias": bias is not None, + "silu": activation is not None, + } + + +def _prefill_static_key(x, weight, bias, conv_states, activation): + return _conv_static_key(x, weight, bias, activation, conv_states) + + +def _prefill_run_key(x, query_start_loc, max_seqlen): + batch_size = query_start_loc.numel() - 1 + max_query_len = x.shape[-1] if max_seqlen is None else max_seqlen + return f"{batch_size}_{max_query_len}" + + +def _prefill_key_distance(run_key, config_key): + batch_size, max_query_len = map(int, run_key.split("_")) + config_batch_size, config_max_query_len = map(int, config_key.split("_")) + return ( + abs(batch_size * max_query_len - config_batch_size * config_max_query_len), + abs(batch_size - config_batch_size), + ) + + +def _update_static_key(x, conv_state, weight, bias, activation, cache_seqlens): + input_3d = x.unsqueeze(-1) if x.ndim == 2 else x + key = _conv_static_key(input_3d, weight, bias, activation, conv_state) + key.update( + tokens=input_3d.shape[-1], + state_len=conv_state.shape[-1], + circular=cache_seqlens is not None, + ) + return key + + +@triton.jit +def _causal_conv1d_fwd_kernel( + x_ptr, + weight_ptr, + bias_ptr, + conv_states_ptr, + query_start_loc_ptr, + cache_indices_ptr, + has_initial_state_ptr, + out_ptr, + num_channels, + stride_x_channel, + stride_x_token, + stride_weight_channel, + stride_weight_tap, + stride_state_cache, + stride_state_channel, + stride_state_token, + stride_out_channel, + stride_out_token, + CONV_WIDTH: tl.constexpr, + HAS_BIAS: tl.constexpr, + APPLY_SILU: tl.constexpr, + PAD_SLOT_ID: tl.constexpr, + BLOCK_TOKENS: tl.constexpr, + BLOCK_CHANNELS: tl.constexpr, +): + HISTORY_LEN: tl.constexpr = CONV_WIDTH - 1 + tl.static_assert(CONV_WIDTH >= 2 and CONV_WIDTH <= 4, "Supported convolution widths are 2, 3 and 4") + tl.static_assert(BLOCK_TOKENS >= HISTORY_LEN, "Only the first token block may read initial state") + + sequence_index = tl.program_id(0) + token_block_index = tl.program_id(1) + channel_indices = tl.program_id(2) * BLOCK_CHANNELS + tl.arange(0, BLOCK_CHANNELS) + valid_channels = channel_indices < num_channels + sequence_start = tl.load(query_start_loc_ptr + sequence_index) + sequence_length = tl.load(query_start_loc_ptr + sequence_index + 1) - sequence_start + if token_block_index * BLOCK_TOKENS >= sequence_length: + return + + token_indices = token_block_index * BLOCK_TOKENS + tl.arange(0, BLOCK_TOKENS) + valid_output = (token_indices[:, None] < sequence_length) & valid_channels[None, :] + input_base = x_ptr + sequence_start * stride_x_token + channel_indices * stride_x_channel + output_base = out_ptr + sequence_start * stride_out_token + channel_indices * stride_out_channel + output_ptrs = output_base[None, :] + token_indices[:, None] * stride_out_token + cache_index = tl.load(cache_indices_ptr + sequence_index) + if cache_index == PAD_SLOT_ID: + padded_input = tl.load(input_base[None, :] + token_indices[:, None] * stride_x_token, valid_output, other=0) + tl.store(output_ptrs, padded_input, valid_output) + return + + state_base = conv_states_ptr + cache_index * stride_state_cache + channel_indices * stride_state_channel + weight_base = weight_ptr + channel_indices * stride_weight_channel + use_initial_state = tl.load(has_initial_state_ptr + sequence_index) + accumulator = tl.full((BLOCK_TOKENS, BLOCK_CHANNELS), 0, tl.float32) + if HAS_BIAS: + bias_values = tl.load(bias_ptr + channel_indices, valid_channels, other=0).to(tl.float32) + accumulator += bias_values[None, :] + + for tap in tl.static_range(CONV_WIDTH): + source_tokens = token_indices[:, None] + tap - HISTORY_LEN + input_values = tl.load( + input_base[None, :] + source_tokens * stride_x_token, valid_output & (source_tokens >= 0), other=0 + ).to(tl.float32) + history_values = tl.load( + state_base[None, :] + (source_tokens + HISTORY_LEN) * stride_state_token, + valid_output & (source_tokens < 0) & use_initial_state, + other=0, + ).to(tl.float32) + window_values = tl.where(source_tokens < 0, history_values, input_values) + conv_weight = tl.load(weight_base + tap * stride_weight_tap, valid_channels, other=0).to(tl.float32) + accumulator = tl.fma(window_values, conv_weight[None, :], accumulator) + if APPLY_SILU: + accumulator = accumulator / (1.0 + tl.exp(-accumulator)) + tl.store(output_ptrs, accumulator, valid_output) + + if token_block_index == 0: + state_offsets = tl.arange(0, triton.next_power_of_2(HISTORY_LEN))[:, None] + tail_tokens = sequence_length - HISTORY_LEN + state_offsets + valid_state = (state_offsets < HISTORY_LEN) & valid_channels[None, :] + tail_inputs = tl.load( + input_base[None, :] + tail_tokens * stride_x_token, valid_state & (tail_tokens >= 0), other=0 + ) + tail_history = tl.load( + state_base[None, :] + (tail_tokens + HISTORY_LEN) * stride_state_token, + valid_state & (tail_tokens < 0) & use_initial_state, + other=0, + ) + updated_state = tl.where(tail_tokens >= 0, tail_inputs, tail_history) + tl.debug_barrier() + tl.store(state_base[None, :] + state_offsets * stride_state_token, updated_state, valid_state) + + +@autotune( + kernel_name="causal_conv1d_prefill:v1", + configs_gen_func=_prefill_configs, + static_key_func=_prefill_static_key, + run_key_func=_prefill_run_key, + run_key_distance_func=_prefill_key_distance, + mutates_args=["conv_states"], +) +def causal_conv1d_fn( + x: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor] = None, + *, + query_start_loc: torch.Tensor, + cache_indices: torch.Tensor, + has_initial_state: torch.Tensor, + conv_states: torch.Tensor, + activation: Optional[str] = "silu", + pad_slot_id: int = -1, + max_seqlen: Optional[int] = None, + run_config: Optional[dict] = None, + **kwargs, +): + if activation not in (None, "silu", "swish"): + raise NotImplementedError(f"Unsupported activation: {activation}") + num_channels, total_tokens = x.shape + batch_size = query_start_loc.numel() - 1 + out = torch.empty_like(x) + if batch_size == 0 or x.numel() == 0: + return out + + bias = bias.contiguous() if bias is not None else None + max_query_len = total_tokens if max_seqlen is None else max_seqlen + if run_config is None: + run_config = {"BLOCK_TOKENS": 16, "BLOCK_CHANNELS": 128, "num_warps": 4} + block_tokens = run_config["BLOCK_TOKENS"] + block_channels = run_config["BLOCK_CHANNELS"] + grid = (batch_size, triton.cdiv(max_query_len, block_tokens), triton.cdiv(num_channels, block_channels)) + _causal_conv1d_fwd_kernel[grid]( + x, + weight, + bias, + conv_states, + query_start_loc, + cache_indices, + has_initial_state, + out, + num_channels, + *x.stride(), + *weight.stride(), + *conv_states.stride(), + *out.stride(), + CONV_WIDTH=weight.shape[1], + HAS_BIAS=bias is not None, + APPLY_SILU=activation is not None, + PAD_SLOT_ID=pad_slot_id, + BLOCK_TOKENS=block_tokens, + BLOCK_CHANNELS=block_channels, + num_warps=run_config["num_warps"], + ) + return out + + +@triton.jit +def _causal_conv1d_update_kernel( + x_ptr, + weight_ptr, + bias_ptr, + conv_state_ptr, + cache_seqlens_ptr, + conv_state_indices_ptr, + out_ptr, + num_channels, + stride_x_batch, + stride_x_channel, + stride_x_token, + stride_weight_channel, + stride_weight_tap, + stride_state_cache, + stride_state_channel, + stride_state_token, + stride_out_batch, + stride_out_channel, + stride_out_token, + CONV_WIDTH: tl.constexpr, + NUM_TOKENS: tl.constexpr, + STATE_LEN: tl.constexpr, + HAS_INDICES: tl.constexpr, + HAS_BIAS: tl.constexpr, + CIRCULAR_BUFFER: tl.constexpr, + APPLY_SILU: tl.constexpr, + PAD_SLOT_ID: tl.constexpr, + BLOCK_CHANNELS: tl.constexpr, +): + tl.static_assert(CONV_WIDTH >= 2 and CONV_WIDTH <= 4, "Supported convolution widths are 2, 3 and 4") + tl.static_assert(STATE_LEN >= CONV_WIDTH - 1, "State must hold the full convolution history") + + batch_index = tl.program_id(0) + channel_indices = tl.program_id(1) * BLOCK_CHANNELS + tl.arange(0, BLOCK_CHANNELS) + valid_channels = channel_indices < num_channels + input_base = x_ptr + batch_index * stride_x_batch + channel_indices * stride_x_channel + output_base = out_ptr + batch_index * stride_out_batch + channel_indices * stride_out_channel + cache_index = tl.load(conv_state_indices_ptr + batch_index) if HAS_INDICES else batch_index + if cache_index == PAD_SLOT_ID: + for token_index in range(NUM_TOKENS): + padded_input = tl.load(input_base + token_index * stride_x_token, valid_channels, other=0) + tl.store(output_base + token_index * stride_out_token, padded_input, valid_channels) + return + state_base = conv_state_ptr + cache_index * stride_state_cache + channel_indices * stride_state_channel + weight_base = weight_ptr + channel_indices * stride_weight_channel + state_cursor = tl.load(cache_seqlens_ptr + batch_index) % STATE_LEN if CIRCULAR_BUFFER else STATE_LEN + history_start = state_cursor + STATE_LEN - (CONV_WIDTH - 1) + + history_0 = tl.load( + state_base + (history_start % STATE_LEN) * stride_state_token, + valid_channels, + other=0, + ).to(tl.float32) + weight_0 = tl.load(weight_base, valid_channels, other=0).to(tl.float32) + weight_1 = tl.load(weight_base + stride_weight_tap, valid_channels, other=0).to(tl.float32) + if CONV_WIDTH >= 3: + history_1 = tl.load( + state_base + ((history_start + 1) % STATE_LEN) * stride_state_token, + valid_channels, + other=0, + ).to(tl.float32) + weight_2 = tl.load(weight_base + 2 * stride_weight_tap, valid_channels, other=0).to(tl.float32) + if CONV_WIDTH == 4: + history_2 = tl.load( + state_base + ((history_start + 2) % STATE_LEN) * stride_state_token, valid_channels, other=0 + ).to(tl.float32) + weight_3 = tl.load(weight_base + 3 * stride_weight_tap, valid_channels, other=0).to(tl.float32) + bias_values = tl.full((BLOCK_CHANNELS,), 0, tl.float32) + if HAS_BIAS: + bias_values = tl.load(bias_ptr + channel_indices, valid_channels, other=0).to(tl.float32) + + for token_index in range(NUM_TOKENS): + input_values = tl.load(input_base + token_index * stride_x_token, valid_channels, other=0).to(tl.float32) + accumulator = tl.fma(weight_0, history_0, bias_values) + if CONV_WIDTH == 2: + accumulator = tl.fma(weight_1, input_values, accumulator) + history_0 = input_values + elif CONV_WIDTH == 3: + accumulator = tl.fma(weight_1, history_1, accumulator) + accumulator = tl.fma(weight_2, input_values, accumulator) + history_0, history_1 = history_1, input_values + else: + accumulator = tl.fma(weight_1, history_1, accumulator) + accumulator = tl.fma(weight_2, history_2, accumulator) + accumulator = tl.fma(weight_3, input_values, accumulator) + history_0, history_1, history_2 = history_1, history_2, input_values + if APPLY_SILU: + accumulator = accumulator / (1.0 + tl.exp(-accumulator)) + tl.store(output_base + token_index * stride_out_token, accumulator, valid_channels) + if CIRCULAR_BUFFER: + tl.store( + state_base + ((state_cursor + token_index) % STATE_LEN) * stride_state_token, + input_values, + valid_channels, + ) + if not CIRCULAR_BUFFER: + state_offsets = tl.arange(0, triton.next_power_of_2(STATE_LEN))[:, None] + valid_state = (state_offsets < STATE_LEN) & valid_channels[None, :] + source_offsets = state_offsets + NUM_TOKENS + old_history = tl.load( + state_base[None, :] + source_offsets * stride_state_token, + valid_state & (source_offsets < STATE_LEN), + other=0, + ) + new_inputs = tl.load( + input_base[None, :] + (source_offsets - STATE_LEN) * stride_x_token, + valid_state & (source_offsets >= STATE_LEN), + other=0, + ) + updated_state = tl.where(source_offsets < STATE_LEN, old_history, new_inputs) + tl.debug_barrier() + tl.store(state_base[None, :] + state_offsets * stride_state_token, updated_state, valid_state) + + +@autotune( + kernel_name="causal_conv1d_update:v1", + kernel_type=AutotuneKernelType.DECODE_ATTENTION, + configs_gen_func=_update_configs, + static_key_func=_update_static_key, + run_key_func=lambda x: x.shape[0], + mutates_args=["conv_state"], +) +def causal_conv1d_update( + x: torch.Tensor, + conv_state: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor] = None, + activation: Optional[str] = None, + cache_seqlens: Optional[torch.Tensor] = None, + conv_state_indices: Optional[torch.Tensor] = None, + pad_slot_id: int = -1, + run_config: Optional[dict] = None, +): + if activation not in (None, "silu", "swish"): + raise NotImplementedError(f"Unsupported activation: {activation}") + out = torch.empty_like(x) + if x.numel() == 0: + return out + + input_3d = x.unsqueeze(-1) if x.ndim == 2 else x + output_3d = out.unsqueeze(-1) if x.ndim == 2 else out + batch_size, num_channels, num_tokens = input_3d.shape + bias = bias.contiguous() if bias is not None else None + if run_config is None: + run_config = {"BLOCK_CHANNELS": 128, "num_warps": 4} + block_channels = run_config["BLOCK_CHANNELS"] + grid = (batch_size, triton.cdiv(num_channels, block_channels)) + _causal_conv1d_update_kernel[grid]( + input_3d, + weight, + bias, + conv_state, + cache_seqlens, + conv_state_indices, + output_3d, + num_channels, + *input_3d.stride(), + *weight.stride(), + *conv_state.stride(), + *output_3d.stride(), + CONV_WIDTH=weight.shape[1], + NUM_TOKENS=num_tokens, + STATE_LEN=conv_state.shape[2], + HAS_INDICES=conv_state_indices is not None, + HAS_BIAS=bias is not None, + CIRCULAR_BUFFER=cache_seqlens is not None, + APPLY_SILU=activation is not None, + PAD_SLOT_ID=pad_slot_id, + BLOCK_CHANNELS=block_channels, + num_warps=run_config["num_warps"], + ) + return out diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/__init__.py b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/__init__.py index cd3b0962a3..1fd7fabe3f 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/__init__.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/__init__.py @@ -7,9 +7,10 @@ # the following copyright notice: # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang from .chunk import chunk_gated_delta_rule -from .fused_recurrent import fused_recurrent_gated_delta_rule +from .fused_recurrent import fused_recurrent_gated_delta_rule, fused_recurrent_gated_delta_rule_packed_decode __all__ = [ "chunk_gated_delta_rule", "fused_recurrent_gated_delta_rule", + "fused_recurrent_gated_delta_rule_packed_decode", ] diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/fused_recurrent.py b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/fused_recurrent.py index 5dfbd6e4ab..83d1bc3325 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/fused_recurrent.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/fused_recurrent.py @@ -380,6 +380,160 @@ def fused_recurrent_gated_delta_rule_fwd( return o, final_state +@triton.jit +def _fused_recurrent_gated_delta_rule_packed_decode_kernel( + mixed_qkv, + a, + b, + A_log, + dt_bias, + o, + h0, + ssm_state_indices, + scale, + stride_mixed_qkv_tok: tl.constexpr, + stride_a_tok: tl.constexpr, + stride_b_tok: tl.constexpr, + stride_state_slot: tl.constexpr, + stride_indices: tl.constexpr, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + SOFTPLUS_THRESHOLD: tl.constexpr, + USE_QK_L2NORM_IN_KERNEL: tl.constexpr, +): + i_v, i_nh = tl.program_id(0), tl.program_id(1) + i_n, i_hv = i_nh // HV, i_nh % HV + i_h = i_hv // (HV // H) + + o_k = tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + mask_k = o_k < K + mask_v = o_v < V + mask_h = mask_k[:, None] & mask_v[None, :] + + state_idx = tl.load(ssm_state_indices + i_n * stride_indices).to(tl.int64) + p_h0 = h0 + state_idx * stride_state_slot + p_h0 += i_hv * K * V + o_k[:, None] * V + o_v[None, :] + b_h = tl.load(p_h0, mask=mask_h, other=0).to(tl.float32) + + p_mixed = mixed_qkv + i_n * stride_mixed_qkv_tok + q_off = i_h * K + o_k + k_off = H * K + i_h * K + o_k + v_off = 2 * H * K + i_hv * V + o_v + b_q = tl.load(p_mixed + q_off, mask=mask_k, other=0).to(tl.float32) + b_k = tl.load(p_mixed + k_off, mask=mask_k, other=0).to(tl.float32) + b_v = tl.load(p_mixed + v_off, mask=mask_v, other=0).to(tl.float32) + + if USE_QK_L2NORM_IN_KERNEL: + b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6) + b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6) + b_q = b_q * scale + + a_val = tl.load(a + i_n * stride_a_tok + i_hv).to(tl.float32) + b_val = tl.load(b + i_n * stride_b_tok + i_hv).to(tl.float32) + A_log_val = tl.load(A_log + i_hv).to(tl.float32) + dt_bias_val = tl.load(dt_bias + i_hv).to(tl.float32) + x = a_val + dt_bias_val + softplus_x = tl.where(x <= SOFTPLUS_THRESHOLD, tl.log(1.0 + tl.exp(x)), x) + g_val = -tl.exp(A_log_val) * softplus_x + beta_val = tl.sigmoid(b_val) + + b_h *= exp(g_val) + b_v -= tl.sum(b_h * b_k[:, None], 0) + b_v *= beta_val + b_h += b_k[:, None] * b_v[None, :] + b_o = tl.sum(b_h * b_q[:, None], 0) + + p_o = o + (i_n * HV + i_hv) * V + o_v + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v) + tl.store(p_h0, b_h.to(p_h0.dtype.element_ty), mask=mask_h) + + +def fused_recurrent_gated_delta_rule_packed_decode( + mixed_qkv: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + ssm_state_indices: torch.Tensor, + use_qk_l2norm_in_kernel: bool = False, + out: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + if mixed_qkv.ndim != 2 or mixed_qkv.stride(-1) != 1: + raise ValueError("mixed_qkv must be a 2D tensor contiguous in its last dimension") + if a.ndim != 2 or b.ndim != 2 or a.stride(-1) != 1 or b.stride(-1) != 1: + raise ValueError("a and b must be 2D tensors contiguous in their last dimension") + if A_log.ndim != 1 or dt_bias.ndim != 1 or A_log.stride(0) != 1 or dt_bias.stride(0) != 1: + raise ValueError("A_log and dt_bias must be contiguous 1D tensors") + if ssm_state_indices.ndim != 1 or not ssm_state_indices.is_contiguous(): + raise ValueError("ssm_state_indices must be a contiguous 1D tensor") + if initial_state.ndim != 4 or initial_state.stride(-1) != 1: + raise ValueError("initial_state must be a 4D tensor contiguous in its last dimension") + + batch = mixed_qkv.shape[0] + num_v_heads, key_dim, value_dim = initial_state.shape[-3:] + if a.shape != (batch, num_v_heads) or b.shape != (batch, num_v_heads): + raise ValueError("a and b must have shape [batch, num_v_heads]") + if ssm_state_indices.shape[0] != batch: + raise ValueError("ssm_state_indices must have one entry per batch row") + if A_log.numel() != num_v_heads or dt_bias.numel() != num_v_heads: + raise ValueError("A_log and dt_bias must match num_v_heads") + + qk_dim = mixed_qkv.shape[1] - num_v_heads * value_dim + if qk_dim <= 0 or qk_dim % 2 != 0: + raise ValueError("mixed_qkv does not contain a valid q/k/v layout") + query_dim = qk_dim // 2 + if query_dim % key_dim != 0: + raise ValueError("query dimension must be divisible by key head dimension") + num_k_heads = query_dim // key_dim + if num_k_heads <= 0 or num_v_heads % num_k_heads != 0: + raise ValueError("invalid GQA head configuration") + + block_k = triton.next_power_of_2(key_dim) + if triton.cdiv(key_dim, block_k) != 1: + raise ValueError("packed decode currently supports one key tile") + block_v = min(triton.next_power_of_2(value_dim), 32) + if out is None: + out = mixed_qkv.new_empty((batch, num_v_heads, value_dim)) + elif out.shape != (batch, num_v_heads, value_dim) or not out.is_contiguous(): + raise ValueError("out must be contiguous with shape [batch, num_v_heads, value_dim]") + + grid = (triton.cdiv(value_dim, block_v), batch * num_v_heads) + _fused_recurrent_gated_delta_rule_packed_decode_kernel[grid]( + mixed_qkv=mixed_qkv, + a=a, + b=b, + A_log=A_log, + dt_bias=dt_bias, + o=out, + h0=initial_state, + ssm_state_indices=ssm_state_indices, + scale=scale, + stride_mixed_qkv_tok=mixed_qkv.stride(0), + stride_a_tok=a.stride(0), + stride_b_tok=b.stride(0), + stride_state_slot=initial_state.stride(0), + stride_indices=ssm_state_indices.stride(0), + H=num_k_heads, + HV=num_v_heads, + K=key_dim, + V=value_dim, + BK=block_k, + BV=block_v, + SOFTPLUS_THRESHOLD=20.0, + USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, + num_warps=1, + num_stages=3, + ) + return out.unsqueeze(1), initial_state + + class FusedRecurrentFunction(torch.autograd.Function): @staticmethod def forward( diff --git a/lightllm/common/basemodel/triton_kernel/linear_att/mtp_state_params.py b/lightllm/common/basemodel/triton_kernel/linear_att/mtp_state_params.py index b81bf9ec16..0413aea898 100644 --- a/lightllm/common/basemodel/triton_kernel/linear_att/mtp_state_params.py +++ b/lightllm/common/basemodel/triton_kernel/linear_att/mtp_state_params.py @@ -3,7 +3,7 @@ import triton.language as tl -@triton.jit +@triton.jit(do_not_specialize=["batch_size"]) def _build_dynamic_mtp_linear_att_state_params_kernel( b_req_idx, b_mtp_index, diff --git a/lightllm/common/basemodel/triton_kernel/mtp_utils.py b/lightllm/common/basemodel/triton_kernel/mtp_utils.py index 7525a754df..85ac266691 100644 --- a/lightllm/common/basemodel/triton_kernel/mtp_utils.py +++ b/lightllm/common/basemodel/triton_kernel/mtp_utils.py @@ -309,12 +309,11 @@ def mtp_scatter_next_token_ids( ) -@triton.jit +@triton.jit(do_not_specialize=["batch_size"]) def _fwd_kernel_gen_b_req_mtp_start_loc( b_mtp_index, b_req_mtp_start_loc, - num_reqs: tl.constexpr, - batch_size: tl.constexpr, + batch_size, BLOCK_SIZE: tl.constexpr, ): offset = tl.arange(0, BLOCK_SIZE) @@ -333,7 +332,6 @@ def gen_b_req_mtp_start_loc(b_mtp_index: torch.Tensor, num_reqs: int): _fwd_kernel_gen_b_req_mtp_start_loc[grid]( b_mtp_index=b_mtp_index, b_req_mtp_start_loc=b_req_mtp_start_loc, - num_reqs=num_reqs, batch_size=batch_size, BLOCK_SIZE=BLOCK_SIZE, num_warps=8, diff --git a/lightllm/common/basemodel/triton_kernel/scalar_linear.py b/lightllm/common/basemodel/triton_kernel/scalar_linear.py new file mode 100644 index 0000000000..e78678c40a --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/scalar_linear.py @@ -0,0 +1,42 @@ +import torch +import triton +import triton.language as tl + + +@triton.jit +def _scalar_linear(X, W, Y, K: tl.constexpr, XS0, XS1, WS, YS, BLOCK: tl.constexpr): + row = tl.program_id(0) + column = tl.arange(0, BLOCK) + value = tl.load(X + row * XS0 + column * XS1, column < K, other=0).to(tl.float64) + weight = tl.load(W + column * WS, column < K, other=0).to(tl.float64) + tl.store(Y + row * YS, tl.sum(value * weight, 0)) + + +def scalar_linear(input: torch.Tensor, weight: torch.Tensor, out=None, alloc_func=torch.empty) -> torch.Tensor: + if input.ndim != 2 or weight.shape != (1, input.shape[1]): + raise ValueError("Scalar projection requires (tokens, hidden) input and (1, hidden) weight") + if not input.is_cuda or weight.device != input.device or weight.dtype != input.dtype: + raise ValueError("Scalar projection requires CUDA input and matching weight dtype/device") + if not input.is_floating_point(): + raise ValueError("Scalar projection requires floating-point activations") + shape = (input.shape[0], 1) + if out is None: + out = alloc_func(shape, dtype=input.dtype, device=input.device) + elif out.shape != shape or out.dtype != input.dtype or out.device != input.device: + raise ValueError("Scalar projection output must match the unpadded shape, dtype and device") + if input.shape[0]: + block = triton.next_power_of_2(max(1, input.shape[1])) + _scalar_linear[(input.shape[0],)]( + input, + weight, + out, + input.shape[1], + input.stride(0), + input.stride(1), + weight.stride(1), + out.stride(0), + block, + num_warps=4 if block <= 8192 else 8, + enable_fp_fusion=False, + ) + return out diff --git a/lightllm/common/gated_residual/__init__.py b/lightllm/common/gated_residual/__init__.py new file mode 100644 index 0000000000..d3003469ec --- /dev/null +++ b/lightllm/common/gated_residual/__init__.py @@ -0,0 +1,93 @@ +from dataclasses import dataclass +from typing import Mapping, Optional + +import torch + +from .kernels import grouped_norm, inject_and_norm, inject_streams +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) + return self._mix(normalized) + + def _mix(self, normalized): + return project_and_mix(normalized, self.down_weight, self.up_weight, self.branches, self.inject_weight) + + @torch.no_grad() + def combine_and_read(self, residual, update, gates): + self._check_write(residual, update, gates) + combined, normalized = inject_and_norm(residual, update, gates, self.norm_weight, self.branches, self.eps) + return combined, self._mix(normalized) + + @torch.no_grad() + def write(self, residual: torch.Tensor, update: torch.Tensor, gates: torch.Tensor) -> torch.Tensor: + self._check_write(residual, update, gates) + return inject_streams(residual, update, gates) + + def _check_write(self, residual, update, gates): + 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") diff --git a/lightllm/common/gated_residual/kernels.py b/lightllm/common/gated_residual/kernels.py new file mode 100644 index 0000000000..8eeb021abe --- /dev/null +++ b/lightllm/common/gated_residual/kernels.py @@ -0,0 +1,146 @@ +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 _gate(X, W, Out, row, branch, D: tl.constexpr, C: tl.constexpr, B: tl.constexpr): + 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 _gates(X, W, Out, D: tl.constexpr, C: tl.constexpr, B: tl.constexpr): + _gate(X, W, Out, tl.program_id(0), tl.program_id(1), D, C, B) + + +@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) + + +@triton.jit +def _inject_norm(R, X, G, W, Y, Z, 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 + update = 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 = (update * gate).to(R.dtype.element_ty).to(tl.float32) + combined = (residual + injection).to(R.dtype.element_ty) + tl.store(Y + offset, combined, col < D) + combined = tl.where(col < D, combined.to(tl.float32), 0.0) + weight = tl.load(W + branch * D + col, col < D, other=0).to(tl.float32) + inv = tl.rsqrt(tl.sum(combined * combined, 0) / D + EPS) + tl.store(Z + offset, (combined * inv) * (1.0 + weight), col < D) + + +def inject_and_norm(residual, update, gates, weight, branches, eps): + output = torch.empty_like(residual) + normalized = torch.empty_like(residual) + dim = update.shape[-1] + rows = update.numel() // dim + if rows: + _inject_norm[(rows, branches)]( + residual, + update, + gates, + weight, + output, + normalized, + dim, + branches, + eps, + triton.next_power_of_2(dim), + enable_fp_fusion=False, + ) + return output, normalized + + +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..c52fcbc945 --- /dev/null +++ b/lightllm/common/gated_residual/projection.py @@ -0,0 +1,208 @@ +import torch +import torch.nn.functional as F +import triton +import triton.language as tl + +from .kernels import _gate, injection_gates, 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(do_not_specialize=["M"]) +def _down( + X, + W, + P, + Y, + M, + N: tl.constexpr, + K: tl.constexpr, + C: tl.constexpr, + BM: tl.constexpr, + BN: tl.constexpr, + BK: tl.constexpr, + S: tl.constexpr, + Inject=None, + Gates=None, + GB: tl.constexpr = 0, +): + if Inject is None: + row_tile = tl.program_id(0) + col_tile = tl.program_id(1) + split = tl.program_id(2) + is_gate: tl.constexpr = False + else: + block = tl.program_id(0) + down_blocks = tl.cdiv(M, BM) * tl.cdiv(N, BN) * S + is_gate = block >= down_blocks + gate_block = block - down_blocks + split = block % S + tile = block // S + row_tile = tile // tl.cdiv(N, BN) + col_tile = tile % tl.cdiv(N, BN) + + if is_gate: + _gate(X, Inject, Gates, gate_block // C, gate_block % C, K, C, GB) + else: + r = row_tile * BM + tl.arange(0, BM) + c = col_tile * BN + tl.arange(0, BN) + 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(do_not_specialize=["M"]) +def _merge(P, Y, M, 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(do_not_specialize=["M"]) +def _up_mix( + X, + W, + Normalized, + Y, + M, + 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, inject_weight=None, gates=None): + 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 + grid = (triton.cdiv(m, bm), triton.cdiv(n, 64), splits) + if inject_weight is not None: + grid = (grid[0] * grid[1] * grid[2] + m * branches,) + _down[grid]( + x, + w, + p, + out, + m, + n, + k, + branches, + bm, + 64, + 64, + splits, + inject_weight, + gates, + triton.next_power_of_2(k) if inject_weight is not None else 0, + num_warps=8 if inject_weight is not None or 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, inject_weight=None +): + gates = None + if normalized.dtype == torch.float32: + low = F.silu(F.linear(normalized, down_weight) / branches) + mixed = mix_streams(normalized, F.linear(low, up_weight), branches) + else: + rows = normalized.reshape(-1, normalized.shape[-1]) + fuse_gates = inject_weight is not None and 0 < rows.shape[0] <= 64 + fuse_gates = fuse_gates and rows.shape[1] >= 8192 and down_weight.shape[0] <= 512 + if fuse_gates: + gates = torch.empty((*normalized.shape[:-1], branches), device=normalized.device, dtype=normalized.dtype) + low = _lowrank(rows, down_weight, branches, inject_weight if fuse_gates else None, gates) + mixed = _mix(rows, low, up_weight, branches) + mixed = mixed.reshape((*normalized.shape[:-1], normalized.shape[-1] // branches)) + if inject_weight is None: + return mixed + if gates is None: + gates = injection_gates(normalized, inject_weight, branches) + return mixed, gates 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..b4f077a7f9 --- /dev/null +++ b/lightllm/common/qsa/attention.py @@ -0,0 +1,127 @@ +from typing import Optional + +import torch +import triton + +from .attention_kernels import _attention, _reduce + + +@torch.no_grad() +def indexed_attention( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + indices: torch.Tensor, + scale: Optional[float] = None, + splits: Optional[int] = None, + selection_lengths: Optional[torch.Tensor] = None, + short_rows_direct: bool = False, + skip_empty_splits: bool = False, + block_size: int = 64, + num_stages: Optional[int] = None, + num_warps: Optional[int] = None, +) -> torch.Tensor: + if query.ndim != 3 or key.ndim != 3 or value.ndim != 3 or indices.ndim != 2: + raise ValueError("Indexed attention requires rank-three Q/K/V and rank-two indices") + rows, q_heads, dim = query.shape + if not q_heads or not dim or not key.shape[1] or q_heads % key.shape[1]: + raise ValueError("Indexed attention query heads must be divisible by KV heads") + if key.shape[0:2] != value.shape[0:2] or key.shape[2] != dim or not value.shape[2]: + raise ValueError("Indexed attention key/value dimensions do not match") + if key.dtype != value.dtype or key.dtype not in (query.dtype, torch.float8_e4m3fn): + raise ValueError("Indexed attention KV must match the query dtype or use FP8 e4m3") + for tensor in (query, key, value): + if not tensor.is_cuda or tensor.device != query.device: + raise ValueError("Indexed attention Q/K/V must be CUDA tensors with matching device") + if not query.is_contiguous(): + raise ValueError("Indexed attention query must be contiguous") + if query.dtype not in (torch.float16, torch.bfloat16, torch.float32): + raise ValueError("Indexed attention requires floating point inputs") + if ( + indices.shape[0] != rows + or indices.dtype != torch.int32 + or indices.device != query.device + or not indices.is_contiguous() + ): + raise ValueError("Indexed attention indices must be a contiguous int32 matrix on the query device") + selected = indices.shape[1] + if selection_lengths is not None and ( + selection_lengths.shape != (rows,) + or selection_lengths.dtype != torch.int32 + or selection_lengths.device != query.device + or not selection_lengths.is_contiguous() + ): + raise ValueError("Selection lengths must be a contiguous int32 vector on the query device") + kv_heads, value_dim = value.shape[1:] + output = torch.empty((rows, q_heads, value_dim), dtype=query.dtype, device=query.device) + if not rows: + return output + if not selected: + return output.zero_() + if block_size < 16 or block_size & (block_size - 1): + raise ValueError("Indexed attention block size must be a power of two >= 16") + if num_stages is not None and num_stages < 1: + raise ValueError("Indexed attention num_stages must be positive") + if num_warps is not None and num_warps not in (1, 2, 4, 8): + raise ValueError("Indexed attention num_warps must be 1, 2, 4 or 8") + if splits is None: + splits = min(32, max(1, triton.cdiv(block_size, rows * kv_heads)), triton.cdiv(selected, block_size)) + if splits < 1: + raise ValueError("Indexed attention split count must be positive") + partial = ( + torch.empty((rows, q_heads, splits, value_dim), dtype=torch.float32, device=query.device) + if splits > 1 + else None + ) + lse = torch.empty((rows, q_heads, splits), dtype=torch.float32, device=query.device) if splits > 1 else None + launch_kwargs = {} + if num_stages is not None: + launch_kwargs["num_stages"] = num_stages + if num_warps is not None: + launch_kwargs["num_warps"] = num_warps + _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 if query.dtype == torch.float32 else 4, triton.next_power_of_2(q_heads // kv_heads)), + max(16, triton.next_power_of_2(dim)), + max(16, triton.next_power_of_2(value_dim)), + block_size, + short_rows_direct, + skip_empty_splits, + "ieee" if query.dtype == torch.float32 else "tf32", + **launch_kwargs, + ) + if splits > 1: + _reduce[(rows, q_heads)]( + selection_lengths, + partial, + lse, + output, + q_heads, + value_dim, + splits, + triton.next_power_of_2(splits), + triton.next_power_of_2(value_dim), + selected, + block_size, + short_rows_direct, + skip_empty_splits, + **launch_kwargs, + ) + return output diff --git a/lightllm/common/qsa/attention_kernels.py b/lightllm/common/qsa/attention_kernels.py new file mode 100644 index 0000000000..ef5683a9e0 --- /dev/null +++ b/lightllm/common/qsa/attention_kernels.py @@ -0,0 +1,139 @@ +import triton +import triton.language as tl + + +@triton.jit +def _attention( + Q, + K, + V, + Indices, + Lengths, + Partial, + LSE, + Output, + Q_HEADS: tl.constexpr, + KV_HEADS: tl.constexpr, + KS: tl.constexpr, + KH: tl.constexpr, + KD: tl.constexpr, + VS: tl.constexpr, + VH: tl.constexpr, + VD: tl.constexpr, + GROUP: tl.constexpr, + D: tl.constexpr, + DV: tl.constexpr, + SELECTED: tl.constexpr, + SPLITS: tl.constexpr, + SCALE: tl.constexpr, + BG: tl.constexpr, + BD: tl.constexpr, + BV: tl.constexpr, + BN: tl.constexpr, + SHORT_DIRECT: tl.constexpr, + SKIP_EMPTY: tl.constexpr, + INPUT_PRECISION: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + kv_head = tl.program_id(1) + split = tl.program_id(2) + limit = SELECTED + if Lengths is not None: + limit = tl.minimum(SELECTED, tl.maximum(0, tl.load(Lengths + row))) + if (not SHORT_DIRECT or limit > BN or split == 0) and ( + not SKIP_EMPTY or SPLITS == 1 or split * BN < limit or (SHORT_DIRECT and limit <= BN and split == 0) + ): + group = tl.arange(0, BG) + head = kv_head * GROUP + group + dim = tl.arange(0, BD) + value_dim = tl.arange(0, BV) + q = tl.load( + Q + (row * Q_HEADS + head[:, None]) * D + dim[None, :], + (group[:, None] < GROUP) & (dim[None, :] < D), + other=0.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.0, + ).to(q.dtype) + score = tl.dot(q, key, input_precision=INPUT_PRECISION) * (SCALE * 1.4426950408889634) + 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.exp2(maximum - safe_max) + prob = tl.exp2(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.0, + ).to(tl.float16 if V.dtype.element_ty == tl.float8e4nv and Q.dtype.element_ty != tl.float32 else q.dtype) + acc = acc * alpha[:, None] + acc = tl.dot(prob.to(value.dtype), value, acc, input_precision=INPUT_PRECISION) + denominator = denominator * alpha + tl.sum(prob, 1) + maximum = new_max + safe_denominator = tl.where(denominator > 0, denominator, 1.0) + result = acc / safe_denominator[:, None] + if SPLITS == 1 or (SHORT_DIRECT and limit <= BN): + tl.store( + Output + (row * Q_HEADS + head[:, None]) * DV + value_dim[None, :], + result, + (group[:, None] < GROUP) & (value_dim[None, :] < DV), + ) + else: + tl.store( + Partial + ((row * Q_HEADS + head[:, None]) * SPLITS + split) * DV + value_dim[None, :], + result, + (group[:, None] < GROUP) & (value_dim[None, :] < DV), + ) + lse = tl.where(denominator > 0, (maximum + tl.log2(safe_denominator)) * 0.6931471805599453, -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..0ac334e3dd --- /dev/null +++ b/lightllm/common/qsa/indexer.py @@ -0,0 +1,302 @@ +from dataclasses import dataclass +from typing import Mapping, Optional + +import torch +import torch.nn.functional as F +import triton + +from lightllm.common.gated_residual.kernels import grouped_norm +from .indexer_kernels import _rotate, _pool, _scores, _expand +from .topk import cuda_topk_indices, cuda_topk_supported, flashinfer_topk_indices, flashinfer_topk_supported + + +@dataclass(frozen=True) +class QSAIndexer: + projection: torch.Tensor + query_norm: torch.Tensor + key_norm: torch.Tensor + query_heads: int = 4 + token_budget: int = 2048 + compress_ratio: int = 4 + eps: float = 1e-6 + use_cuda_topk: bool = False + + def __post_init__(self): + dim = self.key_norm.numel() + if not dim or self.query_heads < 1 or self.compress_ratio < 1 or self.token_budget < 1 or self.eps <= 0: + raise ValueError("QSA dimensions, budget, compression ratio and epsilon must be positive") + if self.token_budget % self.compress_ratio: + raise ValueError("QSA token budget must be divisible by compression ratio") + if self.key_norm.shape != (dim,) or self.query_norm.shape != (dim,): + raise ValueError("QSA normalization weights must be matching vectors") + if ( + self.projection.ndim != 2 + or self.projection.shape[0] != (self.query_heads + 1) * dim + or not self.projection.shape[1] + ): + raise ValueError("QSA projection must contain query heads and one raw key head") + for weight in (self.projection, self.query_norm, self.key_norm): + self._activation(weight) + + @property + def head_dim(self) -> int: + return self.key_norm.numel() + + @classmethod + def from_checkpoint(cls, weights: Mapping[str, torch.Tensor], prefix: str, **kwargs) -> "QSAIndexer": + return cls( + weights[f"{prefix}.index_qk_proj.weight"], + weights[f"{prefix}.q_layernorm.weight"], + weights[f"{prefix}.k_layernorm.weight"], + **kwargs, + ) + + def _activation(self, x: torch.Tensor): + if ( + not x.is_cuda + or not x.is_contiguous() + or x.device != self.projection.device + or x.dtype != self.projection.dtype + ): + raise ValueError("QSA activations/weights must be contiguous CUDA tensors on the same device and dtype") + if x.dtype not in (torch.float16, torch.bfloat16, torch.float32): + raise ValueError("QSA requires floating point activations") + + def _integer(self, x: torch.Tensor): + if x.dtype != torch.int32 or x.device != self.projection.device or not x.is_contiguous(): + raise ValueError("QSA metadata must be contiguous int32 tensors on the input device") + + def _cache(self, x): + if x.dtype == torch.float8_e4m3fn: + if not x.is_cuda or not x.is_contiguous() or x.device != self.projection.device: + raise ValueError("QSA FP8 cache must be contiguous on the indexer device") + else: + self._activation(x) + + 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, projected=None): + 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]) + if projected is None: + 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 :] + + 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._cache(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, + *, + max_visible_length: Optional[int] = None, + physical_stride: int = 1, + ) -> torch.Tensor: + self._cache(query) + self._cache(cache) + if query.dtype != cache.dtype: + raise ValueError("QSA indexer query and pooled cache must have matching dtypes") + 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") + if physical_stride < 1: + raise ValueError("QSA physical cache stride 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] * physical_stride) // self.compress_ratio + if physical + else cache.shape[1] + ) + if max_visible_length is not None: + if max_visible_length < 0: + raise ValueError("QSA maximum visible length must be nonnegative") + visible_blocks = max_visible_length // self.compress_ratio + blocks = min(blocks, triton.next_power_of_2(visible_blocks) if visible_blocks else 0) + block_budget = self.token_budget // self.compress_ratio + topk = min(blocks, block_budget) + cuda_topk = self.use_cuda_topk and cuda_topk_supported(query.device, blocks, topk) + flashinfer_topk = self.use_cuda_topk and not cuda_topk and flashinfer_topk_supported(query.device, blocks, topk) + max_score_elements = 16 * 1024 * 1024 + query_chunk_size = min(query_chunk_size, max(1, max_score_elements // max(blocks, 1))) + score_buffer = torch.empty((min(tokens, query_chunk_size), blocks), dtype=torch.float32, device=query.device) + group_size = 4 if physical and self.query_heads == 4 else 1 + score_block = 128 if group_size == 4 else 64 + for first in range(0, tokens, query_chunk_size): + count = min(query_chunk_size, tokens - first) + scores = score_buffer[:count] + if blocks: + _scores[(triton.cdiv(count, group_size), triton.cdiv(blocks, score_block))]( + query, + cache, + requests, + visible_lengths, + indices, + scores, + first, + self.query_heads, + self.head_dim, + blocks, + indices.shape[1], + physical, + self.compress_ratio, + max(16, triton.next_power_of_2(self.query_heads)), + max(32 if query.dtype == torch.float8_e4m3fn else 16, triton.next_power_of_2(self.head_dim)), + score_block, + block_budget if cuda_topk or flashinfer_topk else 0, + "ieee" if query.dtype == torch.float32 else "tf32", + GROUP_SIZE=group_size, + TOKENS=count, + CACHE_BLOCKS=cache.shape[1] if not physical else 0, + PHYSICAL_STRIDE=physical_stride, + num_warps=4, + ) + if cuda_topk: + selected = cuda_topk_indices( + scores, topk, visible_lengths=visible_lengths, first=first, compress_ratio=self.compress_ratio + ) + elif flashinfer_topk: + selected = flashinfer_topk_indices( + scores, topk, visible_lengths=visible_lengths, first=first, compress_ratio=self.compress_ratio + ) + else: + selected = torch.topk(scores, topk, dim=-1, sorted=False).indices + _expand[(count,)]( + selected, + requests, + visible_lengths, + indices, + output, + first, + topk, + block_budget, + indices.shape[1], + self.compress_ratio, + width, + triton.next_power_of_2(width), + ) + return output diff --git a/lightllm/common/qsa/indexer_kernels.py b/lightllm/common/qsa/indexer_kernels.py new file mode 100644 index 0000000000..c3095fab0a --- /dev/null +++ b/lightllm/common/qsa/indexer_kernels.py @@ -0,0 +1,219 @@ +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.to(Raw.dtype.element_ty), col < D) + + +@triton.jit(do_not_specialize=["QUERY_START", "BLOCKS", "TOKENS"]) +def _scores( + Query, + Cache, + Requests, + Lengths, + Indices, + Output, + QUERY_START, + HEADS: tl.constexpr, + D: tl.constexpr, + BLOCKS, + CAPACITY: tl.constexpr, + PHYSICAL_CACHE: tl.constexpr, + RATIO: tl.constexpr, + BH: tl.constexpr, + BD: tl.constexpr, + BN: tl.constexpr, + MIN_BLOCKS: tl.constexpr = 0, + INPUT_PRECISION: tl.constexpr = "tf32", + GROUP_SIZE: tl.constexpr = 1, + TOKENS=0, + CACHE_BLOCKS: tl.constexpr = 0, + PHYSICAL_STRIDE: tl.constexpr = 1, +): + tl.static_assert(GROUP_SIZE == 1 or (GROUP_SIZE == 4 and HEADS == 4 and PHYSICAL_CACHE)) + group = tl.program_id(0) + if GROUP_SIZE == 4: + rows = group * 4 + tl.arange(0, 4) + valid = rows < TOKENS + requests = tl.load(Requests + QUERY_START + rows, valid, 0).to(tl.int64) + lengths = tl.load(Lengths + QUERY_START + rows, valid, 0) // RATIO + shared = tl.min(tl.where(valid, requests, 0x7FFFFFFF)) == tl.max(tl.where(valid, requests, 0)) + else: + shared = tl.full((), False, tl.int1) + if GROUP_SIZE == 4 and shared: + request = tl.sum(tl.where(tl.arange(0, 4) == 0, requests, 0)) + maximum = tl.max(lengths) + block = tl.program_id(1) * BN + tl.arange(0, BN) + scores = tl.full((4, BN), -float("inf"), tl.float32) + if (tl.program_id(1) * BN < maximum) & (maximum >= MIN_BLOCKS): + head = tl.arange(0, 16) + col = tl.arange(0, BD) + query_rows = QUERY_START + group * 4 + head // HEADS + q = tl.load( + Query + (query_rows[:, None].to(tl.int64) * HEADS + (head % HEADS)[:, None]) * D + col[None, :], + (query_rows[:, None] < QUERY_START + TOKENS) & (col[None, :] < D), + 0.0, + ) + key_rows = tl.load( + Indices + request * CAPACITY + block * RATIO + RATIO - 1, + (block < maximum) & (block < BLOCKS), + 0, + ).to(tl.int64) + key_rows //= PHYSICAL_STRIDE + k = tl.load( + Cache + key_rows[None, :] * D + col[:, None], (block[None, :] < maximum) & (col[:, None] < D), 0.0 + ) + logits = tl.dot(q, k, input_precision=INPUT_PRECISION) + scores = tl.sum(tl.maximum(logits, 0).reshape(4, HEADS, BN), 1) * (D ** -0.5) + scores = tl.where(block[None, :] < lengths[:, None], scores, -float("inf")) + tl.store( + Output + rows[:, None].to(tl.int64) * BLOCKS + block[None, :], + scores, + valid[:, None] & (block[None, :] < BLOCKS) & ((MIN_BLOCKS == 0) | (lengths[:, None] >= MIN_BLOCKS)), + ) + else: + for offset in range(GROUP_SIZE): + row = group * GROUP_SIZE + offset + if GROUP_SIZE == 1 or row < TOKENS: + query_row = row + QUERY_START + request = tl.load(Requests + query_row).to(tl.int64) + visible = tl.load(Lengths + query_row) // RATIO + block = tl.program_id(1) * BN + tl.arange(0, BN) + + scores = tl.full((BN,), -float("inf"), tl.float32) + if (tl.program_id(1) * BN < visible) & (visible >= MIN_BLOCKS): + head = tl.arange(0, BH) + col = tl.arange(0, BD) + q = tl.load( + Query + (query_row.to(tl.int64) * HEADS + head[:, None]) * D + col[None, :], + (head[:, None] < HEADS) & (col[None, :] < D), + other=0.0, + ) + if PHYSICAL_CACHE: + key_rows = tl.load( + Indices + request * CAPACITY + block * RATIO + RATIO - 1, + (block < visible) & (block < BLOCKS), + other=0, + ).to(tl.int64) + key_rows //= PHYSICAL_STRIDE + else: + key_rows = request * (CACHE_BLOCKS if CACHE_BLOCKS else BLOCKS) + block + k = tl.load( + Cache + key_rows[None, :] * D + col[:, None], + (block[None, :] < visible) & (col[:, None] < D), + other=0.0, + ) + logits = tl.dot(q, k, input_precision=INPUT_PRECISION) + logits = tl.where(head[:, None] < HEADS, tl.maximum(logits, 0), 0) + scores = tl.sum(logits, 0) * (D ** -0.5) + scores = tl.where(block < visible, scores, -float("inf")) + tl.store( + Output + row.to(tl.int64) * BLOCKS + block, + scores, + (block < BLOCKS) & ((MIN_BLOCKS == 0) | (visible >= MIN_BLOCKS)), + ) + + +@triton.jit(do_not_specialize=["QUERY_START"]) +def _expand( + Selected, + Requests, + Lengths, + Indices, + Output, + QUERY_START, + 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..6db7b72f48 --- /dev/null +++ b/lightllm/common/qsa/topk.py @@ -0,0 +1,129 @@ +from functools import lru_cache +from inspect import signature + +import torch +import triton +import triton.language as tl +from lightllm.utils.log_utils import init_logger + +logger = init_logger(__name__) + + +_CUDA_TOPK_RUNTIME = str(torch.__version__) == "2.9.1+cu128" and triton.__version__ == "3.5.1" + + +@lru_cache(maxsize=16) +def _supports_device(device): + supported = torch.cuda.get_device_capability(device) == (9, 0) and "H100" in torch.cuda.get_device_name(device) + if supported: + logger.info("QSA CUDA TopK selection enabled") + return supported + + +@triton.jit +def _key(value): + bits = value.to(tl.uint32, bitcast=True) + mask = tl.where((bits & 0x80000000) != 0, 0xFFFFFFFF, 0x80000000).to(tl.uint32) + return tl.where(value == value, bits ^ mask, 0xFFFFFFFF).to(tl.uint32) + + +@triton.jit +def _topk( + Scores, + Selected, + Lengths, + FIRST: tl.constexpr, + RATIO: tl.constexpr, + N: tl.constexpr, + K: tl.constexpr, + B: tl.constexpr, +): + tl.static_assert(N == 1027 and K == 512 and B == 2048) + row = tl.program_id(0).to(tl.int64) + if Lengths is None: + short = tl.full((), False, tl.int1) + else: + short = tl.load(Lengths + row + FIRST) // RATIO < K + if short: + short_col = tl.arange(0, K) + tl.store(Selected + row * K + short_col, short_col.to(tl.int64)) + else: + prefix_col = tl.arange(0, 1024) + prefix = tl.topk(_key(tl.load(Scores + row * N + prefix_col)), K) + merge_col = tl.arange(0, 4) + tail = tl.load(Scores + row * N + 1024 + merge_col, merge_col < 3, other=0.0) + tail = tl.sort(tl.where(merge_col < 3, _key(tail), 0).to(tl.uint32), descending=True) + + prefix_cutoffs = tl.gather(prefix, K - 1 - merge_col, 0) + tail_cutoffs = tl.gather(tail, tl.maximum(merge_col - 1, 0), 0) + tail_cutoffs = tl.where(merge_col == 0, 0xFFFFFFFF, tail_cutoffs).to(tl.uint32) + threshold = tl.max(tl.minimum(prefix_cutoffs, tail_cutoffs)) + col = tl.arange(0, B) + value = tl.load(Scores + row * N + col, col < N, other=-float("inf")) + key = _key(value) + greater = (col < N) & (key > threshold) + equal = (col < N) & (key == threshold) + greater_rank = tl.cumsum(greater.to(tl.int32)) - 1 + equal_rank = tl.cumsum(equal.to(tl.int32)) - 1 + rank = tl.where(greater, greater_rank, tl.sum(greater.to(tl.int32)) + equal_rank) + selected = greater | (equal & (rank < K)) + tl.store(Selected + row * K + rank, col.to(tl.int64), selected) + + +def cuda_topk_supported(device, blocks: int, k: int) -> bool: + return _CUDA_TOPK_RUNTIME and device.type == "cuda" and blocks == 1027 and k == 512 and _supports_device(device) + + +@lru_cache(maxsize=1) +def _flashinfer_topk(): + try: + from flashinfer import top_k_ragged_transform + except ImportError: + return None + if "dsa_graph_safe" not in signature(top_k_ragged_transform).parameters: + return None + return top_k_ragged_transform + + +@lru_cache(maxsize=16) +def flashinfer_topk_supported(device, blocks: int, k: int) -> bool: + return ( + device.type == "cuda" + and k == 512 + and blocks >= k + and torch.cuda.get_device_capability(device) == (9, 0) + and _flashinfer_topk() is not None + ) + + +def flashinfer_topk_indices(scores, k, *, visible_lengths, first=0, compress_ratio=4): + lengths = visible_lengths[first : first + scores.shape[0]] // compress_ratio + return _flashinfer_topk()( + scores, torch.zeros_like(lengths), lengths, k, deterministic=True, tie_break=1, dsa_graph_safe=True + ) + + +def cuda_topk_indices(scores: torch.Tensor, k: int, *, visible_lengths=None, first=0, compress_ratio=4) -> torch.Tensor: + if not ( + scores.is_cuda + and scores.dtype == torch.float32 + and scores.ndim == 2 + and scores.is_contiguous() + and cuda_topk_supported(scores.device, scores.shape[1], k) + ): + return torch.topk(scores, k, dim=-1, sorted=False).indices + if visible_lengths is not None: + if ( + visible_lengths.dtype != torch.int32 + or visible_lengths.device != scores.device + or visible_lengths.ndim != 1 + or not visible_lengths.is_contiguous() + or first < 0 + or first + scores.shape[0] > visible_lengths.numel() + or compress_ratio < 1 + ): + raise ValueError("QSA TopK requires matching contiguous int32 visible lengths and valid bounds") + output = torch.empty((scores.shape[0], k), dtype=torch.int64, device=scores.device) + if scores.shape[0]: + _topk[(scores.shape[0],)](scores, output, visible_lengths, first, compress_ratio, 1027, 512, 2048, num_warps=16) + return output diff --git a/lightllm/common/quantization/aligned_projection.py b/lightllm/common/quantization/aligned_projection.py new file mode 100644 index 0000000000..06e1f9a287 --- /dev/null +++ b/lightllm/common/quantization/aligned_projection.py @@ -0,0 +1,51 @@ +import torch + +from lightllm.common.quantization.no_quant import NoQuantization +from lightllm.common.quantization.quantize_method import WeightPack + + +class OutputAlignedProjection(NoQuantization): + def __init__(self, alignment=8): + super().__init__() + if type(alignment) is not int or alignment < 1 or alignment & (alignment - 1): + raise ValueError("Projection alignment must be a positive power of two") + self.alignment = alignment + self._padded_weight = None + + def _create_weight(self, out_dims, in_dim, dtype, device_id, num_experts=1): + width = sum(out_dims) if isinstance(out_dims, list) else out_dims + physical_width = (width + self.alignment - 1) // self.alignment * self.alignment + self._padded_weight = None + if dtype not in (torch.float16, torch.bfloat16) or physical_width == width or num_experts != 1: + return super()._create_weight(out_dims, in_dim, dtype, device_id, num_experts) + self._padded_weight = torch.zeros((physical_width, in_dim), dtype=dtype, device=device_id) + pack = WeightPack(weight=self._padded_weight[:width]) + parts = self._split_weight_pack(pack, out_dims if isinstance(out_dims, list) else [out_dims], -2) + return pack, parts + + def apply(self, input_tensor, weight_pack, out=None, workspace=None, use_custom_tensor_mananger=True, bias=None): + if self._padded_weight is None or bias is not None or input_tensor.shape[0] == 0: + return super().apply( + input_tensor, + weight_pack, + out, + workspace, + use_custom_tensor_mananger=use_custom_tensor_mananger, + bias=bias, + ) + from lightllm.common.basemodel.layer_infer.cache_tensor_manager import g_cache_manager + + def allocate(shape): + if use_custom_tensor_mananger: + return g_cache_manager.alloc_tensor(shape, input_tensor.dtype, device=input_tensor.device) + return torch.empty(shape, dtype=input_tensor.dtype, device=input_tensor.device) + + rows, width = input_tensor.shape[0], weight_pack.weight.shape[0] + physical_output = allocate((rows, self._padded_weight.shape[0])) + torch.mm(input_tensor, self._padded_weight.T, out=physical_output) + if out is None: + out = allocate((rows, width)) + elif out.shape != (rows, width) or out.dtype != input_tensor.dtype or out.device != input_tensor.device: + raise ValueError("Projection output must match logical shape, dtype and device") + out.copy_(physical_output[:, :width]) + return out diff --git a/lightllm/common/quantization/scalar_projection.py b/lightllm/common/quantization/scalar_projection.py new file mode 100644 index 0000000000..cef317e4a2 --- /dev/null +++ b/lightllm/common/quantization/scalar_projection.py @@ -0,0 +1,21 @@ +import torch + +from lightllm.common.basemodel.triton_kernel.scalar_linear import scalar_linear +from lightllm.common.quantization.no_quant import NoQuantization + + +class StableScalarProjection(NoQuantization): + def apply(self, input_tensor, weight_pack, out=None, workspace=None, use_custom_tensor_mananger=True, bias=None): + if not input_tensor.is_cuda or weight_pack.weight.shape[0] != 1 or bias is not None: + return super().apply( + input_tensor, + weight_pack, + out=out, + workspace=workspace, + use_custom_tensor_mananger=use_custom_tensor_mananger, + bias=bias, + ) + from lightllm.common.basemodel.layer_infer.cache_tensor_manager import g_cache_manager + + allocate = g_cache_manager.empty if use_custom_tensor_mananger else torch.empty + return scalar_linear(input_tensor, weight_pack.weight, out=out, alloc_func=allocate) diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_H100_80GB_HBM3/grouped_matmul:v1/{K=2560,N=160,expert_num=513,mul_routed_weight=false,out_dtype=torch.bfloat16,topk_num=11,use_fp8_w8a8=false}_NVIDIA_H100_80GB_HBM3.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_H100_80GB_HBM3/grouped_matmul:v1/{K=2560,N=160,expert_num=513,mul_routed_weight=false,out_dtype=torch.bfloat16,topk_num=11,use_fp8_w8a8=false}_NVIDIA_H100_80GB_HBM3.json new file mode 100644 index 0000000000..a4e6e1885e --- /dev/null +++ b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_H100_80GB_HBM3/grouped_matmul:v1/{K=2560,N=160,expert_num=513,mul_routed_weight=false,out_dtype=torch.bfloat16,topk_num=11,use_fp8_w8a8=false}_NVIDIA_H100_80GB_HBM3.json @@ -0,0 +1,92 @@ +{ + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 3 + }, + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 5 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 5 + }, + "3": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 5 + }, + "513": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 1 + }, + "514": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 1 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 5 + } +} \ No newline at end of file diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_H100_80GB_HBM3/grouped_matmul:v1/{K=80,N=2560,expert_num=513,mul_routed_weight=true,out_dtype=torch.bfloat16,topk_num=1,use_fp8_w8a8=false}_NVIDIA_H100_80GB_HBM3.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_H100_80GB_HBM3/grouped_matmul:v1/{K=80,N=2560,expert_num=513,mul_routed_weight=true,out_dtype=torch.bfloat16,topk_num=1,use_fp8_w8a8=false}_NVIDIA_H100_80GB_HBM3.json new file mode 100644 index 0000000000..444e3b96fe --- /dev/null +++ b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_H100_80GB_HBM3/grouped_matmul:v1/{K=80,N=2560,expert_num=513,mul_routed_weight=true,out_dtype=torch.bfloat16,topk_num=1,use_fp8_w8a8=false}_NVIDIA_H100_80GB_HBM3.json @@ -0,0 +1,92 @@ +{ + "88": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 4 + }, + "264": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 4 + }, + "5632": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 4 + }, + "11": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 2 + }, + "22": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 2 + }, + "33": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 4 + }, + "44": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 2 + }, + "513": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 2 + }, + "514": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 1 + }, + "5633": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 1 + } +} \ No newline at end of file diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_H200/grouped_matmul:v1/{K=160,N=2560,expert_num=513,mul_routed_weight=true,out_dtype=torch.bfloat16,topk_num=1,use_fp8_w8a8=false}_NVIDIA_H200.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_H200/grouped_matmul:v1/{K=160,N=2560,expert_num=513,mul_routed_weight=true,out_dtype=torch.bfloat16,topk_num=1,use_fp8_w8a8=false}_NVIDIA_H200.json new file mode 100644 index 0000000000..967fd15e61 --- /dev/null +++ b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_H200/grouped_matmul:v1/{K=160,N=2560,expert_num=513,mul_routed_weight=true,out_dtype=torch.bfloat16,topk_num=1,use_fp8_w8a8=false}_NVIDIA_H200.json @@ -0,0 +1,47 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 2 + }, + "1408": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 2 + }, + "1409": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 2 + }, + "22528": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 2 + }, + "22529": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 1 + } +} diff --git a/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_H200/grouped_matmul:v1/{K=2560,N=320,expert_num=513,mul_routed_weight=false,out_dtype=torch.bfloat16,topk_num=11,use_fp8_w8a8=false}_NVIDIA_H200.json b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_H200/grouped_matmul:v1/{K=2560,N=320,expert_num=513,mul_routed_weight=false,out_dtype=torch.bfloat16,topk_num=11,use_fp8_w8a8=false}_NVIDIA_H200.json new file mode 100644 index 0000000000..ace348c88f --- /dev/null +++ b/lightllm/common/triton_utils/autotune_kernel_configs/triton_3.5.1/NVIDIA_H200/grouped_matmul:v1/{K=2560,N=320,expert_num=513,mul_routed_weight=false,out_dtype=torch.bfloat16,topk_num=11,use_fp8_w8a8=false}_NVIDIA_H200.json @@ -0,0 +1,65 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 3 + }, + "33": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 1 + }, + "127": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 1 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 2 + }, + "513": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 2 + }, + "514": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "NEED_TRANS": false, + "num_warps": 4, + "num_stages": 1 + } +} diff --git a/lightllm/distributed/communication_op.py b/lightllm/distributed/communication_op.py index 7e51d8d736..b532e0b5f9 100644 --- a/lightllm/distributed/communication_op.py +++ b/lightllm/distributed/communication_op.py @@ -100,6 +100,21 @@ def all_reduce(self, input_: torch.Tensor) -> None: return dist.all_reduce(input_, group=self.device_group) def all_gather_into_tensor(self, output_: torch.Tensor, input_: torch.Tensor, async_op: bool = False) -> None: + if ( + not async_op + and self.symm_mem_reduce is not None + and output_.is_contiguous() + and output_.device == input_.device + and output_.dtype == input_.dtype + and output_.numel() == input_.numel() * self.dp_world_size + ): + gather_input = input_ + if input_.dtype in (torch.int32, torch.int64) and input_.is_contiguous(): + gather_input = input_.view(-1).view(self.symm_mem_reduce.dtype) + gathered = self.symm_mem_reduce.ordered_all_gather(gather_input) + if gathered is not None: + output_.view(-1).copy_(gathered.view(-1).view(input_.dtype)) + return return dist.all_gather_into_tensor(output_, input_, group=self.device_group, async_op=async_op) @@ -204,7 +219,7 @@ def new_deepep_group( num_max_tokens_per_rank=self.ll_num_tokens, hidden=self.ll_hidden, num_topk=num_experts_per_tok, - use_fp8_dispatch=True, + use_fp8_dispatch="none" not in expert_quant_method_names, allow_multiple_reduction=True, ) self.ep_mega_moe_buffer = None @@ -306,7 +321,7 @@ def all_gather_into_tensor( output_.copy_(input_) return if isinstance(group, CustomProcessGroup): - return group.all_gather_into_tensor(output_, input_) + return group.all_gather_into_tensor(output_, input_, async_op=async_op) else: return dist.all_gather_into_tensor(output_, input_, group, async_op) diff --git a/lightllm/distributed/symm_mem_all_reduce.py b/lightllm/distributed/symm_mem_all_reduce.py index 2256a5093f..727d5a0750 100644 --- a/lightllm/distributed/symm_mem_all_reduce.py +++ b/lightllm/distributed/symm_mem_all_reduce.py @@ -1,5 +1,7 @@ import torch import torch.distributed as dist +import triton +import triton.language as tl from torch.distributed import ProcessGroup from lightllm.utils.log_utils import init_logger @@ -29,6 +31,16 @@ _WORLD_SIZES_MULTIMEM = {"9.0": [4, 6, 8], "10.0": [6, 8], "10.3": [6, 8]} +@triton.jit +def _ordered_peer_sum(Inputs, Output, SIZE: tl.constexpr, START: 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 tl.static_range(len(Inputs)): + value = tl.load(Inputs[rank] + START + index, index < SIZE, other=0) + accumulator += value.to(tl.float64) + tl.store(Output + index, accumulator, index < SIZE) + + class SymmMemAllreduce: """In-place all-reduce via torch symmetric memory (NVLink SHARP / NVLS).""" @@ -65,6 +77,9 @@ def __init__(self, group: ProcessGroup, device, dtype: torch.dtype = torch.bfloa logger.warning("SymmMemAllreduce: multicast pointer unavailable; disabling.") return self.disabled = False + self.handle = handle + self.peer_buffers = None + self.use_peer_reduce = cap_str == "9.0" and self.world_size == 8 and dtype == torch.bfloat16 logger.info( "SymmMemAllreduce enabled: world_size=%d, max_size=%d, multimem=%s", self.world_size, @@ -90,3 +105,48 @@ 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) + + def ordered_reduce_scatter(self, inp: torch.Tensor, output: torch.Tensor, rank: int) -> bool: + if ( + not self.should_use(inp) + or not self.use_peer_reduce + or inp.device != self.device + or inp.numel() < 128 * 2560 + ): + return False + if rank != self.handle.rank: + raise ValueError("Sequence-parallel scatter rank must match the process group rank") + if output.numel() * self.world_size != inp.numel() or output.dtype != inp.dtype: + raise ValueError("Ordered reduce-scatter output must hold one rank's input shard") + if output.device != inp.device or not output.is_contiguous(): + raise ValueError("Ordered reduce-scatter output must be contiguous and on the input device") + if self.peer_buffers is None: + self.peer_buffers = tuple( + self.handle.get_buffer(rank, self.buffer.shape, self.dtype) for rank in range(self.world_size) + ) + self.buffer[: inp.numel()].copy_(inp.view(-1)) + self.handle.barrier(channel=0) + _ordered_peer_sum[(triton.cdiv(output.numel(), 256),)]( + self.peer_buffers, output, output.numel(), rank * output.numel(), 256 + ) + self.handle.barrier(channel=1) + return True diff --git a/lightllm/models/builtin.py b/lightllm/models/builtin.py index d8e091d9aa..d40f653e4a 100644 --- a/lightllm/models/builtin.py +++ b/lightllm/models/builtin.py @@ -127,6 +127,8 @@ def _tarsier_text_model_is(model_type): ModelRegistry.register(["qwen3_vl"], "lightllm.models.qwen3_vl.model:Qwen3VLTpPartModel", is_multimodal=True) ModelRegistry.register(["qwen3_vl_moe"], "lightllm.models.qwen3_vl_moe.model:Qwen3VLMOETpPartModel", is_multimodal=True) ModelRegistry.register("qwen3_next", "lightllm.models.qwen3next.model:Qwen3NextTpPartModel") +ModelRegistry.register("qwen4_exp", "lightllm.models.qwen4_exp.model:Qwen4ExpTpPartModel", is_multimodal=True) +ModelRegistry.register("qwen4_exp_text", "lightllm.models.qwen4_exp.model:Qwen4ExpTpPartModel") ModelRegistry.register( "qwen", "lightllm.models.qwen_vl.model:QWenVLTpPartModel", @@ -200,6 +202,9 @@ def _tarsier_text_model_is(model_type): DraftModelRegistry.register( "qwen3_moe", ("vanilla_no_att", "eagle_no_att"), "lightllm.models.qwen3_moe_mtp.model:Qwen3MOEMTPModel" ) +DraftModelRegistry.register( + ("qwen4_exp", "qwen4_exp_text"), "eagle_with_att", "lightllm.models.qwen4_exp.mtp.model:Qwen4ExpMTPModel" +) # Derive historical package-level class exports from the same registrations. diff --git a/lightllm/models/llama/layer_infer/post_layer_infer.py b/lightllm/models/llama/layer_infer/post_layer_infer.py index 888af5cd35..c3bc37048d 100644 --- a/lightllm/models/llama/layer_infer/post_layer_infer.py +++ b/lightllm/models/llama/layer_infer/post_layer_infer.py @@ -31,6 +31,9 @@ def _norm(self, input, infer_state, layer_weight: LlamaPreAndPostLayerWeight) -> return layer_weight.final_norm_weight_(input=input, eps=self.eps_, alloc_func=self.alloc_tensor) def _slice_get_last_input(self, input_embdings: torch.Tensor, infer_state: LlamaInferStateInfo): + indices = getattr(infer_state, "logits_indices", None) + if indices is not None: + return input_embdings.index_select(0, indices), indices.numel() embed_dim_ = input_embdings.shape[1] if infer_state.is_prefill: # logits 始终只取每个请求最后一个位置的 hidden state，用于正常采样。 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/qwen3_5/layer_infer/transformer_layer_infer.py b/lightllm/models/qwen3_5/layer_infer/transformer_layer_infer.py index 649db03b11..8e4c3c327b 100644 --- a/lightllm/models/qwen3_5/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/qwen3_5/layer_infer/transformer_layer_infer.py @@ -27,8 +27,11 @@ def _get_qkv( ) -> Tuple[torch.Tensor, torch.Tensor]: input = input.view(-1, self.embed_dim_) input = self._tpsp_allgather(input=input, infer_state=infer_state) + return self._project_qkv(input, infer_state, layer_weight) - qkv_gate_out = layer_weight.qkvo_gate_proj.mm(input) + def _project_qkv(self, input, infer_state, layer_weight, qkv_gate_out=None): + if qkv_gate_out is None: + qkv_gate_out = layer_weight.qkvo_gate_proj.mm(input) qkv_out, o_gate = qkv_gate_out.split( [ self.tp_q_head_num_ * self.head_dim_ + (self.tp_k_head_num_ + self.tp_v_head_num_) * self.head_dim_, diff --git a/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py b/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py index 1bc4ef914b..df30288562 100644 --- a/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py @@ -266,9 +266,7 @@ def context_attention_forward( gdn_out = self._linear_post(core_attn_out, z, layer_weight) - if self.tp_world_size_ > 1: - all_reduce(gdn_out, op=dist.ReduceOp.SUM, group=infer_state.dist_group, async_op=False) - return gdn_out + return self._reduce_linear_output(gdn_out, infer_state) def token_attention_forward( self, @@ -296,9 +294,12 @@ def token_attention_forward( ) gdn_out = self._linear_post(core_attn_out, z, layer_weight) + return self._reduce_linear_output(gdn_out, infer_state) + + def _reduce_linear_output(self, output, infer_state): if self.tp_world_size_ > 1: - all_reduce(gdn_out, op=dist.ReduceOp.SUM, group=infer_state.dist_group, async_op=False) - return gdn_out + all_reduce(output, op=dist.ReduceOp.SUM, group=infer_state.dist_group, async_op=False) + return output def _linear_in_proj( self, @@ -316,7 +317,6 @@ def _linear_post( ) -> torch.Tensor: num_tokens = z.shape[0] core_attn_out = core_attn_out.view(-1, core_attn_out.shape[-1]) - z = z.contiguous().view(-1, z.shape[-1]) norm_out = layer_weight.linear_norm(core_attn_out, z, self.eps_) core_attn_out = norm_out.view(num_tokens, -1) output = layer_weight.linear_out_proj.mm(core_attn_out) diff --git a/lightllm/models/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..0dadc06a0e --- /dev/null +++ b/lightllm/models/qwen4_exp/attention_state.py @@ -0,0 +1,198 @@ +from dataclasses import replace + +import torch + +from lightllm.common.attention_state.full_att.gqa.gqa_state import GQAFullAttentionState, GQAAttentionConfig +from .state import QSAExecutionState + + +class QSAIndexerState(GQAFullAttentionState): + def __init__(self, layout_config, page_size, device, pooled_dtype): + super().__init__(layout_config, page_size, device) + self.pooled_dtype = pooled_dtype + + def allocate_storage(self, token_capacity): + self.capacity = token_capacity + self.k_storage = torch.empty(self._storage_shape(), dtype=self.storage_dtype, device=self.device) + self.v_storage = torch.empty(self._storage_shape(), dtype=self.pooled_dtype, device=self.device) + + def get_token_bytes(self): + config = self.config + return ( + len(config.logical_layer_groups) + * config.head_num + * config.head_dim + * (self.storage_dtype.itemsize + self.pooled_dtype.itemsize) + ) + + def _get_transfer_regions(self, page_tokens): + key, value = super()._get_transfer_regions(page_tokens) + value = replace( + value, + dtype=self.pooled_dtype, + byte_end=value.byte_start + + (value.byte_end - value.byte_start) * self.pooled_dtype.itemsize // self.storage_dtype.itemsize, + ) + return key, value + + def copy_full_att_tokens(self, src_indices, dst_indices): + from lightllm.common.attention_state.full_att.dp_transfer import copy_gqa_tokens + + copy_gqa_tokens(self.k_storage, self.v_storage, self.k_storage, self.v_storage, src_indices, dst_indices) + + +class QSAPooledState(GQAFullAttentionState): + def allocate_storage(self, token_capacity): + self.capacity = token_capacity + self.k_storage = torch.empty(self._storage_shape(), dtype=self.storage_dtype, device=self.device) + + def get_token_bytes(self): + return super().get_token_bytes() // 2 + + def _get_transfer_regions(self, page_tokens): + return super()._get_transfer_regions(page_tokens)[:1] + + def copy_full_att_tokens(self, src_indices, dst_indices): + from lightllm.common.attention_state.full_att.dp_transfer import copy_packed_tokens + + copy_packed_tokens(self.k_storage, self.k_storage, src_indices, dst_indices) + + def copy_tokens_from_dp(self, src_token_indices, dst_token_indices, remote_dp_full_att_state): + from lightllm.common.attention_state.full_att.dp_transfer import copy_packed_tokens + + copy_packed_tokens(remote_dp_full_att_state.k_storage, self.k_storage, src_token_indices, dst_token_indices) + + +class Qwen4ExpAttentionState(GQAFullAttentionState): + def __init__(self, layout_config, indexer_dim, rotary_dim, page_size=1, device="meta", fp8=False, compress_ratio=1): + super().__init__(layout_config, page_size, device) + if compress_ratio < 1 or page_size % compress_ratio or indexer_dim % compress_ratio: + raise ValueError("QSA compact cache requires complete aligned blocks and divisible indexer dimensions") + self.compress_ratio = compress_ratio + self.indexer_dim = indexer_dim + if fp8: + self.storage_dtype = torch.float8_e4m3fn + 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, + ) + if compress_ratio > 1: + self.indexer = QSAPooledState( + GQAAttentionConfig( + logical_layer_groups=layout_config.logical_layer_groups, + head_dim=indexer_dim // compress_ratio, + **common, + ), + page_size, + device, + ) + self.indexer.storage_dtype = self.storage_dtype + self.positions = None + return + self.indexer = QSAIndexerState( + GQAAttentionConfig(logical_layer_groups=layout_config.logical_layer_groups, head_dim=indexer_dim, **common), + page_size, + device, + torch.float8_e4m3fn if fp8 else layout_config.dtype, + ) + 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): + suffix = "_fp8" if self.storage_dtype == torch.float8_e4m3fn else "" + if self.compress_ratio > 1: + suffix += f"_c{self.compress_ratio}" + return "qwen4_" + super().get_state_type() + suffix + + def _side_states(self): + return (self.indexer,) if self.positions is None else (self.indexer, self.positions) + + def allocate_storage(self, token_capacity): + super().allocate_storage(token_capacity) + for state in self._side_states(): + state.device = self.device + state.allocate_storage(token_capacity) + + def get_token_bytes(self): + return super().get_token_bytes() + sum(state.get_token_bytes() for state in self._side_states()) + + def qsa_view(self, layer, compressor=None): + if self.compress_ratio > 1: + pooled = self.indexer.k_storage[self.indexer.get_storage_layer_index(layer)].view(-1, self.indexer_dim) + return QSAExecutionState(None, pooled, None, None, compressor) + 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): + if self.positions is None: + return + 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)): + if state is None: + continue + 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)): + if state is None: + continue + 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._side_states(): + 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): + assert ( + self.storage_dtype == remote_dp_full_att_state.storage_dtype + and self.compress_ratio == remote_dp_full_att_state.compress_ratio + ), "QSA DP copy requires the same KV and indexer encoding" + super().copy_tokens_from_dp(src_token_indices, dst_token_indices, remote_dp_full_att_state) + for name in ("indexer", "positions"): + if getattr(self, name) is None: + continue + 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..c1d4c95474 --- /dev/null +++ b/lightllm/models/qwen4_exp/batch.py @@ -0,0 +1,194 @@ +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 _maximum(a, b): + return tl.maximum(a, b) + + +@triton.jit +def _minimum(a, b): + return tl.minimum(a, b) + + +@triton.jit +def _decode_metadata( + Requests, + Lengths, + Starts, + Active, + KeyLengths, + FirstBlocks, + QueryRows, + Visible, + TokenMask, + BATCH: tl.constexpr, + HOLD: tl.constexpr, + RATIO: tl.constexpr, + BLOCK: tl.constexpr, +): + row = tl.arange(0, BLOCK) + valid = row < BATCH + request = tl.load(Requests + row, valid, -1) + previous = tl.load(Requests + row - 1, (row > 0) & valid, -1) + following = tl.load(Requests + row + 1, row + 1 < BATCH, -1) + first = (row == 0) | (request != previous) | (request == HOLD) + last = (row == BATCH - 1) | (request != following) | (request == HOLD) + start = tl.associative_scan(tl.where(first, row, 0), 0, _maximum) + end = tl.associative_scan(tl.where(last, row, BATCH), 0, _minimum, reverse=True) + length = tl.load(Lengths + row, valid, 0) + last_length = tl.load(Lengths + end, valid & (end < BATCH), 0) + active = last & (request != HOLD) + tl.store(Starts + row, start, valid) + tl.store(Starts + BATCH, BATCH) + tl.store(Active + row, tl.where(active, request, -1), valid) + tl.store(KeyLengths + row, tl.where(active, length, 0), valid) + tl.store(FirstBlocks + row, tl.where(active, (length - (row - start + 1)) // RATIO, 0), valid) + tl.store(QueryRows + row, request, valid) + tl.store(Visible + row, tl.where(request != HOLD, last_length - end + row, 0), valid) + tl.store(TokenMask + row, request != HOLD, valid) + + +@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, + decode=False, + ): + 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") + + if decode: + if token_count != batch or capacity != batch: + raise ValueError("Decode metadata requires one token per row and no extra segment capacity") + starts = torch.empty(batch + 1, dtype=torch.int32, device=requests.device) + active, key_lengths, first_blocks, query_rows, visible = (torch.empty_like(requests) for _ in range(5)) + token_mask = torch.empty(batch, dtype=torch.bool, device=requests.device) + if batch: + _decode_metadata[(1,)]( + requests, + lengths, + starts, + active, + key_lengths, + first_blocks, + query_rows, + visible, + token_mask, + batch, + hold_request, + compress_ratio, + triton.next_power_of_2(batch), + ) + else: + starts.zero_() + return cls(starts, active, key_lengths, first_blocks, query_rows, visible, token_mask) + + 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..5bf203b22a --- /dev/null +++ b/lightllm/models/qwen4_exp/infer_struct.py @@ -0,0 +1,96 @@ +import weakref + +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): + self._qwen4_exp_model = weakref.ref(model) + self.qwen4_exp_config = model.config + self.qwen4_exp_fixed_decode = model.fixed_decode_config + self.qwen4_exp_segment_capacity = model.max_req_num + 1 if self.is_prefill else None + self.qwen4_exp_auto_commit = self.is_prefill or model.args.mtp_step == 0 + self.qwen4_exp_graph_preparation = not self.is_prefill and ( + self.is_cuda_graph_capture or self.is_cuda_graph_replay + ) + if not self.qwen4_exp_graph_preparation: + super().init_some_extra_state(model) + self.prepare_qwen_state() + + def prepare_qwen_state(self): + if self.qwen4_exp_graph_preparation: + super().init_some_extra_state(self._qwen4_exp_model()) + config = self.qwen4_exp_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=self.qwen4_exp_segment_capacity, + decode=not self.is_prefill, + ) + 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 + from .qsa_compressor import QSACompressorState + + states = self.att_state_manager.get_full_attention_states(self.att_state_role) + bounded = self.att_state_manager.get_bounded_attention_states(self.att_state_role) + compressor = next((state for state in bounded if isinstance(state, QSACompressorState)), None) + 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: + ring = compressor.get_layer_view(layer) if compressor is not None else None + self.qwen4_exp_qsa[layer] = state.qsa_view(layer, ring) + self.qwen4_exp_lexical_state = next( + (state for state in bounded 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 + ) + ) + + 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) + if self.qwen4_exp_graph_preparation: + if not new_infer_state.qwen4_exp_graph_preparation: + raise ValueError("Qwen graph replay cannot change state preparation mode") + else: + 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..361f603326 --- /dev/null +++ b/lightllm/models/qwen4_exp/layer_infer/pre_layer_infer.py @@ -0,0 +1,25 @@ +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): + if getattr(infer_state, "qwen4_exp_graph_preparation", False): + infer_state.prepare_qwen_state() + 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/replicated_decoder.py b/lightllm/models/qwen4_exp/layer_infer/replicated_decoder.py new file mode 100644 index 0000000000..1dcc1c7b2c --- /dev/null +++ b/lightllm/models/qwen4_exp/layer_infer/replicated_decoder.py @@ -0,0 +1,220 @@ +import torch +import triton +from lightllm.common.basemodel.tp_reduce import ordered_tp_sum +from lightllm.common.basemodel.cuda_graph import CudaGraph +from lightllm.common.basemodel.hidden_collector import FinalHiddenCollector, NoopHiddenCollector +from lightllm.common.basemodel.triton_kernel.fused_moe.topk_select import fused_topk +from lightllm.distributed import dist_group_manager +from lightllm.models.qwen4_exp.layer_infer.transformer_layer_infer import Qwen4ExpTransformerLayerInfer +from ..triton_kernel.replicated_moe import expert_down, expert_up, pack_routes, sum_partials, sum_routes +from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import silu_and_mul_fwd + + +_REPLICATED_DECODE_CONFIG = { + "hidden_size": 2560, + "num_experts": 512, + "moe_intermediate_size": 640, + "shared_expert_intermediate_size": 640, + "num_experts_per_tok": 10, + "norm_topk_prob": True, + "hidden_act": "silu", +} + + +def supports_decode_profile(model): + args = model.args + return ( + model.tp_world_size_ == 8 + and args.dp == 1 + and args.nnodes == 1 + and model.data_type == torch.bfloat16 + and args.llm_kv_type == "fp8kv_spt" + and args.enable_ep_moe + and args.enable_tpsp_mix_mode + and args.mtp_mode == "eagle_with_att" + and args.mtp_step in (2, 3, 4) + and not args.mtp_dynamic_verify + and not args.enable_decode_microbatch_overlap + and args.ssm_state_mode == "replay" + and args.linear_att_ssm_data_type == "float32" + and args.ssm_replay_config.get("cache_len", 16) == 16 + and args.ssm_replay_config.get("compute_mode", "precompute") == "precompute" + and model.quant_cfg.quant_type == "none" + and not model.quant_cfg.quantized_weight + and all(method == "none" for layer in model.quant_cfg.quant_cfg.values() for method in layer.values()) + and all(model.config.get(key) == value for key, value in _REPLICATED_DECODE_CONFIG.items()) + and torch.cuda.get_device_capability() == (9, 0) + ) + + +def supports_replicated_decode(model): + args = model.args + if ( + not supports_decode_profile(model) + or model.is_mtp_draft_model + or args.enable_torch_memory_saver + or args.enable_return_routed_experts + or args.auto_update_redundancy_expert + or type(model.hidden_collector_prototype) not in (FinalHiddenCollector, NoopHiddenCollector) + ): + return False + symmetric = dist_group_manager.get_group(0).symm_mem_reduce + if ( + symmetric is None + or symmetric.disabled + or symmetric.dtype != torch.bfloat16 + or symmetric.world_size != 8 + or symmetric.max_size < 64 * 2560 * 2 * 8 + or not hasattr(torch.ops.symm_mem, "multimem_all_gather_out") + ): + return False + local_ids = list(range(model.tp_rank_ * 64, (model.tp_rank_ + 1) * 64)) + for layer in model.trans_layers_weight: + experts = getattr(layer, "experts", None) + if ( + experts is None + or experts.quant_method.method_name != "none" + or experts.local_logic_expert_ids_list != local_ids + or experts.num_fused_shared_experts + or experts.redundancy_expert_num + or experts.scoring_func != "softmax" + or experts.routed_scaling_factor != 1.0 + or experts.per_expert_scale is not None + or experts.e_score_correction_bias is not None + ): + return False + for pack, shape in ((experts.w13, (64, 1280, 2560)), (experts.w2, (64, 2560, 640))): + if pack.weight.dtype != torch.bfloat16 or pack.weight.shape != shape or not pack.weight.is_contiguous(): + return False + if layer.gate_up_proj.quant_method.method_name != "none" or layer.down_proj.quant_method.method_name != "none": + return False + return True + + +class ReplicatedDecoderLayer(Qwen4ExpTransformerLayerInfer): + def _tpsp_sp_split(self, input, infer_state): + return input + + def _compute_shared_expert(self, input, infer_state, layer_weight): + shape = (input.shape[0], input.shape[1]) + if getattr(self, "shared_shape", None) != shape: + self.shared_shape = shape + self.shared_up = torch.empty(input.shape[0], 1280, dtype=input.dtype, device=input.device) + self.shared_mid = torch.empty(input.shape[0], 640, dtype=input.dtype, device=input.device) + self.shared_output = torch.empty_like(input) + self.shared_gate = torch.empty(input.shape[0], 1, dtype=input.dtype, device=input.device) + layer_weight.gate_up_proj.mm(input, out=self.shared_up) + silu_and_mul_fwd(self.shared_up, self.shared_mid) + layer_weight.down_proj.mm(self.shared_mid, out=self.shared_output) + layer_weight.shared_expert_gate.mm(input, out=self.shared_gate) + return self.shared_output + + def _tpsp_allgather(self, input, infer_state): + return input + + def _tpsp_reduce(self, input, infer_state): + return ordered_tp_sum(input, infer_state.dist_group, self.tp_world_size_) + + def _moe_ffn_edp(self, input, infer_state, layer_weight): + tokens, hidden = input.shape + experts, inter, topk = 64, 640, 10 + if getattr(self, "workspace_tokens", None) != tokens: + self.shared_stream = torch.cuda.Stream() + self.workspace_tokens = tokens + self.routes = torch.empty(experts, tokens * topk, dtype=torch.int32, device="cuda") + self.masked_ids = torch.empty(tokens, topk, dtype=torch.int32, device="cuda") + self.counts = torch.empty(experts, dtype=torch.int32, device="cuda") + self.mid = torch.empty(tokens * topk, inter, dtype=input.dtype, device="cuda") + self.down = torch.empty(tokens * topk, hidden, dtype=input.dtype, device="cuda") + self.partial = torch.empty(tokens, hidden, dtype=input.dtype, device="cuda") + self.shared_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(self.shared_stream): + shared = self._compute_shared_expert(input, infer_state, layer_weight) + weights, ids = fused_topk(input, layer_weight.moe_gate.mm(input), topk, True) + mask = infer_state.qwen4_exp_metadata.token_mask + count = tokens * topk + pack_routes[(experts,)]( + ids, + mask, + self.masked_ids, + self.routes, + self.counts, + tokens, + topk, + self.tp_rank_ * experts, + triton.next_power_of_2(count), + num_warps=4 if count > 256 else 1, + ) + ids = self.masked_ids + expert_up[(experts, triton.cdiv(inter, 32))]( + input, + layer_weight.experts.w13.weight, + self.mid, + self.routes, + self.counts, + hidden, + inter, + tokens, + topk, + 16, + 32, + 256, + num_warps=4, + num_stages=3, + ) + expert_down[(experts, triton.cdiv(hidden, 64))]( + self.mid, + layer_weight.experts.w2.weight, + self.down, + self.routes, + self.counts, + inter, + hidden, + tokens, + topk, + 1, + 16, + 64, + 128, + num_warps=4, + num_stages=3, + ) + sum_routes[(tokens, triton.cdiv(hidden, 256))]( + self.down, + weights, + ids, + self.partial, + hidden, + topk, + self.tp_rank_ * experts, + experts, + triton.next_power_of_2(topk), + num_warps=4, + ) + gathered = infer_state.dist_group.symm_mem_reduce.ordered_all_gather(self.partial) + combined = torch.empty_like(input) + torch.cuda.current_stream().wait_stream(self.shared_stream) + sum_partials[(tokens, triton.cdiv(hidden, 512))]( + gathered, + ids, + combined, + hidden, + tokens, + topk, + experts, + 512, + 10, + num_warps=4, + enable_fp_fusion=False, + Shared=shared, + Gate=self.shared_gate, + ) + return combined + + +class ReplicatedDecodeGraph(CudaGraph): + @staticmethod + def gen_cuda_graph_batch_sizes(batch_step_size_before_split, *args, **kwargs): + sizes = CudaGraph.gen_cuda_graph_batch_sizes(batch_step_size_before_split, *args, **kwargs) + width = batch_step_size_before_split + return sorted(set(sizes) | set(range(width, min(max(sizes), 64) + 1, width))) 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..93bf25a826 --- /dev/null +++ b/lightllm/models/qwen4_exp/layer_infer/transformer_layer_infer.py @@ -0,0 +1,211 @@ +import torch + +from lightllm.common.basemodel.tp_reduce import ordered_tp_sum +from lightllm.utils.envs_utils import get_env_start_args +from lightllm.utils.tensor_utils import tensor_to_no_ref_tensor +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 _moe_ffn_edp(self, input, infer_state, layer_weight): + mask = infer_state.qwen4_exp_metadata.token_mask + mask = self._tpsp_sp_split(mask[:, None], infer_state).view(-1) + return layer_weight.experts.experts( + input, + router_logits=layer_weight.moe_gate.mm(input), + top_k=self.num_experts_per_tok, + renormalize=self.norm_topk_prob, + use_grouped_topk=False, + topk_group=None, + num_expert_group=None, + is_prefill=infer_state.is_prefill, + infer_state=infer_state, + token_mask=mask, + shared_expert_func=lambda: self._compute_shared_expert(input, infer_state, layer_weight), + ) + + 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) + mlp_residual = layer_weight.mlp_hyper_connection.component + residual, (mixed, gates) = mlp_residual.combine_and_read(residual, update, gates) + 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): + hidden = self._tpsp_allgather(input_embdings, infer_state) + index_projection = None + if ( + not infer_state.is_prefill + and infer_state.qwen4_exp_fixed_decode + and hidden.shape[0] <= 128 + and hasattr(layer_weight, "qsa_projection") + ): + projected = torch.mm(hidden, layer_weight.qsa_projection.t()) + width = layer_weight.qkvo_gate_proj.mm_param.weight.shape[0] + qkv_projection, index_projection = projected.split( + (width, layer_weight.indexer.component.projection.shape[0]), dim=-1 + ) + query, kv = self._project_qkv(hidden, infer_state, layer_weight, qkv_gate_out=qkv_projection) + else: + query, kv = self._project_qkv(hidden, 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_) + output = infer_state.qwen4_exp_qsa[self.layer_num_].attend( + layer_weight.indexer.component, + hidden, + query.view(-1, self.tp_q_head_num_, self.head_dim_), + key, + value, + infer_state.qwen4_exp_batch, + decode=not infer_state.is_prefill, + index_projection=index_projection, + attention_splits=None if infer_state.is_prefill else (16 if infer_state.qwen4_exp_fixed_decode else 4), + max_visible_length=getattr(infer_state, "max_kv_seq_len", None) if infer_state.is_prefill else None, + tp_world_size=self.tp_world_size_, + tp_rank=self.tp_rank_, + dist_group=infer_state.dist_group, + ) + return self._get_o(output, infer_state, layer_weight) + + def _linear_prefill_cuda_graph_wrapper(self, mixed_qkvzba, infer_state, layer_weight): + from lightllm.common.basemodel.attention.base_att import AttControl + from lightllm.common.basemodel.attention.linear.flashinfer import FlashInferLinearAttBackend + from lightllm.common.basemodel.attention.linear.replay_gdn import ReplayGDNFlashInferAttBackend + + backend = infer_state.prefill_att_state1.backend + if not isinstance(backend, (FlashInferLinearAttBackend, ReplayGDNFlashInferAttBackend)): + return super()._linear_prefill_cuda_graph_wrapper(mixed_qkvzba, infer_state, layer_weight) + + mixed_qkvzba = mixed_qkvzba.contiguous() + captured_input = tensor_to_no_ref_tensor(mixed_qkvzba) + request_count = infer_state.b_req_idx.numel() + + def prefill(state): + return state.prefill_att_state1.prefill_att( + q=None, + k=None, + v=None, + att_control=AttControl( + linear_att_prefill=True, + linear_att_prefill_dict={ + "mixed_qkvzba": captured_input, + "layer_weight": layer_weight, + "layer_num": self.layer_num_, + }, + ), + alloc_func=self.alloc_tensor, + ) + + previous = infer_state.prefill_cuda_graph_get_current_capture_graph() + previous.__exit__(None, None, None) + infer_state.prefill_cuda_graph_create_graph_obj() + infer_state.prefill_cuda_graph_get_current_capture_graph().__enter__() + infer_state.prefill_att_state1.init_state() + output, z = prefill(infer_state) + output, z = output.contiguous(), z.contiguous() + infer_state.prefill_cuda_graph_get_current_capture_graph().__exit__(None, None, None) + graph, _ = infer_state.prefill_cuda_graph_exe_list.pop() + captured_output = tensor_to_no_ref_tensor(output) + captured_z = tensor_to_no_ref_tensor(z) + + def attention(state): + if state.b_req_idx.numel() == request_count: + graph.replay() + else: + result, result_z = prefill(state) + captured_output.copy_(result) + captured_z.copy_(result_z) + + infer_state.prefill_cuda_graph_add_cpu_runnning_func(func=attention, after_graph=previous) + infer_state.prefill_cuda_graph_create_graph_obj() + infer_state.prefill_cuda_graph_get_current_capture_graph().__enter__() + return output, z + + def _qsa_prefill_cuda_graph_wrapper(self, input_embdings, infer_state, layer_weight): + captured_input = tensor_to_no_ref_tensor(input_embdings) + max_kv_seq_len = infer_state.max_kv_seq_len + previous = infer_state.prefill_cuda_graph_get_current_capture_graph() + previous.__exit__(None, None, None) + infer_state.prefill_cuda_graph_create_graph_obj() + infer_state.prefill_cuda_graph_get_current_capture_graph().__enter__() + output = self._qsa_attention_forward(captured_input, infer_state, layer_weight).contiguous() + infer_state.prefill_cuda_graph_get_current_capture_graph().__exit__(None, None, None) + graph, _ = infer_state.prefill_cuda_graph_exe_list.pop() + captured_output = tensor_to_no_ref_tensor(output) + + def attention(state): + if max_kv_seq_len is None or (state.max_kv_seq_len is not None and state.max_kv_seq_len <= max_kv_seq_len): + graph.replay() + else: + result = self._qsa_attention_forward(captured_input, state, layer_weight) + captured_output.copy_(result) + + infer_state.prefill_cuda_graph_add_cpu_runnning_func(func=attention, after_graph=previous) + infer_state.prefill_cuda_graph_create_graph_obj() + infer_state.prefill_cuda_graph_get_current_capture_graph().__enter__() + return output + + 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) + if torch.cuda.is_current_stream_capturing() and hasattr(infer_state, "prefill_cuda_graph_exe_list"): + return self._qsa_prefill_cuda_graph_wrapper(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..bde6404765 --- /dev/null +++ b/lightllm/models/qwen4_exp/lexical_cache.py @@ -0,0 +1,141 @@ +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 _gather_history( + Rows, + Starts, + Lengths, + Ngram, + Convolution, + OutNgram, + OutConvolution, + NGRAM: tl.constexpr, + WIDTH: tl.constexpr, + EOS: tl.constexpr, + RESET: tl.constexpr, + BLOCK: tl.constexpr, +): + row, part = tl.program_id(0), tl.program_id(1) + request = tl.load(Rows + row).to(tl.int64) + active = request >= 0 + if RESET: + count = tl.load(Starts + row + 1) - tl.load(Starts + row) + active = active & (count != tl.load(Lengths + row)) + col = part * BLOCK + tl.arange(0, BLOCK) + values = tl.load(Convolution + request * WIDTH + col, active & (col < WIDTH), 0) + tl.store(OutConvolution + row.to(tl.int64) * WIDTH + col, values, col < WIDTH) + if part == 0 and NGRAM > 0: + offset = tl.arange(0, triton.next_power_of_2(NGRAM)) + tokens = tl.load(Ngram + request * NGRAM + offset, active & (offset < NGRAM), EOS) + tl.store(OutNgram + row.to(tl.int64) * NGRAM + offset, tokens, offset < NGRAM) + + +@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 + if rows.is_cuda: + ngram = torch.empty((rows.numel(), self.ngram.shape[1]), dtype=self.ngram.dtype, device=rows.device) + convolution = torch.empty( + (rows.numel(), *self.convolution.shape[1:]), dtype=self.convolution.dtype, device=rows.device + ) + width = self.convolution.shape[1] * self.convolution.shape[2] + if rows.numel(): + _gather_history[(rows.numel(), max(1, triton.cdiv(width, 512)))]( + rows, + metadata.starts, + metadata.key_lengths, + self.ngram, + self.convolution, + ngram, + convolution, + self.ngram.shape[1], + width, + self.eos_token_id, + reset_empty_prefix, + 512, + ) + return PLEBatch(metadata.starts, ngram, convolution, metadata.token_mask) + 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..1e0e5fab9b --- /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, req_verify_width): + 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..c534ca8ade --- /dev/null +++ b/lightllm/models/qwen4_exp/model.py @@ -0,0 +1,204 @@ +import json +from pathlib import Path + +import torch + +from lightllm.common.basemodel.cuda_graph import CudaGraph +from .layer_infer.replicated_decoder import ( + ReplicatedDecoderLayer, + ReplicatedDecodeGraph, + supports_decode_profile, + supports_replicated_decode, +) + +from lightllm.models.qwen3_5_moe.model import Qwen3_5MOETpPartModel +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): + _replicated_decode = False + 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 not in ("None", "fp8kv_spt"): + raise ValueError("Qwen4Exp QSA supports BF16 or unit-scale FP8 KV storage") + if args.llm_kv_type == "fp8kv_spt" and args.kv_quant_calibration_config_path: + raise ValueError("Qwen4Exp FP8 QSA uses unit scales and does not support KV calibration files") + 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, + fp8=args.llm_kv_type == "fp8kv_spt", + compress_ratio=config["indexer_compress_ratio"] if args.page_size > 1 else 1, + ) + ] + + @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 = Qwen3_5MOETpPartModel.create_bounded_att_states.__func__( + cls, config, args, model_logical_layer_ids, request_capacity, tp_rank, tp_world_size, device + ) + if args.page_size > 1: + from .qsa_compressor import QSACompressorState + + ratio = config["indexer_compress_ratio"] + if args.page_size % ratio or args.linear_att_hash_page_size % ratio: + raise ValueError("QSA cache pages must align to the indexer compression ratio") + interval = config["full_attention_interval"] + qsa_layers = tuple( + layer for index, layer in enumerate(model_logical_layer_ids) if (index + 1) % interval == 0 + ) + states.append( + QSACompressorState( + qsa_layers, + config["indexer_head_dim"], + int(config["head_dim"] * config["partial_rotary_factor"]), + ratio, + get_torch_dtype(args.data_type), + request_capacity, + args.mtp_step, + tp_rank, + tp_world_size, + device, + ) + ) + layers = tuple(model_logical_layer_ids[layer - 1] for layer in config.get("ple_layer_ids", [])) + if layers: + states.append( + Qwen4ExpLexicalState( + LexicalStateConfig( + logical_layer_ids=layers, + ngram_context=config["ngram_size"] - 1, + conv_context=(config["ple_conv_kernel_size"] - 1) * config["ngram_size"], + width=config["hc_count"] * config["hidden_size"], + eos_token_id=ngram_eos_token_id(config), + dtype=get_torch_dtype(args.data_type), + tp_rank=tp_rank, + ), + request_capacity, + args.mtp_step, + device, + ) + ) + return states + + def prepare_attention_states(self): + super().prepare_attention_states() + from .lexical_state import Qwen4ExpLexicalState + + state = getattr(self.att_state_manager, f"{self.att_state_role}_qwen4_lexical_state", None) + if isinstance(state, Qwen4ExpLexicalState): + state.bind_weights(self.trans_layers_weight) + + def _init_att_backend(self): + self.prefill_att_backend = self.decode_att_backend = TritonAttBackend(model=self) + self.fixed_decode_config = supports_decode_profile(self) and self.args.mtp_step == 3 + + if self.fixed_decode_config and not self.is_mtp_draft_model: + for layer in self.trans_layers_weight: + if not layer.is_linear_attention_layer and not hasattr(layer, "qsa_projection"): + layer.qsa_projection = torch.cat( + (layer.qkvo_gate_proj.mm_param.weight, layer.indexer.component.projection) + ) + + def _init_att_backend1(self): + states = self.att_state_manager.get_bounded_attention_states(self.att_state_role) + if not any(state.get_state_type().startswith("gdn_") for state in states): + self.prefill_att_backend1 = self.decode_att_backend1 = None + return + Qwen3_5MOETpPartModel._init_att_backend1(self) + if self.fixed_decode_config and not self.is_mtp_draft_model: + self.decode_att_backend1.replay_run_config = {"BV": 32, "num_warps": 4, "num_stages": 2} + + def _init_hidden_collector(self): + super()._init_hidden_collector() + self._init_replicated_decode() + + def _init_replicated_decode(self): + self._replicated_decode = supports_replicated_decode(self) + self.decode_graph_class = ReplicatedDecodeGraph if self._replicated_decode else CudaGraph + self._replicated_layers = {} + + def _align_decode_batch_size(self, batch_size): + if self._replicated_decode and 0 < batch_size <= 64: + width = self.mtp_manager.get_decode_batch_alignment(self.is_mtp_draft_model) + aligned = (batch_size + width - 1) // width * width + if aligned <= 64: + return aligned + return super()._align_decode_batch_size(batch_size) + + def _token_forward(self, infer_state): + rows = infer_state.input_ids.numel() + if not self._replicated_decode or not 0 < rows <= 64: + return super()._token_forward(infer_state) + if rows not in self._replicated_layers: + self._replicated_layers[rows] = [ + ReplicatedDecoderLayer(layer.layer_num_, self.config) for layer in self.layers_infer + ] + hidden = self.pre_infer.token_forward(infer_state.input_ids, infer_state, self.pre_post_weight) + for layer, weight in zip(self._replicated_layers[rows], self.trans_layers_weight): + hidden = layer.token_forward(hidden, infer_state, weight) + post = self.post_infer.token_forward(hidden, infer_state, self.pre_post_weight) + infer_state.hidden_collector.add_final_hidden(hidden) + output = self._create_model_output(post, infer_state) + if infer_state.is_cuda_graph_capture: + output.to_no_ref_tensor() + return output 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..0ef960bcd7 --- /dev/null +++ b/lightllm/models/qwen4_exp/mtp/infer.py @@ -0,0 +1,31 @@ +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): + if getattr(infer_state, "qwen4_exp_graph_preparation", False): + infer_state.prepare_qwen_state() + 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..8866ab868c --- /dev/null +++ b/lightllm/models/qwen4_exp/mtp/model.py @@ -0,0 +1,53 @@ +import torch + +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): + supports_selected_logits = True + 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 + _init_att_backend1 = Qwen4ExpTpPartModel._init_att_backend1 + create_full_att_states = classmethod(Qwen4ExpTpPartModel.create_full_att_states.__func__) + create_bounded_att_states = classmethod(Qwen4ExpTpPartModel.create_bounded_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..f9ad0abfd5 --- /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 = 1024 * 1024 * 1024 +_LARGE_BYTES = 512 * 1024 * 1024 + + +class _RegisteredArray(np.ndarray): + def __array_finalize__(self, source): + self.allocation = getattr(source, "allocation", None) + + +class _Registration: + def __init__(self, mapping, device, runtime): + self.mapping = mapping + self.device = device + self.runtime = runtime + self.addresses = [] + + def __del__(self): + if not self.addresses: + return + + try: + with torch.cuda.device(self.device): + torch.cuda.synchronize() + except Exception: + pass + for address in reversed(self.addresses): + try: + self.runtime.cudaHostUnregister(address) + except Exception: + pass + + +def pinned_zeros(shape, *, dtype, device="cpu", pin_memory=True): + size = math.prod(shape) * dtype.itemsize + if device != "cpu" or not pin_memory: + raise ValueError("Qwen recurrent snapshots require pinned CPU storage") + if size < _LARGE_BYTES: + return torch.zeros(shape, dtype=dtype, device="cpu", pin_memory=True) + + cell_bytes = math.prod(shape[1:]) * dtype.itemsize + registration_unit = math.lcm(cell_bytes, mmap.PAGESIZE) + if registration_unit > _CHUNK_BYTES: + raise ValueError("A recurrent snapshot must fit within one bounded host registration") + chunk_bytes = (_CHUNK_BYTES // registration_unit) * registration_unit + runtime = torch.cuda.cudart() + mapping = mmap.mmap(-1, size) + allocation = _Registration(mapping, torch.cuda.current_device(), runtime) + array = np.frombuffer(mapping, dtype=np.uint8).view(_RegisteredArray) + array.allocation = allocation + value = torch.from_numpy(array).view(dtype).reshape(shape) + value.zero_() + for offset in range(0, size, chunk_bytes): + address = value.data_ptr() + offset + status = runtime.cudaHostRegister(address, min(chunk_bytes, size - offset), 0) + if int(status): + raise RuntimeError( + f"Cannot register Qwen pinned snapshot at byte {offset}: {runtime.cudaGetErrorString(status)}" + ) + allocation.addresses.append(address) + + value._qwen4_pinned_allocation = allocation + return value diff --git a/lightllm/models/qwen4_exp/position.py b/lightllm/models/qwen4_exp/position.py new file mode 100644 index 0000000000..9e7d62aac1 --- /dev/null +++ b/lightllm/models/qwen4_exp/position.py @@ -0,0 +1,73 @@ +import torch +import triton +import triton.language as tl + + +@triton.jit +def _combine_mrope( + Cos, + Sin, + OutCos, + OutSin, + ROWS: tl.constexpr, + DIM: tl.constexpr, + C0: tl.constexpr, + C1: tl.constexpr, + C2: tl.constexpr, + S0: tl.constexpr, + S1: tl.constexpr, + S2: tl.constexpr, + HEIGHT: tl.constexpr, + WIDTH: tl.constexpr, + AXES: tl.constexpr, + BLOCK: tl.constexpr, +): + index = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + row, col = index // (2 * DIM), index % DIM + axis = tl.full((BLOCK,), 0, tl.int32) + if AXES: + height = (col % 3 == 1) & (col <= 3 * HEIGHT) + width = (col % 3 == 2) & (col <= 3 * WIDTH) + axis = tl.where(height, 1, tl.where(width, 2, 0)) + cosine = tl.load(Cos + axis * C0 + row * C1 + col * C2, row < ROWS, 0) + sine = tl.load(Sin + axis * S0 + row * S1 + col * S2, row < ROWS, 0) + tl.store(OutCos + index, cosine, row < ROWS) + tl.store(OutSin + index, sine, row < ROWS) + + +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 and (cos.ndim != 3 or cos.shape[0] != 3 or len(sections) != 3): + raise ValueError("Expected three rotary axes and three mRoPE sections") + if cos.is_cuda: + rows, dim = cos.shape[-2:] + out_cos = torch.empty((rows, 2 * dim), dtype=cos.dtype, device=cos.device) + out_sin = torch.empty_like(out_cos) + axes = cos.ndim == 3 + if rows and dim: + _combine_mrope[(triton.cdiv(rows * 2 * dim, 256),)]( + cos, + sin, + out_cos, + out_sin, + rows, + dim, + *(cos.stride() if axes else (0, *cos.stride())), + *(sin.stride() if axes else (0, *sin.stride())), + sections[1] if axes else 0, + sections[2] if axes else 0, + axes, + 256, + ) + return out_cos, out_sin + if cos.ndim == 2: + return cos.repeat(1, 2).contiguous(), sin.repeat(1, 2).contiguous() + 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/qsa_compressor.py b/lightllm/models/qwen4_exp/qsa_compressor.py new file mode 100644 index 0000000000..5a59ae1188 --- /dev/null +++ b/lightllm/models/qwen4_exp/qsa_compressor.py @@ -0,0 +1,185 @@ +import triton +import triton.language as tl + +from lightllm.common.attention_state.bounded_att.window_compressor.window_compressor_state import ( + WindowCompressorState, + WindowCompressorStateConfig, +) + + +class QSACompressorState(WindowCompressorState): + def __init__(self, layers, dim, rotary_dim, ratio, dtype, requests, mtp_step, tp_rank, tp_world_size, device): + width = dim + 2 * rotary_dim + if width % 2: + raise ValueError("QSA compressor fields must have an even total width") + super().__init__( + WindowCompressorStateConfig( + logical_layer_ids=layers, + window_size=ratio, + compress_stride=ratio, + candidate_dim=width // 2, + target="indexer", + dtype=dtype, + tp_rank=tp_rank, + tp_world_size=tp_world_size, + ), + requests, + mtp_step, + device, + ) + self.dim, self.rotary_dim = dim, rotary_dim + + def get_state_type(self): + return "qwen4_qsa_" + super().get_state_type() + + def reset_bounded_att_state(self, req_indices): + for request in req_indices: + self.buffer[:, request].zero_() + + +@triton.jit +def _pool_ring( + Raw, + Cos, + Sin, + Ring, + Indices, + QueryRows, + Visible, + Norm, + Cache, + CAPACITY: tl.constexpr, + RAW_STRIDE: tl.constexpr, + RING_SIZE: tl.constexpr, + WIDTH: tl.constexpr, + D: tl.constexpr, + R: tl.constexpr, + RATIO: tl.constexpr, + EPS: tl.constexpr, + BD: tl.constexpr, + BR: tl.constexpr, +): + token = tl.program_id(0) + request = tl.load(QueryRows + token).to(tl.int64) + visible = tl.load(Visible + token) + if (visible > 0) & (visible % RATIO == 0): + block = visible // RATIO - 1 + position = block * RATIO + tl.arange(0, BR) + col = tl.arange(0, BD) + source = token - RATIO + 1 + tl.arange(0, BR) + source_mask = (source >= 0) & (tl.arange(0, BR) < RATIO) + source_request = tl.load(QueryRows + source, source_mask, other=-1) + source_visible = tl.load(Visible + source, source_mask, other=0) + current = (source >= 0) & (source_request == request) & (source_visible == position + 1) + values = tl.load( + Raw + source[:, None].to(tl.int64) * RAW_STRIDE + col[None, :], + current[:, None] & (tl.arange(0, BR)[:, None] < RATIO) & (col[None, :] < D), + other=0, + ).to(tl.float32) + previous = tl.load( + Ring + (request * RING_SIZE + position[:, None] % RING_SIZE) * WIDTH + col[None, :], + ~current[:, None] & (tl.arange(0, BR)[:, None] < RATIO) & (col[None, :] < D), + other=0, + ).to(tl.float32) + pooled = (tl.sum(values + previous, 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: + first_position = block * RATIO + current_row = (token - RATIO + 1).to(tl.int64) + first_request = tl.load(QueryRows + current_row, current_row >= 0, other=-1) + first_visible = tl.load(Visible + current_row, current_row >= 0, other=0) + first_current = (current_row >= 0) & (first_request == request) & (first_visible == first_position + 1) + ring_row = (request * RING_SIZE + first_position % RING_SIZE) * WIDTH + cos = tl.load(Cos + current_row * R + col, first_current & (col < R), other=0).to(tl.float32) + sin = tl.load(Sin + current_row * R + col, first_current & (col < R), other=0).to(tl.float32) + cos += tl.load(Ring + ring_row + D + col, ~first_current & (col < R), other=0).to(tl.float32) + sin += tl.load(Ring + ring_row + D + R + col, ~first_current & (col < R), other=0).to(tl.float32) + 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) + 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) + destination = tl.load(Indices + request * CAPACITY + block * RATIO + RATIO - 1).to(tl.int64) // RATIO + tl.store(Cache + destination * D + col, normalized.to(Raw.dtype.element_ty), col < D) + + +@triton.jit +def _store_ring( + Raw, + Cos, + Sin, + Ring, + Starts, + Requests, + Lengths, + RAW_STRIDE: tl.constexpr, + RING_SIZE: tl.constexpr, + WIDTH: tl.constexpr, + D: tl.constexpr, + R: tl.constexpr, + BLOCK: tl.constexpr, +): + row = tl.program_id(0) + request = tl.load(Requests + row).to(tl.int64) + start = tl.load(Starts + row) + count = tl.load(Starts + row + 1) - start + length = tl.load(Lengths + row) + offset = tl.program_id(1) + if (request >= 0) & (offset < count): + source_row = start + count - 1 - offset + position = length - 1 - offset + col = tl.arange(0, BLOCK) + raw = tl.load(Raw + source_row.to(tl.int64) * RAW_STRIDE + col, col < D, other=0) + if R > 0: + cos = tl.load(Cos + source_row.to(tl.int64) * R + col - D, (col >= D) & (col < D + R), other=0) + sin = tl.load( + Sin + source_row.to(tl.int64) * R + col - D - R, + (col >= D + R) & (col < D + 2 * R), + other=0, + ) + raw = tl.where(col < D, raw, tl.where(col < D + R, cos, sin)) + tl.store(Ring + (request * RING_SIZE + position % RING_SIZE) * WIDTH + col, raw, col < WIDTH) + + +def pool_with_ring(indexer, raw, ring, batch, cache): + _pool_ring[(raw.shape[0],)]( + raw, + batch.cos, + batch.sin, + ring, + batch.slot_table, + batch.query_rows, + batch.visible_lengths, + indexer.key_norm, + cache, + batch.slot_table.shape[1], + raw.stride(0), + ring.shape[1], + ring.shape[2], + indexer.head_dim, + batch.cos.shape[1], + indexer.compress_ratio, + indexer.eps, + triton.next_power_of_2(indexer.head_dim), + triton.next_power_of_2(indexer.compress_ratio), + enable_fp_fusion=False, + ) + _store_ring[(batch.active_rows.numel(), ring.shape[1])]( + raw, + batch.cos, + batch.sin, + ring, + batch.starts, + batch.active_rows, + batch.key_lengths, + raw.stride(0), + ring.shape[1], + ring.shape[2], + indexer.head_dim, + batch.cos.shape[1], + triton.next_power_of_2(ring.shape[2]), + ) 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..1352f805b6 --- /dev/null +++ b/lightllm/models/qwen4_exp/state.py @@ -0,0 +1,177 @@ +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, all_gather_into_tensor + + +@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 + starts: Optional[torch.Tensor] = None + + @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, + metadata.starts, + ) + + +@dataclass +class QSAExecutionState: + raw_keys: torch.Tensor + + pooled_keys: torch.Tensor + key_cos: torch.Tensor + key_sin: torch.Tensor + compressor: Optional[torch.Tensor] = None + + def select(self, indexer, index_query, batch, decode, max_visible_length, tp_world_size, tp_rank, dist_group): + tokens = index_query.shape[0] + sharded = tp_world_size > 1 and tokens > 0 and tokens % (4 * tp_world_size) == 0 + rows = tokens // tp_world_size if sharded else tokens + first = tp_rank * rows if sharded else 0 + selected = indexer.select( + index_query[first : first + rows], + self.pooled_keys, + batch.query_rows[first : first + rows], + batch.visible_lengths[first : first + rows], + batch.slot_table, + query_chunk_size=max(64, rows) if decode else 1024, + max_visible_length=max_visible_length, + physical_stride=indexer.compress_ratio if self.compressor is not None else 1, + ) + if sharded: + output = torch.empty((tokens, selected.shape[1]), dtype=selected.dtype, device=selected.device) + all_gather_into_tensor(output, selected, group=dist_group) + return output + return selected + + def attend( + self, + indexer, + hidden, + query, + key, + value, + batch: QSABatch, + *, + decode: bool = False, + attention_splits: Optional[int] = None, + max_visible_length: Optional[int] = None, + tp_world_size: int = 1, + tp_rank: int = 0, + dist_group=None, + index_projection=None, + ): + if index_projection is None: + index_query, raw = indexer.project(hidden, batch.cos, batch.sin) + else: + index_query, raw = indexer.project(hidden, batch.cos, batch.sin, projected=index_projection) + if index_query.dtype != self.pooled_keys.dtype: + index_query = index_query.to(self.pooled_keys.dtype) + if self.compressor is not None: + from .qsa_compressor import pool_with_ring + + pool_with_ring(indexer, raw, self.compressor, batch, self.pooled_keys) + else: + 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 = self.select( + indexer, index_query, batch, decode, max_visible_length, tp_world_size, tp_rank, dist_group + ) + + splits = attention_splits + if splits is None and decode: + splits = min(32, triton.cdiv(selected.shape[1], 64)) + wide_block = query.dtype in (torch.float16, torch.bfloat16) + block_size, num_warps = (128 if wide_block else 64), None + if key.dtype == torch.float8_e4m3fn and wide_block: + if not decode and query.shape[0] >= 1024: + block_size, num_warps = 32, 1 + elif decode and query.shape[0] <= 32: + splits = min(16, triton.cdiv(selected.shape[1], block_size)) + return indexed_attention( + query.contiguous(), + key, + value, + selected, + splits=splits, + selection_lengths=batch.visible_lengths, + short_rows_direct=decode or attention_splits is not None, + skip_empty_splits=True, + block_size=block_size, + num_stages=2 if wide_block else None, + num_warps=num_warps, + ) + + +@dataclass +class PLEBatch: + starts: torch.Tensor + ngram_history: torch.Tensor + conv_history: torch.Tensor + mask: Optional[torch.Tensor] = None + conv_inputs: Optional[torch.Tensor] = None + input_tokens: Optional[torch.Tensor] = None + + def apply(self, residual, tokens, layer_weight, tp_world_size: int, dist_group): + hasher = layer_weight.ngram_hasher.component + if self.mask is not None: + tokens = tokens.masked_fill(~self.mask, hasher.eos_token_id) + indices, _ = hasher.hash(tokens, self.starts, self.ngram_history) + embeddings = layer_weight.ngram_embedding(indices.flatten()).view( + tokens.numel(), hasher.heads * layer_weight.ngram_embedding.dim + ) + if tp_world_size > 1: + all_reduce(embeddings, group=dist_group, async_op=False) + update, conv_inputs = layer_weight.ple.component.forward( + residual, embeddings, self.starts, self.conv_history, self.mask + ) + self.input_tokens, self.conv_inputs = tokens, conv_inputs + return update + + def histories_after(self, layer_weight, accepted_lengths=None): + if self.input_tokens is None or self.conv_inputs is None: + raise RuntimeError("PLE candidates must be computed before requesting histories") + ngram = layer_weight.ngram_hasher.component.history_after( + self.input_tokens, self.starts, self.ngram_history, accepted_lengths + ) + conv = layer_weight.ple.component.history_after( + self.conv_inputs, self.starts, self.conv_history, accepted_lengths + ) + return ngram, conv diff --git a/lightllm/models/qwen4_exp/tokenizer.py b/lightllm/models/qwen4_exp/tokenizer.py new file mode 100644 index 0000000000..807cea9ad5 --- /dev/null +++ b/lightllm/models/qwen4_exp/tokenizer.py @@ -0,0 +1,6 @@ +from lightllm.models.qwen2_vl.model import QWen2VLTokenizer + + +class Qwen4ExpTokenizer(QWen2VLTokenizer): + def count_prompt_tokens(self, text_tokens, image_tokens, image_count, audio_tokens, audio_count): + return text_tokens + image_tokens + audio_tokens diff --git a/lightllm/models/qwen4_exp/triton_kernel/__init__.py b/lightllm/models/qwen4_exp/triton_kernel/__init__.py new file mode 100644 index 0000000000..2ae28399f5 --- /dev/null +++ b/lightllm/models/qwen4_exp/triton_kernel/__init__.py @@ -0,0 +1 @@ +pass diff --git a/lightllm/models/qwen4_exp/triton_kernel/cache_pages.py b/lightllm/models/qwen4_exp/triton_kernel/cache_pages.py new file mode 100644 index 0000000000..c541c1b892 --- /dev/null +++ b/lightllm/models/qwen4_exp/triton_kernel/cache_pages.py @@ -0,0 +1,71 @@ +import triton +import triton.language as tl + + +@triton.jit +def copy_slot_pages( + Slots, + Pages, + MemIndices, + PageIndices, + Readies, + SLOT_STRIDE: tl.constexpr, + LAYER_STRIDE: tl.constexpr, + PAGE_STRIDE: tl.constexpr, + TOKENS: tl.constexpr, + LAYERS: tl.constexpr, + DIM: tl.constexpr, + OFFLOAD: tl.constexpr, + BLOCK: tl.constexpr, +): + page = tl.program_id(0) + page_index = tl.load(PageIndices + page).to(tl.int64) + active = page_index >= 0 + if OFFLOAD: + active = active & ~tl.load(Readies + page).to(tl.int1) + if active: + col = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + mask = col < TOKENS * LAYERS * DIM + token = col // (LAYERS * DIM) + layer = ((col // DIM) % LAYERS).to(tl.int64) + slot = tl.load(MemIndices + page * TOKENS + token, mask, other=-1).to(tl.int64) + mask = mask & (slot >= 0) + live = Slots + slot * SLOT_STRIDE + layer * LAYER_STRIDE + col % DIM + cached = Pages + page_index * PAGE_STRIDE + col + if OFFLOAD: + value = tl.load(live, mask, other=0) + tl.store(cached, value, mask) + else: + value = tl.load(cached, mask, other=0) + tl.store(live, value, mask) + + +@triton.jit +def copy_state_pages( + States, + Pages, + StateIndices, + PageIndices, + Readies, + WIDTH: tl.constexpr, + PAGE_STRIDE: tl.constexpr, + OFFLOAD: tl.constexpr, + BLOCK: tl.constexpr, +): + page = tl.program_id(0) + page_index = tl.load(PageIndices + page).to(tl.int64) + active = page_index >= 0 + if OFFLOAD: + active = active & ~tl.load(Readies + page).to(tl.int1) + if active: + state = tl.load(StateIndices + page).to(tl.int64) + col = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + mask = col < WIDTH + live = States + state * WIDTH + col + cached = Pages + page_index * PAGE_STRIDE + col + if OFFLOAD: + value = tl.load(live, mask, other=0) + tl.store(cached, value, mask) + else: + value = tl.load(cached, mask, other=0) + tl.store(live, value, mask) diff --git a/lightllm/models/qwen4_exp/triton_kernel/gated_norm.py b/lightllm/models/qwen4_exp/triton_kernel/gated_norm.py new file mode 100644 index 0000000000..9f50ce3acc --- /dev/null +++ b/lightllm/models/qwen4_exp/triton_kernel/gated_norm.py @@ -0,0 +1,92 @@ +import torch +import triton +import triton.language as tl + + +@triton.jit +def _gated_norm( + X, + Gate, + Weight, + Out, + input_stride_row, + gate_stride_token, + gate_stride_head, + out_stride_row, + GATE_HEADS: tl.constexpr, + D: tl.constexpr, + EPS: tl.constexpr, + SIGMOID: tl.constexpr, + BLOCK: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + col = tl.arange(0, BLOCK) + x = tl.load(X + row * input_stride_row + col, col < D, other=0).to(tl.float32) + inv = tl.rsqrt(tl.sum(x * x, 0) / D + EPS) + + normalized = (x * inv).to(X.dtype.element_ty).to(tl.float32) + weight = tl.load(Weight + col, col < D, other=0).to(tl.float32) + weighted = (normalized * weight).to(X.dtype.element_ty).to(tl.float32) + gate_token = row // GATE_HEADS + gate_head = row % GATE_HEADS + gate = tl.load( + Gate + gate_token * gate_stride_token + gate_head * gate_stride_head + col, + col < D, + other=0, + ).to(tl.float32) + activation = tl.sigmoid(gate) + if not SIGMOID: + activation = gate * activation + tl.store(Out + row * out_stride_row + col, weighted * activation, col < D) + + +@torch.no_grad() +def gated_norm(input: torch.Tensor, gate: torch.Tensor, weight: torch.Tensor, eps: float, activation: str, out=None): + if activation not in ("sigmoid", "silu"): + raise ValueError("Qwen4Exp GDN output gate must be sigmoid or silu") + if ( + input.ndim != 2 + or input.shape[-1] != weight.numel() + or gate.ndim not in (2, 3) + or gate.shape[-1] != input.shape[-1] + or (gate.ndim == 2 and gate.shape != input.shape) + or (gate.ndim == 3 and gate.shape[0] * gate.shape[1] != input.shape[0]) + or not input.is_contiguous() + or gate.stride(-1) != 1 + ): + raise ValueError("GDN norm requires a contiguous input, a packed gate, and a per-head norm weight") + for tensor in (input, gate, weight): + if not tensor.is_cuda or tensor.dtype != input.dtype or tensor.device != input.device: + raise ValueError("GDN norm tensors must be CUDA tensors with matching device and dtype") + if weight.ndim != 1 or not weight.is_contiguous() or input.ndim < 2 or eps <= 0: + raise ValueError("GDN norm requires a weight vector, token/head rows and positive epsilon") + result = torch.empty_like(input) if out is None else out + if ( + result.shape != input.shape + or result.dtype != input.dtype + or result.device != input.device + or not result.is_contiguous() + ): + raise ValueError("GDN norm output must be contiguous and match the input") + rows = input.shape[0] + gate_heads = 1 if gate.ndim == 2 else gate.shape[1] + gate_stride_token = gate.stride(0) + gate_stride_head = 0 if gate.ndim == 2 else gate.stride(1) + if rows: + _gated_norm[(rows,)]( + input, + gate, + weight, + result, + input.stride(0), + gate_stride_token, + gate_stride_head, + result.stride(0), + GATE_HEADS=gate_heads, + D=input.shape[-1], + EPS=eps, + SIGMOID=activation == "sigmoid", + BLOCK=triton.next_power_of_2(input.shape[-1]), + enable_fp_fusion=False, + ) + return result diff --git a/lightllm/models/qwen4_exp/triton_kernel/replicated_moe.py b/lightllm/models/qwen4_exp/triton_kernel/replicated_moe.py new file mode 100644 index 0000000000..4d6ff531b2 --- /dev/null +++ b/lightllm/models/qwen4_exp/triton_kernel/replicated_moe.py @@ -0,0 +1,184 @@ +import triton +import triton.language as tl + + +@triton.jit +def pack_routes( + Ids, + Mask, + Masked, + Routes, + Counts, + TOKENS: tl.constexpr, + TOPK: tl.constexpr, + FIRST: tl.constexpr, + BLOCK: tl.constexpr, +): + expert = tl.program_id(0) + offset = tl.arange(0, BLOCK) + valid = offset < TOKENS * TOPK + live = tl.load(Mask + offset // TOPK, valid, other=0) + selected = tl.load(Ids + offset, valid, other=-1) + selected = tl.where(live, selected, -1) + hit = valid & (selected == FIRST + expert) + position = tl.cumsum(hit.to(tl.int32), 0) - 1 + tl.store(Routes + expert * TOKENS * TOPK + position, offset, hit) + tl.store(Counts + expert, tl.sum(hit.to(tl.int32), 0)) + if expert == 0: + tl.store(Masked + offset, selected, valid) + + +@triton.jit +def expert_up( + Input, + Weight, + Output, + Routes, + Counts, + K: tl.constexpr, + INTER: tl.constexpr, + TOKENS: tl.constexpr, + TOPK: tl.constexpr, + BM: tl.constexpr, + BN: tl.constexpr, + BK: tl.constexpr, +): + expert = tl.program_id(0) + end = tl.load(Counts + expert) + start = 0 + offsets = tl.arange(0, 2 * BN) + col = tl.program_id(1) * BN + offsets % BN + weight_col = col + (offsets // BN) * INTER + kk = tl.arange(0, BK) + for base in range(start, end, BM): + positions = base + tl.arange(0, BM) + rows = tl.load(Routes + expert * TOKENS * TOPK + positions, positions < end, other=0).to(tl.int64) + acc = tl.full((BM, 2 * BN), 0, tl.float32) + for first in range(tl.cdiv(K, BK)): + k = first * BK + kk + x = tl.load( + Input + (rows[:, None] // TOPK) * K + k[None, :], + (positions[:, None] < end) & (k[None, :] < K), + 0, + ) + w = tl.load( + Weight + (expert * 2 * INTER + weight_col[None, :]).to(tl.int64) * K + k[:, None], + (col[None, :] < INTER) & (k[:, None] < K), + 0, + ) + acc = tl.dot(x, w, acc) + gate, value = tl.split(acc.to(Input.dtype.element_ty).reshape((BM, 2, BN)).permute((0, 2, 1))) + gate = gate.to(tl.float32) + gate = (gate / (1.0 + tl.exp(-gate))).to(Input.dtype.element_ty) + mid = value * gate + cols = tl.program_id(1) * BN + tl.arange(0, BN) + tl.store( + Output + rows[:, None].to(tl.int64) * INTER + cols[None, :], + mid, + (positions[:, None] < end) & (cols[None, :] < INTER), + ) + + +@triton.jit +def expert_down( + A, + W, + Y, + Routes, + Counts, + K: tl.constexpr, + N: tl.constexpr, + T: tl.constexpr, + TOPK: tl.constexpr, + INPUT_DIV: tl.constexpr, + BM: tl.constexpr = 16, + BN: tl.constexpr = 64, + BK: tl.constexpr = 128, +): + e = tl.program_id(0) + col = tl.program_id(1) * BN + tl.arange(0, BN) + count = tl.load(Counts + e) + kk = tl.arange(0, BK) + for start in range(0, count, BM): + pos = start + tl.arange(0, BM) + route = tl.load(Routes + e * T * TOPK + pos, pos < count, 0).to(tl.int64) + acc = tl.full((BM, BN), 0, tl.float32) + for first in range(tl.cdiv(K, BK)): + k = first * BK + kk + a = tl.load( + A + (route[:, None] // INPUT_DIV) * K + k[None, :], (pos[:, None] < count) & (k[None, :] < K), 0 + ) + w = tl.load( + W + (e * N + col[None, :]).to(tl.int64) * K + k[:, None], (col[None, :] < N) & (k[:, None] < K), 0 + ) + acc = tl.dot(a, w, acc) + tl.store(Y + route[:, None] * N + col[None, :], acc, (pos[:, None] < count) & (col[None, :] < N)) + + +@triton.jit +def sum_routes( + X, + Weights, + Ids, + Y, + H: tl.constexpr, + TOPK: tl.constexpr, + START: tl.constexpr, + E: tl.constexpr, + BK: tl.constexpr, + BH: tl.constexpr = 256, +): + row = tl.program_id(0) + h = tl.program_id(1) * BH + tl.arange(0, BH) + accumulator = tl.full((BH,), 0, tl.float32) + for route in range(TOPK): + expert = tl.load(Ids + row * TOPK + route) + if expert >= START and expert < START + E: + weight = tl.load(Weights + row * TOPK + route) + value = tl.load(X + (row * TOPK + route).to(tl.int64) * H + h, h < H, 0).to(tl.float32) + accumulator += value * weight + tl.store(Y + row.to(tl.int64) * H + h, accumulator, h < H) + + +@triton.jit +def sum_partials( + Input, + Ids, + Output, + H: tl.constexpr, + TOKENS: tl.constexpr, + TOPK: tl.constexpr, + EXPERTS_PER_RANK: tl.constexpr, + BLOCK: tl.constexpr, + UNROLL: tl.constexpr, + Shared=None, + Gate=None, +): + row = tl.program_id(0) + col = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + accumulator = tl.full((BLOCK,), 0, tl.float32) + seen = tl.full((), 0, tl.int32) + selected = tl.full((), 0, tl.int32) + for reverse_route in tl.static_range(TOPK): + route = TOPK - 1 - reverse_route + expert = tl.load(Ids + row * TOPK + route) + rank = tl.maximum(expert, 0) // EXPERTS_PER_RANK + flag = 1 << rank + last = (expert >= 0) & ((seen & flag) == 0) + selected |= tl.where(last, 1 << route, 0) + seen |= tl.where(last, flag, 0) + for route in tl.range(TOPK, loop_unroll_factor=UNROLL): + expert = tl.load(Ids + row * TOPK + route) + rank = tl.maximum(expert, 0) // EXPERTS_PER_RANK + last = (selected & (1 << route)) != 0 + value = tl.load(Input + (rank * TOKENS + row).to(tl.int64) * H + col, (col < H) & last, 0) + accumulator += value.to(tl.float32) + if Shared is not None: + routed = accumulator.to(Output.dtype.element_ty).to(tl.float32) + shared = tl.load(Shared + row.to(tl.int64) * H + col, col < H, 0).to(tl.float32) + gate = tl.load(Gate + row).to(tl.float32) + gate = tl.sigmoid(gate).to(Gate.dtype.element_ty).to(tl.float32) + shared = (shared * gate).to(Shared.dtype.element_ty).to(tl.float32) + tl.store(Output + row.to(tl.int64) * H + col, routed + shared, col < H) + else: + tl.store(Output + row.to(tl.int64) * H + col, accumulator, col < H) 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 16a8735a7c..749603d582 100644 --- a/lightllm/server/api_cli.py +++ b/lightllm/server/api_cli.py @@ -501,7 +501,8 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: auto: automatically select best backend based on GPU and available packages (priority: fa3 > flashinfer > triton) for hybrid linear-attention models, the second value selects the linear-attention backend - (priority: flashqla > triton); when omitted, it defaults to auto""", + (priority: flashqla > triton); explicit flashinfer requires FlashInfer >= 0.7.2. + When omitted, it defaults to auto""", ) parser.add_argument( "--llm_decode_att_backend", @@ -556,7 +557,8 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: fp8kv_dsa: use dynamic per-token/group FlashMLA FP8 sparse KV cache for DeepSeek-V3.2 or GLM-5.3 Flash. fp8kv_sph and fp8kv_spt requires --kv_quant_calibration_config_path - to load pre-computed FP8 scales. + to load pre-computed FP8 scales, except Qwen4Exp fp8kv_spt, + which uses unit scales and does not accept calibration files. Note: fp8kv_spt requires flashinfer-python>=0.6.5 (default is 0.6.3, may cause runtime issues). Install with: pip install flashinfer-python==0.6.5""", ) @@ -946,6 +948,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/api_start.py b/lightllm/server/api_start.py index e0e6038d60..fd7b64f211 100644 --- a/lightllm/server/api_start.py +++ b/lightllm/server/api_start.py @@ -33,6 +33,7 @@ auto_set_response_parsers, get_running_max_req_size_per_dp, get_default_disable_chunked_mm_input, + get_model_type, ) from lightllm.utils.dist_check_utils import auto_configure_allreduce_flags_from_args @@ -163,7 +164,11 @@ def _launch_subprocesses(args: StartArgs): assert args.disable_dynamic_prompt_cache is True, "need add --disable_dynamic_prompt_cache" assert args.disable_chunked_prefill is True, "need add --disable_chunked_prefill" # FP8 KV cache mode checks - if args.llm_kv_type in ["fp8kv_sph", "fp8kv_spt"]: + unit_scale_qsa = args.llm_kv_type == "fp8kv_spt" and get_model_type(args.model_dir) in ( + "qwen4_exp", + "qwen4_exp_text", + ) + if args.llm_kv_type in ["fp8kv_sph", "fp8kv_spt"] and not unit_scale_qsa: assert ( args.kv_quant_calibration_config_path is not None ), "fp8kv inference mode requires --kv_quant_calibration_config_path. " @@ -185,10 +190,12 @@ def _launch_subprocesses(args: StartArgs): if args.enable_ep_moe: allowed_ep_prefill_att_backends = {"auto", "fa3", "triton", "flashqla"} - for backend in args.llm_prefill_att_backend: + for index, backend in enumerate(args.llm_prefill_att_backend): + if index == 1 and backend == "flashinfer": + continue assert backend in allowed_ep_prefill_att_backends, ( "When --enable_ep_moe is enabled, --llm_prefill_att_backend must be one of " - f"{sorted(allowed_ep_prefill_att_backends)}; flashinfer is not supported." + f"{sorted(allowed_ep_prefill_att_backends)}; flashinfer is only supported for linear prefill." ) allowed_ep_decode_att_backends = {"auto", "fa3", "triton"} for backend in args.llm_decode_att_backend: @@ -782,6 +789,8 @@ def _validate_and_normalize_ssm_config(args: StartArgs): "qwen3_5_moe", "qwen3_5_text", "qwen3_5_moe_text", + "qwen4_exp", + "qwen4_exp_text", "sensenova_loop", ) if is_gdn_or_kda: diff --git a/lightllm/server/core/objs/start_args_type.py b/lightllm/server/core/objs/start_args_type.py index 0609132e62..0f43364d73 100644 --- a/lightllm/server/core/objs/start_args_type.py +++ b/lightllm/server/core/objs/start_args_type.py @@ -246,6 +246,7 @@ class StartArgs: enable_rl: bool = field(default=False) enable_torch_memory_saver: bool = field(default=False) enable_weight_cpu_backup: bool = field(default=False) + offload_ngram_embedding: bool = field(default=False) hardware_platform: str = field(default="cuda", metadata={"choices": ["cuda", "musa"]}) enable_torch_fallback: bool = field(default=False) enable_triton_fallback: bool = field(default=False) diff --git a/lightllm/server/httpserver/manager.py b/lightllm/server/httpserver/manager.py index 95dd057212..bebeedc170 100644 --- a/lightllm/server/httpserver/manager.py +++ b/lightllm/server/httpserver/manager.py @@ -1,6 +1,7 @@ import sys import zmq import zmq.asyncio +from lightllm.server.httpserver.slot_availability import RequestSlotAvailability import asyncio import uvloop import rpyc @@ -27,7 +28,6 @@ from .grammar import create_output_grammar_compiler from lightllm.server.core.objs import Req, FinishStatus, StartArgs from lightllm.server.core.objs import SamplingParams -from lightllm.server.core.objs.out_token_circlequeue import LIGHTLLM_OUT_TOKEN_QUEUE_SIZE from lightllm.server.core.objs.io_objs import GroupReqObjs from lightllm.server.core.objs.shm_req_manager import ShmReqManager from lightllm.server.core.objs.atomic_array_lock import AtomicShmArrayLock, AsyncLock, AtomicLockItem @@ -56,6 +56,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 +345,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 +587,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 +603,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 +767,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 +809,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 +1006,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) # 先保留这个关键得日志，用于方便定位重构中的问题。 @@ -1048,9 +1058,7 @@ async def handle_loop(self): token_list = [] for req in req_status.group_req_objs.shm_req_objs: req_id = req.request_id - read_token_count = 1 - if req.out_tokens_queue.is_full(): - read_token_count = LIGHTLLM_OUT_TOKEN_QUEUE_SIZE + read_token_count = len(req.out_tokens_queue) for _ in range(read_token_count): if not req.out_tokens_queue.is_empty(): diff --git a/lightllm/server/httpserver/slot_availability.py b/lightllm/server/httpserver/slot_availability.py new file mode 100644 index 0000000000..739166fa6d --- /dev/null +++ b/lightllm/server/httpserver/slot_availability.py @@ -0,0 +1,38 @@ +import asyncio + + +class RequestSlotAvailability: + def __init__(self, enabled=True): + self.enabled = enabled + self.generation = 0 + self.event = asyncio.Event() if enabled else None + + @classmethod + def from_args(cls, args): + return cls( + enabled=( + args.run_mode == "normal" + and args.dp == 1 + and args.nnodes == 1 + and not args.enable_rl + and args.httpserver_workers == 1 + ) + ) + + def notify_release(self): + if self.enabled: + self.generation += 1 + self.event.set() + + async def wait(self, observed_generation, timeout): + if not self.enabled: + await asyncio.sleep(timeout) + return False + if observed_generation != self.generation: + return True + self.event.clear() + try: + await asyncio.wait_for(self.event.wait(), timeout=timeout) + return True + except asyncio.TimeoutError: + return False diff --git a/lightllm/server/router/model_infer/mode_backend/base_backend.py b/lightllm/server/router/model_infer/mode_backend/base_backend.py index 208fc46fbc..5348d431bb 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -172,6 +172,9 @@ def init_model(self, kvargs): with self.model.torch_memory_saver.region(tag=MemoryTag.KV_CACHE): manager.finish_state_registration() manager.load_full_att_calibration(args=self.model.args) + if self.args.enable_dp_prompt_cache_fetch: + torch.cuda.synchronize() + dist.barrier() manager.initialize_storage( self.model.args, token_capacity=self.model.max_total_token_num, mem_fraction=self.model.mem_fraction ) diff --git a/lightllm/server/router/model_infer/model_rpc.py b/lightllm/server/router/model_infer/model_rpc.py index c66ea0ec5b..6461e0e93b 100644 --- a/lightllm/server/router/model_infer/model_rpc.py +++ b/lightllm/server/router/model_infer/model_rpc.py @@ -28,7 +28,7 @@ from lightllm.server.core.objs.start_args_type import StartArgs from lightllm.utils.log_utils import init_logger from lightllm.utils.graceful_utils import graceful_registry -from lightllm.utils.process_check import start_parent_check_thread +from lightllm.utils.process_check import install_fatal_thread_excepthook, start_parent_check_thread from lightllm.utils.envs_utils import get_unique_server_name, get_model_infer_recursion_limit from lightllm.utils.torch_memory_saver_utils import MemoryTag from lightllm.server.io_struct import RlOpReq, RlOpRsp @@ -167,6 +167,7 @@ def _init_env( # 注册graceful 退出的处理 graceful_registry(inspect.currentframe().f_code.co_name) + install_fatal_thread_excepthook() setproctitle.setproctitle(f"lightllm::{get_unique_server_name()}::model_infer:RANK{rank}") start_parent_check_thread() 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/model_infer/mtp_speculative/proposers/eagle_with_att.py b/lightllm/server/router/model_infer/mtp_speculative/proposers/eagle_with_att.py index 721af1c27d..967a1972b6 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/proposers/eagle_with_att.py +++ b/lightllm/server/router/model_infer/mtp_speculative/proposers/eagle_with_att.py @@ -73,17 +73,24 @@ def propose_next( verify_draft_input = copy.copy(target_model_input) verify_draft_input.input_ids = target_next_token_ids verify_draft_input.mtp_draft_input_hiddens = target_model_output.mtp_collector.spec_hidden + selected_logits = getattr(draft_model, "supports_selected_logits", False) + if selected_logits: + verify_draft_input.logits_indices = accepted_tail_rows extend_output = draft_model.forward(verify_draft_input) # 只在 req_num 行 logits 上进行 argmax，避免为未接受的 verify 行执行 # vocabulary reduction。第一列 proposal 来自每个请求的 accepted tail。 - accepted_tail_output = ModelOutput( - logits=extend_output.logits.index_select(0, accepted_tail_rows), - logits_token_ids=( - extend_output.logits_token_ids.index_select(0, accepted_tail_rows) - if extend_output.logits_token_ids is not None - else None - ), + accepted_tail_output = ( + extend_output + if selected_logits + else ModelOutput( + logits=extend_output.logits.index_select(0, accepted_tail_rows), + logits_token_ids=( + extend_output.logits_token_ids.index_select(0, accepted_tail_rows) + if extend_output.logits_token_ids is not None + else None + ), + ) ) if self.enable_dynmaic_mtp: draft_token_ids, draft_token_probs = self._gen_argmax_token_ids_and_prob(accepted_tail_output) diff --git a/lightllm/server/router/req_queue/__init__.py b/lightllm/server/router/req_queue/__init__.py index d34a0b2933..94ff216527 100644 --- a/lightllm/server/router/req_queue/__init__.py +++ b/lightllm/server/router/req_queue/__init__.py @@ -13,6 +13,14 @@ def _get_req_queue_class(args, router, dp_size_in_node: int): if args.diverse_mode: return ChunkedBeamContinuesBatchQueue + from .chunked_prefill.cached_prefix_queue import ( + CachedPrefixChunkedPrefillQueue, + supports_cached_prefix_reservations, + ) + from lightllm.utils.config_utils import get_model_type + + if get_model_type(args.model_dir) in {"qwen4_exp", "qwen4_exp_text"} and supports_cached_prefix_reservations(args): + return CachedPrefixChunkedPrefillQueue if args.disable_chunked_prefill: # 虽然也使用chuncked prefill queue 但是由于 args.chunked_prefill_size = args.max_req_total_len # 所以调度的实际行为类似过去的 continues batch 调度，所以将两种调度的实现统一为一种实现，减少代码重复。 diff --git a/lightllm/server/router/req_queue/cached_prefix.py b/lightllm/server/router/req_queue/cached_prefix.py new file mode 100644 index 0000000000..8cd7aa5ca5 --- /dev/null +++ b/lightllm/server/router/req_queue/cached_prefix.py @@ -0,0 +1,71 @@ +import os + +from .peak_reservation import Reservation, estimate_peak + + +class CachedPrefixReservations: + def __init__(self, args, dp_index): + self.page_tokens = int(args.linear_att_hash_page_size) + name = os.environ.get("LIGHTLLM_UNIQUE_SERVICE_NAME_ID") + self.domain = (name + ":" + str(dp_index) + ":" + str(self.page_tokens)).encode() if name else None + self.memo, self.by_tuple = {}, {} + self.shared_prefix_tokens = 0 + + def reset(self, requests): + active = {int(request.request_id) for request in requests} + self.memo = {identity: value for identity, value in self.memo.items() if identity in active} + self.by_tuple.clear() + + def bind(self, reservation, request): + self.by_tuple[id(reservation)] = self.prefix(request) + return reservation + + def prefix(self, request): + try: + mm = getattr(request, "multimodal_params", None) + if ( + self.domain is None + or mm is None + or mm.images + or mm.audios + or request.is_paused + or request.is_aborted + or request.sample_params.disable_prompt_cache + ): + return () + length = min(int(request.prompt_cache_len), int(request.input_len), int(request.shm_cur_kv_len)) + hashes = request.linear_att_token_hash_list.get_all() if length >= self.page_tokens else [] + count = min(length // self.page_tokens, len(hashes), 32) + if count <= 0: + return () + hashes = tuple(int(value) & ((1 << 64) - 1) for value in hashes[:count]) + identity = (int(request.request_id), count, hashes, int(request.input_len)) + saved = self.memo.get(int(request.request_id)) + if saved is not None and saved[0] == identity: + return saved[1] + if not hasattr(request, "shm_prompt_ids"): + request.link_prompt_ids_shm_array() + ids = request.shm_prompt_ids.arr + if ids is None or ids.ndim != 1 or ids.dtype.kind not in "iu" or len(ids) < count * self.page_tokens: + return () + pages = tuple( + hashes[index].to_bytes(8, "little") + + ids[index * self.page_tokens : (index + 1) * self.page_tokens].astype(" 0 + ) + + +class CachedPrefixChunkedPrefillQueue(ChunkedPrefillQueue): + def __init__(self, args, router, dp_index, dp_size_in_node): + super().__init__(args, router, dp_index, dp_size_in_node) + self.prefix_reservations = CachedPrefixReservations(args, dp_index) + + def _init_cache_list(self, current_batch, is_busy): + requests = ( + [] + if current_batch is None + else [ + request for request in current_batch.reqs if request.sample_params.suggested_dp_index == self.dp_index + ] + ) + self.prefix_reservations.reset(requests) + super()._init_cache_list(current_batch, is_busy) + + def _get_req_reservation(self, request, is_busy): + reservation = super()._get_req_reservation(request, is_busy) + return self.prefix_reservations.bind(reservation, request) + + def _estimate_peak_tokens(self): + return self.prefix_reservations.calculate(self.cache_len_list) diff --git a/lightllm/server/router/req_queue/chunked_prefill/impl.py b/lightllm/server/router/req_queue/chunked_prefill/impl.py index f69ab11c5e..ab8bebc95d 100644 --- a/lightllm/server/router/req_queue/chunked_prefill/impl.py +++ b/lightllm/server/router/req_queue/chunked_prefill/impl.py @@ -8,7 +8,7 @@ class ChunkedPrefillQueue(BaseQueue): def _init_cache_list(self, current_batch: Batch, is_busy): if current_batch is not None: self.cache_len_list = [ - req.get_tuple_tokens(is_busy, self.router.router_statics.ema_req_out_len) + self._get_req_reservation(req, is_busy) for req in current_batch.reqs if req.sample_params.suggested_dp_index == self.dp_index ] @@ -18,17 +18,10 @@ def _init_cache_list(self, current_batch: Batch, is_busy): # @calculate_time(show=True, min_cost_ms=0.1) def _can_add_new_req(self, req: Req, is_busy): - self.cache_len_list.append( - req.get_tuple_tokens(is_busy, self.router.router_statics.ema_req_out_len) - ) # hard to analysis + self.cache_len_list.append(self._get_req_reservation(req, is_busy)) # hard to analysis self.cache_len_list.sort(key=lambda x: -x[1]) - left_out_len_array = np.array([e[1] for e in self.cache_len_list]) - has_run_len_array = np.array([e[0] for e in self.cache_len_list]) - cum_run_len_array = np.cumsum(has_run_len_array) - size_array = np.arange(1, len(self.cache_len_list) + 1, 1) - - need_max_token_num = (left_out_len_array * size_array + cum_run_len_array).max() + need_max_token_num = self._estimate_peak_tokens() ok_token_num = need_max_token_num < self.max_total_tokens ok_req_num = len(self.cache_len_list) <= self.running_max_req_size @@ -79,14 +72,18 @@ def generate_new_batch(self, current_batch: Batch): def _calcu_batch_token_load_batch_not_none(self, current_batch: Batch): is_busy = self.is_busy() self._init_cache_list(current_batch, is_busy) - if len(self.cache_len_list) != 0: - self.cache_len_list.sort(key=lambda x: -x[1]) - left_out_len_array = np.array([e[1] for e in self.cache_len_list]) - has_run_len_array = np.array([e[0] for e in self.cache_len_list]) - cum_run_len_array = np.cumsum(has_run_len_array) - size_array = np.arange(1, len(self.cache_len_list) + 1, 1) - need_max_token_num = (left_out_len_array * size_array + cum_run_len_array).max() - else: - need_max_token_num = 0 + need_max_token_num = self._estimate_peak_tokens() return (need_max_token_num, need_max_token_num / self.max_total_tokens) + + def _get_req_reservation(self, req, is_busy): + return req.get_tuple_tokens(is_busy, self.router.router_statics.ema_req_out_len) + + def _estimate_peak_tokens(self): + if not self.cache_len_list: + return 0 + self.cache_len_list.sort(key=lambda entry: -entry[1]) + remaining = np.array([entry[1] for entry in self.cache_len_list]) + resident = np.array([entry[0] for entry in self.cache_len_list]) + count = np.arange(1, len(self.cache_len_list) + 1) + return (remaining * count + np.cumsum(resident)).max() diff --git a/lightllm/server/router/req_queue/peak_reservation.py b/lightllm/server/router/req_queue/peak_reservation.py new file mode 100644 index 0000000000..9df7ab6144 --- /dev/null +++ b/lightllm/server/router/req_queue/peak_reservation.py @@ -0,0 +1,52 @@ +from dataclasses import dataclass +from typing import Optional, Sequence, Tuple + + +@dataclass(frozen=True) +class Reservation: + base_tokens: int + remaining_tokens: int + prefix_pages: Tuple[bytes, ...] = () + cache_domain: Optional[bytes] = None + + +@dataclass(frozen=True) +class Peak: + legacy_tokens: int + shared_tokens: int + max_prefix_credit: int + + +def estimate_peak(entries: Sequence[Reservation], page_tokens: int) -> Peak: + if type(page_tokens) is not int or page_tokens <= 0: + raise ValueError("Page token count must be a positive integer") + for entry in entries: + if ( + type(entry.base_tokens) is not int + or entry.base_tokens < 0 + or type(entry.remaining_tokens) is not int + or entry.remaining_tokens < 0 + ): + raise ValueError("Reservation counts must be nonnegative integers") + if not isinstance(entry.prefix_pages, tuple) or any(type(page) is not bytes for page in entry.prefix_pages): + raise ValueError("Prefix pages must have exact immutable identities") + if entry.cache_domain is not None and type(entry.cache_domain) is not bytes: + raise ValueError("An explicit shared domain must have immutable identity") + if len(entry.prefix_pages) * page_tokens > entry.base_tokens: + raise ValueError("Known cached pages exceed the resident base") + seen = set() + base = credit = legacy = shared = 0 + for count, entry in enumerate(sorted(entries, key=lambda entry: -entry.remaining_tokens), 1): + base += entry.base_tokens + if entry.cache_domain is not None: + prefix = (entry.cache_domain,) + for page in entry.prefix_pages: + prefix += (page,) + if prefix in seen: + credit += page_tokens + else: + seen.add(prefix) + projected = base + count * entry.remaining_tokens + legacy = max(legacy, projected) + shared = max(shared, projected - credit) + return Peak(legacy, shared, credit) diff --git a/lightllm/server/tokenizer.py b/lightllm/server/tokenizer.py index 12ddbb1c68..1fb836998d 100644 --- a/lightllm/server/tokenizer.py +++ b/lightllm/server/tokenizer.py @@ -124,6 +124,15 @@ def get_tokenizer( tokenizer = QWen3VLTokenizer( tokenizer=tokenizer, image_processor=processor.image_processor, model_cfg=model_cfg ) + elif model_type == "qwen4_exp" and "vision_config" in model_cfg: + from ..models.qwen4_exp.image_processing import Qwen4ExpImageProcessor + from ..models.qwen4_exp.tokenizer import Qwen4ExpTokenizer + + tokenizer = Qwen4ExpTokenizer( + tokenizer=tokenizer, + image_processor=Qwen4ExpImageProcessor.from_pretrained(tokenizer_name), + model_cfg=model_cfg, + ) elif model_type in ["qwen3_5", "sensenova_loop", "qwen3_5_moe"] and "vision_config" in model_cfg: from ..models.qwen3_5.model import QWen3_5Tokenizer from transformers import AutoImageProcessor diff --git a/lightllm/server/visualserver/model_infer/model_rpc.py b/lightllm/server/visualserver/model_infer/model_rpc.py index 6451e70008..b8751db58b 100644 --- a/lightllm/server/visualserver/model_infer/model_rpc.py +++ b/lightllm/server/visualserver/model_infer/model_rpc.py @@ -92,6 +92,10 @@ def exposed_init_model(self, kvargs): self.model = ( Qwen2_5_VisionTransformerPretrainedModel(kvargs, **model_cfg["vision_config"]).eval().bfloat16() ) + elif self.model_type == "qwen4_exp": + from lightllm.models.qwen4_exp.vision import Qwen4ExpVisionTransformer + + self.model = Qwen4ExpVisionTransformer(kvargs, **model_cfg["vision_config"]).eval().bfloat16() elif self.model_type in ["qwen3_vl", "qwen3_vl_moe", "qwen3_5", "sensenova_loop", "qwen3_5_moe"]: self.model = ( Qwen3VisionTransformerPretrainedModel(kvargs, **model_cfg["vision_config"]).eval().bfloat16() diff --git a/lightllm/utils/config_utils.py b/lightllm/utils/config_utils.py index 1b0ce53879..187af4706d 100644 --- a/lightllm/utils/config_utils.py +++ b/lightllm/utils/config_utils.py @@ -216,6 +216,8 @@ def auto_set_fused_shared_experts(args) -> None: "qwen3_5_text", "qwen3_5_moe", "qwen3_5_moe_text", + "qwen4_exp", + "qwen4_exp_text", } if model_type not in supported_model_types: logger.info(f"do not enable fused shared experts: unsupported model_type={model_type}") @@ -288,7 +290,15 @@ def get_eos_token_ids(model_path: str) -> Optional[List[int]]: try: config_json = get_config_json(model_path) model_type = config_json.get("model_type") or config_json.get("text_config", {}).get("model_type") - if model_type in {"qwen3_5", "sensenova_loop", "qwen3_5_text", "qwen3_5_moe", "qwen3_5_moe_text"}: + if model_type in { + "qwen3_5", + "sensenova_loop", + "qwen3_5_text", + "qwen3_5_moe", + "qwen3_5_moe_text", + "qwen4_exp", + "qwen4_exp_text", + }: from transformers import AutoTokenizer eos_token_ids = [] @@ -419,6 +429,8 @@ def has_vision_module(model_path: str) -> bool: return True elif model_type == "neo_chat": return True + elif model_type == "qwen4_exp": + return model_cfg.get("vision_config") is not None elif model_type in ["qwen3_5", "sensenova_loop", "qwen3_5_moe"]: return True else: @@ -464,6 +476,8 @@ def is_linear_att_mixed_model(model_path: str) -> bool: "qwen3_5_moe", "qwen3_5_text", "qwen3_5_moe_text", + "qwen4_exp", + "qwen4_exp_text", "glm5_next", "glm5_next_text", ] @@ -559,6 +573,8 @@ def get_reasoning_parser_for_model(model_path: str) -> Optional[str]: "qwen3_5_moe", "qwen3_5_text", "qwen3_5_moe_text", + "qwen4_exp", + "qwen4_exp_text", ]: return "qwen3" diff --git a/lightllm/utils/envs_utils.py b/lightllm/utils/envs_utils.py index 9d8490b69e..a84fe050e1 100644 --- a/lightllm/utils/envs_utils.py +++ b/lightllm/utils/envs_utils.py @@ -2,6 +2,7 @@ import json import torch import uuid +import math from easydict import EasyDict from functools import lru_cache from lightllm.utils.log_utils import init_logger @@ -85,18 +86,22 @@ def is_stop_matching_in_reasoning_enabled() -> bool: @lru_cache(maxsize=None) def get_deepep_num_max_dispatch_tokens_per_rank_prefill(): - # 在单卡最大 prefill batch 之外额外保留 128 个 token，并向上对齐到 8， + # 在 SP 分片后的单卡最大 prefill batch 之外保留 128 个 token，并向上对齐到 8， # 避免 autotune warmup 或调度边界波动使实际输入超过 DeepEP buffer 上限。 # 显式设置 NUM_MAX_DISPATCH_TOKENS_PER_RANK_PREFILL 时仍使用该覆盖值。 configured = os.getenv("NUM_MAX_DISPATCH_TOKENS_PER_RANK_PREFILL", None) if configured is not None: return int(configured) - batch_max_tokens = get_env_start_args().batch_max_tokens or 256 - capacity = ((int(batch_max_tokens) + 128 + 7) // 8) * 8 + args = get_env_start_args() + batch_max_tokens = int(args.batch_max_tokens or 256) + sp_world_size = args.tp // args.dp if getattr(args, "enable_tpsp_mix_mode", False) else 1 + rank_tokens = (batch_max_tokens + sp_world_size - 1) // sp_world_size + capacity = ((rank_tokens + 128 + 7) // 8) * 8 logger.info( - "DeepEP prefill buffer capacity: batch_max_tokens=%s, safety_margin=128, capacity=%s", + "DeepEP prefill buffer capacity: batch_max_tokens=%s, sp_world_size=%s, safety_margin=128, capacity=%s", batch_max_tokens, + sp_world_size, capacity, ) return capacity @@ -113,13 +118,18 @@ def get_deepep_num_max_dispatch_tokens_per_rank_decode(): args = get_env_start_args() required_tokens = int(args.running_max_req_size) * (int(args.mtp_step) + 1) - capacity = ((required_tokens + 12 + 7) // 8) * 8 + sp_world_size = args.tp // args.dp if getattr(args, "enable_tpsp_mix_mode", False) else 1 + alignment = math.lcm(sp_world_size, int(args.mtp_step) + 1) + padded_tokens = (required_tokens + alignment - 1) // alignment * alignment + rank_tokens = padded_tokens // sp_world_size + capacity = ((rank_tokens + 12 + 7) // 8) * 8 logger.info( "DeepEP decode buffer capacity: running_max_req_size=%s, mtp_step=%s, " - "required_tokens=%s, safety_margin=12, capacity=%s", + "required_tokens=%s, sp_world_size=%s, safety_margin=12, capacity=%s", args.running_max_req_size, args.mtp_step, required_tokens, + sp_world_size, capacity, ) return capacity diff --git a/lightllm/utils/net_utils.py b/lightllm/utils/net_utils.py index eab962713d..de972f5e81 100644 --- a/lightllm/utils/net_utils.py +++ b/lightllm/utils/net_utils.py @@ -55,6 +55,7 @@ def validate_ports(ports: list): for port in ports: with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) try: sock.bind(("", port)) except OSError as e: diff --git a/test/__init__.py b/test/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/test/benchmark/qwen38/benchmark_gated_residual.py b/test/benchmark/qwen38/benchmark_gated_residual.py new file mode 100644 index 0000000000..0e3c4fa039 --- /dev/null +++ b/test/benchmark/qwen38/benchmark_gated_residual.py @@ -0,0 +1,93 @@ +import argparse +import json +import os + +import torch +import torch.nn.functional as F +import triton +from triton.testing import do_bench_cudagraph + +from lightllm.common.gated_residual import GatedResidual + + +def eager(component, residual, update, precise_projections=False): + def project(inputs, weight): + if precise_projections: + return F.linear(inputs.double(), weight.double()).to(inputs.dtype) + return F.linear(inputs, weight) + + grouped = residual.float().unflatten(-1, (component.branches, -1)) + normalized = grouped * torch.rsqrt(grouped.square().mean(-1, keepdim=True) + component.eps) + normalized = (normalized.flatten(-2) * (1 + component.norm_weight.float())).to(residual.dtype) + projected = F.silu(project(normalized, component.down_weight) / component.branches) + mix = torch.sigmoid(project(projected, component.up_weight)) + mixed = (normalized * mix).unflatten(-1, (component.branches, -1)).mean(-2) + gates = 2 * torch.sigmoid(project(normalized, component.inject_weight) / component.branches) + return mixed, residual + (update.unsqueeze(-2) * gates.unsqueeze(-1)).flatten(-2) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--tokens", type=int, nargs="+", default=[1, 16, 128, 1024]) + parser.add_argument("--repetitions", type=int, default=3) + args = parser.parse_args() + if not os.environ.get("EXP_ACTIVE"): + parser.error('Wrap this benchmark with exp -m "purpose"') + torch.manual_seed(37) + dim, branches, rank = 2560, 4, 320 + width = dim * branches + + def rand(*shape): + return torch.randn(*shape, device="cuda", dtype=torch.bfloat16) + + component = GatedResidual( + rand(width) * 0.1, + rand(rank, width) / width ** 0.5, + rand(width, rank) / rank ** 0.5, + rand(branches, width) / width ** 0.5, + branches, + ) + print( + json.dumps( + { + "torch": torch.__version__, + "triton": triton.__version__, + "cuda": torch.version.cuda, + "gpu": torch.cuda.get_device_name(), + "dtype": "bfloat16", + "hidden_size": dim, + "branches": branches, + "rank": rank, + "timing": "CUDA graph, median of repetitions", + "validation": "eager equations with FP64 projections quantized at original dtype boundaries", + "eager_timing": "original input-dtype projection; batch-dependent rounding recorded separately", + "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), + } + ) + ) + for tokens in args.tokens: + residual, update = rand(tokens, width), rand(tokens, dim) + + def fused(): + mixed, gates = component.read(residual) + return mixed, component.write(residual, update, gates) + + def reference(): + return eager(component, residual, update) + + actual = fused() + expected = eager(component, residual, update, precise_projections=True) + for value, oracle in zip(actual, expected): + torch.testing.assert_close(value, oracle, atol=2e-3, rtol=torch.finfo(residual.dtype).eps) + results = { + "eager_max_abs": [float((value - old).float().abs().max()) for value, old in zip(actual, reference())] + } + for name, func in [("eager", reference), ("fused", fused)]: + samples = [do_bench_cudagraph(func, rep=100) * 1000 for _ in range(args.repetitions)] + results[name + "_us"] = sorted(samples)[len(samples) // 2] + results[name + "_samples_us"] = samples + print(json.dumps({"tokens": tokens, **results, "speedup": results["eager_us"] / results["fused_us"]})) + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/qwen38/benchmark_lm_head.py b/test/benchmark/qwen38/benchmark_lm_head.py new file mode 100644 index 0000000000..422c9018c2 --- /dev/null +++ b/test/benchmark/qwen38/benchmark_lm_head.py @@ -0,0 +1,111 @@ +import argparse +import importlib.metadata +import json +import os +from pathlib import Path +import subprocess + +import torch +import triton +from safetensors import safe_open + +from check_generation import REVISION, verified_checkpoint + + +@torch.inference_mode() +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model-directory", type=Path, required=True) + parser.add_argument( + "--tp", type=int, default=4, help="Vocabulary partition count; this process runs rank zero only" + ) + parser.add_argument("--rows", type=int, nargs="+", default=[1, 3, 5, 7, 8, 9, 24, 64, 128]) + parser.add_argument("--dtype", choices=["bf16", "fp16"], default="bf16") + parser.add_argument("--samples", type=int, default=3) + parser.add_argument("--rep-ms", type=int, default=40) + args = parser.parse_args() + if not os.environ.get("EXP_ACTIVE"): + raise RuntimeError("Run through exp") + if min(args.rows + [args.tp, args.samples, args.rep_ms]) < 1: + parser.error("Partition count, rows, samples and repeat duration must be positive") + verified = verified_checkpoint(args.model_directory) + checkpoint = json.loads((args.model_directory / "config.json").read_text()) + config = checkpoint.get("text_config", checkpoint) + hidden, vocab = config["hidden_size"], config["vocab_size"] + if args.tp > vocab: + parser.error("Partition count cannot exceed vocabulary size") + + from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.utils.envs_utils import get_env_start_args + + start = StartArgs() + start.hardware_platform, start.data_type, start.tp, start.dp = "cuda", args.dtype, args.tp, 1 + os.environ["LIGHTLLM_START_ARGS"] = json.dumps(vars(start)) + os.environ.update( + LIGHTLLM_DP_WORLD_SIZE=str(args.tp), LIGHTLLM_CURRENT_RANK_IN_DP="0", LIGHTLLM_CURRENT_DEVICE_ID="0" + ) + get_env_start_args.cache_clear() + from lightllm.common.basemodel.layer_weights.meta_weights import LMHeadWeight + from lightllm.models.qwen4_exp.layer_weights.lm_head_weight import Qwen4ExpLMHeadWeight + + dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float16 + baseline = LMHeadWeight(hidden, vocab, "lm_head.weight", dtype) + index = json.loads((args.model_directory / "model.safetensors.index.json").read_text())["weight_map"] + filename = index["lm_head.weight"] + with safe_open(args.model_directory / filename, framework="pt", device="cpu") as source: + baseline.load_hf_weights({"lm_head.weight": source.get_tensor("lm_head.weight")}) + assert baseline.verify_load() + optimized = Qwen4ExpLMHeadWeight(hidden, vocab, "lm_head.weight", dtype, embedding_weight=baseline) + assert optimized.verify_load() and optimized.weight is baseline.weight + torch.manual_seed(3107) + inputs = torch.randn(max(args.rows), hidden, device="cuda", dtype=dtype) + reference = torch.cat([baseline(row[:, None]) for row in inputs], dim=1) + run = Path(os.environ["EXP_RUN_DIR"]) + report = dict( + status="running", + scope="trained rank-zero head component; no distributed serving or vLLM/SGLang performance claim", + revision=REVISION, + configuration=vars(args) | {"model_directory": str(args.model_directory)}, + local_weight_shape=list(baseline.weight.shape), + hardware=subprocess.check_output( + ["nvidia-smi", "--query-gpu=index,name,uuid,memory.used,memory.free", "--format=csv,noheader"], text=True + ), + software={name: importlib.metadata.version(name) for name in ("torch", "triton")}, + cuda=torch.version.cuda, + bf16_reduced_precision=torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction, + source_file=verified["files"][filename], + results=[], + ) + path = run / "results.json" + try: + for rows in args.rows: + tokens = inputs[:rows].T + old, new = baseline(tokens), optimized(tokens) + torch.testing.assert_close(new, reference[:, :rows], atol=0, rtol=0) + timings = {} + for name, weight in (("baseline", baseline), ("optimized", optimized)): + timings[name] = [ + triton.testing.do_bench_cudagraph(lambda: weight(tokens), rep=args.rep_ms) * 1000 + for _ in range(args.samples) + ] + item = dict( + rows=rows, + scalar_matches=True, + baseline_max_abs=float((old - reference[:, :rows]).float().abs().max()), + microseconds=timings, + speedup=sorted(timings["baseline"])[args.samples // 2] + / sorted(timings["optimized"])[args.samples // 2], + ) + report["results"].append(item) + path.write_text(json.dumps(report, indent=2)) + print(json.dumps(item), flush=True) + report["status"] = "passed" + except Exception as error: + report.update(status="failed", error_type=type(error).__name__, error=str(error)) + raise + finally: + path.write_text(json.dumps(report, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/qwen38/benchmark_ngram_hash.py b/test/benchmark/qwen38/benchmark_ngram_hash.py new file mode 100644 index 0000000000..35042cadcb --- /dev/null +++ b/test/benchmark/qwen38/benchmark_ngram_hash.py @@ -0,0 +1,129 @@ +import argparse +import json +import os + +import torch +import triton +from triton.testing import do_bench_cudagraph + +from lightllm.common.ngram_embedding import NGramHasher + + +def eager(hasher, tokens, history): + batch, length = tokens.shape + full = torch.cat((history, tokens.long()), dim=-1) + positions = torch.arange(full.shape[1], device=full.device, dtype=torch.int64) + eos_positions = torch.where(full == hasher.eos_token_id, positions, -1) + inclusive = torch.cummax(eos_positions, dim=1).values + previous_eos = torch.cat((eos_positions.new_full((batch, 1), -1), inclusive[:, :-1]), dim=1) + relative = positions.unsqueeze(0) - previous_eos - 1 + shifted = [full] + for shift in range(1, hasher.context_len + 1): + source = positions - shift + gathered = full.gather(1, source.clamp_min(0).unsqueeze(0).expand(batch, -1)) + valid = (relative >= shift) & (source.unsqueeze(0) >= 0) + shifted.append(torch.where(valid, gathered, hasher.eos_token_id)) + ids = [] + per_order = hasher.heads // hasher.context_len + for order in range(2, hasher.context_len + 2): + mixed = shifted[0] * hasher.multipliers[0] + for index in range(1, order): + mixed = torch.bitwise_xor(mixed, shifted[index] * hasher.multipliers[index]) + begin = (order - 2) * per_order + sizes = hasher.head_vocab_sizes[begin : begin + per_order] + offsets = hasher.head_offsets[begin : begin + per_order] + ids.append(torch.remainder(mixed.unsqueeze(-1), sizes.view(1, 1, -1)) + offsets.view(1, 1, -1)) + return torch.cat(ids, dim=-1)[:, -length:].reshape(-1, hasher.heads), full[:, -hasher.context_len :].contiguous() + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--shapes", nargs="+", default=["1x1", "16x1", "128x1", "1x4096", "16x256"]) + parser.add_argument("--repetitions", type=int, default=3) + args = parser.parse_args() + if not os.environ.get("EXP_ACTIVE"): + parser.error('Wrap this benchmark with exp -m "purpose"') + torch.manual_seed(37) + + sizes = torch.tensor( + [ + 20000003, + 20000023, + 20000033, + 20000047, + 20000059, + 20000063, + 20000069, + 20000077, + 20000081, + 20000093, + 20000107, + 20000147, + 20000153, + 20000159, + 20000161, + 20000171, + ], + device="cuda", + dtype=torch.int64, + ) + offsets = torch.cat([sizes.new_zeros(1), sizes.cumsum(0)[:-1]]) + multipliers = torch.tensor([23703573157769, 20109073645365, 8052911324071], device="cuda", dtype=torch.int64) + hasher = NGramHasher(multipliers, sizes, offsets, 248044) + print( + json.dumps( + { + "torch": torch.__version__, + "triton": triton.__version__, + "cuda": torch.version.cuda, + "gpu": torch.cuda.get_device_name(), + "ngram_size": 3, + "heads_per_ngram": 8, + "multipliers": multipliers.tolist(), + "head_vocab_sizes": sizes.tolist(), + "eos_token_id": hasher.eos_token_id, + "timing": "CUDA graph, median of interleaved repetitions", + "includes_embedding_lookup": False, + "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), + } + ), + flush=True, + ) + for shape in args.shapes: + batch, length = map(int, shape.split("x")) + tokens = torch.randint(0, 248320, (batch, length), dtype=torch.int64, device="cuda") + tokens[:, ::31] = hasher.eos_token_id + starts = torch.arange(batch + 1, device="cuda", dtype=torch.int32) * length + history = hasher.empty_history(batch) + history[:, -1] = 77 + + def fused(): + return hasher.hash(tokens.view(-1), starts, history) + + def reference(): + return eager(hasher, tokens, history) + + for actual, expected in zip(fused(), reference()): + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + samples = {"eager": [], "fused": []} + for repeat in range(args.repetitions): + functions = [("eager", reference), ("fused", fused)] + for name, func in functions[:: 1 if repeat % 2 == 0 else -1]: + samples[name].append(do_bench_cudagraph(func, rep=100) * 1000) + medians = {name + "_us": sorted(values)[len(values) // 2] for name, values in samples.items()} + print( + json.dumps( + { + "batch": batch, + "tokens_per_request": length, + **medians, + "speedup": medians["eager_us"] / medians["fused_us"], + "samples_us": samples, + } + ), + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/qwen38/benchmark_ple.py b/test/benchmark/qwen38/benchmark_ple.py new file mode 100644 index 0000000000..1c98a8481f --- /dev/null +++ b/test/benchmark/qwen38/benchmark_ple.py @@ -0,0 +1,120 @@ +import argparse +import json +import math +import os + +import torch +import torch.nn.functional as F +import triton +from triton.testing import do_bench_cudagraph + +from lightllm.common.ngram_embedding.ple import PLELayer + + +def eager(layer, residual, embeddings, history): + batch, length, width = residual.shape + + def norm(x, weight): + grouped = x.float().unflatten(-1, (layer.branches, -1)) + normalized = grouped * torch.rsqrt(grouped.square().mean(-1, keepdim=True) + layer.eps) + return (normalized.flatten(-2) * (1 + weight.float())).to(x.dtype) + + key = norm(F.linear(embeddings, layer.key_weight), layer.norm_key).unflatten(-1, (layer.branches, -1)) + query = norm(residual, layer.norm_query).unflatten(-1, (layer.branches, -1)) + value = F.linear(embeddings, layer.value_weight) + gate = (key * query).sum(-1, keepdim=True) / math.sqrt(layer.hidden_size) + gate = gate.abs().clamp_min(1e-6).sqrt() * gate.sign() + gated = (torch.sigmoid(gate) * value.unsqueeze(-2)).flatten(-2) + conv_inputs = norm(gated, layer.norm_conv) + full = torch.cat((history, conv_inputs), dim=1) + conv = F.conv1d(full.transpose(1, 2), layer.conv_weight.unsqueeze(1), dilation=layer.dilation, groups=width) + output = gated + F.silu(conv).transpose(1, 2) + return output, full[:, -layer.context_len :].contiguous() + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--shapes", nargs="+", default=["1x1", "16x1", "1x128", "4x128"]) + parser.add_argument("--repetitions", type=int, default=3) + args = parser.parse_args() + if not os.environ.get("EXP_ACTIVE"): + parser.error('Wrap this benchmark with exp -m "purpose"') + torch.manual_seed(137) + dim, branches, embed_dim = 2560, 4, 2560 + width = dim * branches + + def rand(*shape): + return torch.randn(*shape, device="cuda", dtype=torch.bfloat16) + + layer = PLELayer( + rand(width, embed_dim) / embed_dim ** 0.5, + rand(dim, embed_dim) / embed_dim ** 0.5, + rand(width) * 0.1, + rand(width) * 0.1, + rand(width) * 0.1, + rand(width, 4) * 0.1, + branches, + 3, + ) + print( + json.dumps( + { + "torch": torch.__version__, + "triton": triton.__version__, + "cuda": torch.version.cuda, + "gpu": torch.cuda.get_device_name(), + "dtype": "bfloat16", + "hidden_size": dim, + "branches": branches, + "embedding_width": embed_dim, + "kernel_size": 4, + "dilation": 3, + "timing": "CUDA graph, median of interleaved repetitions", + "includes_table_lookup": False, + "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), + } + ), + flush=True, + ) + for shape in args.shapes: + batch, length = map(int, shape.split("x")) + residual = rand(batch, length, width) + embeddings = rand(batch, length, embed_dim) + history = layer.empty_history(batch).normal_() + starts = torch.arange(batch + 1, device="cuda", dtype=torch.int32) * length + + def fused(): + output, conv_inputs = layer.forward( + residual.view(-1, width), embeddings.view(-1, embed_dim), starts, history + ) + return output.view(batch, length, width), layer.history_after(conv_inputs, starts, history) + + def reference(): + return eager(layer, residual, embeddings, history) + + for actual, expected in zip(fused(), reference()): + torch.testing.assert_close( + actual, expected, atol=2 * torch.finfo(residual.dtype).eps, rtol=2 * torch.finfo(residual.dtype).eps + ) + samples = {"eager": [], "fused": []} + for repeat in range(args.repetitions): + functions = [("eager", reference), ("fused", fused)] + for name, func in functions[:: 1 if repeat % 2 == 0 else -1]: + samples[name].append(do_bench_cudagraph(func, rep=100) * 1000) + medians = {name + "_us": sorted(values)[len(values) // 2] for name, values in samples.items()} + print( + json.dumps( + { + "batch": batch, + "tokens_per_request": length, + **medians, + "speedup": medians["eager_us"] / medians["fused_us"], + "samples_us": samples, + } + ), + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/test/benchmark/qwen38/benchmark_qsa.py b/test/benchmark/qwen38/benchmark_qsa.py new file mode 100644 index 0000000000..53eab467dd --- /dev/null +++ b/test/benchmark/qwen38/benchmark_qsa.py @@ -0,0 +1,165 @@ +import argparse +import json +import os + +import torch +import torch.nn.functional as F +import triton +from triton.testing import do_bench_cudagraph + +from lightllm.common.qsa.indexer import QSAIndexer +from lightllm.common.qsa.attention import indexed_attention + + +def eager_select(indexer, query, cache, requests, lengths, table): + blocks = cache.shape[1] + count = lengths // indexer.compress_ratio + + scores = torch.matmul(query.float(), cache[0].float().transpose(-1, -2)) + scores = scores.relu().sum(1) / indexer.head_dim ** 0.5 + scores.masked_fill_(torch.arange(blocks, device=query.device)[None, :] >= count[:, None], -float("inf")) + topk = min(blocks, indexer.token_budget // indexer.compress_ratio) + selected = scores.topk(topk, sorted=False).indices + width = indexer.token_budget + indexer.compress_ratio - 1 + col = torch.arange(width, device=query.device)[None, :].expand(query.shape[0], -1) + chosen = count.clamp_max(indexer.token_budget // indexer.compress_ratio) + full = col < chosen[:, None] * indexer.compress_ratio + rank = (col // indexer.compress_ratio).clamp_max(topk - 1) + position = selected.gather(1, rank) * indexer.compress_ratio + col % indexer.compress_ratio + tail_col = col - chosen[:, None] * indexer.compress_ratio + tail = (tail_col >= 0) & (tail_col < (lengths % indexer.compress_ratio)[:, None]) + position = torch.where(full, position, count[:, None] * indexer.compress_ratio + tail_col) + slots = table[requests.long()[:, None], position.clamp(0, table.shape[1] - 1)] + return slots.masked_fill(~(full | tail), -1) + + +def eager_attention(q, k, v, selected, valid_tokens=2048): + from torch.nn.attention import SDPBackend, sdpa_kernel + + valid_slots = selected[:, :valid_tokens] + gathered_k = k[valid_slots].transpose(1, 2) + gathered_v = v[valid_slots].transpose(1, 2) + with sdpa_kernel(SDPBackend.FLASH_ATTENTION): + return F.scaled_dot_product_attention(q.unsqueeze(2), gathered_k, gathered_v, enable_gqa=True).squeeze(2) + + +def measure(reference, fused, repetitions): + samples = {"torch": [], "kernel": []} + for repeat in range(repetitions): + functions = [("torch", reference), ("kernel", fused)] + for name, function in functions[:: 1 if repeat % 2 == 0 else -1]: + samples[name].append(do_bench_cudagraph(function, rep=100) * 1000) + result = {name + "_us": sorted(values)[len(values) // 2] for name, values in samples.items()} + return {**result, "speedup": result["torch_us"] / result["kernel_us"], "samples_us": samples} + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--contexts", nargs="+", type=int, default=[4099, 65539]) + parser.add_argument("--index-rows", nargs="+", type=int, default=[1, 16]) + parser.add_argument("--attention-rows", nargs="+", type=int, default=[1, 16, 128]) + parser.add_argument("--valid-selections", nargs="+", type=int, default=[2048]) + parser.add_argument("--selection-bounds", action="store_true", help="Use native request visibility bounds") + parser.add_argument("--repetitions", type=int, default=3) + args = parser.parse_args() + if not os.environ.get("EXP_ACTIVE"): + parser.error('Wrap this benchmark with exp -m "purpose"') + if any(count < 1 or count > 2048 for count in args.valid_selections): + parser.error("Valid selection counts must lie in [1, 2048]") + torch.manual_seed(537) + + def rand(*shape): + return torch.randn(*shape, device="cuda", dtype=torch.bfloat16) + + indexer = QSAIndexer(rand(640, 2560), rand(128), rand(128)) + print( + json.dumps( + { + "torch": torch.__version__, + "triton": triton.__version__, + "cuda": torch.version.cuda, + "gpu": torch.cuda.get_device_name(), + "dtype": "bfloat16", + "index_heads": 4, + "index_head_dim": 128, + "token_budget": 2048, + "compress_ratio": 4, + "query_heads": 24, + "kv_heads": 2, + "attention_head_dim": 256, + "timing": "CUDA graph, median of interleaved repetitions", + "attention_baseline": "KV gather plus forced Torch FlashAttention, no padding mask", + "selection_baseline": ( + "Torch matmul with shared cached keys, ReLU head reduction, top-k and slot gather" + ), + "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), + "selection_bounds": args.selection_bounds, + } + ), + flush=True, + ) + for context in args.contexts: + cache = rand(1, context // 4, 128) + table = torch.arange(context, device="cuda", dtype=torch.int32).view(1, -1) + for rows in args.index_rows: + query = rand(rows, 4, 128) + requests = torch.zeros(rows, device="cuda", dtype=torch.int32) + lengths = ( + torch.full((rows,), context, device="cuda", dtype=torch.int32) + - torch.arange(rows, device="cuda", dtype=torch.int32) % 7 + ) + + def reference(): + return eager_select(indexer, query, cache, requests, lengths, table) + + def fused(): + return indexer.select(query, cache, requests, lengths, table) + + torch.testing.assert_close(fused().sort(-1).values, reference().sort(-1).values, atol=0, rtol=0) + print( + json.dumps( + { + "component": "cached_block_selection", + "queries": rows, + "context": context, + **measure(reference, fused, args.repetitions), + } + ), + flush=True, + ) + slots = 65539 + key, value = rand(slots, 2, 256), rand(slots, 2, 256) + for rows in args.attention_rows: + query = rand(rows, 24, 256) + choices = torch.stack([torch.randperm(slots, device="cuda")[:2051] for _ in range(rows)]).int() + for count in args.valid_selections: + selected = choices.clone() + selected[:, count:] = -1 + + bound = slots if count == 2048 else count + lengths = torch.full((rows,), bound, device="cuda", dtype=torch.int32) if args.selection_bounds else None + + def reference(): + return eager_attention(query, key, value, selected, count) + + def fused(): + return indexed_attention(query, key, value, selected, selection_lengths=lengths) + + torch.testing.assert_close(fused(), reference(), atol=0.016, rtol=0.016) + print( + json.dumps( + { + "component": "indexed_gqa", + "queries": rows, + "physical_slots": slots, + "selected": count, + "selection_capacity": selected.shape[1], + **measure(reference, fused, args.repetitions), + } + ), + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/unit_tests/common/att_resources/test_page_pool_storage.py b/unit_tests/common/att_resources/test_page_pool_storage.py index 032f7674eb..5eb5ba52cc 100644 --- a/unit_tests/common/att_resources/test_page_pool_storage.py +++ b/unit_tests/common/att_resources/test_page_pool_storage.py @@ -58,9 +58,9 @@ def test_shared_registration_lives_until_last_view( unregistered = [] unregister = utils._unregister_page_pool_storage - def record_unregister(ptr: int, device_id: int) -> None: - unregistered.append(ptr) - unregister(ptr=ptr, device_id=device_id) + def record_unregister(addresses: tuple[int, ...], device_id: int) -> None: + unregistered.extend(addresses) + unregister(addresses=addresses, device_id=device_id) monkeypatch.setattr(utils, "_unregister_page_pool_storage", record_unregister) tensor = utils.create_page_pool_tensor(size=(4096,), dtype=torch.float32) @@ -75,6 +75,92 @@ def record_unregister(ptr: int, device_id: int) -> None: assert unregistered == [ptr] +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +def test_chunked_shared_pages_keep_dma_boundaries_and_lifetime(cpu_args, monkeypatch, dtype) -> None: + cpu_args.enable_dp_prompt_cache_fetch = True + monkeypatch.setattr(utils, "_HOST_REGISTRATION_BYTES", 16384) + runtime = torch.cuda.cudart() + registered, released = [], [] + + class RecordingRuntime: + def cudaHostRegister(self, address, size, flags): + status = runtime.cudaHostRegister(address, size, flags) + if int(status) == 0: + registered.append((address, size)) + return status + + def cudaHostUnregister(self, address): + released.append(address) + return runtime.cudaHostUnregister(address) + + monkeypatch.setattr(torch.cuda, "cudart", lambda: RecordingRuntime()) + tensor = utils.create_page_pool_tensor(size=(7, 3, 512), dtype=dtype) + cell_bytes = tensor[0].numel() * tensor.element_size() + assert len(registered) > 1 + assert all((address - tensor.data_ptr()) % cell_bytes == 0 for address, _ in registered) + assert all(size <= 16384 for _, size in registered) + destination = torch.empty((3, 5, 512), dtype=dtype, device="cuda") + for index in range(tensor.shape[0]): + tensor[index].fill_(index + 1) + destination[:, 1].copy_(tensor[index], non_blocking=True) + torch.cuda.synchronize() + assert torch.equal(destination[:, 1].cpu(), tensor[index]) + + view = tensor[-1] + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + destination[:, 1].copy_(view, non_blocking=True) + view.fill_(17) + del tensor + gc.collect() + assert not released + assert view.is_shared() and view.is_pinned() + graph.replay() + torch.cuda.synchronize() + assert torch.equal(destination[:, 1].cpu(), view) + graph.reset() + del view + gc.collect() + assert released == [address for address, _ in reversed(registered)] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_private_storage_keeps_whole_tensor_dma(monkeypatch) -> None: + monkeypatch.setattr(utils, "_HOST_REGISTRATION_BYTES", 16384) + tensor = torch.full((7, 3, 512), 7, dtype=torch.float32) + utils.register_page_pool_tensor(tensor) + destination = tensor.to("cuda", non_blocking=True) + torch.cuda.synchronize() + assert torch.equal(destination.cpu(), tensor) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_shared_pool_partial_registration_failure_rolls_back(cpu_args, monkeypatch) -> None: + cpu_args.enable_dp_prompt_cache_fetch = True + monkeypatch.setattr(utils, "_HOST_REGISTRATION_BYTES", 16384) + runtime = torch.cuda.cudart() + registered, released = [], [] + + class FailingRuntime: + def cudaHostRegister(self, address, size, flags): + if registered: + return 1 + status = runtime.cudaHostRegister(address, size, flags) + if int(status) == 0: + registered.append(address) + return status + + def cudaHostUnregister(self, address): + released.append(address) + return runtime.cudaHostUnregister(address) + + monkeypatch.setattr(torch.cuda, "cudart", lambda: FailingRuntime()) + with pytest.raises(torch.cuda.CudaError): + utils.create_page_pool_tensor(size=(7, 3, 512), dtype=torch.float32) + assert registered and released == registered + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") def test_gpu_page_pool_keeps_device_allocation(cpu_args: EasyDict) -> None: cpu_args.enable_dp_prompt_cache_fetch = True diff --git a/unit_tests/common/basemodel/attention/linear/test_flashinfer_gdn.py b/unit_tests/common/basemodel/attention/linear/test_flashinfer_gdn.py new file mode 100644 index 0000000000..5d4fbd0f20 --- /dev/null +++ b/unit_tests/common/basemodel/attention/linear/test_flashinfer_gdn.py @@ -0,0 +1,105 @@ +import inspect + +import pytest +import torch + +from lightllm.common.basemodel.attention.linear.flashinfer import chunk_gated_delta_rule +from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops import chunk_gated_delta_rule as native + + +@pytest.fixture(autouse=True) +def require_flashinfer_gdn(): + if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 9: + pytest.skip("FlashInfer GDN regression requires Hopper") + module = pytest.importorskip("flashinfer.gdn_prefill") + if "use_cp" not in inspect.signature(module.chunk_gated_delta_rule).parameters: + pytest.skip("FlashInfer >= 0.7.2 required") + + +@pytest.mark.parametrize("lengths", [[17, 31], [512], [2048], [8192], [0, 19, 0, 47], [1024] * 8]) +def test_prefill_nonzero_state_and_graph(lengths): + torch.manual_seed(4317) + tokens = sum(lengths) + packed = torch.randn(1, tokens, 10, 128, device="cuda", dtype=torch.bfloat16) + q, k, v = packed[:, :, :2], packed[:, :, 2:4], packed[:, :, 4:] + v.mul_(0.2) + g = -torch.rand(1, tokens, 6, device="cuda") * 0.1 + beta = torch.rand_like(g) + state = torch.randn(len(lengths), 6, 128, 128, device="cuda") * 0.01 + cu = torch.tensor([0, *lengths], device="cuda", dtype=torch.int32).cumsum(0, dtype=torch.int32) + kwargs = dict( + q=q, + k=k, + v=v, + g=g, + beta=beta, + initial_state=state, + output_final_state=True, + cu_seqlens=cu, + use_qk_l2norm_in_kernel=True, + ) + original = state.clone() + if 0 in lengths: + active = [i for i, length in enumerate(lengths) if length] + offsets = torch.tensor([0, *[lengths[i] for i in active]], device="cuda", dtype=torch.int32).cumsum( + 0, dtype=torch.int32 + ) + expected, active_state = native(**dict(kwargs, initial_state=state[active], cu_seqlens=offsets)) + expected_state = state.clone() + expected_state[active] = active_state + else: + expected, expected_state = native(**kwargs) + actual, actual_state = chunk_gated_delta_rule(**kwargs) + torch.testing.assert_close(actual, expected, atol=0.001, rtol=0.03) + torch.testing.assert_close(actual_state, expected_state, atol=0.003, rtol=0.03) + torch.testing.assert_close(state, original, atol=0, rtol=0) + output_only, final = chunk_gated_delta_rule(**dict(kwargs, output_final_state=False)) + assert final is None + torch.testing.assert_close(output_only, actual, atol=0, rtol=0) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output, final = chunk_gated_delta_rule(**kwargs) + try: + for _ in range(3): + packed.mul_(0.7) + state.add_(0.001) + cu.copy_(torch.tensor([0, *reversed(lengths)], device="cuda").cumsum(0)) + graph.replay() + expected, expected_state = chunk_gated_delta_rule(**kwargs) + torch.testing.assert_close(output, expected, atol=0, rtol=0) + torch.testing.assert_close(final, expected_state, atol=0, rtol=0) + finally: + graph.reset() + + +def test_prefill_continuation_uses_native_state_layout(): + torch.manual_seed(801) + q = torch.randn(1, 768, 2, 128, device="cuda", dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn(1, 768, 6, 128, device="cuda", dtype=torch.bfloat16) * 0.2 + g = -torch.rand(1, 768, 6, device="cuda") * 0.1 + beta = torch.rand_like(g) + initial = torch.randn(1, 6, 128, 128, device="cuda") * 0.01 + + def run(kernel, start, end, state): + return kernel( + q=q[:, start:end], + k=k[:, start:end], + v=v[:, start:end], + g=g[:, start:end], + beta=beta[:, start:end], + initial_state=state, + output_final_state=True, + cu_seqlens=torch.tensor([0, end - start], device="cuda", dtype=torch.int32), + use_qk_l2norm_in_kernel=True, + ) + + _, prefix = run(native, 0, 256, initial) + output, final = run(chunk_gated_delta_rule, 256, 768, prefix) + expected, expected_state = run(native, 256, 768, prefix) + torch.testing.assert_close(output, expected, atol=0.001, rtol=0.03) + torch.testing.assert_close(final, expected_state, atol=0.003, rtol=0.03) + expected, expected_state = run(native, 0, 768, initial) + torch.testing.assert_close(output, expected[:, 256:], atol=0.001, rtol=0.03) + torch.testing.assert_close(final, expected_state, atol=0.003, rtol=0.03) diff --git a/unit_tests/common/basemodel/attention/linear/test_gdn.py b/unit_tests/common/basemodel/attention/linear/test_gdn.py index 344bfea57f..9be1de7fb7 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) @@ -178,12 +264,13 @@ def test_missing_linear_backend_arg_uses_auto_selection(monkeypatch, auto_linear assert validate_calls == ["flashqla"] -def test_replay_mode_selects_standalone_gdn_backends(monkeypatch): +@pytest.mark.parametrize("prefill_name", ["flashqla", "flashinfer"]) +def test_replay_mode_selects_standalone_gdn_backends(monkeypatch, prefill_name): monkeypatch.setattr( linear_create_utils, "get_env_start_args", lambda: SimpleNamespace( - llm_prefill_att_backend=["fa3", "flashqla"], + llm_prefill_att_backend=["fa3", prefill_name], llm_decode_att_backend=["flashinfer", "triton"], ssm_state_mode="replay", ), @@ -192,13 +279,15 @@ def test_replay_mode_selects_standalone_gdn_backends(monkeypatch): from lightllm.common.basemodel.attention.linear.gdn import LinearAttBackend from lightllm.common.basemodel.attention.linear.replay_gdn import ( ReplayGDNFlashQlaAttBackend, + ReplayGDNFlashInferAttBackend, ReplayGDNTritonAttBackend, ) prefill_backend = linear_create_utils.get_gdn_linear_prefill_att_backend_class(index=1) decode_backend = linear_create_utils.get_gdn_linear_decode_att_backend_class(index=1) - assert prefill_backend is ReplayGDNFlashQlaAttBackend + expected = {"flashqla": ReplayGDNFlashQlaAttBackend, "flashinfer": ReplayGDNFlashInferAttBackend} + assert prefill_backend is expected[prefill_name] assert decode_backend is ReplayGDNTritonAttBackend assert prefill_backend.__bases__ == (BaseAttBackend,) assert decode_backend.__bases__ == (BaseAttBackend,) @@ -349,7 +438,7 @@ def test_gdn_decode_backend_uses_priority_list(monkeypatch, auto_linear_backend_ monkeypatch.setattr( linear_create_utils, "gdn_decode_backends", - {"optimized": FlashQlaLinearAttBackend, "triton": TritonLinearAttBackend}, + {"optimized": "linear.flashqla:FlashQlaLinearAttBackend", "triton": "linear.triton:TritonLinearAttBackend"}, ) monkeypatch.setattr( linear_create_utils, @@ -378,7 +467,7 @@ def test_gdn_prefill_backend_tries_candidates_in_order(monkeypatch, auto_linear_ monkeypatch.setattr( linear_create_utils, "gdn_prefill_backends", - {"flashqla2": TritonLinearAttBackend, "flashqla3": FlashQlaLinearAttBackend}, + {"flashqla2": "linear.triton:TritonLinearAttBackend", "flashqla3": "linear.flashqla:FlashQlaLinearAttBackend"}, ) monkeypatch.setattr( linear_create_utils, diff --git a/unit_tests/common/basemodel/attention/test_lazy_import.py b/unit_tests/common/basemodel/attention/test_lazy_import.py index e7423d30f9..39eb008c92 100644 --- a/unit_tests/common/basemodel/attention/test_lazy_import.py +++ b/unit_tests/common/basemodel/attention/test_lazy_import.py @@ -90,6 +90,7 @@ def setUp(self): self.attention = importlib.import_module(ATTENTION) def test_importing_selection_apis_loads_no_implementations(self): + importlib.import_module(f"{ATTENTION}.create_linear_utils") importlib.import_module(f"{ATTENTION}.linear") importlib.import_module(f"{ATTENTION}.nsa") self.assertEqual(self.imports.imported, []) @@ -184,25 +185,15 @@ def test_neo_auto_selection_uses_supported_fa3(self): self.assertEqual(self.imports.imported, [f"{ATTENTION}.fa3.neo"]) self.validate.assert_called_once_with("fa3") - def _allow_linear_backend_classes(self): - flashqla = self.imports.allow("linear.flashqla", "FlashQlaLinearAttBackend") - triton = self.imports.allow("linear.triton", "TritonLinearAttBackend") - self.imports.allow("linear.replay_gdn", "ReplayGDNTritonAttBackend") - replay_gdn = self.imports.allowed[f"{ATTENTION}.linear.replay_gdn"] - replay_gdn.ReplayGDNFlashQlaAttBackend = type("ReplayGDNFlashQlaAttBackend", (), {}) - self.imports.allow("linear.kda", "KDALinearAttBackend") - self.imports.allow("linear.replay_kda", "ReplayKDALinearAttBackend") - return flashqla, triton - def test_linear_prefill_auto_selects_validated_flashqla(self): - flashqla, _triton = self._allow_linear_backend_classes() + flashqla = self.imports.allow("linear.flashqla", "FlashQlaLinearAttBackend") selection = importlib.import_module(f"{ATTENTION}.create_linear_utils") self.assertIs(selection.get_gdn_linear_prefill_att_backend_class(), flashqla) self.validate.assert_called_once_with("flashqla") def test_linear_selection_falls_back_to_triton(self): - _flashqla, triton = self._allow_linear_backend_classes() + triton = self.imports.allow("linear.triton", "TritonLinearAttBackend") selection = importlib.import_module(f"{ATTENTION}.create_linear_utils") self.validate.return_value = False @@ -211,6 +202,25 @@ def test_linear_selection_falls_back_to_triton(self): self.args.llm_prefill_att_backend.append("triton") self.assertIs(selection.get_gdn_linear_prefill_att_backend_class(), triton) + def test_linear_flashinfer_selection_is_lazy(self): + backend = self.imports.allow("linear.flashinfer", "FlashInferLinearAttBackend") + self.args.llm_prefill_att_backend = ["triton", "flashinfer"] + selection = importlib.import_module(f"{ATTENTION}.create_linear_utils") + + self.assertIs(selection.get_gdn_linear_prefill_att_backend_class(), backend) + self.assertEqual(self.imports.imported, [f"{ATTENTION}.linear.flashinfer"]) + self.validate.assert_not_called() + + def test_replay_selection_skips_native_backends(self): + backend = self.imports.allow("linear.replay_gdn", "ReplayGDNTritonAttBackend") + self.args.ssm_state_mode = "replay" + selection = importlib.import_module(f"{ATTENTION}.create_linear_utils") + self.validate.return_value = False + + self.assertIs(selection.get_gdn_linear_prefill_att_backend_class(), backend) + self.assertIs(selection.get_gdn_linear_decode_att_backend_class(), backend) + self.assertEqual(self.imports.imported, [f"{ATTENTION}.linear.replay_gdn"]) + def test_vit_selection_loads_only_requested_implementation(self): package = "lightllm.common.basemodel.attention_vit" for name in list(sys.modules): diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_causal_conv1d.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_causal_conv1d.py new file mode 100644 index 0000000000..b5e436a689 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_causal_conv1d.py @@ -0,0 +1,398 @@ +import builtins +from collections import defaultdict + +import pytest +import torch +import torch.nn.functional as F + +from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d import ( + causal_conv1d_fn as dispatch_prefill, + causal_conv1d_update as dispatch_update, +) +from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_triton import ( + _prefill_configs, + _update_configs, + causal_conv1d_fn, + causal_conv1d_update, +) +from lightllm.common.triton_utils import autotuner + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + + +def _prefill_reference(x, weight, bias, states, lengths, indices, initial, activation): + width = weight.shape[-1] + output = x.clone() + start = 0 + for length, slot, use_initial in zip(lengths, indices, initial): + if length and slot != -1: + prefix = states[slot] if use_initial else torch.zeros_like(states[slot]) + extended = torch.cat((prefix, x[:, start : start + length]), dim=-1) + with torch.backends.cudnn.flags(allow_tf32=False): + convolved = F.conv1d( + extended.float().unsqueeze(0), + weight.float().unsqueeze(1), + bias.float() if bias is not None else None, + groups=x.shape[0], + ).squeeze(0) + output[:, start : start + length] = F.silu(convolved) if activation else convolved + states[slot] = extended[:, -(width - 1) :] + start += length + return output + + +def _assert_output(actual, expected): + if actual.dtype == torch.float32: + torch.testing.assert_close(actual, expected, rtol=2e-5, atol=2e-6) + else: + torch.testing.assert_close(actual, expected, rtol=1e-2, atol=2e-3) + + +def _check_prefill_ragged_prefix(dtype, width, layout, run_config=None): + torch.manual_seed(42) + dim = 131 + lengths = [1, 0, 3, 19, 67] + tokens = sum(lengths) + if layout == "projection": + x = torch.randn(tokens, dim + 32, device="cuda", dtype=dtype)[:, :dim].T + else: + x = torch.randn(dim, tokens * 2, device="cuda", dtype=dtype)[:, ::2] + weight = torch.randn(dim, width * 2, device="cuda", dtype=dtype)[:, ::2] * 0.2 + bias = torch.randn(dim * 2, device="cuda", dtype=dtype)[::2] if layout == "projection" else None + activation = "silu" if bias is not None else None + state_storage = torch.randn(6, dim, width + 2, device="cuda", dtype=dtype) + original_storage = state_storage.clone() + states = state_storage[..., : width - 1] + ref_states = states.clone() + indices = [3, 1, -1, 4, 0] + initial = [True, False, False, True, False] + starts = torch.tensor([0, 1, 1, 4, 23, 90], device="cuda", dtype=torch.int32) + expected = _prefill_reference(x, weight, bias, ref_states, lengths, indices, initial, activation) + actual = causal_conv1d_fn( + x, + weight, + bias, + query_start_loc=starts, + cache_indices=torch.tensor(indices, device="cuda"), + has_initial_state=torch.tensor(initial, device="cuda"), + conv_states=states, + activation=activation, + max_seqlen=max(lengths), + run_config=run_config, + ) + _assert_output(actual, expected) + torch.testing.assert_close(states, ref_states, rtol=0, atol=0) + torch.testing.assert_close(state_storage[..., width - 1 :], original_storage[..., width - 1 :], rtol=0, atol=0) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("width", [2, 3, 4]) +@pytest.mark.parametrize("layout", ["projection", "sequence"]) +def test_prefill_ragged_prefix_and_strided_state(dtype, width, layout): + _check_prefill_ragged_prefix(dtype, width, layout) + + +@pytest.mark.parametrize("run_config", _prefill_configs()) +def test_prefill_autotune_configurations(run_config): + _check_prefill_ragged_prefix(torch.bfloat16, 4, "projection", run_config) + + +def test_prefill_new_requests_ignore_stale_state(): + x = torch.randn(2, 65, 33, device="cuda") + weight = torch.randn(65, 4, device="cuda") + states = torch.randn(2, 65, 3, device="cuda") + with torch.backends.cudnn.flags(allow_tf32=False): + expected = F.silu(F.conv1d(F.pad(x, (3, 0)), weight[:, None], groups=65)) + actual = causal_conv1d_fn( + x.transpose(0, 1).reshape(65, -1), + weight, + query_start_loc=torch.tensor([0, 33, 66], device="cuda", dtype=torch.int32), + cache_indices=torch.arange(2, device="cuda", dtype=torch.int32), + has_initial_state=torch.zeros(2, device="cuda", dtype=torch.bool), + conv_states=states, + activation="swish", + ) + _assert_output(actual, expected.transpose(0, 1).reshape(65, -1)) + torch.testing.assert_close(states, x[..., -3:], rtol=0, atol=0) + + +def _update_reference(x, states, weight, bias, indices, cache_lengths, activation): + output = x.clone() + input_3d = x.unsqueeze(-1) if x.ndim == 2 else x + output_3d = output.unsqueeze(-1) if output.ndim == 2 else output + width, state_len = weight.shape[-1], states.shape[-1] + for batch, slot in enumerate(indices): + if slot == -1: + continue + for token in range(input_3d.shape[-1]): + if cache_lengths is None: + history = states[slot, :, -(width - 1) :].clone() + states[slot] = torch.cat((states[slot, :, 1:], input_3d[batch, :, token : token + 1]), dim=-1) + else: + cursor = cache_lengths[batch] + token + history = torch.stack( + [states[slot, :, (cursor - width + 1 + j) % state_len] for j in range(width - 1)], dim=-1 + ) + states[slot, :, cursor % state_len] = input_3d[batch, :, token] + values = torch.cat((history, input_3d[batch, :, token : token + 1]), dim=-1).double() + accumulator = bias.float().clone() if bias is not None else torch.zeros(input_3d.shape[1], device=x.device) + for tap in range(width): + accumulator = (accumulator.double() + values[:, tap] * weight[:, tap].double()).float() + output_3d[batch, :, token] = F.silu(accumulator) if activation else accumulator + return output + + +def _check_update_buffers(dtype, width, circular, tokens, run_config=None): + torch.manual_seed(17) + dim = 131 + x = torch.randn(3, dim + 16, tokens * 2, device="cuda", dtype=dtype)[:, :dim, ::2] + if tokens == 1: + x = x.squeeze(-1) + weight = torch.randn(dim, width, device="cuda", dtype=dtype) * 0.2 + bias = torch.randn(dim * 2, device="cuda", dtype=dtype)[::2] + state_len = width + 1 + storage = torch.randn(5, dim, state_len + 3, device="cuda", dtype=dtype) + original = storage.clone() + states = storage[..., :state_len] + ref_states = states.clone() + indices = [2, -1, 0] + cache_lengths = [0, 3, 13] if circular else None + activation = "silu" if tokens == 1 else None + expected = _update_reference(x, ref_states, weight, bias, indices, cache_lengths, activation) + actual = causal_conv1d_update( + x, + states, + weight, + bias, + activation, + torch.tensor(cache_lengths, device="cuda", dtype=torch.int32) if circular else None, + torch.tensor(indices, device="cuda", dtype=torch.int32), + run_config=run_config, + ) + _assert_output(actual, expected) + torch.testing.assert_close(states, ref_states, rtol=0, atol=0) + torch.testing.assert_close(storage[..., state_len:], original[..., state_len:], rtol=0, atol=0) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("width", [2, 3, 4]) +@pytest.mark.parametrize("circular,tokens", [(False, 1), (False, 7), (True, 1), (True, 7)]) +def test_update_linear_and_ring_buffers(dtype, width, circular, tokens): + _check_update_buffers(dtype, width, circular, tokens) + + +@pytest.mark.parametrize("run_config", _update_configs()) +@pytest.mark.parametrize("circular", [False, True]) +def test_update_autotune_configurations(run_config, circular): + _check_update_buffers(torch.bfloat16, 4, circular, 7, run_config) + + +@pytest.mark.parametrize("mode", ["prefill", "linear", "circular"]) +def test_autotune_and_cached_config_update_state_once(monkeypatch, tmp_path, mode): + torch.manual_seed(13) + tuner = causal_conv1d_fn if mode == "prefill" else causal_conv1d_update + monkeypatch.setattr(autotuner, "get_triton_autotune_level", lambda: autotuner.AutotuneLevel.ADAPTIVE_AUTOTUNE) + monkeypatch.setattr(tuner, "_cache_dir", str(tmp_path), raising=False) + monkeypatch.setattr(tuner, "cached_configs", {}) + monkeypatch.setattr(tuner, "fast_match_configs", defaultdict(dict)) + monkeypatch.setattr(tuner, "warmuped_configs_set", set()) + configs = tuner.configs_gen_func() + monkeypatch.setattr(tuner, "configs_gen_func", lambda: [configs[0], configs[-1]]) + + dim, state_len = 65, 3 if mode == "prefill" else 5 + storage = torch.randn(4, dim, state_len + 2, device="cuda", dtype=torch.bfloat16) + original = storage.clone() + states = storage[..., :state_len] + reference_states = states.clone() + weight = torch.randn(dim, 4, device="cuda", dtype=states.dtype) * 0.2 + indices = torch.tensor([3, 1], device="cuda", dtype=torch.int32) + if mode == "prefill": + x = torch.randn(3, dim + 16, device="cuda", dtype=states.dtype)[:, :dim].T + starts = torch.tensor([0, 1, 3], device="cuda", dtype=torch.int32) + initial = torch.tensor([True, True], device="cuda") + + def run(): + return tuner( + x, + weight, + query_start_loc=starts, + cache_indices=indices, + has_initial_state=initial, + conv_states=states, + max_seqlen=2, + ) + + def reference(): + return _prefill_reference(x, weight, None, reference_states, [1, 2], [3, 1], [True, True], "silu") + + else: + x = torch.randn(2, dim, device="cuda", dtype=states.dtype) + lengths = [0, 13] if mode == "circular" else None + cache_lengths = torch.tensor(lengths, device="cuda", dtype=torch.int32) if lengths is not None else None + + def run(): + return tuner(x, states, weight, activation="silu", cache_seqlens=cache_lengths, conv_state_indices=indices) + + def reference(): + return _update_reference(x, reference_states, weight, None, [3, 1], lengths, "silu") + + clone_mutated_args = tuner._mutate_args_clone + scratch_states = [] + + def check_clones(args, kwargs): + new_args, new_kwargs, original_states, cloned_states = clone_mutated_args(args, kwargs) + assert len(cloned_states) == 1 and original_states[0] is states + assert cloned_states[0].data_ptr() != states.data_ptr() + scratch_states.append(cloned_states[0]) + return new_args, new_kwargs, original_states, cloned_states + + monkeypatch.setattr(tuner, "_mutate_args_clone", check_clones) + expected = reference() + with autotuner.Autotuner.autotune_warmup(tuner.kernel_type): + actual = run() + assert len(scratch_states) == 2 + assert len(list(tmp_path.glob("*.json"))) == 1 + assert all(config is not None for cached in tuner.cached_configs.values() for config in cached.values()) + _assert_output(actual, expected) + torch.testing.assert_close(states, reference_states, rtol=0, atol=0) + + monkeypatch.setattr(autotuner, "get_triton_autotune_level", lambda: autotuner.AutotuneLevel.USE_AUTOTUNE_HIS_CONFIG) + monkeypatch.setattr(tuner, "cached_configs", {}) + monkeypatch.setattr(tuner, "fast_match_configs", defaultdict(dict)) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = run() + expected = reference() + graph.replay() + _assert_output(output, expected) + torch.testing.assert_close(states, reference_states, rtol=0, atol=0) + torch.testing.assert_close(storage[..., state_len:], original[..., state_len:], rtol=0, atol=0) + assert len(scratch_states) > 2 + + +def test_prefill_graph_replays_with_changed_lengths_and_prefixes(): + torch.manual_seed(1) + x = torch.randn(25, 96, device="cuda", dtype=torch.bfloat16).T + weight = torch.randn(96, 4, device="cuda", dtype=x.dtype) * 0.2 + original = torch.randn(2, 96, 3, device="cuda", dtype=x.dtype) + states = original.clone() + starts = torch.tensor([0, 1, 25], dtype=torch.int32, device="cuda") + initial = torch.tensor([True, False], device="cuda") + indices = torch.arange(2, device="cuda", dtype=torch.int32) + + def run(): + return causal_conv1d_fn( + x, + weight, + query_start_loc=starts, + cache_indices=indices, + has_initial_state=initial, + conv_states=states, + max_seqlen=25, + ) + + run() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = run() + for lengths, flags in [([1, 24], [True, False]), ([25, 0], [False, True]), ([2, 23], [True, True])]: + starts.copy_(torch.tensor([0, lengths[0], 25], dtype=starts.dtype, device="cuda")) + initial.copy_(torch.tensor(flags, device="cuda")) + states.copy_(original) + reference_states = original.clone() + expected = _prefill_reference(x, weight, None, reference_states, lengths, [0, 1], flags, "silu") + graph.replay() + _assert_output(output, expected) + torch.testing.assert_close(states, reference_states, rtol=0, atol=0) + + +def test_dispatch_without_sgl_kernel_or_vllm(monkeypatch): + import_fn = builtins.__import__ + + def without_optional_kernels(name, *args, **kwargs): + if name.split(".")[0] in ("sgl_kernel", "vllm"): + raise ModuleNotFoundError(name) + return import_fn(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", without_optional_kernels) + x = torch.randn(13, 65, dtype=torch.bfloat16, device="cuda").T + weight = torch.randn(65, 4, dtype=x.dtype, device="cuda") * 0.2 + states = torch.zeros(1, 65, 3, dtype=x.dtype, device="cuda") + ref_states = states.clone() + starts = torch.tensor([0, 13], dtype=torch.int32, device="cuda") + expected = _prefill_reference(x, weight, None, ref_states, [13], [0], [False], "silu") + _assert_output( + dispatch_prefill( + x, + weight, + query_start_loc=starts, + conv_states=states, + cache_indices=torch.zeros(1, device="cuda", dtype=torch.int32), + has_initial_state=torch.zeros(1, device="cuda", dtype=torch.bool), + ), + expected, + ) + token = torch.randn(1, 65, dtype=x.dtype, device="cuda") + expected = _update_reference(token, ref_states, weight, None, [0], None, "silu") + _assert_output(dispatch_update(token, states, weight, activation="silu"), expected) + torch.testing.assert_close(states, ref_states, rtol=0, atol=0) + + +@pytest.mark.parametrize("tokens", [1, 3]) +def test_sgl_update(tokens): + pytest.importorskip("sgl_kernel") + torch.manual_seed(29) + x = torch.randn(3, 64, tokens, dtype=torch.bfloat16, device="cuda") + if tokens == 1: + x = x.squeeze(-1) + states = torch.randn(5, 64, 3, dtype=x.dtype, device="cuda") + ref_states = states.clone() + weight = torch.randn(64, 4, dtype=x.dtype, device="cuda") * 0.2 + bias = torch.randn(64, dtype=x.dtype, device="cuda") + indices = [2, -1, 0] + expected = _update_reference(x, ref_states, weight, bias, indices, None, "silu") + actual = dispatch_update( + x, + states, + weight, + bias, + activation="silu", + conv_state_indices=torch.tensor(indices, dtype=torch.int32, device="cuda"), + ) + assert actual.shape == x.shape + _assert_output(actual, expected) + torch.testing.assert_close(states, ref_states, rtol=0, atol=0) + + +def test_bfloat16_products_accumulate_in_float32(): + x = torch.tensor([[1.0078125, 1.0]], device="cuda", dtype=torch.bfloat16) + weight = torch.tensor([[1.0078125, -1.0]], device="cuda", dtype=torch.bfloat16) + expected = torch.tensor([[-1.0078125, 0.01568603515625]], device="cuda", dtype=x.dtype) + actual = causal_conv1d_fn( + x, + weight, + activation=None, + query_start_loc=torch.tensor([0, 2], device="cuda", dtype=torch.int32), + cache_indices=torch.zeros(1, device="cuda", dtype=torch.int32), + has_initial_state=torch.zeros(1, device="cuda", dtype=torch.bool), + conv_states=torch.ones(1, 1, 1, device="cuda", dtype=x.dtype), + ) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +def test_empty_inputs_do_not_touch_state(): + state = torch.randn(2, 17, 3, device="cuda") + original = state.clone() + weight = torch.randn(17, 4, device="cuda") + output = causal_conv1d_fn( + torch.empty(17, 0, device="cuda"), + weight, + conv_states=state, + query_start_loc=torch.zeros(3, device="cuda", dtype=torch.int32), + cache_indices=torch.arange(2, device="cuda", dtype=torch.int32), + has_initial_state=torch.ones(2, device="cuda", dtype=torch.bool), + ) + assert output.numel() == 0 + assert causal_conv1d_update(torch.empty(2, 17, 0, device="cuda"), state, weight).numel() == 0 + torch.testing.assert_close(state, original, rtol=0, atol=0) 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..5b73e4ba70 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 @@ -5,6 +5,33 @@ from lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_mtp import causal_conv1d_update +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +@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, diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_fused_recurrent_packed_decode.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_fused_recurrent_packed_decode.py new file mode 100644 index 0000000000..aae63dd572 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_fused_recurrent_packed_decode.py @@ -0,0 +1,62 @@ +import pytest +import torch + +from lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.fused_recurrent import ( + fused_recurrent_gated_delta_rule, + fused_recurrent_gated_delta_rule_packed_decode, +) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_packed_decode_matches_materialized_decode(): + torch.manual_seed(17) + batch, key_heads, value_heads = 2, 2, 4 + key_dim = value_dim = 128 + device = "cuda" + dtype = torch.bfloat16 + + mixed_qkv = torch.randn( + batch, + 2 * key_heads * key_dim + value_heads * value_dim, + device=device, + dtype=dtype, + ) + a_raw = torch.randn(batch, value_heads, device=device, dtype=dtype) + b_raw = torch.randn_like(a_raw) + a_log = torch.randn(value_heads, device=device, dtype=torch.float32) + dt_bias = torch.randn(value_heads, device=device, dtype=torch.float32) + state = torch.randn(batch + 1, value_heads, key_dim, value_dim, device=device, dtype=dtype) + state_reference = state.clone() + state_packed = state.clone() + state_indices = torch.arange(1, batch + 1, device=device, dtype=torch.int32) + + q = mixed_qkv[:, : key_heads * key_dim].view(batch, 1, key_heads, key_dim) + k = mixed_qkv[:, key_heads * key_dim : 2 * key_heads * key_dim].view(batch, 1, key_heads, key_dim) + v = mixed_qkv[:, 2 * key_heads * key_dim :].view(batch, 1, value_heads, value_dim) + + reference, _ = fused_recurrent_gated_delta_rule( + q, + k, + v, + initial_state=state_reference, + ssm_state_indices=state_indices, + use_qk_l2norm_in_kernel=True, + A_log=a_log, + dt_bias=dt_bias, + a_raw=a_raw, + b_raw=b_raw, + ) + packed, _ = fused_recurrent_gated_delta_rule_packed_decode( + mixed_qkv, + a_raw, + b_raw, + a_log, + dt_bias, + key_dim ** -0.5, + state_packed, + state_indices, + use_qk_l2norm_in_kernel=True, + ) + + torch.testing.assert_close(packed, reference, atol=0.02, rtol=0.02) + torch.testing.assert_close(state_packed, state_reference, atol=0.02, rtol=0.02) diff --git a/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_equiv.py b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_equiv.py index 11035c83a3..b280f4fabb 100644 --- a/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_equiv.py +++ b/unit_tests/common/basemodel/triton_kernel/linear_att/test_mtp_fused_recurrent_equiv.py @@ -163,5 +163,57 @@ def test_variable_length_cu_seqlens(seed): ) +@pytest.mark.parametrize("state_dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("steps", [1, 2, 3, 4]) +def test_match_decode_preserves_every_sequential_snapshot(state_dtype, steps): + torch.manual_seed(3107) + batch, heads, value_heads, width = 2, 4, 12, 128 + count = batch * steps + q = torch.randn(batch, steps, heads, width, device="cuda", dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn(batch, steps, value_heads, width, device="cuda", dtype=torch.bfloat16) + a = torch.randn(batch, steps, value_heads, device="cuda", dtype=torch.bfloat16) + b = torch.randn_like(a) + indices = torch.arange(count, device="cuda", dtype=torch.int32).view(batch, steps) + accepted = torch.tensor([1, steps], device="cuda", dtype=torch.int32) + initial = torch.randn(count, value_heads, width, width, device="cuda", dtype=state_dtype) * 0.1 + state = initial[indices.gather(1, accepted.long()[:, None] - 1).flatten().long()].clone() + expected_states = initial.clone() + gates = dict(A_log=torch.zeros(value_heads, device="cuda"), dt_bias=torch.zeros(value_heads, device="cuda")) + outputs = [] + for step in range(steps): + output, _ = fused_recurrent_gated_delta_rule( + q=q[:, step : step + 1], + k=k[:, step : step + 1], + v=v[:, step : step + 1], + initial_state=state, + inplace_final_state=True, + ssm_state_indices=torch.arange(batch, device="cuda", dtype=torch.int32), + use_qk_l2norm_in_kernel=True, + a_raw=a[:, step], + b_raw=b[:, step], + **gates, + ) + outputs.append(output) + expected_states.index_copy_(0, indices[:, step].long(), state) + actual, actual_states = mtp_fused_recurrent_gated_delta_rule( + q=q.reshape(1, count, heads, width), + k=k.reshape(1, count, heads, width), + v=v.reshape(1, count, value_heads, width), + initial_state=initial.clone(), + cu_seqlens=torch.arange(batch + 1, device="cuda", dtype=torch.int64) * steps, + ssm_state_indices=indices, + ssm_state_write_indices=indices, + num_accepted_tokens=accepted, + a_raw=a.reshape(count, value_heads), + b_raw=b.reshape(count, value_heads), + run_config={"BV": 32, "num_warps": 4, "num_stages": 1}, + **gates, + ) + expected = torch.cat(outputs, dim=1).reshape_as(actual) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + torch.testing.assert_close(actual_states, expected_states, atol=0, rtol=0) + + if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/unit_tests/common/basemodel/triton_kernel/test_gated_rmsnorm_strided_gate.py b/unit_tests/common/basemodel/triton_kernel/test_gated_rmsnorm_strided_gate.py new file mode 100644 index 0000000000..64215880e1 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/test_gated_rmsnorm_strided_gate.py @@ -0,0 +1,41 @@ +import pytest +import torch + +from lightllm.common.basemodel.triton_kernel.norm.gated_rmsnorm import gated_rmsnorm_forward + + +@pytest.mark.parametrize("norm_before_gate", [True, False]) +def test_gated_rmsnorm_accepts_strided_3d_gate(norm_before_gate): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for gated RMSNorm test") + + torch.manual_seed(123) + tokens, heads, head_dim = 5, 4, 128 + packed_dim = heads * head_dim + 256 + x = torch.randn((tokens * heads, head_dim), device="cuda", dtype=torch.bfloat16) + weight = torch.randn((head_dim,), device="cuda", dtype=torch.bfloat16) + packed = torch.randn((tokens, packed_dim), device="cuda", dtype=torch.bfloat16) + z = packed[:, 128 : 128 + heads * head_dim].view(tokens, heads, head_dim) + assert not z.is_contiguous() + + run_config = {"BLOCK_N": head_dim, "num_warps": 1} + expected = gated_rmsnorm_forward( + x=x, + weight=weight, + bias=None, + eps=1e-6, + z=z.contiguous().view(-1, head_dim), + norm_before_gate=norm_before_gate, + run_config=run_config, + ) + actual = gated_rmsnorm_forward( + x=x, + weight=weight, + bias=None, + eps=1e-6, + z=z, + norm_before_gate=norm_before_gate, + run_config=run_config, + ) + + torch.testing.assert_close(actual, expected, rtol=0, atol=0) diff --git a/unit_tests/common/basemodel/triton_kernel/test_scalar_linear.py b/unit_tests/common/basemodel/triton_kernel/test_scalar_linear.py new file mode 100644 index 0000000000..2cbd3559d6 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/test_scalar_linear.py @@ -0,0 +1,70 @@ +import pytest +import torch + +from lightllm.common.basemodel.triton_kernel.scalar_linear import scalar_linear +from lightllm.common.quantization.scalar_projection import StableScalarProjection +from lightllm.common.quantization.quantize_method import WeightPack + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("hidden", [65, 2560]) +def test_reference_batching_strides_and_graph(dtype, hidden): + torch.manual_seed(2319) + input = torch.randn((336, hidden * 2), device="cuda", dtype=dtype)[1:335, ::2] + weight = torch.randn((1, hidden * 2), device="cuda", dtype=dtype)[:, ::2] + out = torch.empty((334, 3), device="cuda", dtype=dtype)[:, 1:2] + expected = (input.double() * weight.double()).sum(-1, keepdim=True).to(dtype) + assert scalar_linear(input, weight, out=out) is out + torch.testing.assert_close(out, expected, atol=0, rtol=0) + for size in (1, 3, 8, 24, 64, 128, 256, 334): + pieces = [scalar_linear(chunk, weight) for chunk in input.split(size)] + assert torch.equal(torch.cat(pieces), out) + permutation = torch.randperm(334, device="cuda") + assert torch.equal(scalar_linear(input[permutation], weight), out[permutation]) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + scalar_linear(input, weight, out=out) + input.mul_(0.5) + graph.replay() + torch.cuda.synchronize() + expected = (input.double() * weight.double()).sum(-1, keepdim=True).to(dtype) + torch.testing.assert_close(out, expected, atol=0, rtol=0) + + +def test_small_terms_survive_before_final_dtype_cast(): + input = torch.tensor([[1, 2 ** -8, 2 ** -24, 2 ** -24]], device="cuda", dtype=torch.bfloat16) + weight = torch.ones_like(input) + actual = scalar_linear(input, weight) + torch.testing.assert_close(actual, torch.full_like(actual, 1.0078125), atol=0, rtol=0) + + +def test_dispatch_preserves_weight_layout_output_and_fallbacks(monkeypatch): + monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", str(torch.cuda.current_device())) + method = StableScalarProjection() + input = torch.randn((7, 65), device="cuda", dtype=torch.bfloat16) + weight = torch.randn((1, 65), device="cuda", dtype=input.dtype) + pack = WeightPack(weight=weight, weight_scale=None, weight_zero_point=None) + out = torch.empty((7, 1), device="cuda", dtype=input.dtype) + assert method.method_name == "none" + assert method.apply(input, pack, out=out, use_custom_tensor_mananger=False) is out + expected = (input.double() @ weight.double().t()).to(input.dtype) + torch.testing.assert_close(out, expected, atol=0, rtol=0) + torch.testing.assert_close(method.apply(input, pack), expected, atol=0, rtol=0) + + for outputs, bias in ((2, None), (1, torch.ones((1,), device="cuda", dtype=input.dtype))): + weight = torch.randn((outputs, 65), device="cuda", dtype=input.dtype) + pack = WeightPack(weight=weight, weight_scale=None, weight_zero_point=None) + actual = method.apply(input, pack, bias=bias, use_custom_tensor_mananger=False) + expected = torch.mm(input, weight.t()) if bias is None else torch.addmm(bias, input, weight.t()) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + +def test_empty_input_and_invalid_output(): + input = torch.empty((0, 65), device="cuda", dtype=torch.bfloat16) + weight = torch.ones((1, 65), device="cuda", dtype=input.dtype) + assert scalar_linear(input, weight).shape == (0, 1) + with pytest.raises(ValueError, match="output must match"): + scalar_linear(input, weight, out=torch.empty((0, 2), device="cuda", dtype=input.dtype)) diff --git a/unit_tests/common/fused_moe/test_append_shared_expert_topk.py b/unit_tests/common/fused_moe/test_append_shared_expert_topk.py index 4a197e911f..b603590775 100644 --- a/unit_tests/common/fused_moe/test_append_shared_expert_topk.py +++ b/unit_tests/common/fused_moe/test_append_shared_expert_topk.py @@ -86,6 +86,89 @@ def fused(**kwargs): assert torch.equal(executed["topk_ids"], torch.tensor([[0, 2, 4]], device="cuda")) +@pytest.mark.parametrize("index_dtype", [torch.int32, torch.int64]) +@pytest.mark.parametrize("execution_index_dtype", [None, torch.int64]) +def test_padding_mask_preserves_captured_routes_and_replays(index_dtype, execution_index_dtype): + logical_ids = torch.tensor([[0, 2], [1, 3]], dtype=index_dtype, device="cuda") + logical_weights = torch.tensor([[0.2, 0.8], [0.4, 0.6]], device="cuda") + impl = FuseMoeTriton.__new__(FuseMoeTriton) + impl.n_routed_experts = 4 + impl.num_fused_shared_experts = 1 + impl.execution_index_dtype = execution_index_dtype + impl._select_experts = lambda **kwargs: (logical_weights, logical_ids) + impl._fused_experts = lambda **kwargs: (kwargs["topk_weights"], kwargs["topk_ids"]) + hidden = torch.ones(2, 4, device="cuda") + mask = torch.tensor([True, False], device="cuda") + captured = {} + + def run(): + return impl( + input_tensor=hidden, + router_logits=hidden, + w13=None, + w2=None, + correction_bias=None, + scoring_func="softmax", + top_k=2, + renormalize=False, + use_grouped_topk=False, + topk_group=0, + num_expert_group=0, + is_prefill=False, + moe_capture_callback=lambda ids: captured.setdefault("ids", ids), + token_mask=mask, + ) + + for _ in range(3): + run() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + weights, ids = run() + for values in ([True, False], [False, True], [False, False], [True, True]): + mask.copy_(torch.tensor(values, device="cuda")) + graph.replay() + expected_ids = torch.tensor( + [[0, 2, 4], [1, 3, 4]], dtype=execution_index_dtype or index_dtype, device="cuda" + ).masked_fill(~mask[:, None], -1) + expected_weights = torch.tensor([[0.2, 0.8, 1.0], [0.4, 0.6, 1.0]], device="cuda") + torch.testing.assert_close(ids, expected_ids, atol=0, rtol=0) + torch.testing.assert_close(weights, expected_weights.masked_fill(~mask[:, None], 0), atol=0, rtol=0) + assert captured["ids"] is logical_ids + torch.testing.assert_close( + logical_ids, torch.tensor([[0, 2], [1, 3]], dtype=index_dtype, device="cuda"), atol=0, rtol=0 + ) + torch.testing.assert_close(logical_weights, torch.tensor([[0.2, 0.8], [0.4, 0.6]], device="cuda")) + + +@pytest.mark.parametrize("rows,topk,strided", [(0, 10, False), (17, 3, True), (129, 10, False), (1024, 10, False)]) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("index_dtype", [torch.int32, torch.int64]) +def test_mask_topk_preserves_inputs_and_dynamic_graph(rows, topk, strided, dtype, index_dtype): + from lightllm.common.basemodel.triton_kernel.fused_moe.mask_topk import mask_topk + + stride = 2 if strided else 1 + ids = torch.randint(512, (rows, topk * stride), dtype=index_dtype, device="cuda")[:, ::stride] + weights = torch.randn(rows, topk * stride, dtype=dtype, device="cuda")[:, ::stride] + mask = torch.arange(rows * stride, device="cuda").remainder(3).bool()[::stride] + saved_ids, saved_weights = ids.clone(), weights.clone() + + def run(): + return mask_topk(weights, ids, mask, torch.int64) + + graph = torch.cuda.CUDAGraph() + run() + with torch.cuda.graph(graph): + output_weights, output_ids = run() + for active in (True, False, True): + mask.fill_(active) + graph.replay() + torch.testing.assert_close(output_weights, weights.masked_fill(~mask[:, None], 0), atol=0, rtol=0) + torch.testing.assert_close(output_ids, ids.masked_fill(~mask[:, None], -1).long(), atol=0, rtol=0) + torch.testing.assert_close(ids, saved_ids, atol=0, rtol=0) + torch.testing.assert_close(weights, saved_weights, atol=0, rtol=0) + graph.reset() + + def test_redundancy_repair_does_not_overwrite_logical_ids(monkeypatch): from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.impl.deepgemm_impl import ( FuseMoeDeepGEMM, diff --git a/unit_tests/common/fused_moe/test_bf16_ep.py b/unit_tests/common/fused_moe/test_bf16_ep.py new file mode 100644 index 0000000000..6321a52626 --- /dev/null +++ b/unit_tests/common/fused_moe/test_bf16_ep.py @@ -0,0 +1,173 @@ +import json +import os +from types import SimpleNamespace + +import pytest +import torch +import torch.distributed as dist +import torch.nn.functional as F + + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or int(os.environ.get("WORLD_SIZE", "1")) < 2, + reason="Run with torchrun on multiple CUDA GPUs and DeepEP/DeepGEMM installed", +) + + +def test_bf16_ep(): + import deep_ep + import deep_gemm + from lightllm.distributed import dist_group_manager + from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.utils import envs_utils + from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe_ep.bf16 import fused_experts_bf16 + from lightllm.common.triton_utils.autotuner import Autotuner + + rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(rank) + dist.init_process_group("nccl") + world = dist.get_world_size() + for name, value in { + "LIGHTLLM_GLOBAL_RANK": rank, + "LIGHTLLM_GLOBAL_WORLD_SIZE": world, + "LIGHTLLM_CURRENT_RANK_IN_NODE": rank, + "LIGHTLLM_CURRENT_DEVICE_ID": rank, + "LIGHTLLM_DP_WORLD_SIZE": world, + "LIGHTLLM_CURRENT_RANK_IN_DP": rank, + }.items(): + os.environ[name] = str(value) + args = StartArgs( + batch_max_tokens=8192, running_max_req_size=256, mtp_step=3, tp=world, dp=1, enable_tpsp_mix_mode=True + ) + os.environ["LIGHTLLM_START_ARGS"] = json.dumps(vars(args)) + os.environ.pop("NUM_MAX_DISPATCH_TOKENS_PER_RANK_PREFILL", None) + os.environ.pop("NUM_MAX_DISPATCH_TOKENS_PER_RANK_DECODE", None) + envs_utils.get_deepep_num_max_dispatch_tokens_per_rank_prefill.cache_clear() + envs_utils.get_deepep_num_max_dispatch_tokens_per_rank_decode.cache_clear() + envs_utils.get_deepep_num_max_dispatch_tokens_per_rank.cache_clear() + envs_utils.get_env_start_args.cache_clear() + deep_gemm.set_mk_alignment_for_contiguous_layout(128) + dist.barrier(device_ids=[rank]) + buffer = deep_ep.ElasticBuffer( + dist.group.WORLD, + num_max_tokens_per_rank=envs_utils.get_deepep_num_max_dispatch_tokens_per_rank(), + hidden=256, + num_topk=2, + use_fp8_dispatch=False, + allow_multiple_reduction=True, + ) + dist_group_manager.ep_buffer = buffer + autotune_group = dist.new_group(backend="gloo") + from unittest.mock import patch + + torch.manual_seed(1042) + experts = world * 2 + up = torch.randn(experts, 256, 256, device="cuda", dtype=torch.bfloat16) * 0.05 + down = torch.randn(experts, 256, 128, device="cuda", dtype=torch.bfloat16) * 0.05 + local_up = up[rank * 2 : (rank + 1) * 2].contiguous() + local_down = down[rank * 2 : (rank + 1) * 2].contiguous() + + def reference(x, ids, weights): + total = torch.zeros_like(x, dtype=torch.float32) + for e in range(experts): + rows, routes = torch.where(ids == e) + projected = F.linear(x[rows], up[e]) + gate, value = projected.chunk(2, dim=-1) + activated = (F.silu(gate.float()) * value.float()).to(x.dtype) + output = F.linear(activated, down[e]) + total.index_add_(0, rows, output.float() * weights[rows, routes, None]) + return total.to(x.dtype) + + try: + for tokens, concentrated in ((3, False), (7, True), (17, False), (128, True)): + dummy_x = torch.zeros(256, 256, device="cuda", dtype=torch.bfloat16) + dummy_y = torch.empty_like(dummy_x) + dummy_up = torch.empty_like(dummy_x) + dummy_mid = torch.zeros(256, 128, device="cuda", dtype=torch.bfloat16) + dummy_layout = torch.tensor([128, 256], device="cuda", dtype=torch.int32) + expected_m = tokens + for a, b, out in ((dummy_x, local_up, dummy_up), (dummy_mid, local_down, dummy_y)): + deep_gemm.m_grouped_bf16_gemm_nt_contiguous( + a, + b, + out, + dummy_layout, + use_psum_layout=True, + expected_m_for_psum_layout=expected_m, + ) + torch.cuda.synchronize() + dist.barrier(device_ids=[rank]) + torch.manual_seed(1200 + rank) + x = torch.randn(tokens, 256, device="cuda", dtype=torch.bfloat16) * 0.2 + saved = x.clone() + ids = torch.stack( + [ + torch.arange(tokens, device="cuda") % (2 if concentrated else experts), + (torch.arange(tokens, device="cuda") + 1) % (2 if concentrated else experts), + ], + dim=-1, + ).long() + weights = torch.softmax(torch.randn(tokens, 2, device="cuda"), dim=-1) + expected = reference(x, ids, weights) + + shared_func = None + + def run(prefill): + return fused_experts_bf16( + x, local_up, local_down, weights, ids, experts, prefill, shared_expert_func=shared_func + ) + + for prefill in (True, False): + actual = run(prefill) + torch.testing.assert_close(actual, expected, atol=0.002, rtol=0.03) + torch.testing.assert_close(x, saved, atol=0, rtol=0) + if concentrated: + with patch.object( + dist_group_manager, "get_default_group", return_value=SimpleNamespace(autotune_group=autotune_group) + ), Autotuner.autotune_warmup(): + actual = run(True) + torch.testing.assert_close(actual, expected, atol=0.002, rtol=0.03) + original_ids, original_weights = ids.clone(), weights.clone() + for prefill, overlap in ((False, False), (True, False), (False, True), (True, True)): + ids.copy_(original_ids) + weights.copy_(original_weights) + shared_func = (lambda: x * 0.25) if overlap else None + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + actual = run(prefill) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = run(prefill) + x.mul_(0.5) + expected = reference(x, ids, weights) + if overlap: + expected += shared_func() + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(captured, expected, atol=0.002, rtol=0.03) + ids.fill_(-1) + weights.zero_() + if rank == 0: + ids[0] = torch.tensor([0, experts - 1], device="cuda") + weights[0] = 0.5 + expected = reference(x, ids, weights) + if overlap: + expected += shared_func() + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(captured, expected, atol=0.002, rtol=0.03) + ids.fill_(-1) + weights.zero_() + graph.replay() + torch.cuda.synchronize() + empty_expected = shared_func() if overlap else torch.zeros_like(captured) + torch.testing.assert_close(captured, empty_expected, atol=0, rtol=0) + graph.reset() + dist.barrier(device_ids=[rank]) + finally: + dist_group_manager.ep_buffer = None + del buffer + dist.destroy_process_group() diff --git a/unit_tests/common/fused_moe/test_bf16_ep_activation.py b/unit_tests/common/fused_moe/test_bf16_ep_activation.py new file mode 100644 index 0000000000..ab5629e616 --- /dev/null +++ b/unit_tests/common/fused_moe/test_bf16_ep_activation.py @@ -0,0 +1,102 @@ +import pytest +import torch +import torch.nn.functional as F + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.mark.parametrize("k,n", [(2560, 1280), (640, 2560), (65, 79)]) +def test_small_expert_gemm_matches_linear_and_replays_expert_counts(k, n): + from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe_ep.bf16 import _expert_gemm + + torch.manual_seed(3107) + x = torch.randn(512, k, device="cuda", dtype=torch.bfloat16) * 0.2 + w = torch.randn(4, n, k, device="cuda", dtype=x.dtype) * 0.03 + y = torch.full((512, n), 17, device="cuda", dtype=x.dtype) + ends = torch.tensor([0, 3, 128, 257], device="cuda", dtype=torch.int32) + + def run(): + _expert_gemm(x, w, y, ends, 1, False) + + def check(): + mask = torch.zeros(512, device="cuda", dtype=torch.bool) + start = 0 + for expert, end in enumerate(ends.tolist()): + torch.testing.assert_close(y[start:end], F.linear(x[start:end], w[expert]), rtol=0.008, atol=0.002) + mask[start:end] = True + start = (end + 127) // 128 * 128 + assert (y[~mask] == 17).all() + + run() + check() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + ends.copy_(torch.tensor([17, 128, 256, 399], device="cuda", dtype=ends.dtype)) + x.mul_(0.5) + y.fill_(17) + graph.replay() + check() + + +@pytest.mark.parametrize("width", [79, 128, 640]) +@pytest.mark.parametrize("kind", ["silu", "gelu", "clamp", "clamp_add_one"]) +@pytest.mark.parametrize("shards", [1, 8]) +def test_valid_expert_activation_preserves_padding_and_replays_dynamic_counts(width, kind, shards): + from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe_ep.bf16 import _activate_valid_rows + + x = torch.randn((512, 2 * width), device="cuda", dtype=torch.bfloat16) * 3 + output = torch.full((512, width), 17.0, device="cuda", dtype=x.dtype) + ends = torch.tensor([0, 3, 128, 257], device="cuda", dtype=torch.int32) + clamped = kind.startswith("clamp") + + def run(): + _activate_valid_rows[(4, (width + 127) // 128, shards)]( + x, + output, + ends, + width, + 128, + 1.7 if clamped else None, + 2.0 if clamped else None, + kind == "clamp_add_one", + kind == "gelu", + 16, + 128, + shards, + num_warps=4, + ) + + def check(): + gate, value = x.chunk(2, -1) + gate = gate.float() + if clamped: + gate = gate.clamp(max=2.0) + value = value.clamp(-2.0, 2.0) + gate = gate * torch.sigmoid(gate * 1.7) + if kind == "clamp_add_one": + value = value + 1 + elif kind == "gelu": + gate = F.gelu(gate, approximate="tanh") + else: + gate = F.silu(gate) + expected = (gate.to(x.dtype) * value).to(x.dtype) + mask = torch.zeros(512, device="cuda", dtype=torch.bool) + start = 0 + for end in ends.tolist(): + mask[start:end] = True + start = (end + 127) // 128 * 128 + torch.testing.assert_close(output[mask], expected[mask], rtol=0.008, atol=0.008) + assert (output[~mask] == 17).all() + + run() + check() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + ends.copy_(torch.tensor([5, 128, 259, 512], device="cuda", dtype=ends.dtype)) + x.mul_(0.5) + output.fill_(17) + graph.replay() + check() 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..d8d6d638f0 --- /dev/null +++ b/unit_tests/common/gated_residual/test_gated_residual.py @@ -0,0 +1,287 @@ +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) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test_fused_gate_projection(dtype): + from lightllm.common.gated_residual.kernels import injection_gates + from lightllm.common.gated_residual.projection import project_and_mix + + torch.manual_seed(3816) + down = torch.randn(320, 10240, device="cuda", dtype=dtype) * 0.01 + up = torch.randn(10240, 320, device="cuda", dtype=dtype) * 0.05 + inject = torch.randn(4, 10240, device="cuda", dtype=dtype) * 0.01 + for tokens in (0, 1, 32, 64, 65): + x = torch.randn(1, tokens, 10240, device="cuda", dtype=dtype) + mixed, gates = project_and_mix(x, down, up, 4, inject) + torch.testing.assert_close(mixed, project_and_mix(x, down, up, 4), atol=0, rtol=0) + torch.testing.assert_close(gates, injection_gates(x, inject, 4), atol=0, rtol=0) + + +def test_projection_batch_specialization(): + from lightllm.common.gated_residual.projection import _down, _merge, _up_mix + + dtype = torch.bfloat16 + down = torch.empty(320, 10240, device="cuda", dtype=dtype) + up = torch.empty(10240, 320, device="cuda", dtype=dtype) + inject = torch.empty(4, 10240, device="cuda", dtype=dtype) + for sizes, bm in (((5, 13), 16), ((576, 832), 128)): + previous = None + for rows in sizes: + x = torch.empty(rows, 10240, device="cuda", dtype=dtype) + low = torch.empty(rows, 320, device="cuda", dtype=dtype) + partial = torch.empty(16, rows, 320, device="cuda", dtype=torch.float32) + mixed = torch.empty(rows, 2560, device="cuda", dtype=dtype) + warps = 8 if bm == 128 else 4 + down_args = (x, down, partial, low, rows, 320, 10240, 4, bm, 64, 64, 16) + kernels = [ + _down.warmup(*down_args, grid=(1,), num_warps=warps, enable_fp_fusion=False), + _merge.warmup(partial, low, rows, 320, 4, 16, 128, grid=(1,), num_warps=4, enable_fp_fusion=False), + _up_mix.warmup( + low, + up, + x, + mixed, + rows, + 2560, + 320, + 4, + 4, + bm, + 32, + 64, + grid=(1,), + num_warps=warps, + enable_fp_fusion=False, + ), + ] + if bm == 16: + gates = torch.empty(rows, 4, device="cuda", dtype=dtype) + kernels.append( + _down.warmup(*down_args, inject, gates, 16384, grid=(1,), num_warps=8, enable_fp_fusion=False) + ) + if previous is not None: + assert all(left is right for left, right in zip(previous, kernels)) + previous = kernels + + +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) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("tokens,dim,branches", [(0, 65, 3), (17, 128, 4), (2, 2560, 4)]) +def test_combine_read(dtype, tokens, dim, branches): + from lightllm.common.gated_residual.kernels import grouped_norm, inject_and_norm + + torch.manual_seed(2149) + component = make_component(dim, branches, dtype) + residual = torch.randn(1, tokens, dim * branches, device="cuda", dtype=dtype) + update = torch.randn(1, tokens, dim, device="cuda", dtype=dtype) + gates = torch.sigmoid(torch.randn(1, tokens, branches, device="cuda", dtype=dtype)) * 2 + saved = residual.clone() + + def check(): + expected = component.write(residual, update, gates) + combined, normalized = inject_and_norm(residual, update, gates, component.norm_weight, branches, component.eps) + torch.testing.assert_close(combined, expected, atol=0, rtol=0) + torch.testing.assert_close( + normalized, grouped_norm(expected, component.norm_weight, branches, component.eps), atol=0, rtol=0 + ) + actual, (mixed, injection) = component.combine_and_read(residual, update, gates) + expected_mixed, expected_injection = component.read(expected) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + torch.testing.assert_close(mixed, expected_mixed, atol=0, rtol=0) + torch.testing.assert_close(injection, expected_injection, atol=0, rtol=0) + return actual, mixed, injection + + check() + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + component.combine_and_read(residual, update, gates) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual, (mixed, injection) = component.combine_and_read(residual, update, gates) + update.mul_(0.5) + gates.zero_() + graph.replay() + torch.cuda.synchronize() + for result, expected in zip((actual, mixed, injection), check()): + torch.testing.assert_close(result, expected, atol=0, rtol=0) + torch.testing.assert_close(residual, saved, atol=0, rtol=0) + graph.reset() + + +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..02f7b2539e --- /dev/null +++ b/unit_tests/common/ngram_embedding/test_pinned_embedding.py @@ -0,0 +1,123 @@ +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)) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("width", [37, 160, 513]) +def test_parallel_host_lookup_preserves_strided_output_and_dynamic_ids(dtype, width, monkeypatch): + monkeypatch.setattr(pinned_memory, "_CHUNK_BYTES", 16 * 1024 * 1024) + value = PinnedRowShardedEmbeddingWeight( + width, + 1024, + "table.weight", + dtype, + 4, + cpu_allocator=pinned_memory.pinned_zeros, + tp_rank=1, + tp_world_size=3, + ) + table = ((torch.arange(1024 * width).reshape(1024, width) % 127) - 64).to(dtype) + value.load_hf_weights({"table.weight": table}) + ids = torch.arange(2051, device="cuda", dtype=torch.int64) + backing = torch.full((ids.numel(), width + 7), 91, device="cuda", dtype=dtype) + output = backing[:, :width] + value(ids, out=output) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + value(ids, out=output) + for shift in (0, 341, -700): + ids.copy_(torch.arange(ids.numel(), device="cuda") + shift) + graph.replay() + host_ids = ids.cpu() + expected = table.index_select(0, host_ids.clamp(0, 1023)) + expected[(host_ids < value.tp_vocab_start_id) | (host_ids >= value.tp_vocab_end_id)] = 0 + torch.testing.assert_close(output.cpu(), expected, atol=0, rtol=0) + assert torch.all(backing[:, width:] == 91) + graph.reset() + assert value(ids[:0]).shape == (0, width) diff --git a/unit_tests/common/ngram_embedding/test_ple.py b/unit_tests/common/ngram_embedding/test_ple.py new file mode 100644 index 0000000000..677a923734 --- /dev/null +++ b/unit_tests/common/ngram_embedding/test_ple.py @@ -0,0 +1,259 @@ +import math +from dataclasses import replace + +import pytest +import torch +import torch.nn.functional as F + +from lightllm.common.ngram_embedding.ple import PLELayer + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +def make_layer(dim=65, branches=3, embed_dim=32, kernel=4, dilation=3, dtype=torch.bfloat16): + def rand(*shape): + return torch.randn(*shape, device="cuda", dtype=dtype) + + width = dim * branches + return PLELayer( + rand(width, embed_dim) / embed_dim ** 0.5, + rand(dim, embed_dim) / embed_dim ** 0.5, + rand(width) * 0.1, + rand(width) * 0.1, + rand(width) * 0.1, + rand(width, kernel) * 0.1, + branches, + dilation, + ) + + +def norm(x, weight, branches, eps): + grouped = x.float().unflatten(-1, (branches, -1)) + grouped = grouped * torch.rsqrt(grouped.square().mean(-1, keepdim=True) + eps) + return (grouped.flatten(-2) * (1 + weight.float())).to(x.dtype) + + +def reference_values(layer, residual, embeddings, mask=None): + key = norm(F.linear(embeddings, layer.key_weight), layer.norm_key, layer.branches, layer.eps) + query = norm(residual, layer.norm_query, layer.branches, layer.eps) + key, query = [x.unflatten(-1, (layer.branches, layer.hidden_size)) for x in [key, query]] + value = F.linear(embeddings, layer.value_weight) + gate = (key * query).sum(-1, keepdim=True) / math.sqrt(layer.hidden_size) + gate = gate.abs().clamp_min(1e-6).sqrt() * gate.sign() + gated = (torch.sigmoid(gate) * value.unsqueeze(-2)).flatten(-2) + normalized = norm(gated, layer.norm_conv, layer.branches, layer.eps) + if mask is not None: + gated = gated * mask[:, None] + normalized = normalized * mask[:, None] + return gated, normalized + + +def reference_conv(layer, values, gated, lengths, history): + result, next_state = [], [] + offset = 0 + for request, count in enumerate(lengths): + current = values[offset : offset + count] + full = torch.cat([history[request], current]) + if count: + conv = ( + F.conv1d( + full.T.unsqueeze(0), + layer.conv_weight.unsqueeze(1), + dilation=layer.dilation, + groups=values.shape[1], + ) + .squeeze(0) + .T + ) + result.append(gated[offset : offset + count] + F.silu(conv)) + next_state.append(full[-layer.context_len :] if layer.context_len else full[:0]) + offset += count + output = torch.cat(result) if result else values.new_empty((0, values.shape[1])) + state = torch.stack(next_state) if next_state else history.clone() + return output, state + + +def starts_for(lengths): + return torch.tensor([0] + list(lengths), device="cuda", dtype=torch.int32).cumsum(0).to(torch.int32) + + +def assert_close(actual, expected, dtype): + tolerance = max(2e-6, torch.finfo(dtype).eps * 2) + torch.testing.assert_close(actual, expected, atol=tolerance, rtol=tolerance) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("dim,branches,kernel,dilation", [(65, 3, 4, 3), (128, 4, 3, 2), (32, 2, 1, 1)]) +@pytest.mark.parametrize("masked", [False, True]) +def test_ple_reference_equations(dtype, dim, branches, kernel, dilation, masked): + torch.manual_seed(137) + layer = make_layer(dim, branches, kernel=kernel, dilation=dilation, dtype=dtype) + lengths = [0, 1, 2, 13, 0, 7] + count = sum(lengths) + residual = torch.randn(count, dim * branches, device="cuda", dtype=dtype) + embeddings = torch.randn(count, 32, device="cuda", dtype=dtype) + mask = (torch.arange(count, device="cuda") % 3 != 0) if masked else None + history = layer.empty_history(len(lengths)).normal_() + original = history.clone() + actual, conv_inputs = layer.forward(residual, embeddings, starts_for(lengths), history, mask) + gated, expected_conv_inputs = reference_values(layer, residual, embeddings, mask) + expected, expected_state = reference_conv(layer, expected_conv_inputs, gated, lengths, history) + assert_close(conv_inputs, expected_conv_inputs, dtype) + assert_close(actual, expected, dtype) + assert_close(layer.history_after(conv_inputs, starts_for(lengths), history), expected_state, dtype) + torch.testing.assert_close(history, original, atol=0, rtol=0) + + +def test_production_hidden_width(): + torch.manual_seed(137) + layer = make_layer(dim=2560, branches=4, embed_dim=2560) + residual = torch.randn(2, 10240, device="cuda", dtype=torch.bfloat16) + embeddings = torch.randn(2, 2560, device="cuda", dtype=torch.bfloat16) + history = layer.empty_history(2) + output, inputs = layer.forward(residual, embeddings, starts_for([1, 1]), history) + gated, expected_inputs = reference_values(layer, residual, embeddings) + expected, _ = reference_conv(layer, expected_inputs, gated, [1, 1], history) + assert_close(inputs, expected_inputs, residual.dtype) + assert_close(output, expected, residual.dtype) + + +def test_chunking_speculative_acceptance_and_request_isolation(): + torch.manual_seed(137) + layer = make_layer(dim=32, branches=2, dtype=torch.float32) + lengths = [8, 6, 2] + starts = starts_for(lengths) + residual = torch.randn(sum(lengths), 64, device="cuda") + embeddings = torch.randn(sum(lengths), 32, device="cuda") + history = layer.empty_history(3).normal_() + full, full_inputs = layer.forward(residual, embeddings, starts, history) + full_state = layer.history_after(full_inputs, starts, history) + state = history.clone() + consumed = [0] * 3 + per_request = [[] for _ in range(3)] + for counts in [(1, 0, 1), (2, 3, 1), (1, 1, 0), (4, 2, 0)]: + slices = [slice(int(starts[i]) + consumed[i], int(starts[i]) + consumed[i] + n) for i, n in enumerate(counts)] + chunks = [torch.cat([tensor[s] for s in slices]) for tensor in [residual, embeddings]] + chunk_starts = starts_for(counts) + output, conv_inputs = layer.forward(*chunks, chunk_starts, state) + state = layer.history_after(conv_inputs, chunk_starts, state) + for i, n in enumerate(counts): + per_request[i].append(output[int(chunk_starts[i]) : int(chunk_starts[i + 1])]) + consumed[i] += n + assert_close(torch.cat([torch.cat(parts) for parts in per_request]), full, residual.dtype) + assert_close(state, full_state, residual.dtype) + accepted = torch.tensor([1, 4, 0], device="cuda", dtype=torch.int32) + committed = layer.history_after(full_inputs, starts, history, accepted) + expected = torch.stack( + [ + torch.cat([history[i], full_inputs[int(starts[i]) : int(starts[i]) + n]])[-layer.context_len :] + for i, n in enumerate([1, 4, 0]) + ] + ) + torch.testing.assert_close(committed, expected, atol=0, rtol=0) + + order = torch.tensor([2, 0, 2], device="cuda") + reordered = committed.index_select(0, order) + new_residual, new_embeddings = residual[:3], embeddings[:3] + actual, _ = layer.forward(new_residual, new_embeddings, starts_for([1, 1, 1]), reordered) + gated, normalized = reference_values(layer, new_residual, new_embeddings) + expected, _ = reference_conv(layer, normalized, gated, [1, 1, 1], reordered) + assert_close(actual, expected, residual.dtype) + + +def test_checkpoint_binding_and_graph_replay(): + layer = make_layer(dim=32, branches=2) + prefix = "model.language_model.layers.1.ple" + weights = { + f"{prefix}.{name}": tensor + for name, tensor in [ + ("key_proj.weight", layer.key_weight), + ("value_proj.weight", layer.value_weight), + ("norm_key.weight", layer.norm_key), + ("norm_query.weight", layer.norm_query), + ("norm_conv.weight", layer.norm_conv), + ("conv1d.weight", layer.conv_weight.unsqueeze(1)), + ] + } + bound = PLELayer.from_checkpoint(weights, prefix, branches=2, dilation=3) + assert bound.key_weight is layer.key_weight + assert bound.conv_weight.data_ptr() == layer.conv_weight.data_ptr() + residual = torch.randn(7, 64, device="cuda", dtype=torch.bfloat16) + embeddings = torch.randn(7, 32, device="cuda", dtype=torch.bfloat16) + starts = starts_for([3, 2, 2]) + accepted = torch.tensor([3, 1, 2], device="cuda", dtype=torch.int32) + mask = torch.ones(7, device="cuda", dtype=torch.bool) + history = layer.empty_history(3) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + _, conv_inputs = bound.forward(residual, embeddings, starts, history, mask) + bound.history_after(conv_inputs, starts, history, accepted) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output, conv_inputs = bound.forward(residual, embeddings, starts, history, mask) + next_state = bound.history_after(conv_inputs, starts, history, accepted) + for lengths, counts in [([2, 0, 5], [1, 0, 4]), ([0, 3, 4], [0, 2, 0])]: + residual.normal_() + embeddings.normal_() + history.normal_() + starts.copy_(starts_for(lengths)) + accepted.copy_(torch.tensor(counts, device="cuda", dtype=torch.int32)) + mask[::2] = False + graph.replay() + gated, expected_inputs = reference_values(bound, residual, embeddings, mask) + expected, _ = reference_conv(bound, expected_inputs, gated, lengths, history) + assert_close(output, expected, residual.dtype) + expected_state = [] + for i, count in enumerate(counts): + expected_state.append( + torch.cat([history[i], expected_inputs[int(starts[i]) : int(starts[i]) + count]])[-bound.context_len :] + ) + assert_close(next_state, torch.stack(expected_state), residual.dtype) + + +def test_empty_batches_and_invalid_inputs(): + layer = make_layer(dim=32, branches=2) + for lengths in [[], [0, 0, 0]]: + residual = layer.key_weight.new_empty((0, 64)) + embeddings = layer.key_weight.new_empty((0, 32)) + history = layer.empty_history(len(lengths)) + output, conv_inputs = layer.forward(residual, embeddings, starts_for(lengths), history) + assert output.shape == (0, 64) + torch.testing.assert_close( + layer.history_after(conv_inputs, starts_for(lengths), history), history, atol=0, rtol=0 + ) + with pytest.raises(ValueError, match="projection"): + replace(layer, value_weight=layer.value_weight[:1]) + with pytest.raises(ValueError, match="convolution"): + replace(layer, conv_weight=layer.conv_weight[:, :0]) + residual = layer.key_weight.new_empty((1, 64)) + embeddings = layer.key_weight.new_empty((1, 32)) + with pytest.raises(ValueError, match="history"): + layer.forward(residual, embeddings, starts_for([1]), layer.empty_history(2)) + with pytest.raises(ValueError, match="mask"): + layer.forward(residual, embeddings, starts_for([1]), layer.empty_history(1), torch.ones(1, device="cuda")) + with pytest.raises(ValueError, match="accepted"): + layer.history_after(residual, starts_for([1]), layer.empty_history(1), torch.ones(1, device="cuda")) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +def test_gate_zero_near_zero_and_signed_saturation(dtype): + from lightllm.common.ngram_embedding.ple_kernels import _gate + + dots = torch.tensor([0, 1e-8, -1e-8, 1e-5, -1e-5, 100, -100], device="cuda", dtype=dtype) + key = torch.zeros(7, 4, device="cuda", dtype=dtype) + query = torch.zeros_like(key) + key[:, 0] = dots * 2 + query[:, 0] = 1 + value = torch.tensor([[-2, -1, 1, 2]], device="cuda", dtype=dtype).expand(7, -1).contiguous() + gate = (key * query).sum(-1, keepdim=True) / 2 + gate = gate.abs().clamp_min(1e-6).sqrt() * gate.sign() + expected = torch.sigmoid(gate) * value + actual = torch.empty_like(key) + _gate[(7, 1)](key, query, value, None, actual, 4, 1, False, 4, enable_fp_fusion=False) + if dtype == torch.float32: + torch.testing.assert_close(actual, expected, atol=5e-7, rtol=5e-7) + else: + torch.testing.assert_close(actual, expected, atol=0, rtol=0) diff --git a/unit_tests/common/ngram_embedding/test_row_sharded_embedding.py b/unit_tests/common/ngram_embedding/test_row_sharded_embedding.py new file mode 100644 index 0000000000..cacb940aa4 --- /dev/null +++ b/unit_tests/common/ngram_embedding/test_row_sharded_embedding.py @@ -0,0 +1,161 @@ +from types import SimpleNamespace +import importlib + +import pytest +import torch + +from lightllm.common.basemodel.layer_weights.meta_weights.row_sharded_embedding_weight import RowShardedEmbeddingWeight + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +PREFIX = "model.language_model.layers.1.ple.ple_embedding.ngram_embedding" + + +@pytest.fixture(autouse=True) +def runtime_configuration(monkeypatch): + base = importlib.import_module("lightllm.common.basemodel.layer_weights.meta_weights.base_weight") + platform = importlib.import_module("lightllm.common.basemodel.layer_weights.meta_weights.platform_op") + monkeypatch.setattr(base, "get_current_device_id", lambda: torch.cuda.current_device()) + monkeypatch.setattr( + platform, + "get_env_start_args", + lambda: SimpleNamespace( + hardware_platform="cuda", + enable_torch_fallback=False, + enable_triton_fallback=False, + ), + ) + + +def make_weight(rank=0, world=1): + return RowShardedEmbeddingWeight(5, 24, f"{PREFIX}.weight", torch.bfloat16, 4, rank, world) + + +def shards(table): + return [{f"{PREFIX}.shard_{i}.weight": part} for i, part in enumerate(table.chunk(4))] + + +@pytest.mark.parametrize("world", [1, 2, 3, 8]) +def test_streamed_shards_and_vocab_parallel_lookup(world): + table = torch.arange(120, dtype=torch.bfloat16).view(24, 5) + checkpoint = shards(table) + ids = torch.tensor([0, 5, 6, 7, 8, 11, 12, 15, 17, 18, 23], device="cuda", dtype=torch.int64) + outputs = [] + for rank in range(world): + weight = make_weight(rank, world) + assert not weight.verify_load() + required, loaded = set(), set() + for i in range(4): + if max(i * 6, weight.tp_vocab_start_id) < min((i + 1) * 6, weight.tp_vocab_end_id): + required.add(i) + for index in [2, 0, 3, 1]: + weight.load_hf_weights(checkpoint[index]) + loaded.add(index) + assert weight.verify_load() == (loaded >= required) + torch.testing.assert_close( + weight.weight.cpu(), table[weight.tp_vocab_start_id : weight.tp_vocab_end_id], atol=0, rtol=0 + ) + outputs.append(weight(ids)) + + torch.testing.assert_close( + torch.stack(outputs).sum(0), torch.nn.functional.embedding(ids, table.cuda()), atol=0, rtol=0 + ) + + +def test_missing_shards_bad_shapes_and_other_prefixes(): + weight = make_weight() + table = torch.arange(120, dtype=torch.bfloat16).view(24, 5) + checkpoint = shards(table) + for index in [0, 1, 3]: + weight.load_hf_weights(checkpoint[index]) + assert not weight.verify_load() + weight.load_hf_weights({"other.shard_2.weight": table[:6]}) + assert not weight.verify_load() + with pytest.raises(ValueError, match="shard 2"): + weight.load_hf_weights({f"{PREFIX}.shard_2.weight": table[:5]}) + assert not weight.verify_load() + weight.load_hf_weights(checkpoint[2]) + assert weight.verify_load() + + weight.load_hf_weights({"other.weight": table}) + assert weight.verify_load() + + +def test_full_checkpoint_compatibility_and_reject_mixed_formats(): + table = torch.arange(120, dtype=torch.bfloat16).view(24, 5) + weight = make_weight(1, 3) + weight.load_hf_weights({f"{PREFIX}.weight": table}) + assert weight.verify_load() + weight.load_hf_weights({"other.weight": table}) + assert weight.verify_load() + torch.testing.assert_close(weight.weight.cpu(), table[8:16], atol=0, rtol=0) + with pytest.raises(ValueError, match="mix"): + weight.load_hf_weights(shards(table)[0]) + partial = make_weight() + partial.load_hf_weights(shards(table)[1]) + with pytest.raises(ValueError, match="mix"): + partial.load_hf_weights({f"{PREFIX}.weight": table}) + with pytest.raises(ValueError, match="dimensions"): + make_weight().load_hf_weights({f"{PREFIX}.weight": table[:5]}) + + +def test_row_shard_configuration_and_graph_replay(): + with pytest.raises(ValueError, match="divisible"): + RowShardedEmbeddingWeight(5, 23, f"{PREFIX}.weight", torch.bfloat16, 4, 0, 1) + table = torch.arange(120, dtype=torch.bfloat16).view(24, 5) + weight = make_weight() + for part in shards(table): + weight.load_hf_weights(part) + ids = torch.tensor([0, 6, 12, 23], device="cuda", dtype=torch.int64) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + weight(ids) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = weight(ids) + for values in [[23, 12, 6, 0], [1, 7, 13, 19]]: + ids.copy_(torch.tensor(values, device="cuda")) + graph.replay() + torch.testing.assert_close(output.cpu(), table[values], atol=0, rtol=0) + + +def test_packed_hash_lookup_and_ple_composition(): + from lightllm.common.ngram_embedding import NGramHasher + from lightllm.common.ngram_embedding.ple import PLELayer + + torch.manual_seed(137) + sizes = torch.tensor([2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53], device="cuda", dtype=torch.int64) + offsets = torch.cat([sizes.new_zeros(1), sizes.cumsum(0)[:-1]]) + hasher = NGramHasher( + torch.tensor([23703573157769, 20109073645365, 8052911324071], device="cuda"), sizes, offsets, 248044 + ) + tokens = torch.tensor([1, 2, 248044, 3, 4, 5, 6, 7], device="cuda", dtype=torch.int64) + starts = torch.tensor([0, 3, 4, 4, 8], device="cuda", dtype=torch.int32) + indices, _ = hasher.hash(tokens, starts, hasher.empty_history(4)) + + table = torch.randn(512, 4, dtype=torch.bfloat16) + partition_outputs = [] + for rank in range(3): + weight = RowShardedEmbeddingWeight(4, 512, f"{PREFIX}.weight", torch.bfloat16, 128, rank, 3) + for index, shard in reversed(list(enumerate(table.chunk(128)))): + weight.load_hf_weights({f"{PREFIX}.shard_{index}.weight": shard}) + assert weight.verify_load() + partition_outputs.append(weight(indices.flatten())) + embeddings = torch.stack(partition_outputs).sum(0).view(tokens.numel(), 64) + expected_embeddings = torch.nn.functional.embedding(indices, table.cuda()).flatten(-2) + torch.testing.assert_close(embeddings, expected_embeddings, atol=0, rtol=0) + + def rand(*shape): + return torch.randn(*shape, device="cuda", dtype=torch.bfloat16) + + layer = PLELayer( + rand(64, 64) / 8, rand(32, 64) / 8, rand(64) * 0.1, rand(64) * 0.1, rand(64) * 0.1, rand(64, 4) * 0.1, 2, 3 + ) + history = layer.empty_history(4) + residual = rand(tokens.numel(), 64) + output, conv_inputs = layer.forward(residual, embeddings, starts, history) + expected_output, expected_inputs = layer.forward(residual, expected_embeddings, starts, history) + torch.testing.assert_close(output, expected_output, atol=0, rtol=0) + torch.testing.assert_close(conv_inputs, expected_inputs, atol=0, rtol=0) diff --git a/unit_tests/common/qsa/test_attention.py b/unit_tests/common/qsa/test_attention.py new file mode 100644 index 0000000000..c3cee0766b --- /dev/null +++ b/unit_tests/common/qsa/test_attention.py @@ -0,0 +1,194 @@ +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), (4, 8, 8, 64, 64, 89)], +) +@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)) + + +@pytest.mark.parametrize("splits", [1, 4]) +def test_fp8_cache_matches_dequantized_sdpa(splits): + torch.manual_seed(1437) + q = torch.randn(4, 6, 128, device="cuda", dtype=torch.bfloat16) + k = torch.randn(256, 2, 128, device="cuda", dtype=q.dtype).to(torch.float8_e4m3fn) + v = torch.randn(256, 2, 64, device="cuda", dtype=q.dtype).to(torch.float8_e4m3fn) + ids = torch.stack([torch.randperm(256, device="cuda")[:137] for _ in range(4)]).int() + ids[:, ::7] = -1 + actual = indexed_attention(q, k, v, ids, splits=splits) + expected = reference(q, k.to(q.dtype), v.to(q.dtype), ids) + torch.testing.assert_close(actual, expected, rtol=0.016, atol=0.016) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = indexed_attention(q, k, v, ids, splits=splits) + q.mul_(0.5) + graph.replay() + torch.testing.assert_close(captured, reference(q, k.to(q.dtype), v.to(q.dtype), ids), rtol=0.016, atol=0.016) + + +@pytest.mark.parametrize("rows,block,warps,splits", [(8, 128, 4, 16), (32, 128, 4, 16), (1024, 32, 1, 1)]) +def test_fp8_tiles_match_fp32_reference(rows, block, warps, splits): + torch.manual_seed(9217) + q = torch.randn(rows, 3, 256, device="cuda", dtype=torch.bfloat16) * 0.4 + kv = torch.randn(4096, 2, 256, device="cuda").to(torch.float8_e4m3fn) + k, v = kv[:, :1], kv[:, 1:] + indices = torch.randint(4096, (rows, 2051), device="cuda", dtype=torch.int32) + lengths = torch.arange(rows, device="cuda", dtype=torch.int32).mul_(331).remainder_(2052) + + def run(): + return indexed_attention( + q, + k, + v, + indices, + splits=splits, + selection_lengths=lengths, + short_rows_direct=True, + skip_empty_splits=True, + block_size=block, + num_stages=2, + num_warps=warps, + ) + + def check(output): + for first in range(0, rows, 32): + ids = indices[first : first + 32].clone() + ids.masked_fill_(torch.arange(2051, device="cuda")[None, :] >= lengths[first : first + 32, None], -1) + expected = reference(q[first : first + 32].float(), k.float(), v.float(), ids) + torch.testing.assert_close(output[first : first + 32].float(), expected, atol=0.003, rtol=0.008) + + check(run()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = run() + lengths.add_(257).remainder_(2052) + q.mul_(1.5) + graph.replay() + check(captured) + + +def test_native_interleaved_kv_views_match_contiguous_cache(): + q = torch.randn(3, 24, 256, device="cuda", dtype=torch.bfloat16) + packed = torch.randn(97, 4, 256, device="cuda", dtype=q.dtype) + k, v = packed[:, :2], packed[:, 2:] + assert not k.is_contiguous() and not v.is_contiguous() + indices = torch.stack([torch.randperm(97, device="cuda")[:79] for _ in range(3)]).int() + indices[:, ::7] = -1 + actual = indexed_attention(q, k, v, indices) + expected = indexed_attention(q, k.contiguous(), v.contiguous(), indices) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +def test_direct_short_rows_preserve_fixed_partition_and_graph_replay(dtype): + torch.manual_seed(701) + rows, slots, width = 5, 4096, 2051 + q = torch.randn(rows, 6, 256, device="cuda", dtype=dtype) + k = torch.randn(slots, 1, 256, device="cuda", dtype=dtype) + v = torch.randn_like(k) + ids = torch.stack([torch.randperm(slots, device="cuda")[:width] for _ in range(rows)]).int() + lengths = torch.tensor([0, 16, 64, 65, 2048], device="cuda", dtype=torch.int32) + + def compute(direct): + return indexed_attention(q, k, v, ids, splits=32, selection_lengths=lengths, short_rows_direct=direct) + + torch.testing.assert_close(compute(True), compute(False), atol=0, rtol=0) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + compute(True) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = compute(True) + for bounds in ([2048, 65, 64, 0, 16], [64, 2048, 0, 65, 63]): + lengths.copy_(torch.tensor(bounds, device="cuda", dtype=torch.int32)) + q.normal_() + graph.replay() + torch.testing.assert_close(actual, compute(False), atol=0, rtol=0) diff --git a/unit_tests/common/qsa/test_flashinfer_topk.py b/unit_tests/common/qsa/test_flashinfer_topk.py new file mode 100644 index 0000000000..f3feb6a4b0 --- /dev/null +++ b/unit_tests/common/qsa/test_flashinfer_topk.py @@ -0,0 +1,88 @@ +from dataclasses import replace + +import pytest +import torch + +from lightllm.common.qsa.indexer import QSAIndexer +from lightllm.common.qsa.topk import flashinfer_topk_indices, flashinfer_topk_supported + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +@pytest.fixture(autouse=True) +def require_ragged_topk(): + if not flashinfer_topk_supported(torch.device("cuda", 0), 65542, 512): + pytest.skip("FlashInfer graph-safe TopK on Hopper is required") + + +def test_ragged_graph_uses_visible_lengths_and_chunk_offset(): + rows, columns, k, first = 17, 65542, 512, 5 + scores = torch.randn(rows, columns, device="cuda") + visible = torch.zeros(first + rows + 3, dtype=torch.int32, device="cuda") + lengths = visible[first : first + rows] + limits = torch.tensor([0, 3, 2047, 2048, 2051, 65536, 262144], dtype=torch.int32, device="cuda") + cols = torch.arange(columns, device="cuda") + + def select(): + return flashinfer_topk_indices(scores, k, visible_lengths=visible, first=first) + + select() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + selected = select() + for step in range(12): + lengths.copy_(limits[(torch.arange(rows, device="cuda") + step) % limits.numel()]) + scores.normal_() + if step % 3 == 0: + scores.round_() + scores.masked_fill_(cols[None, :] >= lengths[:, None] // 4, 1e9) + scores.masked_fill_((lengths <= k * 4)[:, None], float("nan")) + graph.replay() + for row, complete in enumerate((lengths // 4).tolist()): + actual = selected[row] + if complete <= k: + expected = torch.arange(k, dtype=torch.int32, device="cuda") + expected[expected >= complete] = -1 + assert torch.equal(actual, expected) + else: + assert bool(((actual >= 0) & (actual < complete)).all()) + assert torch.unique(actual).numel() == k + expected = torch.topk(scores[row, :complete], k).values + values = scores[row, actual.long()].sort(descending=True).values + torch.testing.assert_close(values, expected, atol=0, rtol=0) + ranked = torch.argsort(scores[row, :complete], descending=True, stable=True)[:k] + assert torch.equal(actual.sort().values.long(), ranked.sort().values) + expected_order = selected.clone() + graph.replay() + torch.testing.assert_close(selected, expected_order, atol=0, rtol=0) + + +@pytest.mark.parametrize("chunk_size", [3, 64]) +def test_indexer_ragged_topk_preserves_slots_across_graph_transitions(chunk_size): + rows, capacity = 17, 262169 + norm = torch.ones(128, dtype=torch.bfloat16, device="cuda") + original = QSAIndexer(torch.randn(640, 16, dtype=torch.bfloat16, device="cuda"), norm, norm) + optimized = replace(original, use_cuda_topk=True) + cache = torch.zeros(capacity + 3, 128, dtype=torch.bfloat16, device="cuda") + bits = torch.arange(17, device="cuda") + cache[:, :17] = ((torch.arange(cache.shape[0], device="cuda")[:, None] // 4 >> bits) & 1).to(cache.dtype) + query = torch.zeros(rows, 4, 128, dtype=cache.dtype, device="cuda") + query[:, :, :17] = (2.0 ** bits).to(cache.dtype) + table = torch.arange(capacity, dtype=torch.int32, device="cuda").expand(2, -1).contiguous() + requests = torch.arange(rows, dtype=torch.int32, device="cuda") // 4 % 2 + lengths = torch.full((rows,), 262144, dtype=torch.int32, device="cuda") + data = query, cache, requests, lengths, table + original.select(*data, query_chunk_size=chunk_size) + optimized.select(*data, query_chunk_size=chunk_size) + before, after = torch.cuda.CUDAGraph(), torch.cuda.CUDAGraph() + with torch.cuda.graph(before): + expected = original.select(*data, query_chunk_size=chunk_size) + with torch.cuda.graph(after): + actual = optimized.select(*data, query_chunk_size=chunk_size) + limits = torch.tensor([0, 3, 2047, 2048, 2051, 65536, 262144], dtype=torch.int32, device="cuda") + for step in range(12): + lengths.copy_(limits[(torch.arange(rows, device="cuda") + step) % limits.numel()]) + requests.copy_((torch.arange(rows, device="cuda") + step) % 2) + before.replay() + after.replay() + assert torch.equal(actual.sort(dim=-1).values, expected.sort(dim=-1).values) diff --git a/unit_tests/common/qsa/test_indexer.py b/unit_tests/common/qsa/test_indexer.py new file mode 100644 index 0000000000..6e99d783a0 --- /dev/null +++ b/unit_tests/common/qsa/test_indexer.py @@ -0,0 +1,494 @@ +import math + +import pytest +import torch +import torch.nn.functional as F +import triton + +from lightllm.common.qsa.indexer import QSAIndexer +from lightllm.common.qsa.indexer_kernels import _scores + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float8_e4m3fn]) +@pytest.mark.parametrize("minimum", [0, 2]) +def test_grouped_scores_keep_chunk_offsets_short_rows_and_mixed_requests(dtype, minimum): + query = torch.randn(14, 4, 128, device="cuda").to(dtype) + cache = torch.randn(204, 128, device="cuda").to(dtype) + table = torch.randperm(204, device="cuda").int().view(3, 68) + requests = torch.tensor([0] * 9 + [1, 1, 2, 2, 2], device="cuda", dtype=torch.int32) + lengths = torch.tensor([0] * 5 + [0, 3, 8, 16, 7, 9, 12, 17, 11], device="cuda", dtype=torch.int32) + outputs = [torch.full((9, 17), 42.0, device="cuda") for _ in range(2)] + + def run(): + for group, output in zip((1, 4), outputs): + block = 64 if group == 1 else 128 + _scores[(triton.cdiv(9, group), 1)]( + query, + cache, + requests, + lengths, + table, + output, + 5, + 4, + 128, + 17, + 68, + True, + 4, + 16, + 128, + block, + minimum, + "ieee" if dtype == torch.float32 else "tf32", + GROUP_SIZE=group, + TOKENS=9, + ) + + run() + torch.testing.assert_close(outputs[1], outputs[0], atol=1e-5, rtol=1e-5) + if minimum: + assert torch.all(outputs[1][:2] == 42) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + lengths.zero_() + query.copy_(torch.randn(query.shape, device="cuda").to(dtype)) + graph.replay() + torch.testing.assert_close(outputs[1], outputs[0], atol=1e-5, rtol=1e-5) + + +def make_indexer(dtype=torch.bfloat16, dim=32, heads=4, budget=8, ratio=4): + projection = torch.randn((heads + 1) * dim, 32, device="cuda", dtype=dtype) / math.sqrt(32) + return QSAIndexer( + projection, + torch.randn(dim, device="cuda", dtype=dtype) * 0.1, + torch.randn(dim, device="cuda", dtype=dtype) * 0.1, + heads, + budget, + ratio, + ) + + +def rotary(rows, width, dtype): + angles = torch.randn(rows, width // 2, device="cuda").repeat(1, 2) + return angles.cos().to(dtype), angles.sin().to(dtype) + + +def norm(x, weight, eps): + normalized = x.float() * torch.rsqrt(x.float().square().mean(-1, keepdim=True) + eps) + return (normalized * (1 + weight.float())).to(x.dtype) + + +def rotate(x, cos, sin): + width = cos.shape[-1] + first, second = x[..., :width].chunk(2, -1) + return torch.cat([x[..., :width] * cos + torch.cat([-second, first], -1) * sin, x[..., width:]], -1) + + +def reference_cache(indexer, raw, table, lengths, cos, sin): + result = indexer.empty_cache(table.shape[0], table.shape[1]).fill_(float("nan")) + for row, length in enumerate(lengths): + count = length // indexer.compress_ratio + if not count: + continue + slots = table[row, : count * indexer.compress_ratio].view(count, indexer.compress_ratio).long() + pooled = raw[slots].float().mean(1).to(raw.dtype) + normalized = norm(pooled, indexer.key_norm, indexer.eps) + result[row, :count] = rotate(normalized, cos[slots[:, 0]], sin[slots[:, 0]]) + return result + + +def reference_selection(indexer, query, cache, requests, lengths, table): + rows = [] + for row, (request, length) in enumerate(zip(requests.tolist(), lengths.tolist())): + count = length // indexer.compress_ratio + if count: + scores = (query[row].float() @ cache[request, :count].float().T).relu().sum(0) / math.sqrt(indexer.head_dim) + blocks = scores.topk(min(count, indexer.token_budget // indexer.compress_ratio)).indices + positions = ( + blocks[:, None] * indexer.compress_ratio + torch.arange(indexer.compress_ratio, device=query.device) + ).flatten() + else: + positions = torch.empty(0, device=query.device, dtype=torch.int64) + positions = torch.cat([positions, torch.arange(count * indexer.compress_ratio, length, device=query.device)]) + selected = table[request, positions] + rows.append( + F.pad(selected, (0, indexer.token_budget + indexer.compress_ratio - 1 - selected.numel()), value=-1) + ) + return torch.stack(rows) if rows else table.new_empty((0, indexer.token_budget + indexer.compress_ratio - 1)) + + +def assert_sets(actual, expected): + torch.testing.assert_close(actual.sort(-1).values, expected.sort(-1).values, atol=0, rtol=0) + + +def test_fp8_pool_and_selection(): + torch.manual_seed(1537) + indexer = make_indexer(dim=128) + raw = torch.randn(128, 128, device="cuda", dtype=torch.bfloat16) + cos, sin = rotary(128, 64, raw.dtype) + table = torch.randperm(128, device="cuda").int().view(2, 64) + lengths = torch.tensor([64, 52], device="cuda", dtype=torch.int32) + first = torch.zeros(2, device="cuda", dtype=torch.int32) + cache = torch.zeros(2, 16, 128, device="cuda", dtype=torch.float8_e4m3fn) + expected_cache = reference_cache(indexer, raw, table, lengths, cos, sin).to(cache.dtype) + indexer.pool(raw, table, lengths, cos, sin, cache, first, 16) + for row, length in enumerate(lengths.tolist()): + torch.testing.assert_close( + cache[row, : length // 4].float(), expected_cache[row, : length // 4].float(), rtol=0, atol=0 + ) + query = torch.randn(4, 4, 128, device="cuda").to(cache.dtype) + requests = torch.tensor([0, 0, 1, 1], device="cuda", dtype=torch.int32) + visible = torch.tensor([64, 41, 52, 3], device="cuda", dtype=torch.int32) + expected = reference_selection(indexer, query, cache, requests, visible, table) + assert_sets(indexer.select(query, cache, requests, visible, table), expected) + + +@pytest.mark.parametrize("physical", [False, True]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float8_e4m3fn]) +@pytest.mark.parametrize("maximum", [0, 3, 8, 41, 64]) +def test_prefill_score_bound_preserves_selection_with_short_and_long_rows(physical, dtype, maximum): + torch.manual_seed(984) + indexer = make_indexer(dim=128) + query = torch.randn(9, 4, 128, device="cuda").to(dtype) + cache = torch.randn((128, 128) if physical else (2, 16, 128), device="cuda").to(dtype) + table = torch.randperm(128, device="cuda").int().view(2, 64) + requests = torch.arange(9, device="cuda", dtype=torch.int32) % 2 + lengths = torch.tensor([0, 1, 3, 4, 8, 9, 13, 41, 64], device="cuda", dtype=torch.int32).clamp_max(maximum) + expected = indexer.select(query, cache, requests, lengths, table, query_chunk_size=3) + bounded = indexer.select(query, cache, requests, lengths, table, query_chunk_size=3, max_visible_length=maximum) + assert_sets(bounded, expected) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = indexer.select(query, cache, requests, lengths, table, query_chunk_size=3, max_visible_length=maximum) + lengths.div_(2, rounding_mode="floor") + graph.replay() + assert_sets(actual, indexer.select(query, cache, requests, lengths, table, query_chunk_size=3)) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("dim,rotary_dim,heads,ratio", [(32, 16, 4, 4), (65, 32, 3, 3), (128, 64, 4, 4)]) +def test_projection_pooling_and_causal_selection(dtype, dim, rotary_dim, heads, ratio): + torch.manual_seed(337) + indexer = make_indexer(dtype, dim, heads, budget=2 * ratio, ratio=ratio) + hidden = torch.randn(10, 32, device="cuda", dtype=dtype) + qcos, qsin = rotary(10, rotary_dim, dtype) + query, raw_new = indexer.project(hidden, qcos, qsin) + projected = F.linear(hidden, indexer.projection) + expected_q = projected[:, : heads * dim].reshape(10, heads, dim) + expected_q = rotate(norm(expected_q, indexer.query_norm, indexer.eps), qcos[:, None], qsin[:, None]) + tolerance = max(2e-6, torch.finfo(dtype).eps * 2) + torch.testing.assert_close(query, expected_q, atol=tolerance, rtol=tolerance) + torch.testing.assert_close(raw_new, projected[:, heads * dim :], atol=0, rtol=0) + assert not raw_new.is_contiguous() + raw_copy = torch.empty_like(raw_new) + raw_copy.index_copy_(0, torch.arange(raw_new.shape[0], device="cuda"), raw_new) + torch.testing.assert_close(raw_copy, raw_new.contiguous(), atol=0, rtol=0) + raw = torch.randn(100, dim, device="cuda", dtype=dtype) + cos, sin = rotary(100, rotary_dim, dtype) + table = torch.randperm(100, device="cuda")[:87].view(3, 29).int() + lengths = [23, 17, 0] + cache = indexer.empty_cache(3, 29).fill_(float("nan")) + raw_before = raw.clone() + indexer.pool( + raw, + table, + torch.tensor(lengths, device="cuda", dtype=torch.int32), + cos, + sin, + cache, + torch.zeros(3, device="cuda", dtype=torch.int32), + 29 // ratio, + ) + expected_cache = reference_cache(indexer, raw, table, lengths, cos, sin) + torch.testing.assert_close(cache, expected_cache, atol=tolerance, rtol=tolerance, equal_nan=True) + torch.testing.assert_close(raw, raw_before, atol=0, rtol=0) + requests = torch.tensor([0, 0, 1, 0, 1, 1, 0, 1, 1, 0], device="cuda", dtype=torch.int32) + visible = torch.tensor([0, 1, 3, 4, 5, 8, 9, 15, 17, 23], device="cuda", dtype=torch.int32) + actual = indexer.select(query, cache, requests, visible, table, query_chunk_size=3) + + expected = reference_selection(indexer, query, cache, requests, visible, table) + assert_sets(actual, expected) + for row, (request, length) in enumerate(zip(requests.tolist(), visible.tolist())): + allowed = set(table[request, :length].tolist()) + assert set(actual[row][actual[row] >= 0].tolist()) <= allowed + + +def test_incremental_pool_cache_and_cuda_graph_replay(): + torch.manual_seed(337) + indexer = make_indexer() + raw = torch.randn(48, 32, device="cuda", dtype=torch.bfloat16) + cos, sin = rotary(48, 16, raw.dtype) + table = torch.randperm(48, device="cuda")[:40].view(2, 20).int() + cache = indexer.empty_cache(2, 20).fill_(float("nan")) + lengths = torch.tensor([7, 5], device="cuda", dtype=torch.int32) + first = torch.zeros_like(lengths) + indexer.pool(raw, table, lengths, cos, sin, cache, first, 2) + original = cache[:, :1].clone() + hidden = torch.randn(2, 32, device="cuda", dtype=raw.dtype) + qcos, qsin = rotary(2, 16, raw.dtype) + requests = torch.tensor([0, 1], device="cuda", dtype=torch.int32) + first.fill_(1) + lengths.copy_(torch.tensor([13, 10], device="cuda", dtype=torch.int32)) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + query, _ = indexer.project(hidden, qcos, qsin) + indexer.pool(raw, table, lengths, cos, sin, cache, first, 2) + indexer.select(query, cache, requests, lengths, table, query_chunk_size=1) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + query, _ = indexer.project(hidden, qcos, qsin) + indexer.pool(raw, table, lengths, cos, sin, cache, first, 2) + selected = indexer.select(query, cache, requests, lengths, table, query_chunk_size=1) + for counts, starts in [([13, 10], [1, 1]), ([18, 15], [3, 2])]: + hidden.normal_() + lengths.copy_(torch.tensor(counts, device="cuda", dtype=torch.int32)) + first.copy_(torch.tensor(starts, device="cuda", dtype=torch.int32)) + graph.replay() + expected_cache = reference_cache(indexer, raw, table, counts, cos, sin) + tolerance = 2 * torch.finfo(raw.dtype).eps + torch.testing.assert_close(cache, expected_cache, atol=tolerance, rtol=tolerance, equal_nan=True) + assert_sets(selected, reference_selection(indexer, query, cache, requests, lengths, table)) + torch.testing.assert_close(cache[:, :1], original, atol=0, rtol=0) + + +def test_zero_score_ties_tail_only_empty_and_checkpoint_binding(): + indexer = make_indexer() + prefix = "model.language_model.layers.3.self_attn.indexer" + weights = { + f"{prefix}.index_qk_proj.weight": indexer.projection, + f"{prefix}.q_layernorm.weight": indexer.query_norm, + f"{prefix}.k_layernorm.weight": indexer.key_norm, + } + bound = QSAIndexer.from_checkpoint(weights, prefix, token_budget=8) + assert bound.projection is indexer.projection + query = torch.zeros(3, 4, 32, device="cuda", dtype=torch.bfloat16) + table = torch.arange(29, device="cuda", dtype=torch.int32).view(1, -1) + cache = torch.randn(1, 7, 32, device="cuda", dtype=query.dtype) + requests = torch.zeros(3, device="cuda", dtype=torch.int32) + lengths = torch.tensor([0, 3, 21], device="cuda", dtype=torch.int32) + selected = bound.select(query, cache, requests, lengths, table) + assert torch.all(selected[0] == -1) + assert set(selected[1][selected[1] >= 0].tolist()) == {0, 1, 2} + chosen = selected[2][selected[2] >= 0].tolist() + assert len(chosen) == len(set(chosen)) == 9 + assert chosen[-1] == 20 + for offset in [0, 4]: + block = chosen[offset : offset + 4] + assert block == list(range(block[0], block[0] + 4)) and block[0] % 4 == 0 + empty_query, raw = bound.project(query.new_empty((0, 32)), query.new_empty((0, 0)), query.new_empty((0, 0))) + assert raw.shape == (0, 32) + assert bound.select(empty_query, cache, requests[:0], lengths[:0], table).shape == (0, 11) + + short = bound.select(query[:1], cache[:, :0].contiguous(), requests[:1], lengths[1:2], table[:, :3].contiguous()) + assert set(short[0][short[0] >= 0].tolist()) == {0, 1, 2} + + +def test_invalid_layouts(): + indexer = make_indexer() + with pytest.raises(ValueError, match="budget"): + QSAIndexer(indexer.projection, indexer.query_norm, indexer.key_norm, token_budget=3) + with pytest.raises(ValueError, match="rotary"): + indexer.project( + indexer.projection.new_empty((2, 32)), + indexer.projection.new_empty((2, 3)), + indexer.projection.new_empty((2, 3)), + ) + with pytest.raises(ValueError, match="int32"): + indexer.select( + indexer.projection.new_empty((2, 4, 32)), + indexer.empty_cache(1, 4), + torch.zeros(2, device="cuda", dtype=torch.int64), + torch.ones(2, device="cuda", dtype=torch.int32), + torch.zeros(1, 4, device="cuda", dtype=torch.int32), + ) + + +def test_rejected_suffix_repool_overwrites_candidate_blocks(): + torch.manual_seed(337) + indexer = make_indexer() + table = torch.arange(20, device="cuda", dtype=torch.int32).view(1, -1) + raw = torch.randn(20, 32, device="cuda", dtype=torch.bfloat16) + cos, sin = rotary(20, 16, raw.dtype) + cache = indexer.empty_cache(1, 20).fill_(float("nan")) + indexer.pool( + raw, + table, + torch.tensor([13], device="cuda", dtype=torch.int32), + cos, + sin, + cache, + torch.tensor([0], device="cuda", dtype=torch.int32), + 3, + ) + prefix = cache[:, :1].clone() + + raw[7:12].normal_() + lengths = torch.tensor([12], device="cuda", dtype=torch.int32) + indexer.pool(raw, table, lengths, cos, sin, cache, torch.tensor([7 // 4], device="cuda", dtype=torch.int32), 2) + expected = reference_cache(indexer, raw, table, [12], cos, sin) + torch.testing.assert_close(cache, expected, atol=0.016, rtol=0.016, equal_nan=True) + torch.testing.assert_close(cache[:, :1], prefix, atol=0, rtol=0) + query = torch.randn(1, 4, 32, device="cuda", dtype=raw.dtype) + requests = torch.zeros(1, device="cuda", dtype=torch.int32) + assert_sets( + indexer.select(query, cache, requests, lengths, table), + reference_selection(indexer, query, cache, requests, lengths, table), + ) + + +def test_million_token_context_with_bounded_score_chunks(): + torch.manual_seed(337) + context = 1048579 + indexer = make_indexer(dim=128, budget=2048) + cache = torch.randn(1, context // 4, 128, device="cuda", dtype=torch.bfloat16) + + table = torch.arange(context - 1, -1, -1, device="cuda", dtype=torch.int32).view(1, -1) + query = torch.randn(3, 4, 128, device="cuda", dtype=torch.bfloat16) + requests = torch.zeros(3, device="cuda", dtype=torch.int32) + lengths = torch.tensor([4099, 262147, context], device="cuda", dtype=torch.int32) + actual = indexer.select(query, cache, requests, lengths, table, query_chunk_size=1) + expected = reference_selection(indexer, query, cache, requests, lengths, table) + assert_sets(actual, expected) + + +def test_active_rows_pool_into_shared_request_table(): + torch.manual_seed(337) + indexer = make_indexer() + table = torch.randperm(80, device="cuda").view(4, 20).int() + raw = torch.randn(80, 32, device="cuda", dtype=torch.bfloat16) + cos, sin = rotary(80, 16, raw.dtype) + cache = indexer.empty_cache(4, 20).fill_(float("nan")) + rows = torch.tensor([3, -1, 1], device="cuda", dtype=torch.int32) + lengths = torch.tensor([13, 0, 9], device="cuda", dtype=torch.int32) + indexer.pool(raw, table, lengths, cos, sin, cache, torch.zeros_like(lengths), 4, request_rows=rows) + expected = reference_cache(indexer, raw, table, [0, 9, 0, 13], cos, sin) + torch.testing.assert_close(cache, expected, atol=0.016, rtol=0.016, equal_nan=True) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("queries_per_request", [1, 4]) +@pytest.mark.parametrize("query_chunk_size", [3, 64]) +def test_physical_cache_partial_prefix_forks_and_graph_replay(dtype, queries_per_request, query_chunk_size): + torch.manual_seed(740) + indexer = make_indexer(dtype) + raw = torch.randn(32, 32, device="cuda", dtype=dtype) + cos, sin = rotary(32, 16, dtype) + + table = torch.tensor( + [[0, 1, 2, 3, 4, 5, 6, 7, 9, 10, 11, 12], [0, 1, 2, 3, 4, 5, 6, 8, 13, 14, 15, 16]], + device="cuda", + dtype=torch.int32, + ) + cache = indexer.empty_physical_cache(raw.shape[0]).fill_(float("nan")) + lengths = torch.tensor([7, 7], device="cuda", dtype=torch.int32) + first = torch.zeros_like(lengths) + rows = torch.tensor([1, -1, 0], device="cuda", dtype=torch.int32) + active_lengths = torch.tensor([7, 0, 7], device="cuda", dtype=torch.int32) + active_first = torch.zeros_like(active_lengths) + query = torch.randn(2 * queries_per_request, indexer.query_heads, indexer.head_dim, device="cuda", dtype=dtype) + requests = torch.arange(2, device="cuda", dtype=torch.int32).repeat_interleave(queries_per_request) + + def run(): + indexer.pool(raw, table, active_lengths, cos, sin, cache, active_first, 3, request_rows=rows) + return indexer.select( + query, + cache, + requests, + lengths.repeat_interleave(queries_per_request), + table, + query_chunk_size=query_chunk_size, + ) + + run() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + selected = run() + + for length in (8, 12, 5, 12): + lengths.fill_(length) + active_lengths.copy_(torch.tensor([length, 0, length], device="cuda", dtype=torch.int32)) + graph.replay() + dense = indexer.empty_cache(2, table.shape[1]).fill_(float("nan")) + indexer.pool(raw, table, lengths, cos, sin, dense, first, 3) + for request in range(2): + count = length // indexer.compress_ratio + end_slots = table[ + request, indexer.compress_ratio - 1 : count * indexer.compress_ratio : indexer.compress_ratio + ] + torch.testing.assert_close(cache[end_slots.long()], dense[request, :count], atol=0, rtol=0) + assert_sets( + selected, indexer.select(query, dense, requests, lengths.repeat_interleave(queries_per_request), table) + ) + assert not torch.equal(cache[7], cache[8]) + + untouched = cache[31].clone() + raw[7].add_(3) + run() + dense = indexer.empty_cache(2, table.shape[1]) + indexer.pool(raw, table, lengths, cos, sin, dense, first, 3) + torch.testing.assert_close(cache[7], dense[0, 1], atol=0, rtol=0) + torch.testing.assert_close(cache[31], untouched, equal_nan=True) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("capacity", [4099, 262159]) +def test_sparse_visibility_in_large_cache_and_dynamic_graph(dtype, capacity): + torch.manual_seed(38413) + indexer = make_indexer(dtype, dim=128, heads=4, budget=8) + query = torch.randn(8, 4, 128, device="cuda", dtype=dtype) + table = torch.stack([torch.randperm(capacity, device="cuda") for _ in range(3)]).int() + requests = torch.tensor([0, 0, 0, 0, 1, 1, 2, 2], device="cuda", dtype=torch.int32) + original = torch.tensor([0, 3, 4, 257, 255, 256, capacity - 1, capacity], device="cuda", dtype=torch.int32) + lengths = original.clone() + cache = torch.randn(3, capacity // 4, 128, device="cuda", dtype=dtype) + cache[0, 257 // 4 :] = float("nan") + cache[1, 256 // 4 :] = float("nan") + actual = indexer.select(query, cache, requests, lengths, table) + assert_sets(actual, reference_selection(indexer, query, cache, requests, lengths, table)) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = indexer.select(query, cache, requests, lengths, table) + for pattern in (torch.zeros_like(original), original): + lengths.copy_(pattern) + graph.replay() + assert_sets(output, reference_selection(indexer, query, cache, requests, lengths, table)) + graph.reset() + + +@pytest.mark.parametrize("rows", [4, 36, 40, 64, 128]) +def test_indexer_preprojected_strided_graph(rows): + torch.manual_seed(7912) + projection = torch.randn(640, 2560, dtype=torch.bfloat16, device="cuda") * 0.02 + indexer = QSAIndexer( + projection, + torch.randn(128, dtype=torch.bfloat16, device="cuda") * 0.1, + torch.randn(128, dtype=torch.bfloat16, device="cuda") * 0.1, + ) + hidden = torch.randn(rows, 2560, dtype=torch.bfloat16, device="cuda") + angles = torch.randn(rows, 32, device="cuda").repeat(1, 2) + cos, sin = angles.cos().bfloat16(), angles.sin().bfloat16() + packed = torch.empty(rows, 2560, dtype=hidden.dtype, device="cuda") + projected = packed[:, 1920:] + assert not projected.is_contiguous() + projected.copy_(F.linear(hidden, projection)) + indexer.project(hidden, cos, sin, projected=projected) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = indexer.project(hidden, cos, sin, projected=projected) + for _ in range(8): + hidden.normal_() + expected = indexer.project(hidden, cos, sin) + projected.copy_(F.linear(hidden, projection)) + hidden.zero_() + graph.replay() + for value, reference in zip(actual, expected): + torch.testing.assert_close(value, reference, atol=0, rtol=0) + 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..dfab256d4d --- /dev/null +++ b/unit_tests/common/qsa/test_topk_compatibility.py @@ -0,0 +1,146 @@ +from types import SimpleNamespace + +import pytest +import torch + + +from lightllm.common.qsa import topk +from lightllm.common.qsa.indexer import QSAIndexer + + +@pytest.fixture(autouse=True) +def no_cuda_initialization(): + initially_initialized = torch.cuda.is_initialized() + yield + assert torch.cuda.is_initialized() == initially_initialized + + +@pytest.mark.parametrize("shape,k", [((3, 1027), 512), ((0, 1027), 512), ((2, 11), 0), ((2, 11), 11), ((2, 3, 11), 4)]) +def test_cpu_fallback_keeps_actual_torch_indices(shape, k, monkeypatch): + def unexpected_query(device): + raise AssertionError("CPU fallback queried CUDA") + + monkeypatch.setattr(topk, "_supports_device", unexpected_query) + scores = torch.arange(torch.tensor(shape).prod().item()).reshape(shape).float() + assert torch.equal(topk.cuda_topk_indices(scores, k), torch.topk(scores, k, dim=-1, sorted=False).indices) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64, torch.int64]) +def test_noncontiguous_cpu_fallback(dtype): + scores = torch.arange(22).reshape(11, 2).T.to(dtype) + assert not scores.is_contiguous() + assert torch.equal(topk.cuda_topk_indices(scores, 4), torch.topk(scores, 4, sorted=False).indices) + + +def test_ragged_topk_cpu_and_other_budgets_do_not_query_cuda(monkeypatch): + def unexpected_query(device): + pytest.fail("Unqualified ragged TopK must not initialize CUDA") + + monkeypatch.setattr(torch.cuda, "get_device_capability", unexpected_query) + topk.flashinfer_topk_supported.cache_clear() + assert not topk.flashinfer_topk_supported(torch.device("cpu"), 65542, 512) + assert not topk.flashinfer_topk_supported(torch.device("cuda", 0), 65542, 256) + topk.flashinfer_topk_supported.cache_clear() + + +@pytest.mark.parametrize("k", [-1, 12]) +def test_invalid_budget_keeps_torch_exception(k): + with pytest.raises(RuntimeError): + topk.cuda_topk_indices(torch.ones((2, 11)), k) + + +@pytest.mark.parametrize( + "change", + [ + dict(runtime=False), + dict(dtype=torch.bfloat16), + dict(shape=(3, 1024)), + dict(shape=(3, 8, 1027)), + dict(k=256), + dict(contiguous=False), + dict(device_supported=False), + ], +) +def test_unqualified_cuda_metadata_delegates_without_launch(change, monkeypatch): + config = dict(runtime=True, dtype=torch.float32, shape=(3, 1027), k=512, contiguous=True, device_supported=True) + config.update(change) + scores = SimpleNamespace( + is_cuda=True, + dtype=config["dtype"], + ndim=len(config["shape"]), + shape=config["shape"], + device=torch.device("cuda", 0), + is_contiguous=lambda: config["contiguous"], + ) + sentinel = object() + calls = [] + monkeypatch.setattr(topk, "_CUDA_TOPK_RUNTIME", config["runtime"]) + monkeypatch.setattr(topk, "_supports_device", lambda _: config["device_supported"]) + + def original(value, k, *, dim, sorted): + calls.append((value, k, dim, sorted)) + return SimpleNamespace(indices=sentinel) + + monkeypatch.setattr(topk.torch, "topk", original) + assert topk.cuda_topk_indices(scores, config["k"]) is sentinel + assert calls == [(scores, config["k"], -1, False)] + + +def test_shared_default_and_select_api_remain_compatible(): + import inspect + + assert QSAIndexer.__dataclass_fields__["use_cuda_topk"].default is False + signature = inspect.signature(QSAIndexer.select) + assert list(signature.parameters) == [ + "self", + "query", + "cache", + "requests", + "visible_lengths", + "indices", + "query_chunk_size", + "max_visible_length", + "physical_stride", + ] + assert signature.parameters["query_chunk_size"].default == 64 + assert signature.parameters["max_visible_length"].default is None + assert signature.parameters["max_visible_length"].kind == inspect.Parameter.KEYWORD_ONLY + + +@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..04c84fa580 --- /dev/null +++ b/unit_tests/distributed/test_ordered_tp_sum.py @@ -0,0 +1,192 @@ +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 + from lightllm.distributed.communication_op import all_gather_into_tensor + + graph = None + 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, 128, 512, 1024, 4096, 8192): + 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) + scratch.fill_(rank + 2) + symmetric.all_reduce(scratch) + if gathered is not None: + 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) + owned = torch.empty_like(reference) + all_gather_into_tensor(owned, source, group=group) + symmetric.ordered_all_gather(source + 1) + torch.testing.assert_close(owned, reference, rtol=0, atol=0) + work = all_gather_into_tensor(owned, source, group=group, async_op=True) + assert work is not None + work.wait() + torch.testing.assert_close(owned, reference, 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) + actual = ordered_tp_sum(source, dist.group.WORLD, 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() + + if rows % world == 0: + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + scattered = ordered_tp_sum(source, group, world, scatter_rank=rank) + symmetric.all_reduce(scratch) + symmetric.ordered_all_gather(source[:1]) + for iteration in range(4): + source.fill_((2 ** 40, 1, -(2 ** 40), 1)[rank] if rank < 4 else iteration) + gold = expected() + graph.replay() + torch.testing.assert_close(scattered, gold.chunk(world)[rank], rtol=0, atol=0) + graph.reset() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + scattered = ordered_tp_sum(source, dist.group.WORLD, world, scatter_rank=rank) + for iteration in range(4): + source.fill_((2 ** 40, 1, -(2 ** 40), 1)[rank] if rank < 4 else iteration) + gold = expected() + graph.replay() + torch.testing.assert_close(scattered, gold.chunk(world)[rank], rtol=0, atol=0) + graph.reset() + + owned = torch.empty_like(reference) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + all_gather_into_tensor(owned, source, group=group) + symmetric.ordered_all_gather(source + 1) + for iteration in range(8): + source.fill_(rank + iteration / 8) + expected() + graph.replay() + torch.testing.assert_close(owned, reference, rtol=0, atol=0) + graph.reset() + for dtype in (torch.int32, torch.int64): + for rows in (4, 128, 1024): + source = torch.empty(rows, 2051, dtype=dtype, device=rank) + reference = torch.empty(world * rows, 2051, dtype=dtype, device=rank) + owned = torch.empty_like(reference) + patterns = torch.tensor([-1, 0, 1, 0x7FC00000, -2147483648, 2147483647], device=rank, dtype=dtype) + if dtype == torch.int64: + patterns = torch.cat((patterns, torch.tensor([-(2 ** 63), 2 ** 63 - 1], device=rank))) + + def update(iteration): + source.copy_( + patterns[ + (torch.arange(source.numel(), device=rank) + rank + iteration) % patterns.numel() + ].view_as(source) + ) + dist.all_gather_into_tensor(reference, source) + + update(0) + all_gather_into_tensor(owned, source, group=group) + torch.testing.assert_close(owned, reference, rtol=0, atol=0) + work = all_gather_into_tensor(owned, source, group=group, async_op=True) + work.wait() + torch.testing.assert_close(owned, reference, rtol=0, atol=0) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + all_gather_into_tensor(owned, source, group=group) + symmetric.ordered_all_gather(torch.ones(16, dtype=symmetric.dtype, device=rank)) + for iteration in range(3): + update(iteration + 1) + graph.replay() + torch.testing.assert_close(owned, reference, rtol=0, atol=0) + graph.reset() + + finally: + if graph is not None: + graph.reset() + 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/distributed/test_qsa_selection.py b/unit_tests/distributed/test_qsa_selection.py new file mode 100644 index 0000000000..b8e24fcea9 --- /dev/null +++ b/unit_tests/distributed/test_qsa_selection.py @@ -0,0 +1,98 @@ +import datetime +import os +from types import SimpleNamespace + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + + +def _worker(rank, world, rendezvous): + torch.cuda.set_device(rank) + os.environ["LIGHTLLM_DP_WORLD_SIZE"] = str(world) + dist.init_process_group( + "nccl", + init_method="file://" + rendezvous, + rank=rank, + world_size=world, + timeout=datetime.timedelta(seconds=180), + ) + from lightllm.common.qsa.indexer import QSAIndexer + from lightllm.models.qwen4_exp.state import QSAExecutionState + + graph = None + try: + torch.manual_seed(3519) + indexer = QSAIndexer( + torch.randn(640, 32, dtype=torch.bfloat16, device="cuda"), + torch.zeros(128, dtype=torch.bfloat16, device="cuda"), + torch.zeros(128, dtype=torch.bfloat16, device="cuda"), + 4, + 2048, + 4, + use_cuda_topk=True, + ) + for compact in (False, True): + capacity = 32768 + cache = torch.randn(capacity // 4 if compact else capacity, 128, device="cuda").to(torch.float8_e4m3fn) + pages = torch.randperm(capacity // 4, device="cuda").reshape(4, -1) + table = (pages[:, :, None] * 4 + torch.arange(4, device="cuda")).reshape(4, -1).int() + view = QSAExecutionState(None, cache, None, None, torch.empty(0) if compact else None) + for tokens, decode in ( + (world, True), + (2 * world, True), + (3 * world, True), + (4 * world, True), + (8 * world, True), + (4 * world + 1, True), + (1024, False), + ): + query = torch.randn(tokens, 4, 128, device="cuda").to(cache.dtype) + batch = SimpleNamespace( + query_rows=(torch.arange(tokens, device="cuda") // 4 % 4).int(), + visible_lengths=torch.full((tokens,), table.shape[1], device="cuda", dtype=torch.int32), + slot_table=table, + ) + + def reference(): + return view.select(indexer, query, batch, decode, None, 1, 0, None) + + def run(): + return view.select(indexer, query, batch, decode, None, world, rank, dist.group.WORLD) + + expected = reference() + torch.testing.assert_close(run(), expected, atol=0, rtol=0) + 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) + dist.barrier(device_ids=[rank]) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = run() + for step in range(2): + query.copy_(torch.randn(query.shape, device="cuda").to(query.dtype)) + if step == 1: + query.zero_() + batch.visible_lengths.copy_( + torch.arange(tokens, device="cuda").mul_(137).add_(step * 19).remainder_(table.shape[1]) + ) + batch.query_rows.add_(1).remainder_(4) + expected = reference() + graph.replay() + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + graph.reset() + finally: + if graph is not None: + graph.reset() + dist.destroy_process_group() + + +@pytest.mark.parametrize("world", [2, 8]) +def test_tp_sharded_qsa_selection_matches_replicated_with_graph(world, tmp_path): + if torch.cuda.device_count() < world: + pytest.skip(f"Requires {world} CUDA GPUs") + mp.spawn(_worker, args=(world, str(tmp_path / "qsa-rendezvous")), nprocs=world, join=True) 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..a6f22258ad 100644 --- a/unit_tests/models/qwen2_vl/test_rotary_pos_emb.py +++ b/unit_tests/models/qwen2_vl/test_rotary_pos_emb.py @@ -1,9 +1,10 @@ -import math import torch import pytest from lightllm.models.qwen2_vl.triton_kernel.rotary_pos_emb import apply_rotary_pos_emb_triton +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + def rotate_half(x): """Rotates half the hidden dims of the input.""" @@ -51,5 +52,21 @@ def test_triton_matches_reference(shape): assert torch.allclose(out, ref, rtol=1e-2, atol=1e-2), "Triton 与参考实现不一致" +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_strided_qkv_view_uses_output_strides(dtype): + torch.manual_seed(1907) + qkv = torch.randn(37, 3, 3, 72, device="cuda", dtype=dtype) + tensor = qkv[:, 1] + assert not tensor.is_contiguous() + original = qkv.clone() + angles = torch.randn(37, 36, device="cuda", dtype=torch.float32) + cos, sin = angles.cos(), angles.sin() + expected = apply_rotary_pos_emb_vision(tensor.unsqueeze(0), cos, sin).squeeze(0) + actual = apply_rotary_pos_emb_triton(tensor, cos, sin) + tolerance = 0.016 if dtype == torch.bfloat16 else 2e-5 + torch.testing.assert_close(actual, expected, atol=tolerance, rtol=tolerance) + torch.testing.assert_close(qkv, original, atol=0, rtol=0) + + if __name__ == "__main__": pytest.main() diff --git a/unit_tests/models/qwen3next/test_transformer_layer_infer_strided_gate.py b/unit_tests/models/qwen3next/test_transformer_layer_infer_strided_gate.py new file mode 100644 index 0000000000..5bfae4efa7 --- /dev/null +++ b/unit_tests/models/qwen3next/test_transformer_layer_infer_strided_gate.py @@ -0,0 +1,33 @@ +from types import SimpleNamespace + +import torch + +from lightllm.models.qwen3next.layer_infer.transformer_layer_infer import Qwen3NextTransformerLayerInfer + + +def test_linear_post_passes_strided_gate_without_materializing(): + tokens, heads, head_dim = 3, 4, 8 + packed_dim = heads * head_dim + 16 + packed = torch.randn((tokens, packed_dim)) + z = packed[:, 8 : 8 + heads * head_dim].view(tokens, heads, head_dim) + core_attn_out = torch.randn((1, tokens, heads, head_dim)) + calls = [] + + def linear_norm(input, gate_value, eps): + calls.append((input, gate_value, eps)) + return input + + layer_weight = SimpleNamespace( + linear_norm=linear_norm, + linear_out_proj=SimpleNamespace(mm=lambda input: input), + ) + layer_infer = object.__new__(Qwen3NextTransformerLayerInfer) + layer_infer.eps_ = 1e-6 + + output = layer_infer._linear_post(core_attn_out, z, layer_weight) + + assert len(calls) == 1 + assert calls[0][0].shape == (tokens * heads, head_dim) + assert calls[0][1] is z + assert not calls[0][1].is_contiguous() + assert output.shape == (tokens, heads * head_dim) diff --git a/unit_tests/models/qwen4_exp/__init__.py b/unit_tests/models/qwen4_exp/__init__.py new file mode 100644 index 0000000000..2ae28399f5 --- /dev/null +++ b/unit_tests/models/qwen4_exp/__init__.py @@ -0,0 +1 @@ +pass diff --git a/unit_tests/models/qwen4_exp/conftest.py b/unit_tests/models/qwen4_exp/conftest.py new file mode 100644 index 0000000000..59847d658f --- /dev/null +++ b/unit_tests/models/qwen4_exp/conftest.py @@ -0,0 +1,84 @@ +import json + +import pytest + +from .runtime import runtime + +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.utils import envs_utils + + +@pytest.fixture(autouse=True) +def native_runtime(monkeypatch): + args = StartArgs() + args.hardware_platform = "cuda" + args.mtp_step = 0 + args.enable_ep_moe = False + args.auto_update_redundancy_expert = False + args.enable_tpsp_mix_mode = False + monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(vars(args))) + for name, value in { + "LIGHTLLM_DP_WORLD_SIZE": 1, + "LIGHTLLM_CURRENT_RANK_IN_DP": 0, + "LIGHTLLM_CURRENT_DEVICE_ID": 0, + "LIGHTLLM_GLOBAL_RANK": 0, + "LIGHTLLM_GLOBAL_WORLD_SIZE": 1, + }.items(): + monkeypatch.setenv(name, str(value)) + envs_utils.get_env_start_args.cache_clear() + envs_utils.get_added_mtp_kv_layer_num.cache_clear() + envs_utils.get_redundancy_expert_num.cache_clear() + envs_utils.get_redundancy_expert_ids.cache_clear() + yield + envs_utils.get_env_start_args.cache_clear() + envs_utils.get_added_mtp_kv_layer_num.cache_clear() + envs_utils.get_redundancy_expert_num.cache_clear() + envs_utils.get_redundancy_expert_ids.cache_clear() + + +@pytest.fixture +def config(): + return dict( + hidden_size=32, + n_embed=32, + num_hidden_layers=4, + n_layer=4, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=8, + full_attention_interval=4, + layer_types=["linear_attention", "linear_attention", "linear_attention", "full_attention"], + hc_count=2, + hc_lowrank=8, + rms_norm_eps=1e-6, + hidden_act="silu", + output_gate_type="sigmoid", + linear_num_key_heads=1, + linear_num_value_heads=2, + linear_key_head_dim=128, + linear_value_head_dim=128, + linear_conv_kernel_dim=4, + num_experts=2, + num_experts_per_tok=1, + norm_topk_prob=True, + moe_intermediate_size=32, + shared_expert_intermediate_size=32, + indexer_n_heads=2, + indexer_kv_heads=1, + indexer_head_dim=16, + indexer_budget=8, + indexer_compress_ratio=4, + ple_layer_ids=[2], + ple_embed_dim=32, + ple_conv_kernel_size=4, + ngram_size=3, + heads_per_ngram=2, + ngram_vocab_size_base=5, + make_ngram_vocab_size_divisible_by=4, + split_ngram_parts=4, + eos_token_id=63, + vocab_size=64, + partial_rotary_factor=0.5, + rope_scaling={"mrope_section": [1, 1, 0]}, + tie_word_embeddings=False, + ) diff --git a/unit_tests/models/qwen4_exp/helpers.py b/unit_tests/models/qwen4_exp/helpers.py new file mode 100644 index 0000000000..462a5863ea --- /dev/null +++ b/unit_tests/models/qwen4_exp/helpers.py @@ -0,0 +1,133 @@ +import torch + + +def rand(*shape): + return torch.randn(*shape, dtype=torch.bfloat16) * 0.1 + + +def add_hyper(weights, prefix, config, combine=True): + hidden, branches, rank = config["hidden_size"], config["hc_count"], config["hc_lowrank"] + width = hidden * branches + weights[prefix + ".hc_norm.weight"] = rand(width) + weights[prefix + ".input_mix_weight_down.weight"] = rand(rank, width) + weights[prefix + ".input_mix_weight_up.weight"] = rand(width, rank) + if combine: + weights[prefix + ".block_inject_weight.weight"] = rand(branches, width) + + +def checkpoint(config, layer): + hidden, inter, experts = config["hidden_size"], config["moe_intermediate_size"], config["num_experts"] + prefix = f"model.language_model.layers.{layer}" + weights = {} + for block in ["attn_hyper_connection", "mlp_hyper_connection"]: + add_hyper(weights, prefix + "." + block, config) + weights[prefix + ".mlp.gate.weight"] = rand(experts, hidden) + weights[prefix + ".mlp.experts.gate_up_proj"] = rand(experts, 2 * inter, hidden) + weights[prefix + ".mlp.experts.down_proj"] = rand(experts, hidden, inter) + weights[prefix + ".mlp.shared_expert_gate.weight"] = rand(1, hidden) + for proj, shape in [("gate", (inter, hidden)), ("up", (inter, hidden)), ("down", (hidden, inter))]: + weights[prefix + f".mlp.shared_expert.{proj}_proj.weight"] = rand(*shape) + if config["layer_types"][layer] == "linear_attention": + keydim = config["linear_num_key_heads"] * config["linear_key_head_dim"] + vdim = config["linear_num_value_heads"] * config["linear_value_head_dim"] + channels = 2 * keydim + vdim + weights[prefix + ".linear_attn.conv1d.weight"] = rand(channels, 1, 4) + for proj, rows in [ + ("qkv", channels), + ("z", vdim), + ("b", config["linear_num_value_heads"]), + ("a", config["linear_num_value_heads"]), + ]: + weights[prefix + f".linear_attn.in_proj_{proj}.weight"] = rand(rows, hidden) + weights[prefix + ".linear_attn.out_proj.weight"] = rand(hidden, vdim) + weights[prefix + ".linear_attn.A_log"] = torch.zeros(config["linear_num_value_heads"], dtype=torch.float32) + weights[prefix + ".linear_attn.dt_bias"] = torch.ones(config["linear_num_value_heads"], dtype=torch.float32) + weights[prefix + ".linear_attn.norm.weight"] = 1 + rand(config["linear_value_head_dim"]) + else: + dim, qheads, kvheads = config["head_dim"], config["num_attention_heads"], config["num_key_value_heads"] + for name, shape in [ + ("q_proj", (qheads * dim * 2, hidden)), + ("k_proj", (kvheads * dim, hidden)), + ("v_proj", (kvheads * dim, hidden)), + ("o_proj", (hidden, qheads * dim)), + ]: + weights[prefix + f".self_attn.{name}.weight"] = rand(*shape) + for name in ["q_norm", "k_norm"]: + weights[prefix + f".self_attn.{name}.weight"] = rand(dim) + weights[prefix + ".self_attn.indexer.index_qk_proj.weight"] = rand( + (config["indexer_n_heads"] + 1) * config["indexer_head_dim"], hidden + ) + for name in ["q_layernorm", "k_layernorm"]: + weights[prefix + f".self_attn.indexer.{name}.weight"] = rand(config["indexer_head_dim"]) + if layer + 1 in config["ple_layer_ids"]: + width, embed = config["hc_count"] * hidden, config["ple_embed_dim"] + weights[prefix + ".ple.key_proj.weight"] = rand(width, embed) + weights[prefix + ".ple.value_proj.weight"] = rand(hidden, embed) + weights[prefix + ".ple.conv1d.weight"] = rand(width, 1, 4) + for name in ["norm_key", "norm_query", "norm_conv"]: + weights[prefix + f".ple.{name}.weight"] = rand(width) + weights[prefix + ".ple.ple_embedding.layer_multipliers"] = torch.tensor( + [23703573157769, 20109073645365, 8052911324071], dtype=torch.int64 + ) + weights[prefix + ".ple.ple_embedding.ngram_heads_vocab_sizes"] = torch.tensor([5, 7, 11, 13], dtype=torch.int64) + weights[prefix + ".ple.ple_embedding.ngram_heads_offsets"] = torch.tensor([0, 5, 12, 23], dtype=torch.int64) + table = rand(36, 8) + for shard, values in enumerate(table.chunk(4)): + weights[prefix + f".ple.ple_embedding.ngram_embedding.shard_{shard}.weight"] = values + return weights + + +def prepost_checkpoint(config): + weights = { + "model.language_model.embed_tokens.weight": rand(config["vocab_size"], config["hidden_size"]), + "lm_head.weight": rand(config["vocab_size"], config["hidden_size"]), + } + add_hyper(weights, "model.language_model.hyper_connection_mixer", config, False) + return weights + + +def reference_hyper(x, weights, prefix, config, combine=True): + grouped = x.float().unflatten(-1, (config["hc_count"], config["hidden_size"])) + normalized = grouped * torch.rsqrt(grouped.square().mean(-1, keepdim=True) + config["rms_norm_eps"]) + normalized = (normalized.flatten(-2) * (1 + weights[prefix + ".hc_norm.weight"].cuda().float())).to(x.dtype) + low = torch.nn.functional.silu( + torch.nn.functional.linear(normalized, weights[prefix + ".input_mix_weight_down.weight"].cuda()) + / config["hc_count"] + ) + gates = torch.nn.functional.linear(low, weights[prefix + ".input_mix_weight_up.weight"].cuda()).sigmoid() + mixed = (normalized * gates).unflatten(-1, (config["hc_count"], config["hidden_size"])).mean(-2) + if not combine: + return mixed + injection = 2 * torch.sigmoid( + torch.nn.functional.linear(normalized, weights[prefix + ".block_inject_weight.weight"].cuda()) + / config["hc_count"] + ) + return mixed, injection + + +def reference_moe(x, weights, prefix, config): + logits = torch.nn.functional.linear(x, weights[prefix + ".mlp.gate.weight"].cuda()) + prob = logits.float().softmax(-1) + values, indices = prob.topk(config["num_experts_per_tok"], dim=-1) + if config["norm_topk_prob"]: + values /= values.sum(-1, keepdim=True) + values = values.to(x.dtype) + out = torch.zeros_like(x) + inter = config["moe_intermediate_size"] + + for expert in indices.unique(sorted=True).tolist(): + top_position, token = torch.where((indices == expert).T) + gu = torch.nn.functional.linear(x[token], weights[prefix + ".mlp.experts.gate_up_proj"][expert].cuda()) + output = torch.nn.functional.linear( + torch.nn.functional.silu(gu[:, :inter]) * gu[:, inter:], + weights[prefix + ".mlp.experts.down_proj"][expert].cuda(), + ) + weighted = output * values[token, top_position, None] + out.index_add_(0, token, weighted.to(out.dtype)) + gate = torch.nn.functional.linear(x, weights[prefix + ".mlp.shared_expert_gate.weight"].cuda()).sigmoid() + up = torch.nn.functional.linear(x, weights[prefix + ".mlp.shared_expert.up_proj.weight"].cuda()) + g = torch.nn.functional.linear(x, weights[prefix + ".mlp.shared_expert.gate_proj.weight"].cuda()) + shared = torch.nn.functional.linear( + torch.nn.functional.silu(g) * up, weights[prefix + ".mlp.shared_expert.down_proj.weight"].cuda() + ) + return out + gate * shared diff --git a/unit_tests/models/qwen4_exp/lexical_helpers.py b/unit_tests/models/qwen4_exp/lexical_helpers.py new file mode 100644 index 0000000000..ebd26a3389 --- /dev/null +++ b/unit_tests/models/qwen4_exp/lexical_helpers.py @@ -0,0 +1,32 @@ +import torch + +from lightllm.models.qwen4_exp.layer_weights.transformer_layer_weight import Qwen4ExpTransformerLayerWeight +from lightllm.models.qwen4_exp.batch import PackedBatchMetadata +from lightllm.models.qwen4_exp.lexical_cache import LexicalRequestCache +from lightllm.common.quantization import Quantcfg +from .helpers import checkpoint +from .reference import reference_ple + + +def make_cache(config): + source = checkpoint(config, 1) + weights = Qwen4ExpTransformerLayerWeight(1, torch.bfloat16, config, Quantcfg(config)) + weights.load_hf_weights(dict(source)) + weights.verify_load() + return LexicalRequestCache(weights, 8), source + + +def metadata(requests, lengths, queries): + tensors = [torch.tensor(values, device="cuda", dtype=torch.int32) for values in (requests, lengths, queries)] + return PackedBatchMetadata.prepare(*tensors, sum(queries), 7, 4) + + +def reference_segment(residual, tokens, source, config, ngram, convolution): + updates, states = [], [(ngram.clone(), convolution.clone())] + for row in range(tokens.numel()): + update, ngram, convolution = reference_ple( + residual[row : row + 1], tokens[row : row + 1], source, config, ngram, convolution + ) + updates.append(update) + states.append((ngram.clone(), convolution.clone())) + return torch.cat(updates), states diff --git a/unit_tests/models/qwen4_exp/mtp_reference.py b/unit_tests/models/qwen4_exp/mtp_reference.py new file mode 100644 index 0000000000..e9d4332c42 --- /dev/null +++ b/unit_tests/models/qwen4_exp/mtp_reference.py @@ -0,0 +1,27 @@ +import torch.nn.functional as F + +from .reference import norm, reference_layer +from .helpers import reference_hyper + + +def reference_fusion(embedding, hidden, weights, config): + embedding = norm(embedding, weights["mtp.pre_fc_norm_embedding.weight"], config["rms_norm_eps"]) + hidden = norm(hidden, weights["mtp.pre_fc_norm_hidden.weight"], config["rms_norm_eps"]) + projected_embedding = F.linear(embedding, weights["mtp.fc_embedding.weight"].cuda()) + projected_hidden = F.linear( + hidden.view(-1, config["hc_count"], config["hidden_size"]), weights["mtp.fc_hidden.weight"].cuda() + ) + return (projected_hidden + projected_embedding[:, None]).flatten(1) + + +def reference_draft(tokens, hidden, weights, config, cos, sin, lengths): + embeddings = F.embedding(tokens, weights["model.language_model.embed_tokens.weight"].cuda()) + residual = reference_fusion(embeddings, hidden, weights, config) + mapped = { + name.replace("mtp.layers.0.", "model.language_model.layers.3.", 1): value + for name, value in weights.items() + if name.startswith("mtp.layers.0.") + } + residual = reference_layer(residual, mapped, config, cos, sin, lengths) + mixed = reference_hyper(residual, weights, "mtp.hyper_connection_mixer", config, combine=False) + return residual, F.linear(mixed, weights["lm_head.weight"].cuda()).float() diff --git a/unit_tests/models/qwen4_exp/reference.py b/unit_tests/models/qwen4_exp/reference.py new file mode 100644 index 0000000000..2557c479f5 --- /dev/null +++ b/unit_tests/models/qwen4_exp/reference.py @@ -0,0 +1,235 @@ +import torch +import torch.nn.functional as F + +from .helpers import reference_hyper, reference_moe + + +def grouped_norm(x, weight, config): + values = x.float().unflatten(-1, (config["hc_count"], config["hidden_size"])) + normalized = values * torch.rsqrt(values.square().mean(-1, keepdim=True) + config["rms_norm_eps"]) + return (normalized.flatten(-2) * (1 + weight.cuda().float())).to(x.dtype) + + +def reference_ple(residual, tokens, weights, config, ngram_history, conv_history): + prefix = "model.language_model.layers.1.ple" + history = torch.cat([ngram_history, tokens[:, None]], dim=1) + eos = config["eos_token_id"] + eos = eos[0] if isinstance(eos, list) else eos + context = config["ngram_size"] - 1 + shifted = [tokens] + for distance in range(1, context + 1): + preceding = history[:, context - distance : context] + shifted.append(torch.where((preceding == eos).any(-1), eos, history[:, context - distance])) + multipliers = weights[prefix + ".ple_embedding.layer_multipliers"].cuda() + sizes = weights[prefix + ".ple_embedding.ngram_heads_vocab_sizes"].cuda() + offsets = weights[prefix + ".ple_embedding.ngram_heads_offsets"].cuda() + indices = [] + for head in range(sizes.numel()): + order = head // config["heads_per_ngram"] + 2 + mixed = shifted[0] * multipliers[0] + for position in range(1, order): + mixed = mixed ^ (shifted[position] * multipliers[position]) + indices.append(mixed.remainder(sizes[head]) + offsets[head]) + embedding = reference_ngram_lookup(torch.stack(indices, -1), weights, prefix).flatten(-2) + branches, hidden = config["hc_count"], config["hidden_size"] + key = grouped_norm( + F.linear(embedding, weights[prefix + ".key_proj.weight"].cuda()), weights[prefix + ".norm_key.weight"], config + ).view(-1, branches, hidden) + query = grouped_norm(residual, weights[prefix + ".norm_query.weight"], config).view(-1, branches, hidden) + gate = (key * query).sum(-1, keepdim=True) / (hidden ** 0.5) + gate = gate.abs().clamp_min(1e-6).sqrt() * gate.sign() + value = F.linear(embedding, weights[prefix + ".value_proj.weight"].cuda()) + gated = (gate.sigmoid() * value[:, None]).flatten(-2) + normalized = grouped_norm(gated, weights[prefix + ".norm_conv.weight"], config) + full = torch.cat([conv_history, normalized[:, None]], dim=1) + dilation = config["ngram_size"] + conv = F.conv1d( + full.transpose(1, 2), weights[prefix + ".conv1d.weight"].cuda(), groups=branches * hidden, dilation=dilation + ).squeeze(-1) + conv_context = (config["ple_conv_kernel_size"] - 1) * dilation + return ( + gated + F.silu(conv), + history[:, -context:].contiguous(), + full[:, -conv_context:].contiguous() if conv_context else full[:, :0].contiguous(), + ) + + +def reference_ngram_lookup(indices, weights, prefix): + name = prefix + ".ple_embedding.ngram_embedding" + first = weights[name + ".shard_0.weight"] + shard_rows, dim = first.shape + flat = indices.cpu().flatten() + shards = flat // shard_rows + result = torch.empty(flat.numel(), dim, dtype=first.dtype) + for shard in shards.unique().tolist(): + positions = torch.where(shards == shard)[0] + source = weights[name + f".shard_{shard}.weight"] + rows = (flat[positions] % shard_rows).to(source.device) + result[positions] = source[rows].cpu() + return result.to(indices.device).view(*indices.shape, dim) + + +def reference_gdn(mixed, weights, config, conv_history, recurrent): + prefix = "model.language_model.layers.1.linear_attn" + key_heads, value_heads = config["linear_num_key_heads"], config["linear_num_value_heads"] + key_dim, value_dim = config["linear_key_head_dim"], config["linear_value_head_dim"] + kd, vd = key_heads * key_dim, value_heads * value_dim + qkv = F.linear(mixed, weights[prefix + ".in_proj_qkv.weight"].cuda()) + full = torch.cat([conv_history, qkv[:, :, None]], dim=-1) + conv = (full.float() * weights[prefix + ".conv1d.weight"].squeeze(1).cuda().float()[None]).sum(-1).to(mixed.dtype) + conv = F.silu(conv) + q, k, v = torch.split(conv, [kd, kd, vd], dim=-1) + q, k = q.view(-1, key_heads, key_dim).float(), k.view(-1, key_heads, key_dim).float() + q = q * torch.rsqrt(q.square().sum(-1, keepdim=True) + 1e-6) + k = k * torch.rsqrt(k.square().sum(-1, keepdim=True) + 1e-6) + q = q.repeat_interleave(value_heads // key_heads, dim=1) / (key_dim ** 0.5) + k = k.repeat_interleave(value_heads // key_heads, dim=1) + v = v.view(-1, value_heads, value_dim).float() + a = F.linear(mixed, weights[prefix + ".in_proj_a.weight"].cuda()) + b = F.linear(mixed, weights[prefix + ".in_proj_b.weight"].cuda()) + decay = -weights[prefix + ".A_log"].cuda().float().exp() * F.softplus( + a.float() + weights[prefix + ".dt_bias"].cuda() + ) + beta = b.sigmoid().float() + next_state = recurrent * decay.exp()[:, :, None, None] + memory = (next_state * k[:, :, :, None]).sum(-2) + delta = (v - memory) * beta[:, :, None] + next_state = next_state + k[:, :, :, None] * delta[:, :, None, :] + core = (next_state * q[:, :, :, None]).sum(-2).to(mixed.dtype) + normed = core.float() * torch.rsqrt(core.float().square().mean(-1, keepdim=True) + 1e-6) + weighted = weights[prefix + ".norm.weight"].cuda() * normed.to(mixed.dtype) + z = F.linear(mixed, weights[prefix + ".in_proj_z.weight"].cuda()).view(-1, value_heads, value_dim) + gated = (weighted * z.float().sigmoid()).to(mixed.dtype) + output = F.linear(gated.flatten(1), weights[prefix + ".out_proj.weight"].cuda()) + return output, full[:, :, -3:].contiguous(), next_state + + +def reference_step(residual, tokens, source, config, ngram, ple, conv, recurrent): + lexical, ngram, ple = reference_ple(residual, tokens, source, config, ngram, ple) + expected = residual + lexical + mixed, gates = reference_hyper(expected, source, "model.language_model.layers.1.attn_hyper_connection", config) + output, conv, recurrent = reference_gdn(mixed, source, config, conv, recurrent) + expected = expected + (output[:, None] * gates[:, :, None]).flatten(1) + mixed, gates = reference_hyper(expected, source, "model.language_model.layers.1.mlp_hyper_connection", config) + output = reference_moe(mixed, source, "model.language_model.layers.1", config) + expected = expected + (output[:, None] * gates[:, :, None]).flatten(1) + return expected, ngram, ple, conv, recurrent + + +def norm(x, weight, eps=1e-6): + value = x.float() + return (value * torch.rsqrt(value.square().mean(-1, keepdim=True) + eps) * (1 + weight.cuda().float())).to(x.dtype) + + +def rotate(x, cos, sin): + width = cos.shape[-1] + a, b = x[..., :width].chunk(2, -1) + return torch.cat([x[..., :width] * cos + torch.cat([-b, a], -1) * sin, x[..., width:]], -1) + + +def reference_layer(residual, weights, config, cos, sin, lengths): + prefix = "model.language_model.layers.3" + mixed, gates = reference_hyper(residual, weights, prefix + ".attn_hyper_connection", config) + dim, heads, kvheads = config["head_dim"], config["num_attention_heads"], config["num_key_value_heads"] + qg = F.linear(mixed, weights[prefix + ".self_attn.q_proj.weight"].cuda()).view(-1, heads, dim * 2) + query, gate = qg.chunk(2, -1) + key = F.linear(mixed, weights[prefix + ".self_attn.k_proj.weight"].cuda()).view(-1, kvheads, dim) + value = F.linear(mixed, weights[prefix + ".self_attn.v_proj.weight"].cuda()).view(-1, kvheads, dim) + query = rotate(norm(query, weights[prefix + ".self_attn.q_norm.weight"]), cos[:, None], sin[:, None]) + key = rotate(norm(key, weights[prefix + ".self_attn.k_norm.weight"]), cos[:, None], sin[:, None]) + projected = F.linear(mixed, weights[prefix + ".self_attn.indexer.index_qk_proj.weight"].cuda()) + ih, idim = config["indexer_n_heads"], config["indexer_head_dim"] + iq = projected[:, : ih * idim].view(-1, ih, idim) + iq = rotate(norm(iq, weights[prefix + ".self_attn.indexer.q_layernorm.weight"]), cos[:, None], sin[:, None]) + raw = projected[:, ih * idim :] + ratio, budget = config["indexer_compress_ratio"], config["indexer_budget"] // config["indexer_compress_ratio"] + result = [] + start = 0 + for count in lengths: + for position in range(count): + complete = (position + 1) // ratio + selected = [] + if complete: + groups = raw[start : start + complete * ratio].view(complete, ratio, idim).float().mean(1).to(raw.dtype) + pooled = norm(groups, weights[prefix + ".self_attn.indexer.k_layernorm.weight"]) + pooled = rotate( + pooled, cos[start : start + complete * ratio : ratio], sin[start : start + complete * ratio : ratio] + ) + scores = (iq[start + position].float() @ pooled.float().T).relu().sum(0) / (idim ** 0.5) + blocks = scores.topk(min(complete, budget)).indices + selected = (blocks[:, None] * ratio + torch.arange(ratio, device=residual.device)).flatten().tolist() + selected += list(range(complete * ratio, position + 1)) + ids = torch.tensor(selected, device=residual.device, dtype=torch.int64) + start + out = F.scaled_dot_product_attention( + query[start + position].unsqueeze(0).unsqueeze(2), + key[ids].transpose(0, 1).unsqueeze(0), + value[ids].transpose(0, 1).unsqueeze(0), + enable_gqa=True, + ) + result.append(out[0, :, 0]) + start += count + attention = torch.stack(result) + attention = attention * gate.sigmoid() + output = F.linear(attention.flatten(1), weights[prefix + ".self_attn.o_proj.weight"].cuda()) + residual = residual + (output[:, None, :] * gates[:, :, None]).flatten(1) + mixed, gates = reference_hyper(residual, weights, prefix + ".mlp_hyper_connection", config) + update = reference_moe(mixed, weights, prefix, config) + return residual + (update[:, None, :] * gates[:, :, None]).flatten(1) + + +def reference_model(tokens, cos, sin, lengths, weights, config, input_embeddings=None): + if input_embeddings is None: + input_embeddings = F.embedding(tokens, weights["model.language_model.embed_tokens.weight"].cuda()) + residual = input_embeddings.repeat(1, config["hc_count"]) + for layer in range(config["n_layer"]): + prefix = f"model.language_model.layers.{layer}" + if config["layer_types"][layer] == "full_attention": + residual = reference_layer(residual, weights, config, cos, sin, lengths) + continue + if layer + 1 in config["ple_layer_ids"]: + updates, start = [], 0 + for length in lengths: + ngram = torch.full((1, 2), config["eos_token_id"], device="cuda", dtype=torch.int64) + convolution = torch.zeros(1, 9, 64, device="cuda", dtype=torch.bfloat16) + for row in range(start, start + length): + update, ngram, convolution = reference_ple( + residual[row : row + 1], tokens[row : row + 1], weights, config, ngram, convolution + ) + updates.append(update) + start += length + residual = residual + torch.cat(updates) + mixed, gates = reference_hyper(residual, weights, prefix + ".attn_hyper_connection", config) + + local = { + name.replace(prefix, "model.language_model.layers.1", 1): value + for name, value in weights.items() + if name.startswith(prefix + ".") + } + updates, start = [], 0 + for length in lengths: + channels = ( + 2 * config["linear_num_key_heads"] * config["linear_key_head_dim"] + + config["linear_num_value_heads"] * config["linear_value_head_dim"] + ) + convolution = torch.zeros( + 1, channels, config["linear_conv_kernel_dim"] - 1, device="cuda", dtype=torch.bfloat16 + ) + recurrent = torch.zeros( + 1, + config["linear_num_value_heads"], + config["linear_key_head_dim"], + config["linear_value_head_dim"], + device="cuda", + dtype=torch.float32, + ) + for row in range(start, start + length): + update, convolution, recurrent = reference_gdn( + mixed[row : row + 1], local, config, convolution, recurrent + ) + updates.append(update) + start += length + residual = residual + (torch.cat(updates)[:, None] * gates[:, :, None]).flatten(1) + mixed, gates = reference_hyper(residual, weights, prefix + ".mlp_hyper_connection", config) + residual = residual + (reference_moe(mixed, weights, prefix, config)[:, None] * gates[:, :, None]).flatten(1) + hidden = reference_hyper(residual, weights, "model.language_model.hyper_connection_mixer", config, combine=False) + return F.linear(hidden, weights["lm_head.weight"].cuda()).float() diff --git a/unit_tests/models/qwen4_exp/resource_helpers.py b/unit_tests/models/qwen4_exp/resource_helpers.py new file mode 100644 index 0000000000..c36f8b32e7 --- /dev/null +++ b/unit_tests/models/qwen4_exp/resource_helpers.py @@ -0,0 +1,74 @@ +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 + } + if self.state.positions is not None: + self.qsa_raw_keys = self.state.indexer.k_storage[:, :, 0] + self.qsa_pooled_keys = self.state.indexer.v_storage[:, :, 0] + self.qsa_key_cos = self.state.positions.k_storage[0, :, 0] + self.qsa_key_sin = self.state.positions.v_storage[0, :, 0] + + def alloc(self, count): + return self.allocator.alloc(count) + + def _slot_buffers(self): + return { + f"{state.state_name}.{name}": (tensor, 1) + for state in self.manager.get_full_attention_states() + for name, tensor in ( + ("k", state.k_storage), + ("v", state.v_storage), + ("raw", state.indexer.k_storage), + ("pooled", state.indexer.v_storage), + ("cos", state.positions.k_storage), + ("sin", state.positions.v_storage), + ) + } + + def get_index_kv_buffer(self, indices): + indices = torch.as_tensor(indices, dtype=torch.int64, device=self.state.device) + return {name: tensor.index_select(axis, indices) for name, (tensor, axis) in self._slot_buffers().items()} + + +def bind_states(model): + manager = model.att_state_manager + manager._get_states.cache_clear() + manager.finish_state_registration() + if manager.full_att_allocator is None: + manager.allocate_storage(32, model.args.page_size, 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..91b885507e --- /dev/null +++ b/unit_tests/models/qwen4_exp/runtime.py @@ -0,0 +1,156 @@ +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" + args.llm_prefill_att_backend = ["triton", "triton"] + args.llm_decode_att_backend = ["triton", "triton"] + for name, value in getattr(request, "param", {}).items(): + setattr(args, name, value) + if args.llm_prefill_att_backend[1] == "flashinfer": + import inspect + + module = pytest.importorskip("flashinfer.gdn_prefill") + if "use_cp" not in inspect.signature(module.chunk_gated_delta_rule).parameters: + pytest.skip("FlashInfer >= 0.7.2 required") + 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" + 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, graph_mode=None): + state = Qwen4ExpInferStateInfo() + state.is_cuda_graph_capture = graph_mode == "capture" + state.is_cuda_graph_replay = graph_mode == "replay" + 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) + 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) + if runtime.model.args.page_size > 1: + if not hasattr(runtime.mem, "allocated_lengths"): + runtime.mem.allocated_lengths = {} + parts = [] + page_size = runtime.model.args.page_size + for req, length, count in zip(requests, lengths, query_lengths): + allocated = runtime.mem.allocated_lengths.get(req, 0) + required = (length + page_size - 1) // page_size * page_size + if required > allocated: + slots = runtime.mem.alloc(required - allocated).cuda(non_blocking=True) + runtime.req.req_to_token_indexs[req, allocated:required] = slots + runtime.mem.allocated_lengths[req] = required + parts.append(runtime.req.req_to_token_indexs[req, length - count : length]) + state.mem_index = torch.cat(parts) + else: + state.mem_index = runtime.mem.alloc(len(tokens)).cuda(non_blocking=True) + 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..04c4d95eae --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_attention_state.py @@ -0,0 +1,303 @@ +from dataclasses import replace +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.attention_state.manager import AttentionStateManager +from lightllm.common.attention_state.types import StateTransferLayout +from lightllm.models import get_model_class +from lightllm.models.draft_registry import get_draft_model_class +from lightllm.models.qwen4_exp.model import Qwen4ExpTpPartModel +from lightllm.models.qwen4_exp.mtp.model import Qwen4ExpMTPModel +from lightllm.models.qwen4_exp.attention_state import Qwen4ExpAttentionState +from lightllm.models.qwen4_exp.lexical_state import Qwen4ExpLexicalState, LexicalStateConfig +from lightllm.utils.envs_utils import get_env_start_args +from lightllm.common.attention_state.bounded_att.gdn.gdn_state import GDNState +from lightllm.common.attention_state.bounded_att.gdn.replay_gdn_state import ReplayGDNState + + +@pytest.fixture +def lexical(monkeypatch): + import lightllm.models.qwen4_exp.lexical_state as module + + def copy(*, copies): + for source, destination in copies: + destination.copy_(source) + + monkeypatch.setattr(module, "copy_bounded_att_tensors", copy) + state = Qwen4ExpLexicalState(LexicalStateConfig((1, 5), 2, 6, 8, 63, torch.bfloat16), 4, 2, "cpu") + state.state_name = "main_qwen4_lexical_state" + state.initialize_page_pools(4, 2, pin_memory=False) + return state + + +def test_lazy_registration_and_resource_factories(config): + args = get_env_start_args() + args.data_type, args.linear_att_ssm_data_type = "bf16", "float32" + args.cpu_cache_token_page_size = 8 + args.pd_kv_page_size = 4 + for model_type in ("qwen4_exp", "qwen4_exp_text"): + assert get_model_class({"model_type": model_type}) is Qwen4ExpTpPartModel + assert get_draft_model_class({"model_type": model_type}, "eagle_with_att") is Qwen4ExpMTPModel + manager = AttentionStateManager(request_capacity=4) + for cls, role, layers in ((Qwen4ExpTpPartModel, "main", (0, 1, 2, 3)), (Qwen4ExpMTPModel, "draft0", (4,))): + full = cls.create_full_att_states(config, args, layers) + bounded = cls.create_bounded_att_states(config, args, layers, 4) + for state in (*full, *bounded): + manager.register_state(role, state) + manager.finish_state_registration() + state = manager.get_state("main", "", 3, Qwen4ExpAttentionState) + assert isinstance(state, Qwen4ExpAttentionState) + assert manager.get_state("draft0", "qwen4_gqa", 4) is not state + assert len(manager.get_bounded_attention_states("main")) == 2 + assert not manager.get_bounded_attention_states("draft0") + for state in manager.get_full_attention_states(): + regions = state.get_cpu_transfer_regions() + assert len(regions) == 6 + assert all(left.byte_end <= right.byte_start for left, right in zip(regions, regions[1:])) + assert {region.field for region in regions} == { + "k", + "v", + "indexer.k", + "indexer.v", + "positions.k", + "positions.v", + } + assert state.get_token_bytes() > 0 + + +@pytest.mark.parametrize("mode,state_type", [("native", GDNState), ("replay", ReplayGDNState)]) +def test_mtp_state_modes_keep_qwen4_resources(config, mode, state_type): + args = get_env_start_args() + args.data_type = "bfloat16" + args.mtp_step = 3 + args.ssm_state_mode = mode + args.ssm_replay_config = {"cache_len": 16, "compute_mode": "precompute"} if mode == "replay" else {} + args.linear_att_ssm_data_type = "float32" + layers = (0, 1, 2, 3) + bounded = Qwen4ExpTpPartModel.create_bounded_att_states(config, args, layers, 4, device="cpu") + manager = AttentionStateManager(request_capacity=4) + for state in (*Qwen4ExpTpPartModel.create_full_att_states(config, args, layers), *bounded): + manager.register_state("main", state) + manager.finish_state_registration() + state = manager.get_state("main", "gdn", 1, state_type) + assert state.config.logical_layer_ids == (0, 1, 2) + assert state.get_layer_view(1)[1].dtype == torch.float32 + assert isinstance(manager.get_state("main", "qwen4_gqa", 3), Qwen4ExpAttentionState) + assert any(isinstance(state, Qwen4ExpLexicalState) for state in bounded) + assert not Qwen4ExpMTPModel.create_bounded_att_states(config, args, (4,), 4, device="cpu") + + +@pytest.mark.parametrize("tp", [1, 2, 4, 8]) +def test_replicated_indexer_and_positions_layout(config, tp): + args = get_env_start_args() + args.data_type = "bf16" + for rank in range(tp): + (state,) = Qwen4ExpTpPartModel.create_full_att_states(config, args, (0, 1, 2, 3), rank, tp) + regions = state._get_transfer_regions(8) + for region in regions[2:]: + assert region.write_enabled == (rank == 0) + assert region.shape[2] == 1 + assert sum(region.byte_end - region.byte_start for region in regions) == regions[-1].byte_end + + +@pytest.mark.parametrize("kv_type", ["None", "fp8kv_spt"]) +def test_state_storage_accounts_for_all_fields_and_restores_positions(config, kv_type): + args = get_env_start_args() + args.data_type = "bf16" + args.llm_kv_type = kv_type + (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() + assert view.raw_keys.dtype == torch.bfloat16 + if kv_type == "fp8kv_spt": + assert view.pooled_keys.dtype == state.k_storage.dtype == torch.float8_e4m3fn + fields = {r.field: r for r in state._get_transfer_regions(8)} + assert fields["indexer.k"].dtype == torch.bfloat16 + assert fields["indexer.v"].dtype == torch.float8_e4m3fn + assert fields["positions.k"].dtype == torch.bfloat16 + + +def test_fp8_rejects_calibration_files_and_mixed_encoding_dp_copy(config): + args = get_env_start_args() + args.data_type = "bfloat16" + args.llm_kv_type = "fp8kv_spt" + args.kv_quant_calibration_config_path = "calibration.json" + with pytest.raises(ValueError, match="unit scales"): + Qwen4ExpTpPartModel.create_full_att_states(config, args, (0, 1, 2, 3)) + args.kv_quant_calibration_config_path = None + (fp8,) = Qwen4ExpTpPartModel.create_full_att_states(config, args, (0, 1, 2, 3)) + args.llm_kv_type = "None" + (bf16,) = Qwen4ExpTpPartModel.create_full_att_states(config, args, (0, 1, 2, 3)) + with pytest.raises(AssertionError, match="same KV and indexer encoding"): + fp8.copy_tokens_from_dp(None, None, bf16) + + +def qsa_tensors(state): + return ( + state.k_storage, + state.v_storage, + state.indexer.k_storage, + state.indexer.v_storage, + state.positions.k_storage, + state.positions.v_storage, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("fp8", [False, True]) +@pytest.mark.parametrize("page_device", ["cpu", "cuda"]) +def test_qsa_kv_write_copy_and_page_round_trip(config, fp8, page_device): + args = get_env_start_args() + args.data_type = "bfloat16" + args.cpu_cache_token_page_size = args.pd_kv_page_size = 4 + args.llm_kv_type = "fp8kv_spt" if fp8 else "None" + states = [] + for _ in range(2): + (state,) = Qwen4ExpTpPartModel.create_full_att_states(config, args, (0, 1, 2, 3)) + state.device = torch.device("cuda") + state.allocate_storage(16) + states.append(state) + source, target = states + for tensor in qsa_tensors(source): + tensor.view(torch.uint8).random_(0, 256) + for tensor in qsa_tensors(target): + tensor.view(torch.uint8).fill_(173) + src = torch.tensor([1, 4, 7], device="cuda", dtype=torch.int32) + dst = torch.tensor([2, 6, 8], device="cuda", dtype=torch.int64) + kv = torch.randn((3, 4, 8), device="cuda", dtype=torch.bfloat16) + source.write_full_att_kv(3, src, kv) + k, v = source.get_layer_view(3) + torch.testing.assert_close(k.float()[src], kv[:, :2].to(k.dtype).float(), rtol=0, atol=0) + torch.testing.assert_close(v.float()[src], kv[:, 2:].to(v.dtype).float(), rtol=0, atol=0) + expected = [tensor.view(torch.uint8)[:, src].clone() for tensor in qsa_tensors(source)] + source.copy_full_att_tokens(src, dst) + for tensor, reference in zip(qsa_tensors(source), expected): + assert torch.equal(tensor.view(torch.uint8)[:, dst], reference) + target.copy_tokens_from_dp(src, dst, source) + for tensor, reference in zip(qsa_tensors(target), expected): + assert torch.equal(tensor.view(torch.uint8)[:, dst], reference) + tensor.view(torch.uint8).fill_(173) + regions = source.get_cpu_transfer_regions() if page_device == "cpu" else source.get_pd_transfer_regions() + layout = StateTransferLayout( + source.state_name, + 4, + 16, + 16 + regions[-1].byte_end, + tuple(replace(r, byte_start=r.byte_start + 16, byte_end=r.byte_end + 16) for r in regions), + ) + pages = torch.full( + (3, layout.byte_end + 16), 91, dtype=torch.uint8, device=page_device, pin_memory=page_device == "cpu" + ) + if page_device == "cpu": + source.write_to_cpu(src, pages, [1], layout) + target.read_from_cpu(pages, [1], dst, layout) + else: + source.write_to_pd_page(src, pages, 1, layout) + target.read_from_pd_page(pages, dst, 1, layout) + torch.cuda.synchronize() + for tensor, reference in zip(qsa_tensors(target), expected): + assert torch.equal(tensor.view(torch.uint8)[:, dst], reference) + untouched = [i for i in range(tensor.shape[1]) if i not in (2, 6, 8)] + assert (tensor.view(torch.uint8)[:, untouched] == 173).all() + assert (pages[[0, 2]] == 91).all() + assert (pages[:, :16] == 91).all() and (pages[:, layout.byte_end :] == 91).all() + for region in layout.regions: + tail = pages[1, region.byte_start : region.byte_end].view(region.dtype).view(region.shape) + assert (tail[:, 3:].contiguous().view(torch.uint8) == 91).all() + + +@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, torch.full_like(accepted, 3)) + 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..fbf375c97d --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_batch.py @@ -0,0 +1,121 @@ +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) + + +@pytest.mark.parametrize("size", [0, 1, 8, 32, 128, 512, 1024]) +def test_decode_metadata_matches_variable_groups_and_graph_replay(size): + requests = torch.arange(size, device="cuda", dtype=torch.int32) + lengths = requests + 100 + ones = torch.ones_like(requests) + + def decode(): + return PackedBatchMetadata.prepare(requests, lengths, ones, size, 4095, 4, decode=True) + + decode() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = decode() + for group in (1, 2, 3, 4): + requests.copy_(torch.arange(size, device="cuda", dtype=torch.int32) // group) + requests[max(0, size - group) :] = 4095 + lengths.add_(7) + graph.replay() + expected = PackedBatchMetadata.prepare(requests, lengths, ones, size, 4095, 4) + for name in expected.__dataclass_fields__: + torch.testing.assert_close(getattr(actual, name), getattr(expected, name), atol=0, rtol=0) + graph.reset() 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..6323e01fcc --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_decoder.py @@ -0,0 +1,218 @@ +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_fixed_decode=False, + 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]) +@pytest.mark.parametrize("attention_splits", [None, 4]) +def test_long_decode_attention_matches_single_queries(dtype, prefill, attention_splits): + 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": attention_splits} + expected = torch.cat([execute(row, row + 1, **options) for row in range(rows)]) + if prefill: + torch.testing.assert_close(execute(0, rows, **options), expected, atol=1e-5, rtol=1e-5) + else: + torch.testing.assert_close(execute(0, rows, **options), expected, atol=0, rtol=0) + torch.testing.assert_close( + execute(0, 1, decode=False, attention_splits=attention_splits), 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..25fab5eb06 --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_lexical_cache.py @@ -0,0 +1,125 @@ +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("reset", [False, True]) +def test_history_gather_uses_current_request_and_prefix_on_graph_replay(config, reset): + cache, _ = make_cache(config) + meta = metadata([1, 3, 7, 7], [1, 10, 2, 2], [1] * 4) + cache.prepare(meta, reset) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = cache.prepare(meta, reset) + for request in (2, 4, 5): + meta.active_rows[:2].fill_(request) + meta.key_lengths[0] = 1 if request == 4 else 10 + cache.ngram.add_(1) + cache.convolution.normal_() + graph.replay() + rows = meta.active_rows + empty = rows < 0 + if reset: + empty |= meta.key_lengths == meta.starts[1:] - meta.starts[:-1] + ngram = cache.ngram[rows.clamp_min(0).long()].masked_fill(empty[:, None], cache.eos_token_id) + convolution = cache.convolution[rows.clamp_min(0).long()].masked_fill(empty[:, None, None], 0) + torch.testing.assert_close(actual.ngram_history, ngram, atol=0, rtol=0) + torch.testing.assert_close(actual.conv_history, convolution, atol=0, rtol=0) + graph.reset() + + +@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..4a220834fd --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_linear_decoder.py @@ -0,0 +1,184 @@ +from argparse import Namespace +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.models.qwen4_exp.layer_weights.transformer_layer_weight import Qwen4ExpTransformerLayerWeight +from lightllm.models.qwen4_exp.layer_infer.transformer_layer_infer import Qwen4ExpTransformerLayerInfer +from lightllm.models.qwen4_exp.state import PLEBatch +from lightllm.models.qwen4_exp.batch import PackedBatchMetadata +from lightllm.models.qwen4_exp.lexical_cache import LexicalRequestCache +from lightllm.common.quantization import Quantcfg +from lightllm.common.basemodel.attention.linear.triton import TritonLinearAttBackend +from .helpers import checkpoint +from .reference import reference_step + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +@pytest.mark.parametrize("capture", [False, True]) +def test_native_ple_gdn_and_moe_cached_steps(config, capture): + torch.manual_seed(1037) + source = checkpoint(config, 1) + weights = Qwen4ExpTransformerLayerWeight(1, torch.bfloat16, config, Quantcfg(config)) + weights.load_hf_weights(dict(source)) + weights.verify_load() + infer = Qwen4ExpTransformerLayerInfer(1, config) + channels = 512 + conv = torch.zeros(2, channels, 3, device="cuda", dtype=torch.bfloat16) + recurrent = torch.zeros(2, 2, 128, 128, device="cuda", dtype=torch.float32) + state = Namespace( + is_prefill=False, + need_dp_prefill_balance=False, + dist_group=None, + batch_size=2, + b_req_idx=torch.tensor([0, 1], device="cuda", dtype=torch.int32), + att_state_role="main", + att_state_manager=SimpleNamespace( + get_state=lambda **kwargs: SimpleNamespace(get_layer_view=lambda **kwargs: (conv, recurrent)) + ), + req_manager=SimpleNamespace(get_mamba_cache=lambda layer: (conv, recurrent)), + ) + + class Model: + pass + + model = Model() + model.config, model.tp_world_size_ = config, 1 + model.is_mtp_draft_model = False + model.mtp_manager = SimpleNamespace(get_decode_draft_step=lambda is_draft: 0) + backend = TritonLinearAttBackend(model=model) + state.decode_att_state1 = backend.create_att_decode_state(state) + state.decode_att_state1.init_state() + cache = LexicalRequestCache(weights, 2) + ref_ngram, ref_ple = cache.ngram.clone(), cache.convolution.clone() + ref_conv, ref_recurrent = conv.clone(), recurrent.clone() + ones = torch.ones(2, device="cuda", dtype=torch.int32) + metadata = PackedBatchMetadata.prepare(state.b_req_idx, ones, ones, 2, 7, 4) + + def run(residual, tokens): + batch = cache.prepare(metadata) + state.input_ids, state.qwen4_exp_ple = tokens, {1: batch} + output = infer.token_forward(residual, state, weights) + cache.commit(metadata, batch) + return output + + if capture: + residual = torch.randn(2, 64, device="cuda", dtype=torch.bfloat16) + tokens = torch.zeros(2, device="cuda", dtype=torch.int64) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + run(residual, tokens) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured_output = run(residual, tokens) + conv.zero_() + recurrent.zero_() + for request in range(2): + cache.reset(request) + for ids in [[7, 11], [63, 15], [19, 23]]: + if capture: + residual.normal_() + tokens.copy_(torch.tensor(ids, device="cuda", dtype=torch.int64)) + graph.replay() + actual = captured_output + else: + residual = torch.randn(2, 64, device="cuda", dtype=torch.bfloat16) + tokens = torch.tensor(ids, device="cuda", dtype=torch.int64) + actual = run(residual, tokens) + expected, ref_ngram, ref_ple, ref_conv, ref_recurrent = reference_step( + residual, tokens, source, config, ref_ngram, ref_ple, ref_conv, ref_recurrent + ) + torch.testing.assert_close(actual, expected, atol=0.016, rtol=0.016) + torch.testing.assert_close(cache.ngram, ref_ngram, atol=0, rtol=0) + torch.testing.assert_close(cache.convolution, ref_ple, atol=0.016, rtol=0.016) + torch.testing.assert_close(conv, ref_conv, atol=0.016, rtol=0.016) + + +def test_native_packed_gdn_ple_prefill_matches_recurrent_reference(config): + from lightllm.models.qwen3next.model import Qwen3NextTpPartModel + + torch.manual_seed(1037) + + Qwen3NextTpPartModel._init_triton(None) + source = checkpoint(config, 1) + weights = Qwen4ExpTransformerLayerWeight(1, torch.bfloat16, config, Quantcfg(config)) + weights.load_hf_weights(dict(source)) + weights.verify_load() + infer = Qwen4ExpTransformerLayerInfer(1, config) + conv = torch.zeros(2, 512, 3, device="cuda", dtype=torch.bfloat16) + recurrent = torch.zeros(2, 2, 128, 128, device="cuda", dtype=torch.float32) + starts = torch.tensor([0, 3, 5], device="cuda", dtype=torch.int32) + tokens = torch.tensor([7, 63, 19, 11, 23], device="cuda", dtype=torch.int64) + residual = torch.randn(5, 64, device="cuda", dtype=torch.bfloat16) + batch = PLEBatch(starts, weights.ngram_hasher.component.empty_history(2), weights.ple.component.empty_history(2)) + state = Namespace( + is_prefill=True, + need_dp_prefill_balance=False, + dist_group=None, + batch_size=2, + b_req_idx=torch.tensor([0, 1], device="cuda", dtype=torch.int32), + b1_cu_q_seq_len=starts, + max_q_seq_len=3, + 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..d849bbee4c --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_mtp.py @@ -0,0 +1,311 @@ +import math +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.basemodel.batch_objs import ModelInput +from lightllm.common.basemodel.cuda_graph import CudaGraph +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 + + +@pytest.mark.parametrize("width", [3, 4, 5]) +def test_replicated_decode_graph_layout(monkeypatch, width): + import lightllm.common.basemodel.cuda_graph as graph_module + from lightllm.models.qwen4_exp.model import Qwen4ExpTpPartModel + from lightllm.models.qwen4_exp.layer_infer.replicated_decoder import ReplicatedDecodeGraph + + args = SimpleNamespace(enable_tpsp_mix_mode=True) + monkeypatch.setattr(graph_module, "get_env_start_args", lambda: args) + model = Qwen4ExpTpPartModel.__new__(Qwen4ExpTpPartModel) + model.args, model.tp_world_size_, model.is_mtp_draft_model = args, 8, False + model.mtp_manager = SimpleNamespace(get_decode_batch_alignment=lambda _: width) + model._replicated_decode = True + sizes = ReplicatedDecodeGraph.gen_cuda_graph_batch_sizes( + batch_step_size_before_split=width, + split_batch_size=4 * width, + batch_step_size_after_split=2 * width, + max_batch_size=64 * width, + tp_world_size=8, + ) + assert set(range(width, 65, width)) <= set(sizes) + assert all(size % width == 0 for size in sizes) + assert all(size % 8 == 0 for size in sizes if size > 64) + for rows in (1, width, width + 1, 60, 63, 64, 65, 64 * width): + aligned = model._align_decode_batch_size(rows) + assert aligned >= rows and aligned % width == 0 + if aligned > 64: + assert aligned % 8 == 0 + captured = next(size for size in sizes if size >= aligned) + assert captured % width == 0 + if captured > 64: + assert captured % 8 == 0 + for requests in range(1, 64 // width + 1): + rows = requests * width + assert model._align_decode_batch_size(rows) == rows + model._replicated_decode = False + assert model._align_decode_batch_size(width) == math.lcm(width, 8) + + +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"}, + {"mtp_step": 2, "mtp_mode": "eagle_with_att", "page_size": 4}, + {"mtp_step": 4, "mtp_mode": "eagle_with_att", "page_size": 4}, + ], + indirect=True, +) +@pytest.mark.parametrize("graph_preparation", [False, True]) +def test_native_complete_draft_prefill_cached_decode_and_full_hidden(runtime, config, graph_preparation): + 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], graph_mode="capture" if graph_preparation else None) + 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) + if graph_preparation: + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = draft._token_forward(state) + graph.replay() + 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) + + if graph_preparation: + following = state_for(owner, [1, 3], [5, 4], [31, 37], graph_mode="replay") + following_hidden = torch.randn(2, 64, device="cuda", dtype=torch.bfloat16) + following.mtp_draft_input_hiddens = following_hidden + following.hidden_collector = FinalHiddenCollector() + state.copy_for_cuda_graph(following) + graph.replay() + full_tokens = torch.tensor([7, 11, 13, 29, 31, 17, 19, 23, 37], device="cuda") + full_hidden = torch.cat([full_hidden[:4], following_hidden[:1], full_hidden[4:], following_hidden[1:]]) + positions = torch.tensor([0, 1, 2, 3, 4, 0, 1, 2, 3], 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), + [5, 4], + ) + torch.testing.assert_close(actual.logits, logits[[4, 8]], atol=0.025, rtol=0.025) + torch.testing.assert_close(actual.mtp_collector.spec_hidden, multi[[4, 8]], atol=0.016, rtol=0.016) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@pytest.mark.parametrize("runtime", [{"mtp_step": 3, "mtp_mode": "eagle_with_att", "page_size": 4}], indirect=True) +@pytest.mark.parametrize("prewarm", [False, True]) +def test_selected_draft_logits_keep_hidden_and_kv(runtime, config, monkeypatch, prewarm): + import copy + from lightllm.distributed import dist_group_manager + + torch.manual_seed(2309) + Qwen3NextTpPartModel._init_triton(None) + draft, _ = build_draft(runtime, config) + draft.hidden_collector_prototype = FinalHiddenCollector() + draft.return_all_prompt_logics = False + draft.graph = None + draft._align_decode_batch_size = lambda _: 8 + monkeypatch.setattr(dist_group_manager, "get_group", lambda _: None) + graph = CudaGraph(8, 8, 8, max_batch_size=8, max_len_in_batch=64) + if prewarm: + draft.graph = graph + graph.warmup(draft) + assert set(graph.graph) == {8, (8, 1), (8, 2), (8, 4), (8, 8)} + draft.graph = None + + def unexpected_capture(*args, **kwargs): + raise AssertionError("Selected logits graph was not warmed before serving") + + monkeypatch.setattr(graph, "capture_decode", unexpected_capture) + 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) + prefix = state_for(owner, [1, 3], [3, 3], [7, 11, 13, 17, 19, 23], ready=[0, 0]) + prefix.mtp_draft_input_hiddens = torch.randn(6, 64, device="cuda", dtype=torch.bfloat16) + prefix.hidden_collector = FinalHiddenCollector() + draft._context_forward(prefix) + state = state_for(owner, [1, 1, 3, 3], [4, 5, 4, 5], [29, 31, 37, 41]) + model_input = ModelInput( + batch_size=4, + total_token_num=18, + max_q_seq_len=1, + max_kv_seq_len=5, + input_ids=state.input_ids, + b_req_idx=state.b_req_idx, + b_mtp_index=torch.tensor([0, 1, 0, 1], device="cuda", dtype=torch.int32), + b_seq_len=state.b_seq_len, + b_position_delta=state.b_position_delta, + b_shared_seq_len=torch.zeros(4, device="cuda", dtype=torch.int32), + b_shared_radix_node_id=torch.full((4,), -1, device="cuda", dtype=torch.int64), + multimodal_params=[{"images": [], "audios": []} for _ in range(4)], + mtp_draft_input_hiddens=torch.randn(4, 64, device="cuda", dtype=torch.bfloat16), + ) + full = draft.forward(model_input) + buffers = [] + for resource in draft.att_state_manager.get_full_attention_states(draft.att_state_role): + buffers.extend([resource.k_storage, resource.v_storage, resource.indexer.k_storage]) + assert resource.indexer.v_storage is None + kv = [buffer.index_select(1, state.mem_index.long()) for buffer in buffers] + draft.graph = graph + for selection in ([1, 3], [0, 2], [3], [0, 1, 3]): + selected = copy.copy(model_input) + selected.logits_indices = torch.tensor(selection, device="cuda", dtype=torch.int64) + output = draft.forward(selected) + torch.testing.assert_close(output.logits, full.logits[selection], atol=0.025, rtol=0.025) + torch.testing.assert_close(output.mtp_collector.spec_hidden, full.mtp_collector.spec_hidden, atol=0, rtol=0) + for buffer, values in zip(buffers, kv): + torch.testing.assert_close( + buffer.index_select(1, state.mem_index.long()), values, atol=0, rtol=0, equal_nan=True + ) + assert selected.logits_indices.numel() == len(selection) + assert model_input.logits_indices is None + if not prewarm: + assert set(graph.graph) == {(8, 1), (8, 2), (8, 4)} + model_input.input_ids.add_(2) + model_input.mtp_draft_input_hiddens.normal_() + draft.graph = None + expected = draft.forward(model_input) + draft.graph = graph + selected = copy.copy(model_input) + selected.logits_indices = torch.tensor([0, 3], device="cuda", dtype=torch.int64) + output = draft.forward(selected) + torch.testing.assert_close(output.logits, expected.logits[[0, 3]], atol=0.025, rtol=0.025) + torch.testing.assert_close(output.mtp_collector.spec_hidden, expected.mtp_collector.spec_hidden, atol=0, rtol=0) + complete = draft.forward(model_input) + torch.testing.assert_close(complete.logits, expected.logits, atol=0.025, rtol=0.025) + if not prewarm: + assert set(graph.graph) == {8, (8, 1), (8, 2), (8, 4)} diff --git a/unit_tests/models/qwen4_exp/test_multimodal.py b/unit_tests/models/qwen4_exp/test_multimodal.py new file mode 100644 index 0000000000..6cb7bc2b10 --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_multimodal.py @@ -0,0 +1,213 @@ +from copy import deepcopy +import json +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F + +from lightllm.models.qwen4_exp.configuration import normalize_model_config +from lightllm.models.qwen4_exp.multimodal import canonical_image_ids +from lightllm.server.router.model_infer.infer_batch import g_infer_context +from lightllm.common.basemodel.hidden_collector import NoopHiddenCollector +from .runtime import assemble_model, state_for +from .reference import reference_model + + +def test_model_types_resolve_through_public_registry(): + from lightllm.models.registry import ModelRegistry + from lightllm.models.qwen4_exp.model import Qwen4ExpTpPartModel + + for model_type in ("qwen4_exp", "qwen4_exp_text"): + assert ModelRegistry.get_model_class({"model_type": model_type}) is Qwen4ExpTpPartModel + + +@pytest.mark.parametrize("model_type", ["qwen4_exp", "qwen4_exp_text"]) +def test_startup_detects_hybrid_cache_and_optional_vision(model_type, tmp_path): + from lightllm.utils.config_utils import is_linear_att_mixed_model, has_vision_module, auto_set_fused_shared_experts + + config = {"model_type": model_type, "architectures": ["Qwen4ExpForConditionalGeneration"]} + if model_type == "qwen4_exp": + config["vision_config"] = {"hidden_size": 32} + (tmp_path / "config.json").write_text(json.dumps(config)) + assert is_linear_att_mixed_model(str(tmp_path)) + assert has_vision_module(str(tmp_path)) == (model_type == "qwen4_exp") + args = SimpleNamespace(model_dir=str(tmp_path), enable_ep_moe=False, enable_fused_shared_experts=False) + auto_set_fused_shared_experts(args) + assert args.enable_fused_shared_experts + + +def test_tokenizer_uses_native_processor_and_virtual_image_ranges(monkeypatch): + from lightllm.server import tokenizer as tokenizer_module + from lightllm.models.qwen4_exp.image_processing import Qwen4ExpImageProcessor + from lightllm.models.qwen4_exp.tokenizer import Qwen4ExpTokenizer + + config = dict( + model_type="qwen4_exp", + architectures=["Qwen4ExpForConditionalGeneration"], + vision_config={}, + vision_start_token_id=61, + vision_end_token_id=62, + image_token_id=60, + ) + monkeypatch.setattr(tokenizer_module.AutoTokenizer, "from_pretrained", lambda *args, **kwargs: object()) + monkeypatch.setattr(tokenizer_module.PretrainedConfig, "get_config_dict", lambda *args: (config, {})) + processor = SimpleNamespace(min_pixels=16, max_pixels=64, patch_size=2, merge_size=2) + monkeypatch.setattr(Qwen4ExpImageProcessor, "from_pretrained", lambda *args: processor) + tokenizer = tokenizer_module.get_tokenizer("checkpoint") + assert type(tokenizer) is Qwen4ExpTokenizer + image = SimpleNamespace(token_id=1000000, token_num=2, start_idx=None) + actual = tokenizer.encode([7, 61, 60, 62, 11], SimpleNamespace(images=[image])) + assert actual == [7, 61, 1000000, 1000001, 62, 11] + assert image.start_idx == 2 + assert tokenizer.image_processor is processor + + +@pytest.mark.parametrize("image_count", [0, 1, 2]) +@pytest.mark.parametrize("qwen4", [False, True]) +def test_token_count_hook_matches_qwen4_expansion_and_preserves_legacy(image_count, qwen4): + from lightllm.models.qwen2_vl.model import QWen2VLTokenizer + from lightllm.models.qwen4_exp.tokenizer import Qwen4ExpTokenizer + from lightllm.server.httpserver.manager import HttpServerManager + + cls = Qwen4ExpTokenizer if qwen4 else QWen2VLTokenizer + processor = SimpleNamespace(min_pixels=16, max_pixels=64, patch_size=2, merge_size=2) + tokenizer = cls( + tokenizer=object(), + image_processor=processor, + model_cfg=dict(vision_start_token_id=61, vision_end_token_id=62, image_token_id=60), + ) + images = [] + for index in range(image_count): + image = SimpleNamespace(image_w=4 * (index + 1), image_h=4, token_id=1000 + index * 100) + image.token_num = tokenizer.get_image_token_length(image) + images.append(image) + multimodal = SimpleNamespace(images=images, audios=[]) + prompt = [7] + [61, 60, 62] * image_count + [11] + expanded = tokenizer.encode(prompt, multimodal) + manager = HttpServerManager.__new__(HttpServerManager) + manager.tokenizer = tokenizer + checked = [] + manager._assert_image_token_count = checked.append + actual = manager.tokens(prompt, multimodal, SimpleNamespace()) + assert checked == [image.token_num for image in images] + assert actual == len(expanded) + + +def test_plain_text_counting_fallback_is_unchanged(): + from lightllm.server.httpserver.manager import HttpServerManager + + manager = HttpServerManager.__new__(HttpServerManager) + manager.tokenizer = SimpleNamespace(encode=lambda prompt, *_args, **_kwargs: prompt) + assert manager.tokens([7, 11], SimpleNamespace(images=[], audios=[]), SimpleNamespace()) == 2 + + +def test_token_count_hook_cannot_bypass_image_limits(): + from lightllm.server.httpserver.manager import HttpServerManager + from lightllm.models.qwen4_exp.tokenizer import Qwen4ExpTokenizer + + tokenizer = Qwen4ExpTokenizer( + tokenizer=object(), + image_processor=SimpleNamespace(min_pixels=16, max_pixels=64, patch_size=2, merge_size=2), + model_cfg=dict(vision_start_token_id=61, vision_end_token_id=62, image_token_id=60), + ) + manager = HttpServerManager.__new__(HttpServerManager) + manager.tokenizer = tokenizer + + def reject(_count): + raise ValueError("image token limit") + + manager._assert_image_token_count = reject + with pytest.raises(ValueError, match="image token limit"): + manager.tokens( + [61, 60, 62], + SimpleNamespace(images=[SimpleNamespace(image_w=4, image_h=4)], audios=[]), + SimpleNamespace(), + ) + + +def test_image_processor_loads_real_config_without_runtime_enum_serialization(tmp_path): + from lightllm.models.qwen4_exp.image_processing import Qwen4ExpImageProcessor + + (tmp_path / "preprocessor_config.json").write_text( + json.dumps( + dict( + size={"shortest_edge": 16, "longest_edge": 64}, + patch_size=2, + merge_size=2, + image_mean=[0.5] * 3, + image_std=[0.5] * 3, + image_processor_type="Qwen2VLImageProcessorFast", + ) + ) + ) + processor = Qwen4ExpImageProcessor.from_pretrained(str(tmp_path)) + assert processor.min_pixels == 16 and processor.max_pixels == 64 + assert processor.patch_size == processor.merge_size == 2 + assert processor.image_mean == processor.image_std == [0.5] * 3 + + +def test_normalize_outer_and_text_config_preserves_checkpoint(config): + text = dict(config, rope_parameters={"rope_theta": 1000, "partial_rotary_factor": 0.5}) + text.pop("n_layer") + outer = {"text_config": text, "vision_config": {"hidden_size": 32}, "image_token_id": 60} + original = deepcopy(outer) + actual, vision = normalize_model_config(outer) + standalone, no_vision = normalize_model_config(text) + assert actual["n_layer"] == standalone["n_layer"] == 4 + assert actual["image_token_id"] == 60 + assert actual["rope_theta"] == standalone["rope_theta"] == 1000 + assert vision == outer["vision_config"] and no_vision is None + assert outer == original + customized, _ = normalize_model_config(text, SimpleNamespace(vocab_size=128)) + assert customized["vocab_size"] == 128 and text["vocab_size"] == 64 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@pytest.mark.parametrize("dtype", [torch.int32, torch.int64]) +def test_image_ranges_preserve_text_and_physical_ids(dtype): + second = 2 ** 34 if dtype == torch.int64 else 2000000 + tokens = torch.tensor( + [7, 60, 999999, 1000000, 1000001, 1000002, second + 4, second + 5, 63], device="cuda", dtype=dtype + ) + starts = torch.tensor([1000000, second], device="cuda", dtype=torch.int64) + lengths = torch.tensor([2, 5], device="cuda", dtype=torch.int64) + before = tokens.clone() + actual = canonical_image_ids(tokens, starts, lengths, 60) + expected = torch.tensor([7, 60, 999999, 60, 60, 1000002, 60, second + 5, 63], device="cuda", dtype=dtype) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + torch.testing.assert_close(tokens, before, atol=0, rtol=0) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = canonical_image_ids(tokens, starts, lengths, 60) + tokens[3:5] = 11 + graph.replay() + expected[3:5] = 11 + torch.testing.assert_close(captured, expected, atol=0, rtol=0) + graph.reset() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_image_cache_ids_are_canonicalized_before_ple_commit(runtime, config, monkeypatch): + config["image_token_id"] = 60 + model, source = assemble_model(runtime, config) + runtime.req.init_linear_att_state(SimpleNamespace(req_idx=1)) + image_cache = torch.randn(4, 1, 32, dtype=torch.bfloat16).pin_memory() + monkeypatch.setattr(g_infer_context, "cpu_embed_cache_client", SimpleNamespace(cpu_embed_cache_tensor=image_cache)) + virtual = [7, 1000000, 1000001, 11] + canonical = torch.tensor([7, 60, 60, 11], device="cuda", dtype=torch.int64) + state = state_for(runtime, [1], [4], virtual, ready=[0]) + state.hidden_collector = NoopHiddenCollector() + state.multimodal_params = [ + {"images": [{"token_id": 1000000, "token_num": 2, "start_index_in_embed_cache": 1}], "audios": []} + ] + actual = model._context_forward(state) + torch.testing.assert_close(state.input_ids, torch.tensor(virtual, device="cuda"), atol=0, rtol=0) + torch.testing.assert_close(state.qwen4_exp_ple[1].input_tokens, canonical, atol=0, rtol=0) + torch.testing.assert_close(runtime.req.lexical_caches[1].ngram[1], canonical[-2:], atol=0, rtol=0) + embeddings = F.embedding(canonical.cpu(), source["model.language_model.embed_tokens.weight"]) + embeddings[1:3] = image_cache[1:3, 0] + positions = torch.arange(4, device="cuda") + cos, sin = model._cos_cached[positions].repeat(1, 2), model._sin_cached[positions].repeat(1, 2) + expected = reference_model(canonical, cos, sin, [4], source, config, input_embeddings=embeddings.cuda())[-1:] + torch.testing.assert_close(actual.logits, expected, atol=0.025, rtol=0.025) diff --git a/unit_tests/models/qwen4_exp/test_norm.py b/unit_tests/models/qwen4_exp/test_norm.py new file mode 100644 index 0000000000..50ccc871c0 --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_norm.py @@ -0,0 +1,48 @@ +import pytest +import torch + +from lightllm.models.qwen4_exp.layer_weights.norm_weight import Qwen4ExpGatedNormWeight + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("activation", ["sigmoid", "silu"]) +def test_native_gdn_norm_rounding(dtype, activation): + torch.manual_seed(837) + native = Qwen4ExpGatedNormWeight(128, "linear.norm.weight", dtype, activation) + weight = 1 + torch.randn(128, dtype=dtype) * 0.1 + native.load_hf_weights({"linear.norm.weight": weight}) + values = torch.randn(19, 128, device="cuda", dtype=dtype) + gate = torch.randn_like(values) * 3 + normalized = values.float() * torch.rsqrt(values.float().square().mean(-1, keepdim=True) + 1e-6) + weighted = native.weight * normalized.to(dtype) + gate_activation = gate.float().sigmoid() + if activation == "silu": + gate_activation *= gate.float() + expected = (weighted * gate_activation).to(dtype) + out = torch.empty_like(values) + actual = native(values, gate, 1e-6, out=out) + assert actual is out + tolerance = 2e-6 if dtype == torch.float32 else torch.finfo(dtype).eps * 2 + torch.testing.assert_close(actual, expected, atol=tolerance, rtol=tolerance) + + +@pytest.mark.parametrize("activation", ["sigmoid", "silu"]) +def test_gdn_norm_accepts_strided_gate(activation): + from lightllm.models.qwen4_exp.triton_kernel.gated_norm import gated_norm + + torch.manual_seed(123) + tokens, heads, dim = 5, 4, 128 + values = torch.randn(tokens * heads, dim, device="cuda", dtype=torch.bfloat16) + packed = torch.randn(tokens, heads * dim + 256, device="cuda", dtype=torch.bfloat16) + gate = packed[:, 128 : 128 + heads * dim].view(tokens, heads, dim) + weight = torch.randn(dim, device="cuda", dtype=torch.bfloat16) + + actual = gated_norm(values, gate, weight, 1e-6, activation) + normalized = values.float() * torch.rsqrt(values.float().square().mean(-1, keepdim=True) + 1e-6) + expected = normalized * weight.float() + gate = gate.reshape(-1, dim).float() + gate = gate.sigmoid() if activation == "sigmoid" else torch.nn.functional.silu(gate) + expected = (expected * gate).to(values.dtype) + torch.testing.assert_close(actual, expected, atol=0.008, rtol=0.008) diff --git a/unit_tests/models/qwen4_exp/test_pinned_memory.py b/unit_tests/models/qwen4_exp/test_pinned_memory.py new file mode 100644 index 0000000000..cffc72cf36 --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_pinned_memory.py @@ -0,0 +1,112 @@ +import gc +import weakref + +import pytest +import torch +import triton +import triton.language as tl + +from lightllm.models.qwen4_exp import pinned_memory + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +@triton.jit +def read_ends(Source, Output, SIZE: tl.constexpr): + row = tl.arange(0, 2) + position = tl.where(row == 0, 0, SIZE - 1) + tl.store(Output + row, tl.load(Source + position)) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +def test_registered_snapshot_views_graph_and_lifetime(dtype, monkeypatch): + monkeypatch.setattr(pinned_memory, "_LARGE_BYTES", 1) + monkeypatch.setattr(pinned_memory, "_CHUNK_BYTES", 4096) + value = pinned_memory.pinned_zeros((6144,), dtype=dtype) + assert value.is_pinned() and value.is_contiguous() + assert torch.count_nonzero(value) == 0 + value[0], value[-1] = 3, 7 + output = torch.empty(2, device="cuda", dtype=torch.float32) + read_ends[(1,)](value, output, value.numel()) + torch.testing.assert_close(output, torch.tensor([3, 7], device="cuda", dtype=torch.float32), atol=0, rtol=0) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + read_ends[(1,)](value, output, value.numel()) + value[-1] = -7 + graph.replay() + torch.testing.assert_close(output, torch.tensor([3, -7], device="cuda", dtype=torch.float32), atol=0, rtol=0) + graph.reset() + owner = weakref.ref(value._qwen4_pinned_allocation) + view = value[4096:] + del value + gc.collect() + assert owner() is not None and view.is_pinned() + read_ends[(1,)](view, output, view.numel()) + torch.testing.assert_close(output, torch.tensor([0, -7], device="cuda", dtype=torch.float32), atol=0, rtol=0) + del view + gc.collect() + assert owner() is None + + +def test_partial_registration_failure_releases_registered_ranges(monkeypatch): + monkeypatch.setattr(pinned_memory, "_LARGE_BYTES", 1) + monkeypatch.setattr(pinned_memory, "_CHUNK_BYTES", 4096) + runtime = torch.cuda.cudart() + registered, released = [], [] + + class FailingRuntime: + def cudaHostRegister(self, address, size, flags): + if registered: + return 1 + status = runtime.cudaHostRegister(address, size, flags) + assert int(status) == 0 + registered.append(address) + return status + + def cudaHostUnregister(self, address): + released.append(address) + return runtime.cudaHostUnregister(address) + + def cudaGetErrorString(self, status): + return "intentional registration failure" + + monkeypatch.setattr(torch.cuda, "cudart", lambda: FailingRuntime()) + + def allocate(): + try: + pinned_memory.pinned_zeros((6144,), dtype=torch.float32) + except RuntimeError as error: + assert "intentional registration failure" in str(error) + else: + pytest.fail("Expected registration failure") + + allocate() + gc.collect() + assert registered and released == registered + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +def test_snapshot_dma_restore_never_crosses_registration(dtype, monkeypatch): + monkeypatch.setattr(pinned_memory, "_LARGE_BYTES", 1) + monkeypatch.setattr(pinned_memory, "_CHUNK_BYTES", 16384) + value = pinned_memory.pinned_zeros((7, 3, 512), dtype=dtype) + cell_bytes = value[0].numel() * value.element_size() + starts = value._qwen4_pinned_allocation.addresses + assert len(starts) > 1 + assert all((address - value.data_ptr()) % cell_bytes == 0 for address in starts) + ends = starts[1:] + [value.data_ptr() + value.numel() * value.element_size()] + assert all(end - start <= 16384 for start, end in zip(starts, ends)) + destination = torch.empty((3, 5, 512), dtype=dtype, device="cuda") + for index in range(value.shape[0]): + value[index].fill_(index + 1) + + destination[:, 1, ...] = value[index] + torch.cuda.synchronize() + torch.testing.assert_close(destination[:, 1].cpu(), value[index], atol=0, rtol=0) + + +def test_snapshot_larger_than_registration_bound_is_rejected(monkeypatch): + monkeypatch.setattr(pinned_memory, "_LARGE_BYTES", 1) + monkeypatch.setattr(pinned_memory, "_CHUNK_BYTES", 16384) + with pytest.raises(ValueError, match="one bounded host registration"): + pinned_memory.pinned_zeros((2, 8192), dtype=torch.float32) diff --git a/unit_tests/models/qwen4_exp/test_position.py b/unit_tests/models/qwen4_exp/test_position.py new file mode 100644 index 0000000000..603e0ccfb1 --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_position.py @@ -0,0 +1,28 @@ +import pytest +import torch + +from lightllm.models.qwen4_exp.position import combine_mrope + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@pytest.mark.parametrize("rows", [0, 1, 32, 512]) +@pytest.mark.parametrize("axes", [False, True]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) +def test_rotary_axis_selection_matches_cpu_on_graph_replay(rows, axes, dtype): + shape = (3, rows, 66) if axes else (rows, 66) + cos = torch.randn(shape, device="cuda", dtype=dtype)[..., ::2] + sin = torch.randn(shape, device="cuda", dtype=dtype)[..., ::2] + sections = [8, 8, 0] + combine_mrope(cos, sin, sections) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = combine_mrope(cos, sin, sections) + for _ in range(3): + cos.add_(1) + sin.mul_(2) + graph.replay() + expected = combine_mrope(cos.cpu(), sin.cpu(), sections) + for result, reference in zip(actual, expected): + assert result.is_contiguous() + torch.testing.assert_close(result.cpu(), reference, atol=0, rtol=0) + graph.reset() diff --git a/unit_tests/models/qwen4_exp/test_qsa_compressor.py b/unit_tests/models/qwen4_exp/test_qsa_compressor.py new file mode 100644 index 0000000000..646449e87f --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_qsa_compressor.py @@ -0,0 +1,218 @@ +from dataclasses import replace + +import pytest +import torch + +from lightllm.common.attention_state.types import StateTransferLayout +from lightllm.common.qsa.indexer import QSAIndexer +from lightllm.models.qwen4_exp.batch import PackedBatchMetadata +from lightllm.models.qwen4_exp.qsa_compressor import QSACompressorState, pool_with_ring +from lightllm.models.qwen4_exp.state import QSABatch +from lightllm.models.qwen4_exp.model import Qwen4ExpTpPartModel +from lightllm.models.qwen4_exp.mtp.model import Qwen4ExpMTPModel +from lightllm.utils.envs_utils import get_env_start_args + + +@pytest.mark.parametrize("fp8", [False, True]) +def test_compact_factory_accounts_for_storage_and_draft(config, fp8): + args = get_env_start_args() + args.data_type = "bfloat16" + args.page_size = 4 + args.llm_kv_type = "fp8kv_spt" if fp8 else "None" + args.mtp_step = 3 + (state,) = Qwen4ExpTpPartModel.create_full_att_states(config, args, (0, 1, 2, 3)) + state.device = torch.device("cpu") + state.allocate_storage(16) + assert state.positions is None + assert state.indexer.v_storage is None + assert state.compress_ratio == 4 + assert state.qsa_view(3).pooled_keys.shape == (5, 16) + total = sum(t.numel() * t.element_size() for t in (state.k_storage, state.v_storage, state.indexer.k_storage)) + assert total == 20 * state.get_token_bytes() + assert len(state._get_transfer_regions(8)) == 3 + for cls, layers in ((Qwen4ExpTpPartModel, (0, 1, 2, 3)), (Qwen4ExpMTPModel, (4,))): + bounded = cls.create_bounded_att_states(config, args, layers, 4, device="cpu") + compressor = next(s for s in bounded if isinstance(s, QSACompressorState)) + assert compressor.buffer.shape == (1, 4, 7, 24) + assert not compressor.get_cpu_transfer_regions() + assert compressor.get_pd_transfer_regions()[0].shape == (1, 4, 24) + if cls is Qwen4ExpMTPModel: + assert len(bounded) == 1 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("fp8", [False, True]) +@pytest.mark.parametrize("graph", [False, True]) +@pytest.mark.parametrize("segment_capacity", [None, 257]) +def test_ring_pool_matches_full_history_after_chunking_and_rejection(fp8, graph, segment_capacity): + torch.manual_seed(2437) + device, dtype = "cuda", torch.bfloat16 + dim, rotary, capacity = 16, 4, 64 + indexer = QSAIndexer( + torch.randn(48, 32, device=device, dtype=dtype), + torch.randn(dim, device=device, dtype=dtype), + torch.randn(dim, device=device, dtype=dtype), + query_heads=2, + token_budget=8, + ) + storage_dtype = torch.float8_e4m3fn if fp8 else dtype + raw = torch.zeros(capacity, dim, device=device, dtype=dtype) + cos = torch.zeros(capacity, rotary, device=device, dtype=dtype) + sin = torch.zeros_like(cos) + dense = torch.empty(4, 8, dim, device=device, dtype=storage_dtype) + compact = torch.empty(capacity // 4, dim, device=device, dtype=storage_dtype) + ring = QSACompressorState((3,), dim, rotary, 4, dtype, 4, 3, 0, 1, device).get_layer_view(3) + table = torch.full((4, 32), -1, device=device, dtype=torch.int32) + slots = torch.randperm(16, device=device).view(2, 8) * 4 + for row, pages in zip((1, 2), slots): + table[row] = (pages[:, None] + torch.arange(4, device=device)).flatten().int() + + def phase(prefixes, counts, captured=False): + lengths = [a + b for a, b in zip(prefixes, counts)] + metadata = PackedBatchMetadata.prepare( + torch.tensor([1, 2], device=device, dtype=torch.int32), + torch.tensor(lengths, device=device, dtype=torch.int32), + torch.tensor(counts, device=device, dtype=torch.int32), + sum(counts), + 3, + 4, + segment_capacity=segment_capacity, + ) + ids = torch.cat([table[r, p : p + n] for r, p, n in zip((1, 2), prefixes, counts)]).long() + projected = torch.randn(sum(counts), 48, device=device, dtype=dtype) + current = projected[:, 32:] + angles = torch.randn(sum(counts), rotary, device=device) + c, s = angles.cos().to(dtype), angles.sin().to(dtype) + batch = QSABatch.from_metadata(metadata, table, ids, c, s, 4) + if captured: + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + saved = ring.clone() + with torch.cuda.stream(stream): + for _ in range(3): + pool_with_ring(indexer, current, ring, batch, compact) + torch.cuda.current_stream().wait_stream(stream) + ring.copy_(saved) + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + pool_with_ring(indexer, current, ring, batch, compact) + ring.copy_(saved) + g.replay() + else: + pool_with_ring(indexer, current, ring, batch, compact) + raw.index_copy_(0, ids, current) + cos.index_copy_(0, ids, c) + sin.index_copy_(0, ids, s) + indexer.pool( + raw, + table, + batch.key_lengths, + cos, + sin, + dense, + batch.first_blocks, + batch.max_new_blocks, + request_rows=batch.active_rows, + ) + for r, length in zip((1, 2), lengths): + completed = length // 4 + physical = table[r, 3 : completed * 4 : 4].long() // 4 + torch.testing.assert_close(compact.float()[physical], dense.float()[r, :completed], atol=0, rtol=0) + query = torch.randn(sum(counts), 2, dim, device=device, dtype=dtype).to(storage_dtype) + actual = indexer.select(query, compact, batch.query_rows, batch.visible_lengths, table, physical_stride=4) + expected = indexer.select(query, dense, batch.query_rows, batch.visible_lengths, table) + torch.testing.assert_close(actual.sort(-1).values, expected.sort(-1).values, atol=0, rtol=0) + + phase([0, 0], [3, 2]) + phase([3, 2], [6, 7]) + phase([9, 9], [4, 4], graph) + phase([10, 11], [4, 4], graph) + phase([14, 15], [3, 2], graph) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("fp8", [False, True]) +@pytest.mark.parametrize("page_device", ["cpu", "cuda"]) +def test_compact_token_storage_round_trip(config, fp8, page_device): + args = get_env_start_args() + args.data_type, args.page_size = "bfloat16", 4 + args.llm_kv_type = "fp8kv_spt" if fp8 else "None" + args.cpu_cache_token_page_size = args.pd_kv_page_size = 8 + states = [] + for _ in range(2): + (state,) = Qwen4ExpTpPartModel.create_full_att_states(config, args, (0, 1, 2, 3)) + state.device = torch.device("cuda") + state.allocate_storage(16) + states.append(state) + source, target = states + + def tensors(state): + return state.k_storage, state.v_storage, state.indexer.k_storage + + for tensor in tensors(source): + tensor.view(torch.uint8).random_(0, 256) + for tensor in tensors(target): + tensor.view(torch.uint8).fill_(173) + src = torch.arange(4, 12, device="cuda", dtype=torch.int32) + dst = torch.cat([torch.arange(12, 16, device="cuda"), torch.arange(0, 4, device="cuda")]).long() + expected = [tensor.view(torch.uint8)[:, src].clone() for tensor in tensors(source)] + source.copy_full_att_tokens(src, dst) + target.copy_tokens_from_dp(src, dst, source) + for tensor, gold in zip(tensors(target), expected): + assert torch.equal(tensor.view(torch.uint8)[:, dst], gold) + tensor.view(torch.uint8).fill_(173) + regions = source.get_cpu_transfer_regions() if page_device == "cpu" else source.get_pd_transfer_regions() + layout = StateTransferLayout( + source.state_name, + 8, + 16, + 16 + regions[-1].byte_end, + tuple(replace(r, byte_start=r.byte_start + 16, byte_end=r.byte_end + 16) for r in regions), + ) + pages = torch.full( + (2, layout.byte_end + 16), 91, dtype=torch.uint8, device=page_device, pin_memory=page_device == "cpu" + ) + if page_device == "cpu": + source.write_to_cpu(src, pages, [1], layout) + target.read_from_cpu(pages, [1], dst, layout) + else: + source.write_to_pd_page(src, pages, 1, layout) + target.read_from_pd_page(pages, dst, 1, layout) + torch.cuda.synchronize() + for tensor, gold in zip(tensors(target), expected): + assert torch.equal(tensor.view(torch.uint8)[:, dst], gold) + assert (pages[0] == 91).all() + assert (pages[:, :16] == 91).all() and (pages[:, layout.byte_end :] == 91).all() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("length", [0, 3, 4, 13]) +def test_compressor_pd_restores_logical_window_across_mtp_capacities(length): + source = QSACompressorState((3, 7), 16, 4, 4, torch.bfloat16, 4, 3, 0, 1, "cuda") + target = QSACompressorState((3, 7), 16, 4, 4, torch.bfloat16, 4, 1, 0, 1, "cuda") + source.buffer.normal_() + target.buffer.fill_(173) + source.state_name = target.state_name = "main_qsa" + regions = source.get_pd_transfer_regions() + layout = StateTransferLayout( + source.state_name, + None, + 16, + 16 + regions[-1].byte_end, + tuple(replace(r, byte_start=r.byte_start + 16, byte_end=r.byte_end + 16) for r in regions), + ) + pages = torch.full((2, layout.byte_end + 16), 91, dtype=torch.uint8, device="cuda") + source.write_to_pd_page(1, pages, 1, length, layout) + target.read_from_pd_page(pages, 2, 1, length, layout) + for position in range(length - 4, length): + actual = target.buffer[:, 2, position % 5] + expected = source.buffer[:, 1, position % 7] if position >= 0 else torch.zeros_like(actual) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + assert (pages[0] == 91).all() + assert (pages[:, :16] == 91).all() and (pages[:, layout.byte_end :] == 91).all() + before = source.buffer.clone() + source.save_bounded_att_runtime_to_pages([1], [512], "small", [0]) + source.restore_bounded_att_runtime_from_pages([2], [512], "small", [0]) + torch.testing.assert_close(source.buffer, before, atol=0, rtol=0) + with pytest.raises(AssertionError, match="close a compressor stride"): + source.restore_bounded_att_runtime_from_pages([2], [511], "small", [0]) diff --git a/unit_tests/models/qwen4_exp/test_reduction.py b/unit_tests/models/qwen4_exp/test_reduction.py new file mode 100644 index 0000000000..2eaf564f06 --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_reduction.py @@ -0,0 +1,104 @@ +import datetime +import json +from types import SimpleNamespace + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + + +def _check_reductions(rank, rendezvous, cuda=False): + import os + + from lightllm.models.qwen4_exp.layer_infer.transformer_layer_infer import Qwen4ExpTransformerLayerInfer + from lightllm.server.core.objs.start_args_type import StartArgs + from lightllm.utils.envs_utils import get_env_start_args + + args = StartArgs() + args.enable_tpsp_mix_mode = False + os.environ["LIGHTLLM_START_ARGS"] = json.dumps(vars(args)) + get_env_start_args.cache_clear() + device = torch.device("cuda", rank) if cuda else torch.device("cpu") + if cuda: + torch.cuda.set_device(rank) + dist.init_process_group( + "nccl" if cuda else "gloo", + init_method="file://" + rendezvous, + rank=rank, + world_size=4, + timeout=datetime.timedelta(seconds=60), + ) + try: + layer = Qwen4ExpTransformerLayerInfer.__new__(Qwen4ExpTransformerLayerInfer) + layer.tp_world_size_ = 4 + layer.tp_rank_ = rank + state = SimpleNamespace(dist_group=dist.group.WORLD) + + contribution = (256.0, 1.0, -256.0, 0.25)[rank] + for dtype in (torch.bfloat16, torch.float16, torch.float32): + for rows in (1, 3, 8): + value = torch.full((rows, 16), contribution, dtype=dtype, device=device) + value[:, 1::2].neg_() + expected = torch.full_like(value, 1.25) + expected[:, 1::2].neg_() + for reduce in (layer._reduce_linear_output, layer._tpsp_reduce): + actual = reduce(value.clone(), state) + assert actual.dtype == dtype + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + midpoint_parts = (1.0, 2 ** -8, 2 ** -24, 2 ** -24) + for rows in (1, 3, 256, 512): + value = torch.full((rows, 16), midpoint_parts[rank], dtype=torch.bfloat16, device=device) + expected = torch.full_like(value, 1.0078125) + for reduce in (layer._reduce_linear_output, layer._tpsp_reduce): + actual = reduce(value.clone(), state) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + value = torch.full((16, 3), contribution, dtype=torch.bfloat16, device=device).t() + assert layer._reduce_linear_output(value, state) is value + torch.testing.assert_close(value, torch.full_like(value, 1.25), atol=0, rtol=0) + + args.enable_tpsp_mix_mode = True + os.environ["LIGHTLLM_START_ARGS"] = json.dumps(vars(args)) + get_env_start_args.cache_clear() + for rows in (4, 12): + common = torch.arange(rows * 16, device=device, dtype=torch.float32).reshape(rows, 16) / 16 + value = (common + rank / 4).to(torch.bfloat16) + expected = (common * 4 + 1.5).to(value.dtype).chunk(4)[rank] + actual = layer._tpsp_reduce(value, state) + assert actual.shape == (rows // 4, 16) and actual.dtype == value.dtype + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + if cuda: + for _ in range(3): + layer._tpsp_reduce(value, state) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + result = layer._tpsp_reduce(value, state) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(result, expected, atol=0, rtol=0) + graph.reset() + finally: + dist.destroy_process_group() + + +@pytest.mark.skipif(not dist.is_available() or not dist.is_gloo_available(), reason="Gloo is required") +def test_tp_and_sequence_parallel_accumulation(tmp_path): + mp.spawn(_check_reductions, args=(str(tmp_path / "rendezvous"),), nprocs=4, join=True) + + +@pytest.mark.skipif(torch.cuda.device_count() < 4 or not dist.is_nccl_available(), reason="Four CUDA GPUs are required") +def test_cuda_tp_and_sequence_parallel_accumulation(tmp_path): + mp.spawn(_check_reductions, args=(str(tmp_path / "nccl-rendezvous"), True), nprocs=4, join=True) + + +def test_single_rank_preserves_tensor(): + from lightllm.models.qwen4_exp.layer_infer.transformer_layer_infer import Qwen4ExpTransformerLayerInfer + + layer = Qwen4ExpTransformerLayerInfer.__new__(Qwen4ExpTransformerLayerInfer) + layer.tp_world_size_ = 1 + value = torch.tensor([[1.25]], dtype=torch.bfloat16) + assert layer._reduce_linear_output(value, None) is value + assert layer._tpsp_reduce(value, None) is value diff --git a/unit_tests/models/qwen4_exp/test_replicated_moe.py b/unit_tests/models/qwen4_exp/test_replicated_moe.py new file mode 100644 index 0000000000..f3e55d2fb7 --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_replicated_moe.py @@ -0,0 +1,231 @@ +import pytest +import torch +import triton +import torch.nn.functional as F + +from lightllm.models.qwen3next.triton_kernel.shared_expert_gate import sigmoid_mul_ +from lightllm.models.qwen4_exp.triton_kernel.replicated_moe import ( + expert_down, + expert_up, + pack_routes, + sum_partials, + sum_routes, +) + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.mark.parametrize("rows", [3, 4, 5, 28, 63, 64]) +def test_routes_graph(rows): + topk, experts, first = 10, 64, 192 + ids = torch.zeros(rows, topk, dtype=torch.int32, device="cuda") + mask = torch.ones(rows, dtype=torch.bool, device="cuda") + masked = torch.empty_like(ids) + routes = torch.empty(experts, rows * topk, dtype=torch.int32, device="cuda") + counts = torch.empty(experts, dtype=torch.int32, device="cuda") + + def run(): + pack_routes[(experts,)]( + ids, + mask, + masked, + routes, + counts, + rows, + topk, + first, + triton.next_power_of_2(rows * topk), + num_warps=4 if rows * topk > 256 else 1, + ) + + run() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + for change in range(4): + torch.manual_seed(947 + change) + ids.copy_(torch.rand(rows, 512, device="cuda").argsort(dim=1)[:, :topk]) + if change == 2: + ids.copy_(first + torch.arange(topk, device="cuda")[None, :]) + mask.copy_(torch.arange(rows, device="cuda") % 3 != 0) + if change == 0: + mask.zero_() + graph.replay() + expected = ids.masked_fill(~mask[:, None], -1) + torch.testing.assert_close(masked, expected, rtol=0, atol=0) + for expert in range(experts): + positions = (expected.flatten() == first + expert).nonzero().flatten().to(torch.int32) + assert counts[expert].item() == positions.numel() + torch.testing.assert_close(routes[expert, : positions.numel()], positions, rtol=0, atol=0) + graph.reset() + + +def test_partials_order_graph(): + rows, hidden, topk = 3, 513, 10 + ids = torch.full((rows, topk), -1, dtype=torch.int32, device="cuda") + partials = torch.zeros(8, rows, hidden, dtype=torch.bfloat16, device="cuda") + partials[0].fill_(2 ** 25) + partials[1].fill_(-(2 ** 25)) + partials[2].fill_(1) + output = torch.empty(rows, hidden, dtype=torch.bfloat16, device="cuda") + + def run(): + sum_partials[(rows, triton.cdiv(hidden, 512))]( + partials, + ids, + output, + hidden, + rows, + topk, + 64, + 512, + 10, + num_warps=4, + enable_fp_fusion=False, + ) + + run() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + for routes, expected in (([0, 64, 128], 1), ([0, 64, 128, 0], 0), ([], 0)): + ids.fill_(-1) + if routes: + ids[:, : len(routes)] = torch.tensor(routes, device="cuda", dtype=ids.dtype) + graph.replay() + torch.testing.assert_close(output, torch.full_like(output, expected), rtol=0, atol=0) + graph.reset() + + +@pytest.mark.parametrize("rows", [3, 5, 28]) +def test_experts_graph(rows): + torch.manual_seed(1337) + experts, hidden, intermediate, topk = 4, 2560, 640, 3 + inputs = torch.randn(rows, hidden, dtype=torch.bfloat16, device="cuda") * 0.1 + up = torch.randn(experts, 2 * intermediate, hidden, dtype=inputs.dtype, device="cuda") * 0.02 + down = torch.randn(experts, hidden, intermediate, dtype=inputs.dtype, device="cuda") * 0.02 + ids = torch.zeros(rows, topk, dtype=torch.int32, device="cuda") + weights = torch.full((rows, topk), 1 / topk, device="cuda") + mask = torch.ones(rows, dtype=torch.bool, device="cuda") + masked = torch.empty_like(ids) + routes = torch.empty(experts, rows * topk, dtype=torch.int32, device="cuda") + counts = torch.empty(experts, dtype=torch.int32, device="cuda") + mid = torch.empty(rows * topk, intermediate, dtype=inputs.dtype, device="cuda") + projected = torch.empty(rows * topk, hidden, dtype=inputs.dtype, device="cuda") + output = torch.empty_like(inputs) + + def run(): + pack_routes[(experts,)]( + ids, mask, masked, routes, counts, rows, topk, 0, triton.next_power_of_2(rows * topk), num_warps=1 + ) + expert_up[(experts, triton.cdiv(intermediate, 32))]( + inputs, + up, + mid, + routes, + counts, + hidden, + intermediate, + rows, + topk, + 16, + 32, + 256, + num_warps=4, + num_stages=3, + ) + expert_down[(experts, triton.cdiv(hidden, 64))]( + mid, + down, + projected, + routes, + counts, + intermediate, + hidden, + rows, + topk, + 1, + 16, + 64, + 128, + num_warps=4, + num_stages=3, + ) + sum_routes[(rows, triton.cdiv(hidden, 256))]( + projected, weights, masked, output, hidden, topk, 0, experts, triton.next_power_of_2(topk), num_warps=4 + ) + + run() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + for change in range(4): + inputs.mul_(0.9) + ids.copy_(torch.rand(rows, experts, device="cuda").argsort(dim=1)[:, :topk]) + if change == 2: + ids.copy_(torch.arange(topk, device="cuda")[None, :]) + mask.copy_(torch.arange(rows, device="cuda") % 3 != 0) + if change == 0: + mask.zero_() + graph.replay() + expected = torch.zeros(rows, hidden, dtype=torch.float32, device="cuda") + for route in range(topk): + for expert in range(experts): + active = mask & (ids[:, route] == expert) + gate, value = F.linear(inputs[active], up[expert]).chunk(2, dim=-1) + activated = F.silu(gate.float()).to(inputs.dtype) * value + result = F.linear(activated, down[expert]).float() + expected[active] += result * weights[active, route, None] + torch.testing.assert_close(output, expected.to(output.dtype), rtol=0.008, atol=0.002) + assert torch.count_nonzero(output[~mask]) == 0 + graph.reset() + + +@pytest.mark.parametrize("rows,hidden", [(4, 513), (28, 2560), (64, 2560)]) +def test_shared_gate_and_partials_graph(rows, hidden): + torch.manual_seed(7911) + ids = torch.zeros(rows, 10, dtype=torch.int32, device="cuda") + partials = torch.randn(8, rows, hidden, dtype=torch.bfloat16, device="cuda") + shared = torch.randn(rows, hidden, dtype=torch.bfloat16, device="cuda") + gate = torch.empty(rows, 1, dtype=torch.bfloat16, device="cuda") + actual = torch.empty_like(shared) + expected = torch.empty_like(shared) + + def reduce(output, **kwargs): + sum_partials[(rows, triton.cdiv(hidden, 512))]( + partials, + ids, + output, + hidden, + rows, + 10, + 64, + 512, + 10, + num_warps=4, + enable_fp_fusion=False, + **kwargs, + ) + + gate.zero_() + reduce(actual, Shared=shared, Gate=gate) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + reduce(actual, Shared=shared, Gate=gate) + for change in range(8): + partials.normal_() + shared.normal_() + gate.copy_(torch.linspace(-12, 12, rows, device="cuda")[:, None]) + ids.copy_(torch.rand(rows, 512, device="cuda").argsort(dim=1)[:, :10]) + ids[::3].fill_(-1) + if change % 2: + partials[0].fill_(2 ** 25) + partials[1].fill_(-(2 ** 25)) + partials[2].fill_(1) + ids[1] = torch.tensor([0, 64, 128, 0, -1, -1, -1, -1, -1, -1], device="cuda") + reduce(expected) + expected.add_(sigmoid_mul_(shared.clone(), gate)) + graph.replay() + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + graph.reset() 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..c5f5fda1fd --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_runtime.py @@ -0,0 +1,397 @@ +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_decode_graph_preparation_reads_current_positions_and_lexical_history(runtime): + old = state_for(runtime, [1, 3], [3, 2], [7, 11], graph_mode="capture") + old.prepare_qwen_state() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + old.prepare_qwen_state() + cache = runtime.req.lexical_caches[1] + for requests, delta in (([3, 1], [-1, 7]), ([1, 3], [4, 0])): + new = state_for(runtime, requests, [4, 4], [13, 17], graph_mode="replay") + new.b_position_delta.copy_(torch.tensor(delta, device="cuda")) + cache.ngram.random_(0, 63) + cache.convolution.normal_() + before = cache.snapshot(1) + old.copy_for_cuda_graph(new) + graph.replay() + positions = torch.tensor(delta, device="cuda") + 3 + torch.testing.assert_close(old.position_ids, positions.expand(3, -1).to(torch.int32), atol=0, rtol=0) + torch.testing.assert_close(old.qwen4_exp_batch.cos, runtime.model._cos_cached[positions].repeat(1, 2)) + torch.testing.assert_close(old.qwen4_exp_batch.sin, runtime.model._sin_cached[positions].repeat(1, 2)) + batch = old.qwen4_exp_ple[1] + torch.testing.assert_close(batch.ngram_history, cache.ngram[requests], atol=0, rtol=0) + torch.testing.assert_close(batch.conv_history, cache.convolution[requests], atol=0, rtol=0) + torch.testing.assert_close(cache.ngram[1], before.ngram, atol=0, rtol=0) + torch.testing.assert_close(cache.convolution[1], before.convolution, atol=0, rtol=0) + + +@pytest.mark.parametrize("runtime", [{}, {"ssm_state_mode": "replay"}], ids=["native", "replay"], indirect=True) +def test_inference_state_prefill_decode_and_automatic_commit(runtime, config): + Qwen3NextTpPartModel._init_triton(None) + req, weight = runtime.req, runtime.weight + for request in (1, 3): + req.init_linear_att_state(SimpleNamespace(req_idx=request)) + state = state_for(runtime, [1, 3], [3, 2], [7, 63, 19, 11, 23], ready=[0, 0]) + assert state.qwen4_exp_metadata.starts.tolist() == [0, 3, 5, 5, 5, 5] + assert state.qwen4_exp_batch.slot_table.data_ptr() == req.req_to_token_indexs.data_ptr() + torch.testing.assert_close( + runtime.mem.qsa_key_cos[state.mem_index.long()], state.qwen4_exp_batch.cos, atol=0, rtol=0 + ) + infer = Qwen4ExpTransformerLayerInfer(1, config) + residual = torch.randn(5, 64, device="cuda", dtype=torch.bfloat16) + reference = {} + expected = [] + for request, start, end in ((1, 0, 3), (3, 3, 5)): + ngram = weight.ngram_hasher.component.empty_history(1) + ple = weight.ple.component.empty_history(1) + conv = torch.zeros(1, 512, 3, device="cuda", dtype=torch.bfloat16) + ssm = torch.zeros(1, 2, 128, 128, device="cuda", dtype=torch.float32) + for row in range(start, end): + output, ngram, ple, conv, ssm = reference_step( + residual[row : row + 1], state.input_ids[row : row + 1], runtime.source, config, ngram, ple, conv, ssm + ) + expected.append(output) + reference[request] = (ngram, ple, conv, ssm) + actual = infer.context_forward(residual, state, weight) + torch.testing.assert_close(actual, torch.cat(expected), atol=0.016, rtol=0.016) + + state = state_for(runtime, [1, 3], [4, 3], [29, 31]) + residual = torch.randn(2, 64, device="cuda", dtype=torch.bfloat16) + actual = infer.token_forward(residual, state, weight) + for row, request in enumerate((1, 3)): + expected, ngram, ple, _, _ = reference_step( + residual[row : row + 1], state.input_ids[row : row + 1], runtime.source, config, *reference[request] + ) + torch.testing.assert_close(actual[row : row + 1], expected, atol=0.016, rtol=0.016) + torch.testing.assert_close(req.lexical_caches[1].ngram[request], ngram[0], atol=0, rtol=0) + torch.testing.assert_close(req.lexical_caches[1].convolution[request], ple[0], atol=0.016, rtol=0.016) + + +def test_native_prefill_graph_state_copy_changes_request_count(runtime, config): + source = checkpoint(config, 3) + weight = Qwen4ExpTransformerLayerWeight(3, torch.bfloat16, config, Quantcfg(config)) + weight.load_hf_weights(dict(source)) + weight.verify_load() + infer = Qwen4ExpTransformerLayerInfer(3, config) + old = state_for(runtime, [1], [4], [7, 11, 13, 17], ready=[0]) + residual = torch.randn(4, 64, device="cuda", dtype=torch.bfloat16) + for _ in range(3): + infer.context_forward(residual, old, weight) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = infer.context_forward(residual, old, weight) + pointers = [ + getattr(old.qwen4_exp_metadata, name).data_ptr() for name in old.qwen4_exp_metadata.__dataclass_fields__ + ] + before = runtime.mem.get_index_kv_buffer(old.mem_index) + new = state_for(runtime, [2, 3], [2, 2], [19, 23, 29, 31], ready=[0, 0]) + old.copy_for_prefill_cuda_graph(new) + residual.normal_() + graph.replay() + expected = reference_layer(residual, source, config, new.qwen4_exp_batch.cos, new.qwen4_exp_batch.sin, [2, 2]) + torch.testing.assert_close(output, expected, atol=0.016, rtol=0.016) + assert pointers == [ + getattr(old.qwen4_exp_metadata, name).data_ptr() for name in old.qwen4_exp_metadata.__dataclass_fields__ + ] + for name, values in runtime.mem.get_index_kv_buffer(torch.arange(4, device="cuda")).items(): + torch.testing.assert_close(values, before[name], atol=0, rtol=0, equal_nan=True) + + +@pytest.mark.parametrize( + "runtime", + [ + {"page_size": 4, "ssm_state_mode": "replay"}, + {"page_size": 4, "llm_prefill_att_backend": ["triton", "flashinfer"]}, + {"page_size": 4, "ssm_state_mode": "replay", "llm_prefill_att_backend": ["triton", "flashinfer"]}, + ], + indirect=True, +) +def test_prefill_graph_reuses_current_attention_state(runtime, config, monkeypatch): + from lightllm.common.basemodel.prefill_cuda_graph import PrefillCudaGraph + + Qwen3NextTpPartModel._init_triton(None) + model, source = assemble_model(runtime, config) + for request in (1, 2, 3): + runtime.req.init_linear_att_state(SimpleNamespace(req_idx=request)) + + def run(state): + state.hidden_collector = NoopHiddenCollector() + return model._context_forward(state).logits.clone() + + state = state_for(runtime, [1], [4], [7, 11, 13, 17], ready=[0]) + state.max_kv_seq_len = 4 + expected = run(state) + runtime.req.init_linear_att_state(SimpleNamespace(req_idx=1)) + qsa_calls = [] + qsa_forward = Qwen4ExpTransformerLayerInfer._qsa_attention_forward + + def tracked(self, *args, **kwargs): + qsa_calls.append(torch.cuda.is_current_stream_capturing()) + return qsa_forward(self, *args, **kwargs) + + monkeypatch.setattr(Qwen4ExpTransformerLayerInfer, "_qsa_attention_forward", tracked) + model.prefill_graph = PrefillCudaGraph(None, 1) + actual = run(state) + torch.testing.assert_close(actual, expected, atol=0.025, rtol=0.025) + assert qsa_calls == [True] + + state = state_for(runtime, [2, 3], [2, 2], [19, 23, 29, 31], ready=[0, 0]) + state.max_kv_seq_len = 2 + actual = run(state) + expected = reference_model( + state.input_ids, state.qwen4_exp_batch.cos, state.qwen4_exp_batch.sin, [2, 2], source, config + ) + torch.testing.assert_close(actual, expected[[1, 3]], atol=0.025, rtol=0.025) + assert qsa_calls == [True] + prefix = state + state = state_for(runtime, [3], [6], [37, 41, 43, 47], ready=[2]) + state.max_kv_seq_len = 6 + actual = run(state) + expected = reference_model( + torch.tensor([29, 31, 37, 41, 43, 47], device="cuda"), + torch.cat((prefix.qwen4_exp_batch.cos[2:], state.qwen4_exp_batch.cos)), + torch.cat((prefix.qwen4_exp_batch.sin[2:], state.qwen4_exp_batch.sin)), + [6], + source, + config, + ) + torch.testing.assert_close(actual, expected[[-1]], atol=0.025, rtol=0.025) + assert qsa_calls == [True, False] + + +@pytest.mark.parametrize( + "runtime", + [ + {"page_size": 4, "llm_prefill_att_backend": ["triton", "flashinfer"]}, + { + "page_size": 4, + "ssm_state_mode": "replay", + "mtp_step": 3, + "llm_prefill_att_backend": ["triton", "flashinfer"], + }, + ], + indirect=True, +) +def test_flashinfer_prefill_graph_replays_and_falls_back_for_changed_batch(runtime, config, monkeypatch): + from lightllm.common.basemodel.prefill_cuda_graph import PrefillCudaGraph + + Qwen3NextTpPartModel._init_triton(None) + model, _ = assemble_model(runtime, config) + requests = (1, 2, 3) + cases = [ + ([1], [4], [7, 11, 13, 17], [0]), + ([1], [8], [19, 23, 29, 31], [4]), + ([3, 2], [2, 2], [37, 41, 43, 47], [0, 0]), + ([2], [6], [11, 13, 17, 19], [2]), + ] + bounded = model.att_state_manager.get_bounded_attention_states(model.att_state_role) + linear = next(state for state in bounded if state.get_state_type().startswith("gdn_")) + + def run(): + results = [] + for ids, lengths, tokens, ready in cases: + state = state_for(runtime, ids, lengths, tokens, ready=ready) + state.hidden_collector = NoopHiddenCollector() + logits = model._context_forward(state).logits.clone() + buffers = [ + tensor[list(requests)].clone() + for layer in linear.config.logical_layer_ids + for tensor in linear.get_layer_view(logical_layer_id=layer) + ] + results.append((logits, buffers)) + return results + + for request in requests: + runtime.req.init_linear_att_state(SimpleNamespace(req_idx=request)) + expected = run() + for request in requests: + runtime.req.init_linear_att_state(SimpleNamespace(req_idx=request)) + + calls = [] + kernel = model.prefill_att_backend1.prefill_kernel + + def tracked(**kwargs): + calls.append(torch.cuda.is_current_stream_capturing()) + return kernel(**kwargs) + + monkeypatch.setattr(model.prefill_att_backend1, "prefill_kernel", tracked) + model.prefill_graph = PrefillCudaGraph(None, 1) + actual = run() + linear_layers = len(linear.config.logical_layer_ids) + assert calls.count(True) == linear_layers + assert calls.count(False) == linear_layers + for (logits, buffers), (expected_logits, expected_buffers) in zip(actual, expected): + torch.testing.assert_close(logits, expected_logits, atol=0.025, rtol=0.025) + for value, reference in zip(buffers, expected_buffers): + torch.testing.assert_close(value, reference, atol=0.003, rtol=0.03) + + +def test_native_config_normalization_and_periodic_layout(runtime, config, tmp_path): + outer = dict(text_config=dict(config), vision_config={"hidden_size": 16}) + outer["text_config"]["rope_parameters"] = dict( + rope_type="default", rope_theta=10000, partial_rotary_factor=0.5, mrope_section=[1, 1, 0] + ) + (tmp_path / "config.json").write_text(json.dumps(outer)) + model = Qwen4ExpTpPartModel.__new__(Qwen4ExpTpPartModel) + model.weight_dir_, model.finetune_config, model.tp_world_size_ = str(tmp_path), None, 1 + model._init_config() + assert model.config["n_layer"] == 4 and model.config["n_embed"] == 32 + assert model.config["rope_theta"] == 10000 + assert model.num_kv_heads == 2 + invalid = dict(config, layer_types=["full_attention"] * 4) + with pytest.raises(ValueError, match="periodic"): + Qwen4ExpTpPartModel.normalize_model_config(invalid) + + +def test_native_decode_graph_replay_updates_positions_slots_and_request_order(runtime, config): + source = checkpoint(config, 3) + weight = Qwen4ExpTransformerLayerWeight(3, torch.bfloat16, config, Quantcfg(config)) + weight.load_hf_weights(dict(source)) + weight.verify_load() + infer = Qwen4ExpTransformerLayerInfer(3, config) + prefix = state_for(runtime, [1, 3], [4, 4], [7, 11, 13, 17, 19, 23, 29, 31], ready=[0, 0]) + prefix_residual = torch.randn(8, 64, device="cuda", dtype=torch.bfloat16) + infer.context_forward(prefix_residual, prefix, weight) + old = state_for(runtime, [1, 3], [5, 5], [37, 41]) + residual = torch.randn(2, 64, device="cuda", dtype=torch.bfloat16) + for _ in range(3): + infer.token_forward(residual, old, weight) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = infer.token_forward(residual, old, weight) + previous_residual = residual.clone() + previous_cos, previous_sin = old.qwen4_exp_batch.cos.clone(), old.qwen4_exp_batch.sin.clone() + new = state_for(runtime, [3, 1], [6, 6], [43, 47]) + old.copy_for_cuda_graph(new) + residual.normal_() + graph.replay() + all_residual, all_cos, all_sin = [], [], [] + for row, (start, end) in enumerate(((0, 4), (4, 8))): + incoming = 1 - row + all_residual.extend( + [prefix_residual[start:end], previous_residual[row : row + 1], residual[incoming : incoming + 1]] + ) + all_cos.extend( + [ + prefix.qwen4_exp_batch.cos[start:end], + previous_cos[row : row + 1], + new.qwen4_exp_batch.cos[incoming : incoming + 1], + ] + ) + all_sin.extend( + [ + prefix.qwen4_exp_batch.sin[start:end], + previous_sin[row : row + 1], + new.qwen4_exp_batch.sin[incoming : incoming + 1], + ] + ) + expected = reference_layer(torch.cat(all_residual), source, config, torch.cat(all_cos), torch.cat(all_sin), [6, 6]) + torch.testing.assert_close(output, expected[[11, 5]], atol=0.016, rtol=0.016) + + +@pytest.mark.parametrize( + "runtime", + [ + {}, + {"ssm_state_mode": "replay"}, + {"page_size": 4}, + {"page_size": 4, "ssm_state_mode": "replay"}, + {"page_size": 4, "llm_prefill_att_backend": ["triton", "flashinfer"]}, + {"page_size": 4, "ssm_state_mode": "replay", "llm_prefill_att_backend": ["triton", "flashinfer"]}, + ], + ids=["native", "replay", "compact-native", "compact-replay", "flashinfer-native", "flashinfer-replay"], + indirect=True, +) +@pytest.mark.parametrize("graph_preparation", [False, True]) +def test_complete_small_model_prefill_and_cached_logits(runtime, config, graph_preparation): + 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], graph_mode="capture" if graph_preparation else None) + 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], graph_mode="replay" if graph_preparation else None) + 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) + current = third if graph_preparation else fourth + all_cos = torch.cat([all_cos[:5], current.qwen4_exp_batch.cos[1:], all_cos[5:], current.qwen4_exp_batch.cos[:1]]) + all_sin = torch.cat([all_sin[:5], current.qwen4_exp_batch.sin[1:], all_sin[5:], current.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..c75f75e418 --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_sequence_parallel.py @@ -0,0 +1,133 @@ +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, + combine_and_read=lambda value, update, gates: (value + update, (value + update, None)), + ) + 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"]) + + +@pytest.mark.parametrize("enabled", [False, True]) +@pytest.mark.parametrize("prefill", [False, True]) +def test_qsa_shares_gathered_rows_between_projection_and_indexer(enabled, prefill): + full = torch.arange(64, dtype=torch.float32).reshape(8, 8) + layer = partitioned_layer(2, 1, enabled, full) + layer.tp_q_head_num_, layer.head_dim_ = 1, 8 + gathers = [] + + def gather(value, state): + gathers.append(value) + return full + + def project(value, state, weight): + assert value is full + return value, value + + def attend(indexer, hidden, query, key, value, batch, **kwargs): + assert hidden is full + assert kwargs["decode"] is not prefill + return query + + layer._tpsp_allgather = gather + layer._project_qkv = project + layer._post_cache_kv = lambda *args: None + layer._get_o = lambda output, *args: output + manager = SimpleNamespace(get_layer_view=lambda number: (None, None)) + state = SimpleNamespace( + is_prefill=prefill, + att_state_role="main", + att_state_manager=SimpleNamespace(get_state=lambda *args: manager), + qwen4_exp_qsa={1: SimpleNamespace(attend=attend)}, + qwen4_exp_fixed_decode=False, + qwen4_exp_batch=None, + dist_group=None, + ) + weight = SimpleNamespace(indexer=SimpleNamespace(component=None)) + value = full.chunk(2)[1].contiguous() if enabled else full + actual = layer._qsa_attention_forward(value, state, weight) + torch.testing.assert_close(actual.view_as(full), full) + assert len(gathers) == 1 and gathers[0] is value 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..3e1e584e7c --- /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, req_verify_width, role): + assert role == "main" + lexical.commit_bounded_att_state(requests, accepted, req_verify_width) + req.req_to_mtp_state_index[requests] = accepted - 1 + + backend = SimpleNamespace( + model=SimpleNamespace(req_manager=req, att_state_manager=SimpleNamespace(commit_bounded_att_state=commit)), + is_linear_att_mixed_model=True, + max_draft_step=3, + ) + meta = metadata([5, 5, 5, 2, 2, 7], [11, 12, 13, 11, 12, 1], [1] * 6) + tokens = torch.empty(6, device="cuda", dtype=torch.int64) + residual = torch.empty(6, 64, device="cuda", dtype=torch.bfloat16) + + def run(): + batch = cache.prepare(meta) + result = batch.apply(residual, tokens, cache.layer_weight, 1, None) + speculative.stage(meta, batch) + return result + + tokens.fill_(19) + residual.normal_() + if capture: + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + run() + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + result = run() + + cases = ( + ([5, 5, 5, 2, 2, 7], [2, 1], [19, 63, 23, 29, 31, 37]), + ([2, 2, 5, 5, 5, 7], [2, 3], [41, 43, 47, 63, 5, 7]), + ([5, 5, 5, 5, 2, 7], [1, 1], [11, 13, 17, 19, 23, 29]), + ) + for requests, accepts, ids in cases: + starts = [0] + [i for i in range(1, len(requests)) if requests[i] != requests[i - 1]] + counts = [end - start for start, end in zip(starts, starts[1:] + [len(requests)])] + lengths = [] + for count in counts: + lengths.extend(range(11, 11 + count)) + meta.copy_from(metadata(requests, lengths, [1] * 6)) + tokens.copy_(torch.tensor(ids, device="cuda")) + residual.normal_() + before_ngram, before_conv = cache.ngram.clone(), cache.convolution.clone() + if capture: + graph.replay() + actual = result + else: + actual = run() + torch.testing.assert_close(cache.ngram, before_ngram, atol=0, rtol=0) + torch.testing.assert_close(cache.convolution, before_conv, atol=0, rtol=0) + + next_ids = torch.arange(101, 107, device="cuda", dtype=torch.int64) + proposals = req.req_sampling_params_manager.req_to_next_token_ids + proposals.fill_(-1) + expected_mask = [] + for start, count, accepted in zip(starts, counts, accepts + [1]): + request = requests[start] + proposals[request, 1:accepted] = next_ids[start : start + accepted - 1] + expected_mask.extend([1] * accepted + [0] * (count - accepted)) + accept_lengths, accepted_index = verify_mtp_tokens( + backend, + next_ids, + torch.tensor(requests, device="cuda", dtype=torch.int32), + torch.tensor(starts, device="cuda", dtype=torch.int32), + torch.tensor(sum((list(range(count)) for count in counts), []), device="cuda", dtype=torch.int32), + ) + assert accept_lengths.tolist() == accepts + [1] + assert accepted_index.tolist() == expected_mask + for start, count, accepted in zip(starts[:2], counts[:2], accepts): + request = requests[start] + expected, states = reference_segment( + residual[start : start + count], + tokens[start : start + count], + source, + config, + before_ngram[request : request + 1], + before_conv[request : request + 1], + ) + torch.testing.assert_close(actual[start : start + count], expected, atol=0.016, rtol=0.016) + ngram, convolution = states[accepted] + torch.testing.assert_close(cache.ngram[request], ngram[0], atol=0, rtol=0) + torch.testing.assert_close(cache.convolution[request], convolution[0], atol=0.016, rtol=0.016) + assert req.req_to_mtp_state_index[request] == accepted - 1 + torch.testing.assert_close(cache.ngram[7], untouched.ngram, atol=0, rtol=0) + torch.testing.assert_close(cache.convolution[7], untouched.convolution, atol=0, rtol=0) + if capture: + graph.reset() + + +def test_microbatch_candidates_survive_reverse_verification(config): + cache, source = make_cache(config) + speculative = LexicalSpeculativeCache(cache, 3) + batches = [] + for request, ids in ((5, [19, 63, 23]), (2, [29, 31, 37])): + meta = metadata([request] * 3, [11, 12, 13], [1] * 3) + tokens = torch.tensor(ids, device="cuda") + residual = torch.randn(3, 64, device="cuda", dtype=torch.bfloat16) + batch = cache.prepare(meta) + batch.apply(residual, tokens, cache.layer_weight, 1, None) + speculative.stage(meta, batch) + batches.append((request, residual, tokens)) + for request, residual, tokens in reversed(batches): + _, states = reference_segment( + residual, + tokens, + source, + config, + cache.ngram[request : request + 1].clone(), + cache.convolution[request : request + 1].clone(), + ) + speculative.commit( + torch.tensor([request] * 3, device="cuda", dtype=torch.int32), + torch.tensor([0], device="cuda", dtype=torch.int32), + torch.tensor([2], device="cuda", dtype=torch.int32), + 7, + ) + torch.testing.assert_close(cache.ngram[request], states[2][0][0], atol=0, rtol=0) + torch.testing.assert_close(cache.convolution[request], states[2][1][0], atol=0.016, rtol=0.016) diff --git a/unit_tests/models/qwen4_exp/test_vision.py b/unit_tests/models/qwen4_exp/test_vision.py new file mode 100644 index 0000000000..ccf34a0521 --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_vision.py @@ -0,0 +1,49 @@ +import pytest +import torch + +from lightllm.models.qwen4_exp.vision import Qwen4ExpVisionTransformer +from lightllm.server import visualserver +from lightllm.common.basemodel.attention_vit.create_utils import get_vit_att_backend_class +from .vision_reference import interpolate_positions, vision_reference + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +def test_vision_encoder_matches_independent_packed_frames(dtype, monkeypatch): + monkeypatch.setattr(visualserver, "VIT_ATTN_BACKEND", get_vit_att_backend_class("sdpa")) + config = dict( + depth=2, + hidden_size=32, + out_hidden_size=16, + intermediate_size=48, + hidden_act="gelu_pytorch_tanh", + num_heads=4, + in_channels=3, + patch_size=2, + temporal_patch_size=2, + spatial_merge_size=2, + num_position_embeddings=16, + deepstack_visual_indexes=[], + ) + torch.manual_seed(1907) + model = Qwen4ExpVisionTransformer({"data_type": dtype}, **config).eval().to(device="cuda", dtype=dtype) + grid = torch.tensor([[1, 4, 6], [2, 2, 4]], device="cuda") + pixels = torch.randn(40, 24, device="cuda", dtype=dtype) + source = {name: value.detach().clone() for name, value in model.state_dict().items()} + tolerance = 0.016 if dtype == torch.bfloat16 else 2e-5 + with torch.no_grad(): + actual_positions = model.fast_pos_embed_interpolate(grid) + expected_positions = interpolate_positions(source["pos_embed.weight"], grid, 2) + torch.testing.assert_close(actual_positions, expected_positions, atol=tolerance, rtol=tolerance) + actual, deepstack = model(pixels, grid) + expected = vision_reference(pixels, grid, source, config) + torch.testing.assert_close(actual_positions, expected_positions, atol=tolerance, rtol=tolerance) + torch.testing.assert_close(actual, expected, atol=tolerance, rtol=tolerance) + assert not deepstack + assert actual.shape == (10, 16) + + +def test_vision_rejects_unsupported_deepstack(): + with pytest.raises(ValueError, match="deepstack"): + Qwen4ExpVisionTransformer({"data_type": "bf16"}, deepstack_visual_indexes=[1]) diff --git a/unit_tests/models/qwen4_exp/test_weights.py b/unit_tests/models/qwen4_exp/test_weights.py new file mode 100644 index 0000000000..51bebb72ee --- /dev/null +++ b/unit_tests/models/qwen4_exp/test_weights.py @@ -0,0 +1,99 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.models.qwen4_exp.checkpoint import normalize_text_weight_names, ngram_table_shape +from lightllm.models.qwen4_exp.layer_weights.transformer_layer_weight import Qwen4ExpTransformerLayerWeight +from lightllm.models.qwen4_exp.layer_weights.pre_and_post_layer_weight import Qwen4ExpPreAndPostLayerWeight +from lightllm.models.qwen4_exp.layer_infer.pre_layer_infer import Qwen4ExpPreLayerInfer +from lightllm.models.qwen4_exp.layer_infer.post_layer_infer import Qwen4ExpPostLayerInfer +from lightllm.common.quantization import Quantcfg +from .helpers import checkpoint, prepost_checkpoint, reference_hyper + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +def test_scalar_gate_dispatch_is_model_owned(config): + from lightllm.common.quantization.no_quant import NoQuantization + from lightllm.common.quantization.scalar_projection import StableScalarProjection + from lightllm.models.qwen3_5_moe.layer_weights.transformer_layer_weight import Qwen35MOETransformerLayerWeight + + native = Qwen4ExpTransformerLayerWeight(0, torch.bfloat16, config, Quantcfg(config)) + legacy = Qwen35MOETransformerLayerWeight(0, torch.bfloat16, config, Quantcfg(config)) + assert isinstance(native.shared_expert_gate.quant_method, StableScalarProjection) + assert type(legacy.shared_expert_gate.quant_method) is NoQuantization + + +@pytest.mark.parametrize("layer", [0, 1, 3]) +def test_incremental_native_checkpoint_loading(config, layer): + torch.manual_seed(737) + weights = checkpoint(config, layer) + native = Qwen4ExpTransformerLayerWeight(layer, torch.bfloat16, config, Quantcfg(config)) + assert not native.attn_hyper_connection.verify_load() + with pytest.raises(RuntimeError, match="not fully loaded"): + _ = native.attn_hyper_connection.component + + for key in reversed(list(weights)): + native.load_hf_weights({key: weights[key]}) + native.verify_load() + assert not hasattr(native, "att_norm_weight_") + assert not hasattr(native, "ffn_norm_weight_") + if layer == 3: + assert native.qk_norm_weight_.q_weight.dtype == torch.float32 + expected = 1 + weights[f"model.language_model.layers.{layer}.self_attn.q_norm.weight"].float().cuda() + torch.testing.assert_close(native.qk_norm_weight_.q_weight, expected, atol=0, rtol=0) + assert native.indexer.component.head_dim == 16 + else: + assert native.linear_norm.activation == "sigmoid" + hidden = torch.randn(3, 32, device="cuda", dtype=torch.bfloat16) + actual = native.linear_in_proj.mm(hidden, use_custom_tensor_mananger=False) + p = f"model.language_model.layers.{layer}.linear_attn" + expected = torch.cat( + [ + torch.nn.functional.linear(hidden, weights[p + f".in_proj_{name}.weight"].cuda()) + for name in ["qkv", "z", "b", "a"] + ], + -1, + ) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + if layer == 1: + assert native.ngram_embedding.weight.shape == (36, 8) + assert native.ngram_hasher.component.heads == 4 + assert native.ple.component.context_len == 9 + + +def test_embedding_widening_and_final_native_logits(config): + torch.manual_seed(737) + source = prepost_checkpoint(config) + native = Qwen4ExpPreAndPostLayerWeight(torch.bfloat16, config) + native.load_hf_weights(dict(source)) + native.verify_load() + assert not hasattr(native, "final_norm_weight_") + pre, post = Qwen4ExpPreLayerInfer(config), Qwen4ExpPostLayerInfer(config) + state = SimpleNamespace(is_prefill=False, is_draft_model=False, batch_size=2, prompt_logics=None, dist_group=None) + ids = torch.tensor([2, 17], device="cuda", dtype=torch.int64) + widened = pre.token_forward(ids, state, native) + torch.testing.assert_close( + widened, source["model.language_model.embed_tokens.weight"].cuda()[ids].repeat(1, 2), atol=0, rtol=0 + ) + widened = widened + torch.randn_like(widened) * 0.1 + actual = post.token_forward(widened, state, native).logits + mixed = reference_hyper(widened, source, "model.language_model.hyper_connection_mixer", config, False) + expected = torch.nn.functional.linear(mixed, source["lm_head.weight"].cuda()).float() + torch.testing.assert_close(actual, expected, atol=0.016, rtol=0.016) + + +def test_namespace_and_allocation_metadata(): + value = torch.zeros(1) + weights = { + "model.language_model.x": value, + "model.visual.x": value, + "mtp.x": value, + "adapter.model.language_model.x": value, + } + normalize_text_weight_names(weights) + assert set(weights) == {"model.x", "model.visual.x", "mtp.x", "adapter.model.language_model.x"} + with pytest.raises(ValueError, match="Conflicting"): + normalize_text_weight_names({"model.x": value, "model.language_model.x": value.clone()}) + assert ngram_table_shape(20000000, 16, 0, 2560, 128) == (320001536, 160) diff --git a/unit_tests/models/qwen4_exp/vision_reference.py b/unit_tests/models/qwen4_exp/vision_reference.py new file mode 100644 index 0000000000..bd0f2b68b1 --- /dev/null +++ b/unit_tests/models/qwen4_exp/vision_reference.py @@ -0,0 +1,93 @@ +import torch +import torch.nn.functional as F + + +def patch_coordinates(grid, merge, device): + coordinates = [] + for frames, height, width in grid.cpu().tolist(): + coordinates.extend( + (y + dy, x + dx) + for _ in range(frames) + for y in range(0, height, merge) + for x in range(0, width, merge) + for dy in range(merge) + for dx in range(merge) + ) + return torch.tensor(coordinates, dtype=torch.int64, device=device) + + +def interpolate_positions(table, grid, merge): + side = int(table.shape[0] ** 0.5) + result = [] + for frames, height, width in grid.cpu().tolist(): + coords = patch_coordinates(torch.tensor([[frames, height, width]]), merge, table.device) + ys = torch.linspace(0, side - 1, height, device=table.device)[coords[:, 0]] + xs = torch.linspace(0, side - 1, width, device=table.device)[coords[:, 1]] + y0, x0 = ys.floor().long(), xs.floor().long() + y1, x1 = (y0 + 1).clamp_max(side - 1), (x0 + 1).clamp_max(side - 1) + dy, dx = ys - y0, xs - x0 + indices = torch.stack([y0 * side + x0, y0 * side + x1, y1 * side + x0, y1 * side + x1], -1) + coefficients = torch.stack([(1 - dy) * (1 - dx), (1 - dy) * dx, dy * (1 - dx), dy * dx], -1) + result.append((table[indices] * coefficients[:, :, None]).sum(1).to(table.dtype)) + return torch.cat(result) + + +def vision_reference(pixels, grid, weights, config): + hidden, heads = config["hidden_size"], config["num_heads"] + patch, temporal, merge = config["patch_size"], config["temporal_patch_size"], config["spatial_merge_size"] + values = F.conv3d( + pixels.view(-1, config["in_channels"], temporal, patch, patch), + weights["patch_embed.proj.weight"], + weights["patch_embed.proj.bias"], + stride=(temporal, patch, patch), + ).view(-1, hidden) + values = values + interpolate_positions(weights["pos_embed.weight"], grid, merge) + dim = hidden // heads + theta = config.get("rope_parameters", {}).get("rope_theta", 10000.0) + invfreq = theta ** (-torch.arange(0, dim // 2, 2, device=pixels.device).float() / (dim // 2)) + coordinates = patch_coordinates(grid, merge, pixels.device) + angles = (coordinates[:, :, None] * invfreq).flatten(1).repeat(1, 2) + cos, sin = angles.cos()[:, None], angles.sin()[:, None] + + def norm(value, prefix): + return F.layer_norm(value, (value.shape[-1],), weights[prefix + ".weight"], weights[prefix + ".bias"], 1e-6) + + def rotate(value): + value = value.float() + a, b = value.chunk(2, -1) + return (value * cos + torch.cat([-b, a], -1) * sin).to(pixels.dtype) + + lengths = [height * width for frames, height, width in grid.cpu().tolist() for _ in range(frames)] + for layer in range(config["depth"]): + prefix = f"blocks.{layer}" + qkv = F.linear( + norm(values, prefix + ".norm1"), weights[prefix + ".attn.qkv.weight"], weights[prefix + ".attn.qkv.bias"] + ) + q, k, v = qkv.view(-1, 3, heads, dim).unbind(1) + q, k = rotate(q), rotate(k) + outputs, offset = [], 0 + for count in lengths: + outputs.append( + F.scaled_dot_product_attention( + q[offset : offset + count].transpose(0, 1)[None], + k[offset : offset + count].transpose(0, 1)[None], + v[offset : offset + count].transpose(0, 1)[None], + )[0].transpose(0, 1) + ) + offset += count + attention = torch.cat(outputs).reshape(-1, hidden) + values = values + F.linear( + attention, weights[prefix + ".attn.proj.weight"], weights[prefix + ".attn.proj.bias"] + ) + mlp = F.linear( + norm(values, prefix + ".norm2"), + weights[prefix + ".mlp.linear_fc1.weight"], + weights[prefix + ".mlp.linear_fc1.bias"], + ) + mlp = F.gelu(mlp, approximate="tanh" if config["hidden_act"] == "gelu_pytorch_tanh" else "none") + values = values + F.linear( + mlp, weights[prefix + ".mlp.linear_fc2.weight"], weights[prefix + ".mlp.linear_fc2.bias"] + ) + values = norm(values, "merger.norm").view(-1, hidden * merge * merge) + values = F.gelu(F.linear(values, weights["merger.linear_fc1.weight"], weights["merger.linear_fc1.bias"])) + return F.linear(values, weights["merger.linear_fc2.weight"], weights["merger.linear_fc2.bias"]) diff --git a/unit_tests/models/sensenova_loop/test_gdn_state.py b/unit_tests/models/sensenova_loop/test_gdn_state.py index c25576a19a..3db796daab 100644 --- a/unit_tests/models/sensenova_loop/test_gdn_state.py +++ b/unit_tests/models/sensenova_loop/test_gdn_state.py @@ -79,10 +79,6 @@ def test_gdn_decode_keeps_logical_layer_state( if device == "cpu": - def conv_pack(mixed, z, a, b, conv, *args): - conv[1].add_(1) - return (*backend._rearrange_mixed_qkv(mixed, decode=True), z, a, b) - def conv_mtp(mixed, conv, *args, **kwargs): conv[1].add_(1) return mixed @@ -91,12 +87,19 @@ def recurrent(*, initial_state, v, **kwargs): initial_state[mtp_step + 1 : 2 * (mtp_step + 1)].add_(1) return v, None - monkeypatch.setattr("lightllm.common.basemodel.attention.linear.gdn.conv_pack_gdn_decode_inputs", conv_pack) + def packed_recurrent(*, mixed_qkv, initial_state, **kwargs): + _, _, value = backend._rearrange_mixed_qkv(mixed_qkv, decode=True) + return recurrent(initial_state=initial_state, v=value) + + monkeypatch.setattr( + "lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d.causal_conv1d_update", conv_mtp + ) monkeypatch.setattr( "lightllm.common.basemodel.triton_kernel.linear_att.causal_conv1d_mtp.causal_conv1d_update", conv_mtp ) monkeypatch.setattr( - "lightllm.common.basemodel.attention.linear.gdn.fused_recurrent_gated_delta_rule", recurrent + "lightllm.common.basemodel.attention.linear.gdn.fused_recurrent_gated_delta_rule_packed_decode", + packed_recurrent, ) monkeypatch.setattr( "lightllm.common.basemodel.attention.linear.gdn.mtp_fused_recurrent_gated_delta_rule", diff --git a/unit_tests/server/core/objs/test_stop_sequence_output.py b/unit_tests/server/core/objs/test_stop_sequence_output.py index 4e0d514c63..358aff8f73 100644 --- a/unit_tests/server/core/objs/test_stop_sequence_output.py +++ b/unit_tests/server/core/objs/test_stop_sequence_output.py @@ -126,6 +126,30 @@ async def _read_http_tokens(manager, req): return status.out_token_info_list +@pytest.mark.parametrize("count", [1, 3, 7]) +@pytest.mark.parametrize("start", [0, 6]) +def test_http_notification_drains_ready_tokens(count, start): + async def run(): + manager = _http_manager() + req = _req(list(range(10, 10 + count)), finish=FinishStatus.FINISHED_LENGTH) + req.out_tokens_queue.head = req.out_tokens_queue.tail = start + for index in range(count): + req.out_tokens_queue.push(str(index), 2 + index, False, index + 1) + + results = await _read_http_tokens(manager, req) + + assert req.out_tokens_queue.is_empty() + assert [result[1] for result in results] == [str(index) for index in range(count)] + assert [result[2]["id"] for result in results] == list(range(10, 10 + count)) + assert [result[2]["count_output_tokens"] for result in results] == list(range(1, count + 1)) + assert [result[2]["logprob"] for result in results] == [-0.5] * count + assert all(not result[3].is_finished() for result in results[:-1]) + assert results[-1][3].is_finished_length() + req.merge_final_token_metadata.assert_awaited_once() + + asyncio.run(run()) + + def _match_http_stream(manager, generator, params, prompt_ids, reqs): for req in reqs: manager.req_id_to_out_inf[convert_sub_id_to_group_id(req.request_id)] = SimpleNamespace( 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_eagle_with_att.py b/unit_tests/server/router/model_infer/mtp_speculative/test_eagle_with_att.py index 365acb30b2..311012f546 100644 --- a/unit_tests/server/router/model_infer/mtp_speculative/test_eagle_with_att.py +++ b/unit_tests/server/router/model_infer/mtp_speculative/test_eagle_with_att.py @@ -125,7 +125,8 @@ def forward(model_input): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") -def test_eagle_with_att_commits_verify_kv_then_recurrently_decodes(monkeypatch): +@pytest.mark.parametrize("selected_logits", [False, True]) +def test_eagle_with_att_commits_verify_kv_then_recurrently_decodes(monkeypatch, selected_logits): device = "cuda" original_input_ids = torch.tensor([10, 11, 12, 20, 21, 22], dtype=torch.int64, device=device) target_hidden = torch.arange(12, dtype=torch.float32, device=device).reshape(6, 2) @@ -148,8 +149,11 @@ def forward(model_input): } ) if len(draft_calls) == 1: + logits = torch.arange(30, 36, dtype=torch.float32, device=device).unsqueeze(1) + if selected_logits: + logits = logits.index_select(0, model_input.logits_indices) return ModelOutput( - logits=torch.arange(30, 36, dtype=torch.float32, device=device).unsqueeze(1), + logits=logits, mtp_collector=ModelMtpOutputCollector(spec_hidden=extend_hidden), ) return ModelOutput( @@ -158,7 +162,7 @@ def forward(model_input): ) backend = SimpleNamespace( - draft_models=[SimpleNamespace(forward=forward)], + draft_models=[SimpleNamespace(forward=forward, supports_selected_logits=selected_logits)], _gen_argmax_token_ids=lambda output: output.logits[:, 0].long(), _gen_argmax_token_ids_and_prob=lambda output: ( output.logits[:, 0].long(), 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_api_start_deepep.py b/unit_tests/server/test_api_start_deepep.py index 949021cd34..1972fff244 100644 --- a/unit_tests/server/test_api_start_deepep.py +++ b/unit_tests/server/test_api_start_deepep.py @@ -6,7 +6,8 @@ from lightllm.server.core.objs.start_args_type import StartArgs -def test_decode_node_batch_max_tokens_covers_mtp_rows(monkeypatch, tmp_path): +@pytest.mark.parametrize("linear_backend", ["triton", "flashinfer"]) +def test_decode_node_batch_max_tokens_covers_mtp_rows(monkeypatch, tmp_path, linear_backend): for name in ( "_set_envs_and_config", "auto_set_max_req_total_len", @@ -25,6 +26,7 @@ def stop(args): run_mode="decode", model_dir=str(tmp_path), enable_ep_moe=True, + llm_prefill_att_backend=["triton", linear_backend], running_max_req_size=256, mtp_mode="vanilla_no_att", mtp_step=2, @@ -43,3 +45,48 @@ def stop(args): _launch_subprocesses(args) assert args.batch_max_tokens == 768 + + +@pytest.mark.parametrize( + "model_type,kv_type,requires_calibration", + [ + ("qwen4_exp", "fp8kv_spt", False), + ("qwen4_exp_text", "fp8kv_spt", False), + ("qwen4_exp", "fp8kv_sph", True), + ("llama", "fp8kv_spt", True), + ], +) +def test_qsa_unit_scale_fp8_startup(monkeypatch, tmp_path, model_type, kv_type, requires_calibration): + for name in ( + "_set_envs_and_config", + "auto_set_max_req_total_len", + "auto_set_fused_shared_experts", + "set_unique_server_name", + "auto_set_page_size", + ): + monkeypatch.setattr(f"lightllm.server.api_start.{name}", lambda args: None) + + def stop(args): + raise RuntimeError("startup validation passed") + + monkeypatch.setattr("lightllm.server.api_start.auto_set_response_parsers", stop) + (tmp_path / "config.json").write_text(json.dumps({"model_type": model_type})) + args = StartArgs( + model_dir=str(tmp_path), + llm_kv_type=kv_type, + running_max_req_size=256, + max_req_total_len=4096, + batch_max_tokens=4096, + chunked_prefill_size=2048, + eos_id=[0], + data_type="bfloat16", + disable_vision=True, + disable_audio=True, + disable_shm_warning=True, + ) + if requires_calibration: + with pytest.raises(AssertionError, match="requires --kv_quant_calibration_config_path"): + _launch_subprocesses(args) + else: + with pytest.raises(RuntimeError, match="startup validation passed"): + _launch_subprocesses(args) diff --git a/unit_tests/server/test_native_stream_details.py b/unit_tests/server/test_native_stream_details.py new file mode 100644 index 0000000000..0d983258a0 --- /dev/null +++ b/unit_tests/server/test_native_stream_details.py @@ -0,0 +1,69 @@ +import asyncio +import copy +import json +from types import SimpleNamespace + +import pytest + +from lightllm.server import api_lightllm + + +@pytest.mark.parametrize("details", [None, False, True]) +@pytest.mark.parametrize("reported", [False, True]) +def test_stream_details_are_opt_in_and_preserve_metadata(monkeypatch, details, reported): + class Sampling: + best_of = 1 + + def init(self, **kwargs): + assert "return_details" not in kwargs + + async def verify_cpu_async(self, compiler): + pass + + monkeypatch.setattr(api_lightllm, "SamplingParams", Sampling) + parameters = {"do_sample": False} + if details is not None: + parameters["return_details"] = details + + class Request: + async def json(self): + return {"inputs": [11, 12], "parameters": parameters} + + metadata = dict( + id=7, + prompt_tokens=2, + count_output_tokens=1, + logprob=-0.2, + logprobs={}, + prompt_cache_len=1, + mtp_accepted_token_num=2, + mtp_verify_token_num=3, + mtp_verify_step_num=1, + input_usage=dict(input_text_tokens=2), + ) + fields = ("prompt_cache_len", "mtp_accepted_token_num", "mtp_verify_token_num", "mtp_verify_step_num") + if not reported: + for name in fields: + metadata.pop(name) + before = copy.deepcopy(metadata) + + async def generate(*args, **kwargs): + yield 0, "token", metadata, SimpleNamespace(is_finished=lambda: True, get_finish_reason=lambda: "length") + + async def collect(): + response = await api_lightllm.lightllm_generate_stream( + Request(), SimpleNamespace(tokenizer=None, generate=generate, output_grammar_compiler=None) + ) + return [json.loads(chunk.decode().removeprefix("data:")) async for chunk in response.body_iterator] + + bodies = asyncio.run(collect()) + assert len(bodies) == 1 + body = bodies[0] + for name in fields: + if details and reported: + assert body["token"][name] == metadata[name] + else: + assert name not in body["token"] + assert body["input_usage"] == metadata["input_usage"] and body["details"] is None + assert body["token"]["id"] == 7 and body["finish_reason"] == "length" + assert metadata == before diff --git a/unit_tests/server/test_ssm_state_mode.py b/unit_tests/server/test_ssm_state_mode.py index e60e71028d..da786bddef 100644 --- a/unit_tests/server/test_ssm_state_mode.py +++ b/unit_tests/server/test_ssm_state_mode.py @@ -21,6 +21,8 @@ ("replay", "bfloat16", 2, "qwen3_5", 16, None), ("replay", "bfloat16", 2, "qwen3_5", 4, None), ("replay", "bfloat16", 0, "qwen3_5", 16, None), + ("replay", "float32", 3, "qwen4_exp", 16, None), + ("replay", "bfloat16", 3, "qwen4_exp_text", 8, None), ("replay", "float32", 4, "qwen3_5", 4, "capacity"), ("replay", "float32", 2, "llama", 16, "require a GDN"), ], @@ -46,7 +48,7 @@ def test_ssm_mode_validation(monkeypatch, mode, dtype, step, model, capacity, er _validate_and_normalize_ssm_config(args) -@pytest.mark.parametrize("model", ["qwen3_5", "glm5_next"]) +@pytest.mark.parametrize("model", ["qwen3_5", "glm5_next", "qwen4_exp", "qwen4_exp_text"]) @pytest.mark.parametrize("mode", ["native", "replay"]) @pytest.mark.parametrize("step", [15, 16]) def test_gdn_kda_mtp_width_limit(monkeypatch, model, mode, step): @@ -59,7 +61,7 @@ def test_gdn_kda_mtp_width_limit(monkeypatch, model, mode, step): _validate_and_normalize_ssm_config(args) -@pytest.mark.parametrize("model", ["qwen3_5", "glm5_next"]) +@pytest.mark.parametrize("model", ["qwen3_5", "glm5_next", "qwen4_exp", "qwen4_exp_text"]) @pytest.mark.parametrize("compute_mode", ["inline", "precompute"]) @pytest.mark.parametrize("cache_len", [32, 64]) def test_replay_large_cache_does_not_bypass_mtp_width_limit(monkeypatch, model, compute_mode, cache_len): diff --git a/unit_tests/utils/test_envs_utils.py b/unit_tests/utils/test_envs_utils.py index 75f48cb174..f8bec308db 100644 --- a/unit_tests/utils/test_envs_utils.py +++ b/unit_tests/utils/test_envs_utils.py @@ -1,5 +1,7 @@ from types import SimpleNamespace +import pytest + from lightllm.utils import envs_utils from lightllm.utils.envs_utils import ( get_deepep_num_max_dispatch_tokens_per_rank, @@ -177,3 +179,42 @@ def test_deepep_buffer_capacity_keeps_explicit_environment_override(monkeypatch) get_deepep_num_max_dispatch_tokens_per_rank_prefill.cache_clear() get_deepep_num_max_dispatch_tokens_per_rank_decode.cache_clear() + + +@pytest.mark.parametrize("tp,dp", [(8, 1), (8, 2), (8, 8)]) +@pytest.mark.parametrize("mtp_step", [0, 2, 3]) +@pytest.mark.parametrize("tokens,requests", [(8192, 256), (8195, 257), (1, 1)]) +def test_deepep_capacity_covers_sequence_parallel_padding(monkeypatch, tp, dp, mtp_step, tokens, requests): + import math + + monkeypatch.delenv("NUM_MAX_DISPATCH_TOKENS_PER_RANK_PREFILL", raising=False) + monkeypatch.delenv("NUM_MAX_DISPATCH_TOKENS_PER_RANK_DECODE", raising=False) + monkeypatch.setattr( + envs_utils, + "get_env_start_args", + lambda: SimpleNamespace( + batch_max_tokens=tokens, + running_max_req_size=requests, + mtp_step=mtp_step, + enable_tpsp_mix_mode=True, + tp=tp, + dp=dp, + ), + ) + get_deepep_num_max_dispatch_tokens_per_rank_prefill.cache_clear() + get_deepep_num_max_dispatch_tokens_per_rank_decode.cache_clear() + try: + prefill = get_deepep_num_max_dispatch_tokens_per_rank_prefill() + decode = get_deepep_num_max_dispatch_tokens_per_rank_decode() + world = tp // dp + alignment = math.lcm(world, mtp_step + 1) + verify_rows = requests * (mtp_step + 1) + padded_rows = (verify_rows + alignment - 1) // alignment * alignment + assert prefill >= (tokens + world - 1) // world + 128 + assert decode >= padded_rows // world + assert prefill % 8 == decode % 8 == 0 + if tp == 8 and dp == 1 and tokens == 8192 and mtp_step == 3: + assert (prefill, decode) == (1152, 144) + finally: + get_deepep_num_max_dispatch_tokens_per_rank_prefill.cache_clear() + get_deepep_num_max_dispatch_tokens_per_rank_decode.cache_clear() diff --git a/unit_tests/utils/test_net_utils.py b/unit_tests/utils/test_net_utils.py new file mode 100644 index 0000000000..5ca317a94a --- /dev/null +++ b/unit_tests/utils/test_net_utils.py @@ -0,0 +1,45 @@ +import errno +import socket +import sys + +import pytest + +from lightllm.utils.net_utils import validate_ports + + +pytestmark = pytest.mark.skipif(sys.platform != "linux", reason="requires Linux TCP port reuse semantics") + + +def test_validate_ports_rejects_duplicates(): + with pytest.raises(RuntimeError, match="conflicting ports"): + validate_ports([0, 0]) + + +def test_validate_ports_rejects_active_listener(): + with socket.socket() as listener: + listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + listener.bind(("127.0.0.1", 0)) + listener.listen() + port = listener.getsockname()[1] + with pytest.raises(RuntimeError, match=f"port {port} has been used"): + validate_ports([port]) + + +def test_validate_ports_allows_restart_after_closed_connection(): + with socket.socket() as listener: + listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + listener.bind(("127.0.0.1", 0)) + listener.listen() + listener.settimeout(2) + address = listener.getsockname() + with socket.create_connection(address, timeout=2) as client: + connection, _ = listener.accept() + with connection: + connection.shutdown(socket.SHUT_WR) + assert client.recv(1) == b"" + + with socket.socket() as probe: + with pytest.raises(OSError) as error: + probe.bind(("", address[1])) + assert error.value.errno == errno.EADDRINUSE + validate_ports([address[1]]) diff --git a/unit_tests/utils/test_startup_process_lifecycle.py b/unit_tests/utils/test_startup_process_lifecycle.py index 9b9d38fb60..ad5131bd8a 100644 --- a/unit_tests/utils/test_startup_process_lifecycle.py +++ b/unit_tests/utils/test_startup_process_lifecycle.py @@ -4,6 +4,7 @@ import os import signal import sys +import threading import time import traceback from types import SimpleNamespace @@ -133,3 +134,48 @@ def test_metrics_cleanup_on_launcher_failure(failure_mode): controller.close() result_reader.close() result_writer.close() + + +def _run_model_thread(fail, started): + model_rpc = importlib.import_module("lightllm.server.router.model_infer.model_rpc") + threading.excepthook = threading.__excepthook__ + model_rpc.graceful_registry = lambda _: None + model_rpc.start_parent_check_thread = lambda: None + model_rpc.get_unique_server_name = lambda: "model-thread-regression" + model_rpc.get_model_infer_recursion_limit = lambda: 3000 + + def infer(): + started.set() + if fail: + raise RuntimeError("injected inference thread failure") + + class Server: + def __init__(self, *args, **kwargs): + pass + + def start(self): + thread = threading.Thread(target=infer, daemon=True) + thread.start() + thread.join() + + model_rpc.ThreadedServer = Server + model_rpc._init_env(SimpleNamespace(), 0, 0, 1, None, "unused", mp.get_context("fork").Event()) + + +@pytest.mark.parametrize("fail", [False, True]) +def test_model_worker_exits_on_inference_thread_failure(fail): + importlib.import_module("lightllm.server.router.model_infer.model_rpc") + ctx = mp.get_context("fork") + started = ctx.Event() + worker = ctx.Process(target=_run_model_thread, args=(fail, started)) + worker.start() + try: + worker.join(timeout=5) + assert started.is_set(), "inference thread did not start" + assert not worker.is_alive(), "model worker survived a failed inference thread" + assert worker.exitcode == (1 if fail else 0) + finally: + if worker.is_alive(): + worker.kill() + worker.join(timeout=5) + worker.close()