sglang 0.5.2rc2__py3-none-any.whl → 0.5.3rc2__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- sglang/bench_one_batch.py +7 -9
- sglang/bench_one_batch_server.py +330 -31
- sglang/bench_serving.py +267 -32
- sglang/global_config.py +2 -2
- sglang/lang/backend/runtime_endpoint.py +1 -1
- sglang/launch_server.py +14 -0
- sglang/profiler.py +2 -2
- sglang/srt/batch_invariant_ops/__init__.py +27 -0
- sglang/srt/batch_invariant_ops/batch_invariant_ops.py +549 -0
- sglang/srt/configs/__init__.py +8 -0
- sglang/srt/configs/device_config.py +3 -1
- sglang/srt/configs/dots_ocr.py +64 -0
- sglang/srt/configs/dots_vlm.py +139 -0
- sglang/srt/configs/falcon_h1.py +360 -0
- sglang/srt/configs/load_config.py +9 -0
- sglang/srt/configs/model_config.py +181 -82
- sglang/srt/configs/qwen3_next.py +326 -0
- sglang/srt/configs/qwen3_vl.py +586 -0
- sglang/srt/connector/__init__.py +8 -1
- sglang/srt/connector/remote_instance.py +82 -0
- sglang/srt/constrained/base_grammar_backend.py +49 -12
- sglang/srt/constrained/llguidance_backend.py +0 -1
- sglang/srt/constrained/outlines_backend.py +0 -1
- sglang/srt/constrained/outlines_jump_forward.py +1 -1
- sglang/srt/constrained/xgrammar_backend.py +30 -9
- sglang/srt/custom_op.py +11 -1
- sglang/srt/debug_utils/dump_comparator.py +81 -44
- sglang/srt/debug_utils/dump_loader.py +97 -0
- sglang/srt/debug_utils/dumper.py +21 -6
- sglang/srt/debug_utils/text_comparator.py +73 -11
- sglang/srt/disaggregation/ascend/conn.py +2 -2
- sglang/srt/disaggregation/ascend/transfer_engine.py +47 -9
- sglang/srt/disaggregation/base/conn.py +1 -1
- sglang/srt/disaggregation/common/conn.py +279 -108
- sglang/srt/disaggregation/decode.py +71 -19
- sglang/srt/disaggregation/decode_kvcache_offload_manager.py +185 -0
- sglang/srt/disaggregation/decode_schedule_batch_mixin.py +29 -17
- sglang/srt/disaggregation/fake/conn.py +1 -1
- sglang/srt/disaggregation/mini_lb.py +6 -445
- sglang/srt/disaggregation/mooncake/conn.py +55 -537
- sglang/srt/disaggregation/nixl/conn.py +326 -53
- sglang/srt/disaggregation/prefill.py +36 -17
- sglang/srt/disaggregation/utils.py +40 -54
- sglang/srt/distributed/device_communicators/all_reduce_utils.py +16 -0
- sglang/srt/distributed/device_communicators/shm_broadcast.py +4 -2
- sglang/srt/distributed/device_communicators/symm_mem.py +164 -0
- sglang/srt/distributed/parallel_state.py +156 -80
- sglang/srt/entrypoints/engine.py +59 -18
- sglang/srt/entrypoints/grpc_request_manager.py +855 -0
- sglang/srt/entrypoints/grpc_server.py +810 -0
- sglang/srt/entrypoints/http_server.py +130 -59
- sglang/srt/entrypoints/openai/protocol.py +112 -4
- sglang/srt/entrypoints/openai/serving_base.py +65 -3
- sglang/srt/entrypoints/openai/serving_chat.py +204 -55
- sglang/srt/entrypoints/openai/serving_completions.py +14 -3
- sglang/srt/entrypoints/openai/serving_embedding.py +9 -3
- sglang/srt/entrypoints/openai/serving_rerank.py +3 -1
- sglang/srt/entrypoints/openai/serving_responses.py +48 -3
- sglang/srt/entrypoints/openai/serving_score.py +1 -0
- sglang/srt/environ.py +285 -0
- sglang/srt/eplb/eplb_manager.py +2 -2
- sglang/srt/eplb/expert_distribution.py +26 -13
- sglang/srt/eplb/expert_location.py +38 -8
- sglang/srt/eplb/expert_location_updater.py +1 -1
- sglang/srt/function_call/base_format_detector.py +3 -6
- sglang/srt/function_call/ebnf_composer.py +11 -9
- sglang/srt/function_call/function_call_parser.py +9 -2
- sglang/srt/function_call/glm4_moe_detector.py +4 -4
- sglang/srt/function_call/gpt_oss_detector.py +23 -0
- sglang/srt/function_call/json_array_parser.py +63 -0
- sglang/srt/function_call/kimik2_detector.py +17 -4
- sglang/srt/function_call/qwen3_coder_detector.py +1 -1
- sglang/srt/function_call/utils.py +96 -5
- sglang/srt/grpc/__init__.py +1 -0
- sglang/srt/grpc/compile_proto.py +245 -0
- sglang/srt/grpc/sglang_scheduler_pb2.py +111 -0
- sglang/srt/grpc/sglang_scheduler_pb2.pyi +434 -0
- sglang/srt/grpc/sglang_scheduler_pb2_grpc.py +239 -0
- sglang/srt/layers/activation.py +143 -9
- sglang/srt/layers/attention/aiter_backend.py +14 -15
- sglang/srt/layers/attention/ascend_backend.py +115 -9
- sglang/srt/layers/attention/attention_registry.py +206 -0
- sglang/srt/layers/attention/base_attn_backend.py +12 -3
- sglang/srt/layers/attention/cutlass_mla_backend.py +3 -3
- sglang/srt/layers/attention/dual_chunk_flashattention_backend.py +1 -1
- sglang/srt/layers/attention/fla/chunk.py +242 -0
- sglang/srt/layers/attention/fla/chunk_delta_h.py +314 -0
- sglang/srt/layers/attention/fla/chunk_o.py +178 -0
- sglang/srt/layers/attention/fla/chunk_scaled_dot_kkt.py +151 -0
- sglang/srt/layers/attention/fla/cumsum.py +300 -0
- sglang/srt/layers/attention/fla/fused_recurrent.py +640 -0
- sglang/srt/layers/attention/fla/fused_sigmoid_gating_recurrent.py +232 -0
- sglang/srt/layers/attention/fla/index.py +37 -0
- sglang/srt/layers/attention/fla/l2norm.py +150 -0
- sglang/srt/layers/attention/fla/layernorm_gated.py +326 -0
- sglang/srt/layers/attention/fla/op.py +66 -0
- sglang/srt/layers/attention/fla/solve_tril.py +465 -0
- sglang/srt/layers/attention/fla/utils.py +331 -0
- sglang/srt/layers/attention/fla/wy_fast.py +158 -0
- sglang/srt/layers/attention/flashattention_backend.py +41 -8
- sglang/srt/layers/attention/flashinfer_backend.py +118 -198
- sglang/srt/layers/attention/flashinfer_mla_backend.py +27 -27
- sglang/srt/layers/attention/flashmla_backend.py +7 -5
- sglang/srt/layers/attention/hybrid_attn_backend.py +68 -53
- sglang/srt/layers/attention/hybrid_linear_attn_backend.py +602 -0
- sglang/srt/layers/attention/intel_amx_backend.py +3 -0
- sglang/srt/layers/attention/mamba/causal_conv1d.py +129 -0
- sglang/srt/layers/attention/mamba/causal_conv1d_triton.py +969 -0
- sglang/srt/layers/attention/mamba/mamba.py +629 -0
- sglang/srt/layers/attention/mamba/mamba_utils.py +81 -0
- sglang/srt/layers/attention/mamba/ops/__init__.py +2 -0
- sglang/srt/layers/attention/mamba/ops/layernorm_gated.py +172 -0
- sglang/srt/layers/attention/mamba/ops/mamba_ssm.py +442 -0
- sglang/srt/layers/attention/mamba/ops/ssd_bmm.py +264 -0
- sglang/srt/layers/attention/mamba/ops/ssd_chunk_scan.py +622 -0
- sglang/srt/layers/attention/mamba/ops/ssd_chunk_state.py +757 -0
- sglang/srt/layers/attention/mamba/ops/ssd_combined.py +262 -0
- sglang/srt/layers/attention/mamba/ops/ssd_state_passing.py +275 -0
- sglang/srt/layers/attention/npu_ops/mla_preprocess.py +393 -0
- sglang/srt/layers/attention/nsa/dequant_k_cache.py +163 -0
- sglang/srt/layers/attention/nsa/index_buf_accessor.py +354 -0
- sglang/srt/layers/attention/nsa/nsa_indexer.py +761 -0
- sglang/srt/layers/attention/nsa/quant_k_cache.py +255 -0
- sglang/srt/layers/attention/nsa/tilelang_kernel.py +785 -0
- sglang/srt/layers/attention/nsa/transform_index.py +144 -0
- sglang/srt/layers/attention/nsa/utils.py +24 -0
- sglang/srt/layers/attention/nsa_backend.py +887 -0
- sglang/srt/layers/attention/tbo_backend.py +6 -6
- sglang/srt/layers/attention/torch_flex_backend.py +325 -0
- sglang/srt/layers/attention/torch_native_backend.py +12 -6
- sglang/srt/layers/attention/triton_backend.py +57 -7
- sglang/srt/layers/attention/trtllm_mha_backend.py +5 -7
- sglang/srt/layers/attention/trtllm_mla_backend.py +276 -39
- sglang/srt/layers/attention/vision.py +58 -0
- sglang/srt/layers/attention/wave_backend.py +4 -4
- sglang/srt/layers/attention/wave_ops/decode_attention.py +2 -4
- sglang/srt/layers/attention/wave_ops/extend_attention.py +1 -3
- sglang/srt/layers/communicator.py +8 -0
- sglang/srt/layers/dp_attention.py +41 -2
- sglang/srt/layers/elementwise.py +3 -1
- sglang/srt/layers/layernorm.py +34 -15
- sglang/srt/layers/linear.py +55 -7
- sglang/srt/layers/logits_processor.py +44 -12
- sglang/srt/layers/moe/__init__.py +2 -1
- sglang/srt/layers/moe/cutlass_w4a8_moe.py +3 -3
- sglang/srt/layers/moe/ep_moe/kernels.py +2 -2
- sglang/srt/layers/moe/ep_moe/layer.py +256 -63
- sglang/srt/layers/moe/flashinfer_cutedsl_moe.py +183 -0
- sglang/srt/layers/moe/fused_moe_native.py +5 -3
- sglang/srt/layers/moe/fused_moe_triton/configs/{triton_3_4_0/E=128,N=768,device_name=NVIDIA_B200,dtype=fp8_w8a8,block_shape=[128, 128].json → triton_3_3_1/E=128,N=768,device_name=NVIDIA_H20,dtype=fp8_w8a8,block_shape=[128, 128].json } +35 -35
- sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=128,N=352,device_name=NVIDIA_RTX_5880_Ada_Generation,dtype=fp8_w8a8.json +146 -0
- sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=256,N=256,device_name=NVIDIA_H800,dtype=fp8_w8a8,block_shape=[128, 128].json +146 -0
- sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=256,N=512,device_name=NVIDIA_H20.json +146 -0
- sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=512,N=128,device_name=NVIDIA_H100_80GB_HBM3.json +146 -0
- sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=512,N=128,device_name=NVIDIA_H20-3e.json +146 -0
- sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=512,N=128,device_name=NVIDIA_H200.json +146 -0
- sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=512,N=128,device_name=NVIDIA_H800,dtype=fp8_w8a8,block_shape=[128, 128].json +146 -0
- sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=512,N=256,device_name=NVIDIA_B200.json +146 -0
- sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=512,N=256,device_name=NVIDIA_H20-3e.json +146 -0
- sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=512,N=256,device_name=NVIDIA_H200.json +146 -0
- sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=512,N=64,device_name=NVIDIA_H100_80GB_HBM3.json +146 -0
- sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=512,N=64,device_name=NVIDIA_H200.json +146 -0
- sglang/srt/layers/moe/fused_moe_triton/fused_moe.py +5 -2
- sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_config.py +7 -3
- sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_kernels.py +23 -20
- sglang/srt/layers/moe/fused_moe_triton/layer.py +71 -70
- sglang/srt/layers/moe/moe_runner/__init__.py +2 -1
- sglang/srt/layers/moe/moe_runner/base.py +274 -1
- sglang/srt/layers/moe/moe_runner/runner.py +80 -0
- sglang/srt/layers/moe/moe_runner/triton.py +448 -0
- sglang/srt/layers/moe/token_dispatcher/__init__.py +16 -4
- sglang/srt/layers/moe/token_dispatcher/{base_dispatcher.py → base.py} +67 -17
- sglang/srt/layers/moe/token_dispatcher/deepep.py +118 -56
- sglang/srt/layers/moe/token_dispatcher/standard.py +44 -2
- sglang/srt/layers/moe/topk.py +30 -9
- sglang/srt/layers/moe/utils.py +22 -6
- sglang/srt/layers/parameter.py +23 -6
- sglang/srt/layers/quantization/awq.py +19 -7
- sglang/srt/layers/quantization/base_config.py +11 -6
- sglang/srt/layers/quantization/blockwise_int8.py +38 -27
- sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py +1 -0
- sglang/srt/layers/quantization/compressed_tensors/compressed_tensors_moe.py +50 -30
- sglang/srt/layers/quantization/compressed_tensors/schemes/__init__.py +2 -0
- sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py +13 -1
- sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_int8.py +173 -0
- sglang/srt/layers/quantization/deep_gemm_wrapper/configurer.py +2 -10
- sglang/srt/layers/quantization/deep_gemm_wrapper/entrypoint.py +27 -0
- sglang/srt/layers/quantization/fp8.py +78 -49
- sglang/srt/layers/quantization/fp8_utils.py +51 -32
- sglang/srt/layers/quantization/gptq.py +25 -17
- sglang/srt/layers/quantization/modelopt_quant.py +190 -55
- sglang/srt/layers/quantization/moe_wna16.py +21 -18
- sglang/srt/layers/quantization/mxfp4.py +74 -42
- sglang/srt/layers/quantization/quark/quark_moe.py +48 -30
- sglang/srt/layers/quantization/unquant.py +135 -47
- sglang/srt/layers/quantization/w4afp8.py +26 -17
- sglang/srt/layers/quantization/w8a8_fp8.py +35 -20
- sglang/srt/layers/quantization/w8a8_int8.py +91 -41
- sglang/srt/layers/rotary_embedding.py +78 -31
- sglang/srt/layers/sampler.py +213 -21
- sglang/srt/layers/utils.py +23 -0
- sglang/srt/lora/backend/base_backend.py +50 -8
- sglang/srt/lora/backend/chunked_backend.py +348 -0
- sglang/srt/lora/backend/triton_backend.py +99 -5
- sglang/srt/lora/layers.py +32 -0
- sglang/srt/lora/lora.py +8 -3
- sglang/srt/lora/lora_manager.py +52 -118
- sglang/srt/lora/mem_pool.py +25 -11
- sglang/srt/lora/triton_ops/__init__.py +4 -0
- sglang/srt/lora/triton_ops/chunked_sgmv_expand.py +214 -0
- sglang/srt/lora/triton_ops/chunked_sgmv_shrink.py +174 -0
- sglang/srt/lora/utils.py +22 -11
- sglang/srt/managers/async_dynamic_batch_tokenizer.py +170 -0
- sglang/srt/managers/cache_controller.py +199 -301
- sglang/srt/managers/data_parallel_controller.py +115 -80
- sglang/srt/managers/detokenizer_manager.py +19 -15
- sglang/srt/managers/disagg_service.py +46 -0
- sglang/srt/managers/io_struct.py +340 -109
- sglang/srt/managers/mm_utils.py +44 -6
- sglang/srt/managers/multi_tokenizer_mixin.py +357 -407
- sglang/srt/managers/multimodal_processor.py +1 -2
- sglang/srt/managers/overlap_utils.py +53 -0
- sglang/srt/managers/schedule_batch.py +240 -138
- sglang/srt/managers/schedule_policy.py +144 -17
- sglang/srt/managers/scheduler.py +502 -209
- sglang/srt/managers/scheduler_input_blocker.py +1 -1
- sglang/srt/managers/scheduler_metrics_mixin.py +99 -126
- sglang/srt/managers/scheduler_output_processor_mixin.py +75 -22
- sglang/srt/managers/scheduler_profiler_mixin.py +6 -6
- sglang/srt/managers/scheduler_update_weights_mixin.py +7 -0
- sglang/srt/managers/tokenizer_communicator_mixin.py +675 -0
- sglang/srt/managers/tokenizer_manager.py +320 -632
- sglang/srt/managers/tp_worker.py +81 -22
- sglang/srt/managers/tp_worker_overlap_thread.py +71 -56
- sglang/srt/managers/utils.py +1 -45
- sglang/srt/mem_cache/allocator.py +14 -20
- sglang/srt/mem_cache/allocator_ascend.py +41 -27
- sglang/srt/mem_cache/base_prefix_cache.py +1 -1
- sglang/srt/mem_cache/chunk_cache.py +8 -1
- sglang/srt/mem_cache/evict_policy.py +23 -0
- sglang/srt/mem_cache/hicache_storage.py +43 -24
- sglang/srt/mem_cache/hiradix_cache.py +222 -75
- sglang/srt/mem_cache/memory_pool.py +535 -58
- sglang/srt/mem_cache/memory_pool_host.py +239 -228
- sglang/srt/mem_cache/radix_cache.py +222 -73
- sglang/srt/mem_cache/radix_cache_cpp.py +11 -8
- sglang/srt/mem_cache/storage/__init__.py +10 -0
- sglang/srt/mem_cache/storage/aibrix_kvcache/aibrix_kvcache_storage.py +151 -0
- sglang/srt/mem_cache/storage/aibrix_kvcache/unit_test.py +109 -0
- sglang/srt/mem_cache/storage/backend_factory.py +223 -0
- sglang/srt/mem_cache/storage/eic/eic_storage.py +778 -0
- sglang/srt/mem_cache/storage/eic/test_unit.py +115 -0
- sglang/srt/mem_cache/storage/hf3fs/hf3fs_client.py +164 -0
- sglang/srt/mem_cache/storage/hf3fs/{client_hf3fs.py → hf3fs_usrbio_client.py} +5 -1
- sglang/srt/mem_cache/storage/hf3fs/storage_hf3fs.py +259 -62
- sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py +284 -0
- sglang/srt/mem_cache/storage/lmcache/unit_test.py +121 -0
- sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py +166 -17
- sglang/srt/mem_cache/swa_radix_cache.py +25 -36
- sglang/srt/metrics/collector.py +511 -132
- sglang/srt/metrics/func_timer.py +2 -7
- sglang/srt/metrics/startup_func_log_and_timer.py +150 -0
- sglang/srt/metrics/utils.py +8 -1
- sglang/srt/model_executor/cpu_graph_runner.py +640 -0
- sglang/srt/model_executor/cuda_graph_runner.py +52 -37
- sglang/srt/model_executor/forward_batch_info.py +82 -40
- sglang/srt/model_executor/model_runner.py +432 -157
- sglang/srt/model_executor/npu_graph_runner.py +12 -5
- sglang/srt/model_loader/__init__.py +9 -3
- sglang/srt/model_loader/loader.py +133 -5
- sglang/srt/model_loader/remote_instance_weight_loader_utils.py +69 -0
- sglang/srt/model_loader/weight_utils.py +158 -3
- sglang/srt/models/apertus.py +686 -0
- sglang/srt/models/bailing_moe.py +820 -217
- sglang/srt/models/bailing_moe_nextn.py +168 -0
- sglang/srt/models/deepseek_nextn.py +6 -1
- sglang/srt/models/deepseek_v2.py +607 -130
- sglang/srt/models/dots_ocr.py +173 -0
- sglang/srt/models/dots_vlm.py +174 -0
- sglang/srt/models/dots_vlm_vit.py +337 -0
- sglang/srt/models/ernie4.py +1 -1
- sglang/srt/models/falcon_h1.py +576 -0
- sglang/srt/models/gemma3_causal.py +0 -2
- sglang/srt/models/gemma3_mm.py +1 -1
- sglang/srt/models/gemma3n_mm.py +2 -2
- sglang/srt/models/glm4_moe.py +4 -4
- sglang/srt/models/glm4_moe_nextn.py +2 -2
- sglang/srt/models/glm4v.py +5 -3
- sglang/srt/models/glm4v_moe.py +4 -1
- sglang/srt/models/gpt_oss.py +8 -31
- sglang/srt/models/kimi_vl_moonvit.py +2 -2
- sglang/srt/models/llama.py +4 -0
- sglang/srt/models/llama4.py +9 -0
- sglang/srt/models/llama_eagle3.py +13 -0
- sglang/srt/models/longcat_flash.py +3 -3
- sglang/srt/models/longcat_flash_nextn.py +1 -1
- sglang/srt/models/mllama4.py +40 -4
- sglang/srt/models/opt.py +637 -0
- sglang/srt/models/qwen2_5_vl.py +29 -5
- sglang/srt/models/qwen2_audio.py +1 -1
- sglang/srt/models/qwen2_moe.py +120 -13
- sglang/srt/models/qwen2_vl.py +1 -1
- sglang/srt/models/qwen3.py +18 -3
- sglang/srt/models/qwen3_moe.py +32 -4
- sglang/srt/models/qwen3_next.py +1069 -0
- sglang/srt/models/qwen3_next_mtp.py +112 -0
- sglang/srt/models/qwen3_vl.py +787 -0
- sglang/srt/models/qwen3_vl_moe.py +471 -0
- sglang/srt/models/registry.py +15 -3
- sglang/srt/models/sarashina2_vision.py +269 -0
- sglang/srt/models/solar.py +505 -0
- sglang/srt/models/starcoder2.py +357 -0
- sglang/srt/models/step3_vl.py +1 -1
- sglang/srt/models/torch_native_llama.py +9 -2
- sglang/srt/models/utils.py +51 -0
- sglang/srt/multimodal/processors/base_processor.py +15 -7
- sglang/srt/multimodal/processors/dots_vlm.py +98 -0
- sglang/srt/multimodal/processors/glm4v.py +9 -9
- sglang/srt/multimodal/processors/internvl.py +153 -129
- sglang/srt/multimodal/processors/qwen_vl.py +23 -6
- sglang/srt/multimodal/processors/sarashina2_vision.py +81 -0
- sglang/srt/offloader.py +27 -3
- sglang/srt/parser/jinja_template_utils.py +6 -0
- sglang/srt/sampling/sampling_batch_info.py +38 -17
- sglang/srt/sampling/sampling_params.py +7 -0
- sglang/srt/server_args.py +966 -267
- sglang/srt/server_args_config_parser.py +146 -0
- sglang/srt/single_batch_overlap.py +151 -0
- sglang/srt/speculative/cpp_ngram/ngram.cpp +374 -0
- sglang/srt/speculative/cpp_ngram/ngram.h +110 -0
- sglang/srt/speculative/cpp_ngram/ngram_cache.py +138 -0
- sglang/srt/speculative/cpp_ngram/ngram_cache_binding.cpp +43 -0
- sglang/srt/speculative/cpp_ngram/param.h +125 -0
- sglang/srt/speculative/cpp_ngram/queue.h +71 -0
- sglang/srt/speculative/eagle_draft_cuda_graph_runner.py +7 -1
- sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py +13 -2
- sglang/srt/speculative/{eagle_utils.py → eagle_info.py} +207 -757
- sglang/srt/speculative/eagle_worker.py +99 -28
- sglang/srt/speculative/ngram_utils.py +428 -0
- sglang/srt/speculative/ngram_worker.py +245 -0
- sglang/srt/speculative/spec_info.py +52 -0
- sglang/srt/speculative/spec_utils.py +606 -0
- sglang/srt/speculative/standalone_worker.py +109 -0
- sglang/srt/torch_memory_saver_adapter.py +5 -7
- sglang/srt/tracing/trace.py +578 -0
- sglang/srt/two_batch_overlap.py +8 -5
- sglang/srt/utils/__init__.py +2 -0
- sglang/srt/{utils.py → utils/common.py} +433 -77
- sglang/srt/{hf_transformers_utils.py → utils/hf_transformers_utils.py} +53 -5
- sglang/srt/{patch_torch.py → utils/patch_torch.py} +8 -0
- sglang/srt/utils/rpd_utils.py +452 -0
- sglang/srt/utils/slow_rank_detector.py +71 -0
- sglang/srt/warmup.py +8 -4
- sglang/srt/weight_sync/utils.py +2 -2
- sglang/test/attention/test_trtllm_mla_backend.py +169 -5
- sglang/test/get_logits_ut.py +57 -0
- sglang/test/run_eval.py +79 -11
- sglang/test/runners.py +5 -1
- sglang/test/simple_eval_common.py +5 -2
- sglang/test/simple_eval_mmmu_vlm.py +441 -0
- sglang/test/test_block_fp8.py +2 -2
- sglang/test/test_cutlass_moe.py +24 -6
- sglang/test/test_deterministic.py +297 -0
- sglang/test/test_disaggregation_utils.py +77 -0
- sglang/test/test_fp4_moe.py +370 -1
- sglang/test/test_programs.py +1 -1
- sglang/test/test_utils.py +383 -5
- sglang/utils.py +21 -1
- sglang/version.py +1 -1
- {sglang-0.5.2rc2.dist-info → sglang-0.5.3rc2.dist-info}/METADATA +69 -124
- {sglang-0.5.2rc2.dist-info → sglang-0.5.3rc2.dist-info}/RECORD +375 -245
- sglang/srt/disaggregation/launch_lb.py +0 -118
- sglang/srt/mem_cache/lora_radix_cache.py +0 -421
- /sglang/srt/{poll_based_barrier.py → utils/poll_based_barrier.py} +0 -0
- {sglang-0.5.2rc2.dist-info → sglang-0.5.3rc2.dist-info}/WHEEL +0 -0
- {sglang-0.5.2rc2.dist-info → sglang-0.5.3rc2.dist-info}/licenses/LICENSE +0 -0
- {sglang-0.5.2rc2.dist-info → sglang-0.5.3rc2.dist-info}/top_level.txt +0 -0
@@ -0,0 +1,139 @@
|
|
1
|
+
from typing import Any, List, Optional, Union
|
2
|
+
|
3
|
+
from transformers import AutoProcessor, LlamaTokenizerFast, PretrainedConfig
|
4
|
+
from transformers.feature_extraction_utils import BatchFeature
|
5
|
+
from transformers.image_utils import ImageInput
|
6
|
+
from transformers.processing_utils import ProcessingKwargs, Unpack
|
7
|
+
from transformers.tokenization_utils_base import PreTokenizedInput, TextInput
|
8
|
+
|
9
|
+
try:
|
10
|
+
from transformers import Qwen2_5_VLProcessor
|
11
|
+
except ImportError:
|
12
|
+
raise ImportError(
|
13
|
+
"Qwen2_5_VLProcessor can not be found. Please upgrade your transformers version."
|
14
|
+
)
|
15
|
+
|
16
|
+
from sglang.srt.configs.deepseekvl2 import DeepseekV2Config
|
17
|
+
|
18
|
+
|
19
|
+
class DotsVisionConfig(PretrainedConfig):
|
20
|
+
model_type: str = "dots_vit"
|
21
|
+
|
22
|
+
def __init__(
|
23
|
+
self,
|
24
|
+
embed_dim: int = 1536, # vision encoder embed size
|
25
|
+
hidden_size: int = 1536, # after merger hidden size
|
26
|
+
intermediate_size: int = 4224,
|
27
|
+
num_hidden_layers: int = 42,
|
28
|
+
num_attention_heads: int = 12,
|
29
|
+
num_channels: int = 3,
|
30
|
+
patch_size: int = 14,
|
31
|
+
spatial_merge_size: int = 2,
|
32
|
+
temporal_patch_size: int = 1,
|
33
|
+
rms_norm_eps: float = 1e-5,
|
34
|
+
use_bias: bool = False,
|
35
|
+
attn_implementation="flash_attention_2", # "eager","sdpa","flash_attention_2"
|
36
|
+
initializer_range=0.02,
|
37
|
+
init_merger_std=0.02,
|
38
|
+
is_causal=False, # ve causal forward
|
39
|
+
post_norm=True,
|
40
|
+
gradient_checkpointing=False,
|
41
|
+
**kwargs,
|
42
|
+
):
|
43
|
+
super().__init__(**kwargs)
|
44
|
+
self.embed_dim = embed_dim
|
45
|
+
self.hidden_size = hidden_size
|
46
|
+
self.intermediate_size = intermediate_size
|
47
|
+
self.num_hidden_layers = num_hidden_layers
|
48
|
+
self.num_attention_heads = num_attention_heads
|
49
|
+
self.num_channels = num_channels
|
50
|
+
self.patch_size = patch_size
|
51
|
+
self.spatial_merge_size = spatial_merge_size
|
52
|
+
self.temporal_patch_size = temporal_patch_size
|
53
|
+
self.rms_norm_eps = rms_norm_eps
|
54
|
+
self.use_bias = use_bias
|
55
|
+
self.attn_implementation = attn_implementation
|
56
|
+
self.initializer_range = initializer_range
|
57
|
+
self.init_merger_std = init_merger_std
|
58
|
+
self.is_causal = is_causal
|
59
|
+
self.post_norm = post_norm
|
60
|
+
self.gradient_checkpointing = gradient_checkpointing
|
61
|
+
|
62
|
+
|
63
|
+
class DotsVLMConfig(PretrainedConfig):
|
64
|
+
model_type = "dots_vlm"
|
65
|
+
|
66
|
+
def __init__(self, **kwargs):
|
67
|
+
super().__init__(**kwargs)
|
68
|
+
vision_config = kwargs.get("vision_config", {})
|
69
|
+
self.im_span_id = kwargs.get("image_token_id", 128815)
|
70
|
+
self.video_span_id = kwargs.get("video_token_id", 128836)
|
71
|
+
self.vision_config = DotsVisionConfig(**vision_config)
|
72
|
+
self.language_config = DeepseekV2Config(**kwargs)
|
73
|
+
self.architectures = ["DotsVLMForCausalLM"]
|
74
|
+
|
75
|
+
|
76
|
+
class DotsVLMProcessorKwargs(ProcessingKwargs, total=False):
|
77
|
+
_defaults = {
|
78
|
+
"text_kwargs": {
|
79
|
+
"padding": False,
|
80
|
+
},
|
81
|
+
}
|
82
|
+
|
83
|
+
|
84
|
+
class DotsVLMProcessor(Qwen2_5_VLProcessor):
|
85
|
+
r"""
|
86
|
+
Constructs a DotsVLM processor which derives from Qwen2_5_VLProcessor, but overrides the image and video token ids.
|
87
|
+
Besides, its tokenizer is a LlamaTokenizerFast instead of Qwen2TokenizerFast.
|
88
|
+
[`DotsVLMProcessor`] offers all the functionalities of [`DotsVisionConfig`] and [`LlamaTokenizerFast`]. See the
|
89
|
+
[`~DotsVLMProcessor.__call__`] and [`~DotsVLMProcessor.decode`] for more information.
|
90
|
+
Args:
|
91
|
+
image_processor ([`Qwen2VLImageProcessor`], *optional*):
|
92
|
+
The image processor is a required input.
|
93
|
+
tokenizer ([`LlamaTokenizerFast`], *optional*):
|
94
|
+
The tokenizer is a required input.
|
95
|
+
chat_template (`str`, *optional*): A Jinja template which will be used to convert lists of messages
|
96
|
+
in a chat into a tokenizable string.
|
97
|
+
"""
|
98
|
+
|
99
|
+
attributes = ["image_processor", "tokenizer"]
|
100
|
+
|
101
|
+
valid_kwargs = ["chat_template"]
|
102
|
+
|
103
|
+
tokenizer_class = ("LlamaTokenizer", "LlamaTokenizerFast")
|
104
|
+
|
105
|
+
def __init__(
|
106
|
+
self, image_processor=None, tokenizer=None, chat_template=None, **kwargs
|
107
|
+
):
|
108
|
+
super().__init__(image_processor, tokenizer, chat_template=chat_template)
|
109
|
+
self.image_token = (
|
110
|
+
"<|imgpad|>"
|
111
|
+
if not hasattr(tokenizer, "image_token")
|
112
|
+
else tokenizer.image_token
|
113
|
+
)
|
114
|
+
self.video_token = (
|
115
|
+
"<|video_pad|>"
|
116
|
+
if not hasattr(tokenizer, "video_token")
|
117
|
+
else tokenizer.video_token
|
118
|
+
)
|
119
|
+
self.img_token = (
|
120
|
+
"<|img|>" if not hasattr(tokenizer, "img_token") else tokenizer.img_token
|
121
|
+
)
|
122
|
+
self.endofimg_token = (
|
123
|
+
"<|endofimg|>"
|
124
|
+
if not hasattr(tokenizer, "endofimg_token")
|
125
|
+
else tokenizer.endofimg_token
|
126
|
+
)
|
127
|
+
self.image_token_id = (
|
128
|
+
tokenizer.image_token_id
|
129
|
+
if getattr(tokenizer, "image_token_id", None)
|
130
|
+
else tokenizer.encode(self.image_token)[0]
|
131
|
+
)
|
132
|
+
self.video_token_id = (
|
133
|
+
tokenizer.video_token_id
|
134
|
+
if getattr(tokenizer, "video_token_id", None)
|
135
|
+
else tokenizer.encode(self.video_token)[0]
|
136
|
+
)
|
137
|
+
|
138
|
+
|
139
|
+
AutoProcessor.register(DotsVLMConfig, DotsVLMProcessor)
|
@@ -0,0 +1,360 @@
|
|
1
|
+
# coding=utf-8
|
2
|
+
# Copyright 2024 TII and the HuggingFace Inc. team. All rights reserved.
|
3
|
+
#
|
4
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
5
|
+
# you may not use this file except in compliance with the License.
|
6
|
+
# You may obtain a copy of the License at
|
7
|
+
#
|
8
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
9
|
+
#
|
10
|
+
# Unless required by applicable law or agreed to in writing, software
|
11
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
12
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
13
|
+
# See the License for the specific language governing permissions and
|
14
|
+
# limitations under the License.
|
15
|
+
"""Falcon-H1 model configuration"""
|
16
|
+
|
17
|
+
import enum
|
18
|
+
import os
|
19
|
+
|
20
|
+
import numpy as np
|
21
|
+
import torch
|
22
|
+
from transformers.configuration_utils import PretrainedConfig
|
23
|
+
from transformers.modeling_rope_utils import rope_config_validation
|
24
|
+
from transformers.utils import logging
|
25
|
+
|
26
|
+
from sglang.srt.distributed.utils import divide
|
27
|
+
from sglang.srt.layers.attention.mamba.mamba_utils import MambaStateShapeCalculator
|
28
|
+
from sglang.srt.layers.dp_attention import (
|
29
|
+
get_attention_tp_size,
|
30
|
+
get_tensor_model_parallel_world_size,
|
31
|
+
)
|
32
|
+
|
33
|
+
logger = logging.get_logger(__name__)
|
34
|
+
|
35
|
+
|
36
|
+
class FalconH1Config(PretrainedConfig):
|
37
|
+
r"""
|
38
|
+
This is the configuration class to store the configuration of a [`FalconH1Model`]. It is used to instantiate a
|
39
|
+
FalconH1Model model according to the specified arguments, defining the model architecture. Instantiating a configuration
|
40
|
+
with defaults taken from [ibm-fms/FalconH1-9.8b-2.2T-hf](https://huggingface.co/ibm-fms/FalconH1-9.8b-2.2T-hf).
|
41
|
+
The FalconH1Model is a hybrid [mamba2](https://github.com/state-spaces/mamba) architecture with SwiGLU.
|
42
|
+
The checkpoints are jointly trained by IBM, Princeton, and UIUC.
|
43
|
+
Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
|
44
|
+
documentation from [`PretrainedConfig`] for more information.
|
45
|
+
Args:
|
46
|
+
vocab_size (`int`, *optional*, defaults to 128000):
|
47
|
+
Vocabulary size of the FalconH1 model. Defines the number of different tokens that can be represented by the
|
48
|
+
`inputs_ids` passed when calling [`FalconH1Model`]
|
49
|
+
tie_word_embeddings (`bool`, *optional*, defaults to `False`):
|
50
|
+
Whether the model's input and output word embeddings should be tied. Note that this is only relevant if the
|
51
|
+
model has a output word embedding layer.
|
52
|
+
hidden_size (`int`, *optional*, defaults to 4096):
|
53
|
+
Dimension of the hidden representations.
|
54
|
+
intermediate_size (`int`, *optional*, defaults to 14336):
|
55
|
+
Dimension of the MLP representations.
|
56
|
+
num_hidden_layers (`int`, *optional*, defaults to 32):
|
57
|
+
Number of hidden layers in the Transformer encoder.
|
58
|
+
num_attention_heads (`int`, *optional*, defaults to 32):
|
59
|
+
Number of attention heads for each attention layer in the Transformer encoder.
|
60
|
+
num_key_value_heads (`int`, *optional*, defaults to 8):
|
61
|
+
This is the number of key_value heads that should be used to implement Grouped Query Attention. If
|
62
|
+
`num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
|
63
|
+
`num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When
|
64
|
+
converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
|
65
|
+
by meanpooling all the original heads within that group. For more details, check out [this
|
66
|
+
paper](https://huggingface.co/papers/2305.13245). If it is not specified, will default to `8`.
|
67
|
+
hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
|
68
|
+
The non-linear activation function (function or string) in the decoder.
|
69
|
+
initializer_range (`float`, *optional*, defaults to 0.02):
|
70
|
+
The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
|
71
|
+
rms_norm_eps (`float`, *optional*, defaults to 1e-05):
|
72
|
+
The epsilon used by the rms normalization layers.
|
73
|
+
use_cache (`bool`, *optional*, defaults to `True`):
|
74
|
+
Whether or not the model should return the last key/values attentions (not used by all models). Only
|
75
|
+
relevant if `config.is_decoder=True`.
|
76
|
+
num_logits_to_keep (`int` or `None`, *optional*, defaults to 1):
|
77
|
+
Number of prompt logits to calculate during generation. If `None`, all logits will be calculated. If an
|
78
|
+
integer value, only last `num_logits_to_keep` logits will be calculated. Default is 1 because only the
|
79
|
+
logits of the last prompt token are needed for generation. For long sequences, the logits for the entire
|
80
|
+
sequence may use a lot of memory so, setting `num_logits_to_keep=1` will reduce memory footprint
|
81
|
+
significantly.
|
82
|
+
pad_token_id (`int`, *optional*, defaults to 0):
|
83
|
+
The id of the padding token.
|
84
|
+
bos_token_id (`int`, *optional*, defaults to 1):
|
85
|
+
The id of the "beginning-of-sequence" token.
|
86
|
+
eos_token_id (`int`, *optional*, defaults to 2):
|
87
|
+
The id of the "end-of-sequence" token.
|
88
|
+
max_position_embeddings (`int`, *optional*, defaults to 8192):
|
89
|
+
Max cached sequence length for the model
|
90
|
+
attention_dropout (`float`, *optional*, defaults to 0.0):
|
91
|
+
The dropout ratio for the attention probabilities.
|
92
|
+
mamba_d_ssm (`int`, *optional*, defaults to 1024):
|
93
|
+
The dimension of the SSM state space latents.
|
94
|
+
mamba_n_heads (`int`, *optional*, defaults to 128):
|
95
|
+
The number of mamba heads used in the v2 implementation.
|
96
|
+
mamba_d_head (`int`, *optional*, defaults to `"auto"`):
|
97
|
+
Head embedding dimension size
|
98
|
+
mamba_n_groups (`int`, *optional*, defaults to 1):
|
99
|
+
The number of the mamba groups used in the v2 implementation.
|
100
|
+
mamba_d_state (`int`, *optional*, defaults to 256):
|
101
|
+
The dimension the mamba state space latents
|
102
|
+
mamba_d_conv (`int`, *optional*, defaults to 4):
|
103
|
+
The size of the mamba convolution kernel
|
104
|
+
mamba_expand (`int`, *optional*, defaults to 2):
|
105
|
+
Expanding factor (relative to hidden_size) used to determine the mamba intermediate size
|
106
|
+
mamba_chunk_size (`int`, *optional*, defaults to 256):
|
107
|
+
The chunks in which to break the sequence when doing prefill/training
|
108
|
+
mamba_conv_bias (`bool`, *optional*, defaults to `True`):
|
109
|
+
Flag indicating whether or not to use bias in the convolution layer of the mamba mixer block.
|
110
|
+
mamba_proj_bias (`bool`, *optional*, defaults to `False`):
|
111
|
+
Flag indicating whether or not to use bias in the input and output projections (["in_proj", "out_proj"]) of the mamba mixer block
|
112
|
+
mamba_norm_before_gate (`bool`, *optional*, defaults to `True`):
|
113
|
+
Whether to use RMSNorm before the gate in the Mamba block
|
114
|
+
mamba_rms_norm (`bool`, *optional*, defaults to `False`):
|
115
|
+
Whether to use RMSNorm instead of LayerNorm in the Mamba block
|
116
|
+
projectors_bias (`bool`, *optional*, defaults to `False`):
|
117
|
+
Flag indicating whether or not to use bias in the input and output projections (["in_proj", "out_proj"]) of the attention block
|
118
|
+
rope_theta (`float`, *optional*, defaults to 100000.0):
|
119
|
+
The theta value used for the RoPE embeddings.
|
120
|
+
rope_scaling (`float`, *optional*):
|
121
|
+
The scaling value used for the RoPE embeddings. If `None`, no scaling is applied.
|
122
|
+
lm_head_multiplier (`float`, *optional*, defaults to 1.0):
|
123
|
+
The multiplier for the LM head. This is used to scale the output of the LM head.
|
124
|
+
embedding_multiplier (`float`, *optional*, defaults to 1.0):
|
125
|
+
The multiplier for the embedding layer. This is used to scale the output of the embedding layer.
|
126
|
+
mlp_multipliers (`list[float]`, *optional*):
|
127
|
+
The multipliers for the MLP layers. This is used to scale the output of the MLP layers. The first value is
|
128
|
+
the multiplier of gate layer, the second value is the multiplier of the down_proj layer.
|
129
|
+
key_multiplier (`float`, *optional*):
|
130
|
+
The multiplier for the key layer. This is used to scale the output of the key layer.
|
131
|
+
attention_out_multiplier (`float`, *optional*):
|
132
|
+
The multiplier for the attention output layer. This is used to scale the output of the attention output
|
133
|
+
attention_in_multiplier (`float`, *optional*):
|
134
|
+
The multiplier for the attention input layer. This is used to scale the output of the attention input layer.
|
135
|
+
ssm_multipliers (`list[float]`, *optional*):
|
136
|
+
The multipliers for the SSM layers. This is used to scale the output of the SSM layers.
|
137
|
+
ssm_in_multiplier (`float`, *optional*):
|
138
|
+
The multiplier for the SSM input layer. This is used to scale the output of the SSM input layer.
|
139
|
+
ssm_out_multiplier (`float`, *optional*):
|
140
|
+
The multiplier for the SSM output layer. This is used to scale the output of the SSM output layer.
|
141
|
+
"""
|
142
|
+
|
143
|
+
model_type = "falcon_h1"
|
144
|
+
keys_to_ignore_at_inference = ["past_key_values"]
|
145
|
+
|
146
|
+
def __init__(
|
147
|
+
self,
|
148
|
+
vocab_size=128000,
|
149
|
+
tie_word_embeddings=False,
|
150
|
+
hidden_size=4096,
|
151
|
+
intermediate_size=14336,
|
152
|
+
num_hidden_layers=32,
|
153
|
+
num_attention_heads=32,
|
154
|
+
num_key_value_heads=8,
|
155
|
+
hidden_act="silu",
|
156
|
+
initializer_range=0.02,
|
157
|
+
rms_norm_eps=1e-5,
|
158
|
+
use_cache=True,
|
159
|
+
num_logits_to_keep=1,
|
160
|
+
pad_token_id=0,
|
161
|
+
bos_token_id=1,
|
162
|
+
eos_token_id=2,
|
163
|
+
max_position_embeddings=8192,
|
164
|
+
attention_dropout=0.0,
|
165
|
+
mamba_d_ssm=1024,
|
166
|
+
mamba_n_heads=128,
|
167
|
+
mamba_d_head="auto",
|
168
|
+
mamba_n_groups=1,
|
169
|
+
mamba_d_state=256,
|
170
|
+
mamba_d_conv=4,
|
171
|
+
mamba_expand=2,
|
172
|
+
mamba_chunk_size=256,
|
173
|
+
mamba_conv_bias=True,
|
174
|
+
mamba_proj_bias=False,
|
175
|
+
mamba_norm_before_gate=True,
|
176
|
+
mamba_rms_norm=False,
|
177
|
+
projectors_bias=False,
|
178
|
+
rope_theta=100000.0,
|
179
|
+
rope_scaling=None,
|
180
|
+
lm_head_multiplier=1.0,
|
181
|
+
embedding_multiplier=1.0,
|
182
|
+
mlp_multipliers=None,
|
183
|
+
key_multiplier=None,
|
184
|
+
attention_out_multiplier=None,
|
185
|
+
attention_in_multiplier=None,
|
186
|
+
ssm_multipliers=None,
|
187
|
+
ssm_in_multiplier=None,
|
188
|
+
ssm_out_multiplier=None,
|
189
|
+
**kwargs,
|
190
|
+
):
|
191
|
+
self.vocab_size = vocab_size
|
192
|
+
self.hidden_size = hidden_size
|
193
|
+
self.intermediate_size = intermediate_size
|
194
|
+
self.num_hidden_layers = num_hidden_layers
|
195
|
+
self.num_attention_heads = num_attention_heads
|
196
|
+
self.max_position_embeddings = max_position_embeddings
|
197
|
+
self.attention_dropout = attention_dropout
|
198
|
+
self.attention_bias = False
|
199
|
+
self.mlp_bias = False
|
200
|
+
|
201
|
+
# for backward compatibility
|
202
|
+
if num_key_value_heads is None:
|
203
|
+
num_key_value_heads = num_attention_heads
|
204
|
+
|
205
|
+
self.num_key_value_heads = num_key_value_heads
|
206
|
+
self.hidden_act = hidden_act
|
207
|
+
self.initializer_range = initializer_range
|
208
|
+
self.rms_norm_eps = rms_norm_eps
|
209
|
+
|
210
|
+
self.use_cache = use_cache
|
211
|
+
self.num_logits_to_keep = num_logits_to_keep
|
212
|
+
|
213
|
+
self.rope_theta = rope_theta
|
214
|
+
self.rope_scaling = None
|
215
|
+
self.rope_scaling = rope_scaling
|
216
|
+
self.projectors_bias = projectors_bias
|
217
|
+
mamba_intermediate = (
|
218
|
+
mamba_expand * hidden_size if mamba_d_ssm is None else mamba_d_ssm
|
219
|
+
)
|
220
|
+
|
221
|
+
if mamba_intermediate % mamba_n_heads != 0:
|
222
|
+
raise ValueError("mamba_n_heads must divide mamba_expand * hidden_size")
|
223
|
+
|
224
|
+
# for the mamba_v2, must satisfy the following
|
225
|
+
if mamba_d_head == "auto":
|
226
|
+
mamba_d_head = mamba_intermediate // mamba_n_heads
|
227
|
+
|
228
|
+
if mamba_d_head * mamba_n_heads != mamba_intermediate:
|
229
|
+
raise ValueError(
|
230
|
+
"The dimensions for the Mamba head state do not match the model intermediate_size"
|
231
|
+
)
|
232
|
+
|
233
|
+
self.mamba_d_ssm = mamba_d_ssm
|
234
|
+
self.mamba_n_heads = mamba_n_heads
|
235
|
+
self.mamba_d_head = mamba_d_head
|
236
|
+
self.mamba_n_groups = mamba_n_groups
|
237
|
+
self.mamba_d_state = mamba_d_state
|
238
|
+
self.mamba_d_conv = mamba_d_conv
|
239
|
+
self.mamba_expand = mamba_expand
|
240
|
+
self.mamba_chunk_size = mamba_chunk_size
|
241
|
+
self.mamba_conv_bias = mamba_conv_bias
|
242
|
+
self.mamba_proj_bias = mamba_proj_bias
|
243
|
+
|
244
|
+
self.mamba_norm_before_gate = mamba_norm_before_gate
|
245
|
+
self.mamba_rms_norm = mamba_rms_norm
|
246
|
+
|
247
|
+
self.lm_head_multiplier = lm_head_multiplier
|
248
|
+
self.embedding_multiplier = embedding_multiplier
|
249
|
+
|
250
|
+
if mlp_multipliers is not None:
|
251
|
+
self.mlp_multipliers = mlp_multipliers
|
252
|
+
else:
|
253
|
+
self.mlp_multipliers = [1.0, 1.0]
|
254
|
+
|
255
|
+
if attention_out_multiplier is not None:
|
256
|
+
self.attention_out_multiplier = attention_out_multiplier
|
257
|
+
else:
|
258
|
+
self.attention_out_multiplier = 1.0
|
259
|
+
|
260
|
+
if attention_in_multiplier is not None:
|
261
|
+
self.attention_in_multiplier = attention_in_multiplier
|
262
|
+
else:
|
263
|
+
self.attention_in_multiplier = 1.0
|
264
|
+
|
265
|
+
if key_multiplier is not None:
|
266
|
+
self.key_multiplier = key_multiplier
|
267
|
+
else:
|
268
|
+
self.key_multiplier = 1.0
|
269
|
+
|
270
|
+
if ssm_multipliers is not None:
|
271
|
+
self.ssm_multipliers = ssm_multipliers
|
272
|
+
else:
|
273
|
+
self.ssm_multipliers = [1.0, 1.0, 1.0, 1.0, 1.0]
|
274
|
+
|
275
|
+
if ssm_in_multiplier is not None:
|
276
|
+
self.ssm_in_multiplier = ssm_in_multiplier
|
277
|
+
else:
|
278
|
+
self.ssm_in_multiplier = 1.0
|
279
|
+
|
280
|
+
if ssm_out_multiplier is not None:
|
281
|
+
self.ssm_out_multiplier = ssm_out_multiplier
|
282
|
+
else:
|
283
|
+
self.ssm_out_multiplier = 1.0
|
284
|
+
|
285
|
+
super().__init__(
|
286
|
+
pad_token_id=pad_token_id,
|
287
|
+
bos_token_id=bos_token_id,
|
288
|
+
eos_token_id=eos_token_id,
|
289
|
+
tie_word_embeddings=tie_word_embeddings,
|
290
|
+
**kwargs,
|
291
|
+
)
|
292
|
+
|
293
|
+
@property
|
294
|
+
def layers_block_type(self):
|
295
|
+
return ["falcon_h1" for i in range(self.num_hidden_layers)]
|
296
|
+
|
297
|
+
@property
|
298
|
+
def mamba_cache_per_req(self):
|
299
|
+
conv_state_shape, temporal_state_shape, conv_dtype, ssm_dtype, mamba_layers = (
|
300
|
+
self.hybrid_gdn_params
|
301
|
+
)
|
302
|
+
mamba_layers_len = len(mamba_layers)
|
303
|
+
|
304
|
+
return (
|
305
|
+
int(np.prod(conv_state_shape)) * conv_dtype.itemsize
|
306
|
+
+ int(np.prod(temporal_state_shape)) * ssm_dtype.itemsize
|
307
|
+
) * mamba_layers_len
|
308
|
+
|
309
|
+
@property
|
310
|
+
def full_attention_layer_ids(self):
|
311
|
+
# For Falcon-H1, we do have attention on all layers
|
312
|
+
return range(self.num_hidden_layers)
|
313
|
+
|
314
|
+
@property
|
315
|
+
def linear_layer_ids(self):
|
316
|
+
# For Falcon-H1, we do have mamba on all layers
|
317
|
+
return range(self.num_hidden_layers)
|
318
|
+
|
319
|
+
@property
|
320
|
+
def hybrid_gdn_params(self):
|
321
|
+
world_size = get_tensor_model_parallel_world_size()
|
322
|
+
|
323
|
+
n_groups = self.mamba_n_groups
|
324
|
+
if self.mamba_n_groups % world_size != 0:
|
325
|
+
# - for TP we shard conv_dim by sharding on n_groups,
|
326
|
+
# - but if n_groups cannot divide tp_size, we need to
|
327
|
+
# extend some extra groups
|
328
|
+
extra_groups = MambaStateShapeCalculator.extra_groups_for_head_shards(
|
329
|
+
self.mamba_n_groups, world_size
|
330
|
+
)
|
331
|
+
n_groups += extra_groups
|
332
|
+
|
333
|
+
conv_dim = self.mamba_d_ssm + 2 * n_groups * self.mamba_d_state
|
334
|
+
|
335
|
+
conv_state_shape = (
|
336
|
+
divide(conv_dim, world_size),
|
337
|
+
self.mamba_d_conv - 1,
|
338
|
+
)
|
339
|
+
|
340
|
+
# we TP-ize on the heads dimension
|
341
|
+
temporal_state_shape = (
|
342
|
+
self.mamba_d_state,
|
343
|
+
self.mamba_d_head,
|
344
|
+
divide(self.mamba_n_heads, world_size),
|
345
|
+
)
|
346
|
+
conv_dtype = torch.bfloat16
|
347
|
+
dtype_map = {
|
348
|
+
"float32": torch.float32,
|
349
|
+
"bfloat16": torch.bfloat16,
|
350
|
+
}
|
351
|
+
ssm_dtype = dtype_map[os.environ["SGLANG_MAMBA_SSM_DTYPE"]]
|
352
|
+
mamba_layers = self.linear_layer_ids
|
353
|
+
|
354
|
+
return (
|
355
|
+
conv_state_shape,
|
356
|
+
temporal_state_shape,
|
357
|
+
conv_dtype,
|
358
|
+
ssm_dtype,
|
359
|
+
mamba_layers,
|
360
|
+
)
|
@@ -23,6 +23,9 @@ class LoadFormat(str, enum.Enum):
|
|
23
23
|
LAYERED = "layered"
|
24
24
|
JAX = "jax"
|
25
25
|
REMOTE = "remote"
|
26
|
+
REMOTE_INSTANCE = "remote_instance"
|
27
|
+
RDMA = "rdma"
|
28
|
+
LOCAL_CACHED = "local_cached"
|
26
29
|
|
27
30
|
|
28
31
|
@dataclass
|
@@ -46,6 +49,7 @@ class LoadConfig:
|
|
46
49
|
checkpoints.
|
47
50
|
decryption_key_file: If set, decrypts the output files with a password read
|
48
51
|
from this file (after PBKDF2).
|
52
|
+
decrypt_max_concurrency: The maximum number of concurrent processes to decrypt the safetensor files. -1 means no limit.
|
49
53
|
"""
|
50
54
|
|
51
55
|
load_format: Union[str, LoadFormat] = LoadFormat.AUTO
|
@@ -53,6 +57,11 @@ class LoadConfig:
|
|
53
57
|
model_loader_extra_config: Optional[Union[str, dict]] = field(default_factory=dict)
|
54
58
|
ignore_patterns: Optional[Union[List[str], str]] = None
|
55
59
|
decryption_key_file: Optional[str] = None
|
60
|
+
decrypt_max_concurrency: int = -1
|
61
|
+
tp_rank: Optional[int] = None
|
62
|
+
remote_instance_weight_loader_seed_instance_ip: Optional[str] = None
|
63
|
+
remote_instance_weight_loader_seed_instance_service_port: Optional[int] = None
|
64
|
+
remote_instance_weight_loader_send_weights_group_ports: Optional[List[int]] = None
|
56
65
|
|
57
66
|
def __post_init__(self):
|
58
67
|
model_loader_extra_config = self.model_loader_extra_config or {}
|