sglang 0.5.2rc2__py3-none-any.whl → 0.5.3__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.
Files changed (377) hide show
  1. sglang/bench_one_batch.py +7 -9
  2. sglang/bench_one_batch_server.py +330 -31
  3. sglang/bench_serving.py +267 -32
  4. sglang/global_config.py +2 -2
  5. sglang/lang/backend/runtime_endpoint.py +1 -1
  6. sglang/launch_server.py +14 -0
  7. sglang/profiler.py +2 -2
  8. sglang/srt/batch_invariant_ops/__init__.py +27 -0
  9. sglang/srt/batch_invariant_ops/batch_invariant_ops.py +549 -0
  10. sglang/srt/configs/__init__.py +8 -0
  11. sglang/srt/configs/device_config.py +3 -1
  12. sglang/srt/configs/dots_ocr.py +64 -0
  13. sglang/srt/configs/dots_vlm.py +139 -0
  14. sglang/srt/configs/falcon_h1.py +360 -0
  15. sglang/srt/configs/load_config.py +9 -0
  16. sglang/srt/configs/model_config.py +181 -82
  17. sglang/srt/configs/qwen3_next.py +326 -0
  18. sglang/srt/configs/qwen3_vl.py +586 -0
  19. sglang/srt/connector/__init__.py +8 -1
  20. sglang/srt/connector/remote_instance.py +82 -0
  21. sglang/srt/constrained/base_grammar_backend.py +49 -12
  22. sglang/srt/constrained/llguidance_backend.py +0 -1
  23. sglang/srt/constrained/outlines_backend.py +0 -1
  24. sglang/srt/constrained/outlines_jump_forward.py +1 -1
  25. sglang/srt/constrained/xgrammar_backend.py +30 -9
  26. sglang/srt/custom_op.py +11 -1
  27. sglang/srt/debug_utils/dump_comparator.py +81 -44
  28. sglang/srt/debug_utils/dump_loader.py +97 -0
  29. sglang/srt/debug_utils/dumper.py +21 -6
  30. sglang/srt/debug_utils/text_comparator.py +73 -11
  31. sglang/srt/disaggregation/ascend/conn.py +2 -2
  32. sglang/srt/disaggregation/ascend/transfer_engine.py +47 -9
  33. sglang/srt/disaggregation/base/conn.py +1 -1
  34. sglang/srt/disaggregation/common/conn.py +279 -108
  35. sglang/srt/disaggregation/decode.py +71 -19
  36. sglang/srt/disaggregation/decode_kvcache_offload_manager.py +185 -0
  37. sglang/srt/disaggregation/decode_schedule_batch_mixin.py +29 -17
  38. sglang/srt/disaggregation/fake/conn.py +1 -1
  39. sglang/srt/disaggregation/mini_lb.py +6 -445
  40. sglang/srt/disaggregation/mooncake/conn.py +55 -537
  41. sglang/srt/disaggregation/nixl/conn.py +326 -53
  42. sglang/srt/disaggregation/prefill.py +36 -17
  43. sglang/srt/disaggregation/utils.py +40 -54
  44. sglang/srt/distributed/device_communicators/all_reduce_utils.py +16 -0
  45. sglang/srt/distributed/device_communicators/shm_broadcast.py +4 -2
  46. sglang/srt/distributed/device_communicators/symm_mem.py +164 -0
  47. sglang/srt/distributed/parallel_state.py +156 -80
  48. sglang/srt/entrypoints/engine.py +59 -18
  49. sglang/srt/entrypoints/grpc_request_manager.py +855 -0
  50. sglang/srt/entrypoints/grpc_server.py +810 -0
  51. sglang/srt/entrypoints/http_server.py +130 -59
  52. sglang/srt/entrypoints/openai/protocol.py +112 -4
  53. sglang/srt/entrypoints/openai/serving_base.py +65 -3
  54. sglang/srt/entrypoints/openai/serving_chat.py +204 -55
  55. sglang/srt/entrypoints/openai/serving_completions.py +14 -3
  56. sglang/srt/entrypoints/openai/serving_embedding.py +9 -3
  57. sglang/srt/entrypoints/openai/serving_rerank.py +3 -1
  58. sglang/srt/entrypoints/openai/serving_responses.py +48 -3
  59. sglang/srt/entrypoints/openai/serving_score.py +1 -0
  60. sglang/srt/environ.py +285 -0
  61. sglang/srt/eplb/eplb_manager.py +2 -2
  62. sglang/srt/eplb/expert_distribution.py +26 -13
  63. sglang/srt/eplb/expert_location.py +38 -8
  64. sglang/srt/eplb/expert_location_updater.py +1 -1
  65. sglang/srt/function_call/base_format_detector.py +3 -6
  66. sglang/srt/function_call/ebnf_composer.py +11 -9
  67. sglang/srt/function_call/function_call_parser.py +9 -2
  68. sglang/srt/function_call/glm4_moe_detector.py +4 -4
  69. sglang/srt/function_call/gpt_oss_detector.py +23 -0
  70. sglang/srt/function_call/json_array_parser.py +63 -0
  71. sglang/srt/function_call/kimik2_detector.py +17 -4
  72. sglang/srt/function_call/qwen3_coder_detector.py +1 -1
  73. sglang/srt/function_call/utils.py +96 -5
  74. sglang/srt/grpc/__init__.py +1 -0
  75. sglang/srt/grpc/compile_proto.py +245 -0
  76. sglang/srt/grpc/sglang_scheduler_pb2.py +111 -0
  77. sglang/srt/grpc/sglang_scheduler_pb2.pyi +434 -0
  78. sglang/srt/grpc/sglang_scheduler_pb2_grpc.py +239 -0
  79. sglang/srt/layers/activation.py +143 -9
  80. sglang/srt/layers/attention/aiter_backend.py +14 -15
  81. sglang/srt/layers/attention/ascend_backend.py +115 -9
  82. sglang/srt/layers/attention/attention_registry.py +206 -0
  83. sglang/srt/layers/attention/base_attn_backend.py +12 -3
  84. sglang/srt/layers/attention/cutlass_mla_backend.py +3 -3
  85. sglang/srt/layers/attention/dual_chunk_flashattention_backend.py +1 -1
  86. sglang/srt/layers/attention/fla/chunk.py +242 -0
  87. sglang/srt/layers/attention/fla/chunk_delta_h.py +314 -0
  88. sglang/srt/layers/attention/fla/chunk_o.py +178 -0
  89. sglang/srt/layers/attention/fla/chunk_scaled_dot_kkt.py +151 -0
  90. sglang/srt/layers/attention/fla/cumsum.py +300 -0
  91. sglang/srt/layers/attention/fla/fused_recurrent.py +640 -0
  92. sglang/srt/layers/attention/fla/fused_sigmoid_gating_recurrent.py +232 -0
  93. sglang/srt/layers/attention/fla/index.py +37 -0
  94. sglang/srt/layers/attention/fla/l2norm.py +150 -0
  95. sglang/srt/layers/attention/fla/layernorm_gated.py +326 -0
  96. sglang/srt/layers/attention/fla/op.py +66 -0
  97. sglang/srt/layers/attention/fla/solve_tril.py +465 -0
  98. sglang/srt/layers/attention/fla/utils.py +331 -0
  99. sglang/srt/layers/attention/fla/wy_fast.py +158 -0
  100. sglang/srt/layers/attention/flashattention_backend.py +41 -8
  101. sglang/srt/layers/attention/flashinfer_backend.py +118 -198
  102. sglang/srt/layers/attention/flashinfer_mla_backend.py +27 -27
  103. sglang/srt/layers/attention/flashmla_backend.py +7 -5
  104. sglang/srt/layers/attention/hybrid_attn_backend.py +68 -53
  105. sglang/srt/layers/attention/hybrid_linear_attn_backend.py +602 -0
  106. sglang/srt/layers/attention/intel_amx_backend.py +3 -0
  107. sglang/srt/layers/attention/mamba/causal_conv1d.py +129 -0
  108. sglang/srt/layers/attention/mamba/causal_conv1d_triton.py +969 -0
  109. sglang/srt/layers/attention/mamba/mamba.py +629 -0
  110. sglang/srt/layers/attention/mamba/mamba_utils.py +81 -0
  111. sglang/srt/layers/attention/mamba/ops/__init__.py +2 -0
  112. sglang/srt/layers/attention/mamba/ops/layernorm_gated.py +172 -0
  113. sglang/srt/layers/attention/mamba/ops/mamba_ssm.py +442 -0
  114. sglang/srt/layers/attention/mamba/ops/ssd_bmm.py +264 -0
  115. sglang/srt/layers/attention/mamba/ops/ssd_chunk_scan.py +622 -0
  116. sglang/srt/layers/attention/mamba/ops/ssd_chunk_state.py +757 -0
  117. sglang/srt/layers/attention/mamba/ops/ssd_combined.py +262 -0
  118. sglang/srt/layers/attention/mamba/ops/ssd_state_passing.py +275 -0
  119. sglang/srt/layers/attention/npu_ops/mla_preprocess.py +393 -0
  120. sglang/srt/layers/attention/nsa/dequant_k_cache.py +163 -0
  121. sglang/srt/layers/attention/nsa/index_buf_accessor.py +354 -0
  122. sglang/srt/layers/attention/nsa/nsa_indexer.py +761 -0
  123. sglang/srt/layers/attention/nsa/quant_k_cache.py +255 -0
  124. sglang/srt/layers/attention/nsa/tilelang_kernel.py +785 -0
  125. sglang/srt/layers/attention/nsa/transform_index.py +144 -0
  126. sglang/srt/layers/attention/nsa/utils.py +24 -0
  127. sglang/srt/layers/attention/nsa_backend.py +887 -0
  128. sglang/srt/layers/attention/tbo_backend.py +6 -6
  129. sglang/srt/layers/attention/torch_flex_backend.py +325 -0
  130. sglang/srt/layers/attention/torch_native_backend.py +12 -6
  131. sglang/srt/layers/attention/triton_backend.py +57 -7
  132. sglang/srt/layers/attention/trtllm_mha_backend.py +5 -7
  133. sglang/srt/layers/attention/trtllm_mla_backend.py +276 -39
  134. sglang/srt/layers/attention/vision.py +58 -0
  135. sglang/srt/layers/attention/wave_backend.py +4 -4
  136. sglang/srt/layers/attention/wave_ops/decode_attention.py +2 -4
  137. sglang/srt/layers/attention/wave_ops/extend_attention.py +1 -3
  138. sglang/srt/layers/communicator.py +8 -0
  139. sglang/srt/layers/dp_attention.py +41 -2
  140. sglang/srt/layers/elementwise.py +3 -1
  141. sglang/srt/layers/layernorm.py +34 -15
  142. sglang/srt/layers/linear.py +55 -7
  143. sglang/srt/layers/logits_processor.py +44 -12
  144. sglang/srt/layers/moe/__init__.py +2 -1
  145. sglang/srt/layers/moe/cutlass_w4a8_moe.py +3 -3
  146. sglang/srt/layers/moe/ep_moe/kernels.py +2 -2
  147. sglang/srt/layers/moe/ep_moe/layer.py +256 -63
  148. sglang/srt/layers/moe/flashinfer_cutedsl_moe.py +183 -0
  149. sglang/srt/layers/moe/fused_moe_native.py +5 -3
  150. 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
  151. 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
  152. 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
  153. sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=256,N=512,device_name=NVIDIA_H20.json +146 -0
  154. 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
  155. sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=512,N=128,device_name=NVIDIA_H20-3e.json +146 -0
  156. sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=512,N=128,device_name=NVIDIA_H200.json +146 -0
  157. 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
  158. sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=512,N=256,device_name=NVIDIA_B200.json +146 -0
  159. sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=512,N=256,device_name=NVIDIA_H20-3e.json +146 -0
  160. sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=512,N=256,device_name=NVIDIA_H200.json +146 -0
  161. 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
  162. sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=512,N=64,device_name=NVIDIA_H200.json +146 -0
  163. sglang/srt/layers/moe/fused_moe_triton/fused_moe.py +5 -2
  164. sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_config.py +7 -3
  165. sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_kernels.py +23 -20
  166. sglang/srt/layers/moe/fused_moe_triton/layer.py +71 -70
  167. sglang/srt/layers/moe/moe_runner/__init__.py +2 -1
  168. sglang/srt/layers/moe/moe_runner/base.py +274 -1
  169. sglang/srt/layers/moe/moe_runner/runner.py +80 -0
  170. sglang/srt/layers/moe/moe_runner/triton.py +448 -0
  171. sglang/srt/layers/moe/token_dispatcher/__init__.py +16 -4
  172. sglang/srt/layers/moe/token_dispatcher/{base_dispatcher.py → base.py} +67 -17
  173. sglang/srt/layers/moe/token_dispatcher/deepep.py +118 -56
  174. sglang/srt/layers/moe/token_dispatcher/standard.py +44 -2
  175. sglang/srt/layers/moe/topk.py +30 -9
  176. sglang/srt/layers/moe/utils.py +22 -6
  177. sglang/srt/layers/parameter.py +23 -6
  178. sglang/srt/layers/quantization/awq.py +19 -7
  179. sglang/srt/layers/quantization/base_config.py +11 -6
  180. sglang/srt/layers/quantization/blockwise_int8.py +38 -27
  181. sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py +1 -0
  182. sglang/srt/layers/quantization/compressed_tensors/compressed_tensors_moe.py +50 -30
  183. sglang/srt/layers/quantization/compressed_tensors/schemes/__init__.py +2 -0
  184. sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py +13 -1
  185. sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_int8.py +173 -0
  186. sglang/srt/layers/quantization/deep_gemm_wrapper/configurer.py +2 -10
  187. sglang/srt/layers/quantization/deep_gemm_wrapper/entrypoint.py +27 -0
  188. sglang/srt/layers/quantization/fp8.py +78 -49
  189. sglang/srt/layers/quantization/fp8_utils.py +51 -32
  190. sglang/srt/layers/quantization/gptq.py +25 -17
  191. sglang/srt/layers/quantization/modelopt_quant.py +190 -55
  192. sglang/srt/layers/quantization/moe_wna16.py +21 -18
  193. sglang/srt/layers/quantization/mxfp4.py +74 -42
  194. sglang/srt/layers/quantization/quark/quark_moe.py +48 -30
  195. sglang/srt/layers/quantization/unquant.py +135 -47
  196. sglang/srt/layers/quantization/w4afp8.py +26 -17
  197. sglang/srt/layers/quantization/w8a8_fp8.py +35 -20
  198. sglang/srt/layers/quantization/w8a8_int8.py +91 -41
  199. sglang/srt/layers/rotary_embedding.py +78 -31
  200. sglang/srt/layers/sampler.py +213 -21
  201. sglang/srt/layers/utils.py +23 -0
  202. sglang/srt/lora/backend/base_backend.py +50 -8
  203. sglang/srt/lora/backend/chunked_backend.py +348 -0
  204. sglang/srt/lora/backend/triton_backend.py +99 -5
  205. sglang/srt/lora/layers.py +32 -0
  206. sglang/srt/lora/lora.py +8 -3
  207. sglang/srt/lora/lora_manager.py +52 -118
  208. sglang/srt/lora/mem_pool.py +25 -11
  209. sglang/srt/lora/triton_ops/__init__.py +4 -0
  210. sglang/srt/lora/triton_ops/chunked_sgmv_expand.py +214 -0
  211. sglang/srt/lora/triton_ops/chunked_sgmv_shrink.py +174 -0
  212. sglang/srt/lora/utils.py +22 -11
  213. sglang/srt/managers/async_dynamic_batch_tokenizer.py +170 -0
  214. sglang/srt/managers/cache_controller.py +199 -301
  215. sglang/srt/managers/data_parallel_controller.py +115 -80
  216. sglang/srt/managers/detokenizer_manager.py +19 -15
  217. sglang/srt/managers/disagg_service.py +46 -0
  218. sglang/srt/managers/io_struct.py +340 -109
  219. sglang/srt/managers/mm_utils.py +44 -6
  220. sglang/srt/managers/multi_tokenizer_mixin.py +357 -407
  221. sglang/srt/managers/multimodal_processor.py +1 -2
  222. sglang/srt/managers/overlap_utils.py +53 -0
  223. sglang/srt/managers/schedule_batch.py +240 -138
  224. sglang/srt/managers/schedule_policy.py +144 -17
  225. sglang/srt/managers/scheduler.py +502 -209
  226. sglang/srt/managers/scheduler_input_blocker.py +1 -1
  227. sglang/srt/managers/scheduler_metrics_mixin.py +99 -126
  228. sglang/srt/managers/scheduler_output_processor_mixin.py +75 -22
  229. sglang/srt/managers/scheduler_profiler_mixin.py +6 -6
  230. sglang/srt/managers/scheduler_update_weights_mixin.py +7 -0
  231. sglang/srt/managers/tokenizer_communicator_mixin.py +675 -0
  232. sglang/srt/managers/tokenizer_manager.py +320 -632
  233. sglang/srt/managers/tp_worker.py +81 -22
  234. sglang/srt/managers/tp_worker_overlap_thread.py +71 -56
  235. sglang/srt/managers/utils.py +1 -45
  236. sglang/srt/mem_cache/allocator.py +14 -20
  237. sglang/srt/mem_cache/allocator_ascend.py +41 -27
  238. sglang/srt/mem_cache/base_prefix_cache.py +1 -1
  239. sglang/srt/mem_cache/chunk_cache.py +8 -1
  240. sglang/srt/mem_cache/evict_policy.py +23 -0
  241. sglang/srt/mem_cache/hicache_storage.py +43 -24
  242. sglang/srt/mem_cache/hiradix_cache.py +222 -75
  243. sglang/srt/mem_cache/memory_pool.py +535 -58
  244. sglang/srt/mem_cache/memory_pool_host.py +239 -228
  245. sglang/srt/mem_cache/radix_cache.py +222 -73
  246. sglang/srt/mem_cache/radix_cache_cpp.py +11 -8
  247. sglang/srt/mem_cache/storage/__init__.py +10 -0
  248. sglang/srt/mem_cache/storage/aibrix_kvcache/aibrix_kvcache_storage.py +151 -0
  249. sglang/srt/mem_cache/storage/aibrix_kvcache/unit_test.py +109 -0
  250. sglang/srt/mem_cache/storage/backend_factory.py +223 -0
  251. sglang/srt/mem_cache/storage/eic/eic_storage.py +778 -0
  252. sglang/srt/mem_cache/storage/eic/test_unit.py +115 -0
  253. sglang/srt/mem_cache/storage/hf3fs/hf3fs_client.py +164 -0
  254. sglang/srt/mem_cache/storage/hf3fs/{client_hf3fs.py → hf3fs_usrbio_client.py} +5 -1
  255. sglang/srt/mem_cache/storage/hf3fs/storage_hf3fs.py +259 -62
  256. sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py +284 -0
  257. sglang/srt/mem_cache/storage/lmcache/unit_test.py +121 -0
  258. sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py +166 -17
  259. sglang/srt/mem_cache/swa_radix_cache.py +25 -36
  260. sglang/srt/metrics/collector.py +511 -132
  261. sglang/srt/metrics/func_timer.py +2 -7
  262. sglang/srt/metrics/startup_func_log_and_timer.py +150 -0
  263. sglang/srt/metrics/utils.py +8 -1
  264. sglang/srt/model_executor/cpu_graph_runner.py +640 -0
  265. sglang/srt/model_executor/cuda_graph_runner.py +52 -37
  266. sglang/srt/model_executor/forward_batch_info.py +82 -40
  267. sglang/srt/model_executor/model_runner.py +432 -157
  268. sglang/srt/model_executor/npu_graph_runner.py +12 -5
  269. sglang/srt/model_loader/__init__.py +9 -3
  270. sglang/srt/model_loader/loader.py +133 -5
  271. sglang/srt/model_loader/remote_instance_weight_loader_utils.py +69 -0
  272. sglang/srt/model_loader/weight_utils.py +158 -3
  273. sglang/srt/models/apertus.py +686 -0
  274. sglang/srt/models/bailing_moe.py +820 -217
  275. sglang/srt/models/bailing_moe_nextn.py +168 -0
  276. sglang/srt/models/deepseek_nextn.py +6 -1
  277. sglang/srt/models/deepseek_v2.py +607 -130
  278. sglang/srt/models/dots_ocr.py +173 -0
  279. sglang/srt/models/dots_vlm.py +174 -0
  280. sglang/srt/models/dots_vlm_vit.py +337 -0
  281. sglang/srt/models/ernie4.py +1 -1
  282. sglang/srt/models/falcon_h1.py +576 -0
  283. sglang/srt/models/gemma3_causal.py +0 -2
  284. sglang/srt/models/gemma3_mm.py +1 -1
  285. sglang/srt/models/gemma3n_mm.py +2 -2
  286. sglang/srt/models/glm4_moe.py +4 -4
  287. sglang/srt/models/glm4_moe_nextn.py +2 -2
  288. sglang/srt/models/glm4v.py +5 -3
  289. sglang/srt/models/glm4v_moe.py +4 -1
  290. sglang/srt/models/gpt_oss.py +8 -31
  291. sglang/srt/models/kimi_vl_moonvit.py +2 -2
  292. sglang/srt/models/llama.py +4 -0
  293. sglang/srt/models/llama4.py +9 -0
  294. sglang/srt/models/llama_eagle3.py +13 -0
  295. sglang/srt/models/longcat_flash.py +3 -3
  296. sglang/srt/models/longcat_flash_nextn.py +1 -1
  297. sglang/srt/models/mllama4.py +40 -4
  298. sglang/srt/models/opt.py +637 -0
  299. sglang/srt/models/qwen2_5_vl.py +29 -5
  300. sglang/srt/models/qwen2_audio.py +1 -1
  301. sglang/srt/models/qwen2_moe.py +120 -13
  302. sglang/srt/models/qwen2_vl.py +1 -1
  303. sglang/srt/models/qwen3.py +18 -3
  304. sglang/srt/models/qwen3_moe.py +32 -4
  305. sglang/srt/models/qwen3_next.py +1069 -0
  306. sglang/srt/models/qwen3_next_mtp.py +112 -0
  307. sglang/srt/models/qwen3_vl.py +787 -0
  308. sglang/srt/models/qwen3_vl_moe.py +471 -0
  309. sglang/srt/models/registry.py +15 -3
  310. sglang/srt/models/sarashina2_vision.py +269 -0
  311. sglang/srt/models/solar.py +505 -0
  312. sglang/srt/models/starcoder2.py +357 -0
  313. sglang/srt/models/step3_vl.py +1 -1
  314. sglang/srt/models/torch_native_llama.py +9 -2
  315. sglang/srt/models/utils.py +51 -0
  316. sglang/srt/multimodal/processors/base_processor.py +15 -7
  317. sglang/srt/multimodal/processors/dots_vlm.py +98 -0
  318. sglang/srt/multimodal/processors/glm4v.py +9 -9
  319. sglang/srt/multimodal/processors/internvl.py +153 -129
  320. sglang/srt/multimodal/processors/qwen_vl.py +23 -6
  321. sglang/srt/multimodal/processors/sarashina2_vision.py +81 -0
  322. sglang/srt/offloader.py +27 -3
  323. sglang/srt/parser/jinja_template_utils.py +6 -0
  324. sglang/srt/sampling/sampling_batch_info.py +38 -17
  325. sglang/srt/sampling/sampling_params.py +7 -0
  326. sglang/srt/server_args.py +966 -267
  327. sglang/srt/server_args_config_parser.py +146 -0
  328. sglang/srt/single_batch_overlap.py +151 -0
  329. sglang/srt/speculative/cpp_ngram/ngram.cpp +374 -0
  330. sglang/srt/speculative/cpp_ngram/ngram.h +110 -0
  331. sglang/srt/speculative/cpp_ngram/ngram_cache.py +138 -0
  332. sglang/srt/speculative/cpp_ngram/ngram_cache_binding.cpp +43 -0
  333. sglang/srt/speculative/cpp_ngram/param.h +125 -0
  334. sglang/srt/speculative/cpp_ngram/queue.h +71 -0
  335. sglang/srt/speculative/eagle_draft_cuda_graph_runner.py +7 -1
  336. sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py +13 -2
  337. sglang/srt/speculative/{eagle_utils.py → eagle_info.py} +207 -757
  338. sglang/srt/speculative/eagle_worker.py +99 -28
  339. sglang/srt/speculative/ngram_utils.py +428 -0
  340. sglang/srt/speculative/ngram_worker.py +245 -0
  341. sglang/srt/speculative/spec_info.py +52 -0
  342. sglang/srt/speculative/spec_utils.py +606 -0
  343. sglang/srt/speculative/standalone_worker.py +109 -0
  344. sglang/srt/torch_memory_saver_adapter.py +5 -7
  345. sglang/srt/tracing/trace.py +578 -0
  346. sglang/srt/two_batch_overlap.py +8 -5
  347. sglang/srt/utils/__init__.py +2 -0
  348. sglang/srt/{utils.py → utils/common.py} +433 -77
  349. sglang/srt/{hf_transformers_utils.py → utils/hf_transformers_utils.py} +53 -5
  350. sglang/srt/{patch_torch.py → utils/patch_torch.py} +8 -0
  351. sglang/srt/utils/rpd_utils.py +452 -0
  352. sglang/srt/utils/slow_rank_detector.py +71 -0
  353. sglang/srt/warmup.py +8 -4
  354. sglang/srt/weight_sync/utils.py +2 -2
  355. sglang/test/attention/test_trtllm_mla_backend.py +169 -5
  356. sglang/test/get_logits_ut.py +57 -0
  357. sglang/test/run_eval.py +79 -11
  358. sglang/test/runners.py +5 -1
  359. sglang/test/simple_eval_common.py +5 -2
  360. sglang/test/simple_eval_mmmu_vlm.py +441 -0
  361. sglang/test/test_block_fp8.py +2 -2
  362. sglang/test/test_cutlass_moe.py +24 -6
  363. sglang/test/test_deterministic.py +297 -0
  364. sglang/test/test_disaggregation_utils.py +77 -0
  365. sglang/test/test_fp4_moe.py +370 -1
  366. sglang/test/test_programs.py +1 -1
  367. sglang/test/test_utils.py +383 -5
  368. sglang/utils.py +21 -1
  369. sglang/version.py +1 -1
  370. {sglang-0.5.2rc2.dist-info → sglang-0.5.3.dist-info}/METADATA +69 -124
  371. {sglang-0.5.2rc2.dist-info → sglang-0.5.3.dist-info}/RECORD +375 -245
  372. sglang/srt/disaggregation/launch_lb.py +0 -118
  373. sglang/srt/mem_cache/lora_radix_cache.py +0 -421
  374. /sglang/srt/{poll_based_barrier.py → utils/poll_based_barrier.py} +0 -0
  375. {sglang-0.5.2rc2.dist-info → sglang-0.5.3.dist-info}/WHEEL +0 -0
  376. {sglang-0.5.2rc2.dist-info → sglang-0.5.3.dist-info}/licenses/LICENSE +0 -0
  377. {sglang-0.5.2rc2.dist-info → sglang-0.5.3.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 {}